Merge pull request python-quantities#235 from zm711/numpy-2.0-compat

apdavison · web-flow · commit bb96cbb826eb · 2024-08-27T10:22:37.000+02:00
Add numpy 2.0 compatibility
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -34,6 +34,9 @@ jobs:
           - python-version: "3.12"
             numpy-version: "1.26"
             os: ubuntu-latest
+          - python-version: "3.12"
+            numpy-version: "2.0"
+            os: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
 
@@ -128,4 +131,4 @@ jobs:
 
       - name: Check type information
         run: |
-          mypy quantities
+          mypy quantities
diff --git a/quantities/__init__.py b/quantities/__init__.py
@@ -265,6 +265,8 @@
 
 """
 
+class QuantitiesDeprecationWarning(DeprecationWarning):
+    pass
 
 from ._version import __version__
 
diff --git a/quantities/dimensionality.py b/quantities/dimensionality.py
@@ -9,6 +9,8 @@
 from .registry import unit_registry
 from .decorators import memoize
 
+_np_version = tuple(map(int, np.__version__.split('.')))
+
 def assert_isinstance(obj, types):
     try:
         assert isinstance(obj, types)
@@ -329,10 +331,11 @@ def _d_copy(q1, out=None):
 
 def _d_clip(a1, a2, a3, q):
     return q.dimensionality
-try:
+
+if _np_version < (2, 0, 0):
     p_dict[np.core.umath.clip] = _d_clip
-except AttributeError:
-    pass # For compatibility with Numpy < 1.17 when clip wasn't a ufunc yet
+else:
+    p_dict[np.clip] = _d_clip
 
 def _d_sqrt(q1, out=None):
     return q1._dimensionality**0.5
diff --git a/quantities/quantity.py b/quantities/quantity.py
@@ -3,10 +3,11 @@
 
 import copy
 from functools import wraps
+import warnings
 
 import numpy as np
 
-from . import markup
+from . import markup, QuantitiesDeprecationWarning
 from .dimensionality import Dimensionality, p_dict
 from .registry import unit_registry
 from .decorators import with_doc
@@ -114,15 +115,19 @@ class Quantity(np.ndarray):
     # TODO: what is an appropriate value?
     __array_priority__ = 21
 
-    def __new__(cls, data, units='', dtype=None, copy=True):
+    def __new__(cls, data, units='', dtype=None, copy=None):
+        if copy is not None:
+            warnings.warn(("The 'copy' argument in Quantity is deprecated and will be removed in the future. "
+                           "The argument has no effect since quantities-0.16.0 (to aid numpy-2.0 support)."),
+                           QuantitiesDeprecationWarning, stacklevel=2)
         if isinstance(data, Quantity):
             if units:
                 data = data.rescale(units)
             if isinstance(data, unit_registry['UnitQuantity']):
                 return 1*data
-            return np.array(data, dtype=dtype, copy=copy, subok=True).view(cls)
+            return np.asanyarray(data, dtype=dtype).view(cls)
 
-        ret = np.array(data, dtype=dtype, copy=copy).view(cls)
+        ret = np.asarray(data, dtype=dtype).view(cls)
         ret._dimensionality.update(validate_dimensionality(units))
         return ret
 
@@ -210,15 +215,17 @@ def rescale(self, units=None, dtype=None):
             dtype = self.dtype
         if self.dimensionality == to_dims:
             return self.astype(dtype)
-        to_u = Quantity(1.0, to_dims)
-        from_u = Quantity(1.0, self.dimensionality)
+        to_u = Quantity(1.0, to_dims, dtype=dtype)
+        from_u = Quantity(1.0, self.dimensionality, dtype=dtype)
         try:
             cf = get_conversion_factor(from_u, to_u)
         except AssertionError:
             raise ValueError(
                 'Unable to convert between units of "%s" and "%s"'
                 %(from_u._dimensionality, to_u._dimensionality)
             )
+        if np.dtype(dtype).kind in 'fc':
+            cf = np.array(cf, dtype=dtype)
         new_magnitude = cf*self.magnitude
         dtype = np.result_type(dtype, new_magnitude)
         return Quantity(new_magnitude, to_u, dtype=dtype)
@@ -272,7 +279,7 @@ def __array_prepare__(self, obj, context=None):
         uf, objs, huh = context
         if uf.__name__.startswith('is'):
             return obj
-        #print self, obj, res, uf, objs
+
         try:
             res._dimensionality = p_dict[uf](*objs)
         except KeyError:
@@ -283,11 +290,21 @@ def __array_prepare__(self, obj, context=None):
                 )
         return res
 
-    def __array_wrap__(self, obj, context=None):
-        if not isinstance(obj, Quantity):
-            # backwards compatibility with numpy-1.3
-            obj = self.__array_prepare__(obj, context)
-        return obj
+    def __array_wrap__(self, obj, context=None, return_scalar=False):
+        _np_version = tuple(map(int, np.__version__.split('.')))
+        # For NumPy < 2.0 we do old behavior
+        if _np_version < (2, 0, 0):
+            if not isinstance(obj, Quantity):
+                return self.__array_prepare__(obj, context)
+            else:
+                return obj
+        # For NumPy > 2.0 we either do the prepare or the wrap
+        else:
+            if not isinstance(obj, Quantity):
+                return self.__array_prepare__(obj, context)
+            else:
+                return super().__array_wrap__(obj, context, return_scalar)
+
 
     @with_doc(np.ndarray.__add__)
     @scale_other_units
@@ -476,7 +493,7 @@ def sum(self, axis=None, dtype=None, out=None):
         ret = self.magnitude.sum(axis, dtype, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -487,8 +504,7 @@ def nansum(self, axis=None, dtype=None, out=None):
         import numpy as np
         return Quantity(
             np.nansum(self.magnitude, axis, dtype, out),
-            self.dimensionality,
-            copy=False
+            self.dimensionality
         )
 
     @with_doc(np.ndarray.fill)
@@ -523,7 +539,7 @@ def argsort(self, axis=-1, kind='quick', order=None):
     @with_doc(np.ndarray.searchsorted)
     def searchsorted(self,values, side='left'):
         if not isinstance (values, Quantity):
-            values = Quantity(values, copy=False)
+            values = Quantity(values)
 
         if values._dimensionality != self._dimensionality:
             raise ValueError("values does not have the same units as self")
@@ -539,7 +555,7 @@ def max(self, axis=None, out=None):
         ret = self.magnitude.max(axis, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -553,16 +569,15 @@ def argmax(self, axis=None, out=None):
     def nanmax(self, axis=None, out=None):
         return Quantity(
             np.nanmax(self.magnitude),
-            self.dimensionality,
-            copy=False
+            self.dimensionality
         )
 
     @with_doc(np.ndarray.min)
     def min(self, axis=None, out=None):
         ret = self.magnitude.min(axis, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -572,8 +587,7 @@ def min(self, axis=None, out=None):
     def nanmin(self, axis=None, out=None):
         return Quantity(
             np.nanmin(self.magnitude),
-            self.dimensionality,
-            copy=False
+            self.dimensionality
         )
 
     @with_doc(np.ndarray.argmin)
@@ -590,10 +604,10 @@ def nanargmax(self,axis=None, out=None):
 
     @with_doc(np.ndarray.ptp)
     def ptp(self, axis=None, out=None):
-        ret = self.magnitude.ptp(axis, None if out is None else out.magnitude)
+        ret = np.ptp(self.magnitude, axis, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -620,7 +634,7 @@ def clip(self, min=None, max=None, out=None):
         )
         dim = self.dimensionality
         if out is None:
-            return Quantity(clipped, dim, copy=False)
+            return Quantity(clipped, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -631,7 +645,7 @@ def round(self, decimals=0, out=None):
         ret = self.magnitude.round(decimals, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -642,7 +656,7 @@ def trace(self, offset=0, axis1=0, axis2=1, dtype=None, out=None):
         ret = self.magnitude.trace(offset, axis1, axis2, dtype, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -652,16 +666,15 @@ def trace(self, offset=0, axis1=0, axis2=1, dtype=None, out=None):
     def squeeze(self, axis=None):
         return Quantity(
             self.magnitude.squeeze(axis),
-            self.dimensionality,
-            copy=False
+            self.dimensionality
         )
 
     @with_doc(np.ndarray.mean)
     def mean(self, axis=None, dtype=None, out=None):
         ret = self.magnitude.mean(axis, dtype, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -672,15 +685,14 @@ def nanmean(self, axis=None, dtype=None, out=None):
         import numpy as np
         return Quantity(
             np.nanmean(self.magnitude, axis, dtype, out),
-            self.dimensionality,
-            copy=False)
+            self.dimensionality)
 
     @with_doc(np.ndarray.var)
     def var(self, axis=None, dtype=None, out=None, ddof=0):
         ret = self.magnitude.var(axis, dtype, out, ddof)
         dim = self._dimensionality**2
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -691,7 +703,7 @@ def std(self, axis=None, dtype=None, out=None, ddof=0):
         ret = self.magnitude.std(axis, dtype, out, ddof)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -701,8 +713,7 @@ def std(self, axis=None, dtype=None, out=None, ddof=0):
     def nanstd(self, axis=None, dtype=None, out=None, ddof=0):
         return Quantity(
             np.nanstd(self.magnitude, axis, dtype, out, ddof),
-            self._dimensionality,
-            copy=False
+            self._dimensionality
         )
 
     @with_doc(np.ndarray.prod)
@@ -715,7 +726,7 @@ def prod(self, axis=None, dtype=None, out=None):
         ret = self.magnitude.prod(axis, dtype, None if out is None else out.magnitude)
         dim = self._dimensionality**power
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -726,7 +737,7 @@ def cumsum(self, axis=None, dtype=None, out=None):
         ret = self.magnitude.cumsum(axis, dtype, None if out is None else out.magnitude)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if not isinstance(out, Quantity):
             raise TypeError("out parameter must be a Quantity")
         out._dimensionality = dim
@@ -743,7 +754,7 @@ def cumprod(self, axis=None, dtype=None, out=None):
         ret = self.magnitude.cumprod(axis, dtype, out)
         dim = self.dimensionality
         if out is None:
-            return Quantity(ret, dim, copy=False)
+            return Quantity(ret, dim)
         if isinstance(out, Quantity):
             out._dimensionality = dim
         return out
diff --git a/quantities/tests/common.py b/quantities/tests/common.py
@@ -19,7 +19,11 @@ def assertQuantityEqual(self, q1, q2, msg=None, delta=None):
         Make sure q1 and q2 are the same quantities to within the given
         precision.
         """
