Explicitly wrap np.ptp for numpy 2.0 compatibility

brian-team · Sep 26, 2023 · decab78 · decab78
1 parent 76545b0
commit decab78
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/brian2/tests/test_units.py b/brian2/tests/test_units.py
@@ -1005,6 +1005,8 @@ def test_numpy_functions_same_dimensions():
     values = [np.array([1, 2]), np.ones((3, 3))]
     units = [volt, second, siemens, mV, kHz]
 
+    from brian2.units.unitsafefunctions import ptp
+
     # numpy functions
     keep_dim_funcs = [
         np.abs,
@@ -1013,7 +1015,7 @@ def test_numpy_functions_same_dimensions():
         np.mean,
         np.min,
         np.negative,
-        np.ptp,
+        ptp,
         np.round,
         np.squeeze,
         np.std,

diff --git a/brian2/units/unitsafefunctions.py b/brian2/units/unitsafefunctions.py
@@ -13,6 +13,7 @@
     fail_for_dimension_mismatch,
     is_dimensionless,
     wrap_function_dimensionless,
+    wrap_function_keep_dimensions,
     wrap_function_remove_dimensions,
 )
 
@@ -44,6 +45,7 @@
     "zeros_like",
     "arange",
     "linspace",
+    "ptp",
 ]
 
 
@@ -93,6 +95,8 @@ def where(condition, *args, **kwds):  # pylint: disable=C0111
 expm1 = wrap_function_dimensionless(np.expm1)
 log1p = wrap_function_dimensionless(np.log1p)
 
+ptp = wrap_function_keep_dimensions(np.ptp)
+
 
 @check_units(x=1, result=1)
 def exprel(x):