Add support for NumPy 2.0

bjodah · bjodah · commit ed430211844f · 2024-06-18T10:21:48.000+02:00
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/quantity.py b/quantities/quantity.py
@@ -120,9 +120,9 @@ def __new__(cls, data, units='', dtype=None, copy=True):
                 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 +210,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 +274,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:
@@ -590,7 +592,7 @@ 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)
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/uncertainquantity.py b/quantities/uncertainquantity.py
@@ -118,7 +118,7 @@ def __mul__(self, other):
             ru = (sru**2+oru**2)**0.5
             u = res.view(Quantity) * ru
         except AttributeError:
-            other = np.array(other, copy=False, subok=True)
+            other = np.asanyarray(other)
             u = (self.uncertainty**2*other**2)**0.5
 
         res._uncertainty = u
@@ -140,7 +140,7 @@ def __truediv__(self, other):
             ru = (sru**2+oru**2)**0.5
             u = res.view(Quantity) * ru
         except AttributeError:
-            other = np.array(other, copy=False, subok=True)
+            other = np.asanyarray(other)
             u = (self.uncertainty**2/other**2)**0.5
 
         res._uncertainty = u