-        delta = 1e-5 if delta is None else delta
+        if delta is None:
+            # NumPy 2 introduced float16, so we base tolerance on machine epsilon
+            delta1 = np.finfo(q1.dtype).eps if isinstance(q1, np.ndarray) and q1.dtype.kind in 'fc' else 1e-15
+            delta2 = np.finfo(q2.dtype).eps if isinstance(q2, np.ndarray) and q2.dtype.kind in 'fc' else 1e-15
+            delta = max(delta1, delta2)**0.3
         msg = '' if msg is None else ' (%s)' % msg
 
         q1 = Quantity(q1)
@@ -28,6 +32,7 @@ def assertQuantityEqual(self, q1, q2, msg=None, delta=None):
             raise self.failureException(
                 f"Shape mismatch ({q1.shape} vs {q2.shape}){msg}"
                 )
+
         if not np.all(np.abs(q1.magnitude - q2.magnitude) < delta):
             raise self.failureException(
                 "Magnitudes differ by more than %g (%s vs %s)%s"
diff --git a/quantities/tests/test_arithmetic.py b/quantities/tests/test_arithmetic.py
@@ -48,32 +48,9 @@ def check(f, *args, **kwargs):
     return (new, )
 
 
-class iter_dtypes:
-
-    def __init__(self):
-        self._i = 1
-        self._typeDict = np.sctypeDict.copy()
-        self._typeDict[17] = int
-        self._typeDict[18] = long
-        self._typeDict[19] = float
-        self._typeDict[20] = complex
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        if self._i > 20:
-            raise StopIteration
-
-        i = self._i
-        self._i += 1
-        return self._typeDict[i]
-
-    def next(self):
-        return self.__next__()
-
 def get_dtypes():
-    return list(iter_dtypes())
+    numeric_dtypes = 'iufc'  # https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html
+    return [v for v in np.sctypeDict.values() if np.dtype(v).kind in numeric_dtypes] + [int, long, float, complex]
 
 
 class iter_types:
diff --git a/quantities/tests/test_methods.py b/quantities/tests/test_methods.py
diff --git a/quantities/tests/test_umath.py b/quantities/tests/test_umath.py
diff --git a/quantities/uncertainquantity.py b/quantities/uncertainquantity.py
