Support Numpy 1 and Numpy 2

pyproject.toml

@@ -36,7 +36,7 @@ dependencies = [
     "joblib",
     "matplotlib >= 3.2.0",
     "mrcfile",
-    "numpy>=1.21.5, <2.0.0",
+    "numpy>=1.21.5",
     "packaging",
     "pooch>=1.7.0",
     "pillow",

src/aspire/__init__.py

@@ -2,11 +2,14 @@
 import logging.config
 import os
 import pkgutil
+import warnings
 from datetime import datetime
 from pathlib import Path
 
 import confuse
+import numpy as np
 import pooch
+from packaging.version import parse as parse_version
 
 import aspire
 from aspire.exceptions import handle_exception
@@ -84,3 +87,16 @@ def __getattr__(attr):
         return importlib.import_module(f"aspire.{attr}")
     else:
         raise AttributeError(f"module `{__name__}` has no attribute `{attr}`.")
+
+
+if parse_version(np.version.version) >= parse_version("2.0.0"):
+    # ImportWarnings are generally filtered, but raise this one for now.
+    with warnings.catch_warnings():
+        warnings.simplefilter("default")
+        warnings.warn(
+            "ASPIRE's Numpy 2 support is in beta.  If you experience a runtime"
+            " crash relating to mismatched dtypes or a Numpy call please try"
+            ' `pip install "numpy<2"` and report to ASPIRE developers.',
+            ImportWarning,
+            stacklevel=1,
+        )

src/aspire/abinitio/commonline_sync3n.py

@@ -266,7 +266,7 @@ def _sync3n_S_to_rot(self, S, W=None, n_eigs=4):
         # Enforce we are returning actual rotations
         rotations = nearest_rotations(rotations, allow_reflection=True)
 
-        return rotations.astype(self.dtype)
+        return rotations.astype(self.dtype, copy=False)
 
     def _construct_sync3n_matrix(self, Rij):
         """

src/aspire/basis/basis.py

@@ -373,9 +373,9 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
         :param size: The size of the vectors for which to define the basis.
             Currently only square images and cubic volumes are supported.
         :param ell_max: The maximum order ell of the basis elements. If no input
-            (= None), it will be set to np.Inf and the basis includes all
+            (= None), it will be set to np.inf and the basis includes all
             ell such that the resulting basis vectors are concentrated
-            below the Nyquist frequency (default Inf).
+            below the Nyquist frequency (default inf).
         """
         if ell_max is None:
             ell_max = np.inf

src/aspire/basis/fb_2d.py

@@ -31,9 +31,9 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
             May be a 2-tuple or an integer, in which case a square basis is assumed.
             Currently only square images are supported.
         :ell_max: The maximum order ell of the basis elements. If no input
-            (= None), it will be set to np.Inf and the basis includes all
+            (= None), it will be set to np.inf and the basis includes all
             ell such that the resulting basis vectors are concentrated
-            below the Nyquist frequency (default Inf).
+            below the Nyquist frequency (default inf).
         """
 
         if isinstance(size, int):

src/aspire/basis/fb_3d.py

@@ -26,9 +26,9 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
             May be a 3-tuple or an integer, in which case a cubic basis is assumed.
             Currently only cubic images are supported.
         :param ell_max: The maximum order ell of the basis elements. If no input
-            (= None), it will be set to np.Inf and the basis includes all
+            (= None), it will be set to np.inf and the basis includes all
             ell such that the resulting basis vectors are concentrated
-            below the Nyquist frequency (default Inf).
+            below the Nyquist frequency (default inf).
         """
         if isinstance(size, int):
             size = (size, size, size)

src/aspire/basis/fle_2d.py

@@ -314,7 +314,7 @@ def _compute_nufft_points(self):
         )
         grid_xy[0] = xp.cos(phi)  # x
         grid_xy[1] = xp.sin(phi)  # y
-        grid_xy = grid_xy * nodes * h
+        grid_xy[:] = grid_xy * nodes * h
         self.grid_xy = grid_xy.reshape(2, -1)
 
     def _build_interpolation_matrix(self):
@@ -358,7 +358,7 @@ def _lap_eig_disk(self):
         num_ells = 1 + 2 * max_ell
         self._ells = np.zeros((num_ells, max_k), dtype=int)
         self._ks = np.zeros((num_ells, max_k), dtype=int)
-        self.bessel_zeros = np.ones((num_ells, max_k), dtype=np.float64) * np.Inf
+        self.bessel_zeros = np.ones((num_ells, max_k), dtype=np.float64) * np.inf
 
         # keep track of which order Bessel function we're on
         self._ells[0, :] = 0

src/aspire/basis/fpswf_2d.py

@@ -87,7 +87,7 @@ def _precomp(self):
         self.quad_rule_radial_wts = e
         self.num_angular_pts = f
 
-        us_fft_pts = np.row_stack((self.quad_rule_pts_x, self.quad_rule_pts_y))
+        us_fft_pts = np.vstack((self.quad_rule_pts_x, self.quad_rule_pts_y))
         us_fft_pts = self.bandlimit / self.rcut * us_fft_pts
         (
             blk_r,

src/aspire/covariance/covar2d.py

@@ -607,7 +607,7 @@ def _calc_op(self):
             ctf_basis_k_sq = ctf_basis_k_t @ ctf_basis_k
             A_mean_k = weight * ctf_basis_k_sq
             A_mean += A_mean_k
-            A_covar_k = np.sqrt(weight) * ctf_basis_k_sq
+            A_covar_k = np.sqrt(weight).astype(self.dtype) * ctf_basis_k_sq
             A_covar[k] = A_covar_k
 
             M_covar += A_covar_k

src/aspire/image/image.py

@@ -72,7 +72,7 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True):
 
     # Apply mask images and calculate mean and std values of background
     imgs_masked = imgs * mask
-    denominator = np.sum(mask)
+    denominator = np.count_nonzero(mask)  # scalar int
     first_moment = np.sum(imgs_masked, axis=(1, 2)) / denominator
     second_moment = np.sum(imgs_masked**2, axis=(1, 2)) / denominator
     mean = first_moment.reshape(-1, 1, 1)

src/aspire/image/xform.py

@@ -145,7 +145,7 @@ def __init__(self, factor):
 
     def _forward(self, im, indices):
         if self.multipliers.size == 1:  # if we have a scalar multiplier
-            im_new = im * self.multipliers
+            im_new = im * self.multipliers.astype(im.dtype)
         else:
             im_new = im * self.multipliers[indices]
 

src/aspire/nufft/finufft.py

@@ -99,7 +99,7 @@ def transform(self, signal):
         ), f"Signal frame to be transformed must have shape {self.sz}"
 
         # FINUFFT was designed for a complex input array
-        signal = np.array(signal, copy=False, dtype=self.complex_dtype, order="C")
+        signal = np.asarray(signal, dtype=self.complex_dtype, order="C")
 
         result = self._transform_plan.execute(signal)
 
@@ -130,7 +130,7 @@ def adjoint(self, signal):
                 signal = signal.reshape(self.num_pts)
 
         # FINUFFT was designed for a complex input array
-        signal = np.array(signal, copy=False, dtype=self.complex_dtype, order="C")
+        signal = np.asarray(signal, dtype=self.complex_dtype, order="C")
 
         result = self._adjoint_plan.execute(signal)
 

src/aspire/reconstruction/kernel.py

@@ -32,7 +32,7 @@
             to be able to use it within optimization loops. This operator allows one to use the FourierKernel object
             with the underlying 'kernel' attribute tweaked with a regularization parameter.
         """
-        new_kernel = self.kernel + delta
+        new_kernel = self.kernel + float(delta)
         return FourierKernel(new_kernel)
 
     def circularize(self):
@@ -191,7 +191,7 @@
 
     def circularize(self):
         _L = self.M // 2
-        xx = np.empty((self.r, self.r, _L, _L, _L))
+        xx = np.empty((self.r, self.r, _L, _L, _L), self.dtype)
         for k in range(self.r):
             for j in range(self.r):
                 xx[k, j] = FourierKernel(self.kermat[k, j]).circularize().real

src/aspire/reconstruction/mean.py

@@ -139,7 +139,7 @@ def _compute_kernel(self):
                     if j != k:
                         kernel[j, k] += batch_kernel
 
-        kermat_f = np.zeros((self.r, self.r, _2L, _2L, _2L))
+        kermat_f = np.zeros((self.r, self.r, _2L, _2L, _2L), dtype=self.dtype)
         logger.info("Computing non-centered Fourier Transform Kernel Mat")
         for k in range(self.r):
             for j in range(self.r):
@@ -184,10 +184,12 @@ def src_backward(self):
                     i,
                     self.weights[:, k],
                     symmetry_group=symmetry_group,
-                ) / (self.src.n * sym_order)
+                ) / float(self.src.n * sym_order)
                 vol_rhs[k] += batch_vol_rhs.astype(self.dtype)
 
-        res = np.sqrt(self.src.n * sym_order) * self.basis.evaluate_t(vol_rhs)
+        res = np.sqrt(self.src.n * sym_order, dtype=self.dtype) * self.basis.evaluate_t(
+            vol_rhs
+        )
         logger.info(f"Determined weighted adjoint mappings. Shape = {res.shape}")
 
         return res
@@ -262,7 +264,7 @@ def cb(xk):
                 f"Conjugate gradient unable to converge after {info} iterations."
             )
 
-        return x.reshape(self.r, self.basis.count)
+        return x.reshape(self.r, self.basis.count).astype(self.dtype, copy=False)
 
     def apply_kernel(self, vol_coef, kernel=None):
         """
@@ -283,8 +285,7 @@ def apply_kernel(self, vol_coef, kernel=None):
         vols_out = Volume(
             np.zeros((self.r, self.src.L, self.src.L, self.src.L), dtype=self.dtype)
         )
-
-        vol = Coef(self.basis, vol_coef).evaluate()
+        vol = Coef(self.basis, vol_coef.astype(self.dtype, copy=False)).evaluate()
 
         for k in range(self.r):
             for j in range(self.r):

src/aspire/utils/matrix.py

@@ -1,5 +1,5 @@
 """
-Utilties for arrays/n-dimensional matrices.
+Utilities for arrays/n-dimensional matrices.
 """
 
 import numpy as np
@@ -493,7 +493,7 @@ def fix_signs(u):
     # Now we only care about the sign +1/-1.
     #  The following corrects for any numerical division noise,
     #  and also remaps 0 to +1.
-    signs = np.sign(signs * 2 + 1)
+    signs = np.sign(signs.real * 2 + 1)
 
     # Apply signs elementwise to matrix
     return u * signs

src/aspire/utils/misc.py

@@ -1,5 +1,5 @@
 """
-Miscellaneous Utilities that have no better place (yet).
+Miscellaneous utilities that have no better place (yet).
 """
 
 import hashlib

tests/test_mrc.py

@@ -81,15 +81,13 @@ def testUpdate(self):
             with mrcfile.new_mmap(
                 files[1], shape=(self.n, self.n), mrc_mode=2, overwrite=True
             ) as mrc:
-                mrc.set_data(self.a.astype(np.float32))
+                mrc.set_data(self.a)
+                self.stats.update_header(mrc)
                 mrc.header.time = epoch
                 mrc.header.label[0] = label
 
             # Our homebrew and mrcfile files should now match to the bit.
             comparison = sha256sum(files[0]) == sha256sum(files[1])
-            # Expected hash:
-            # 71355fa0bcd5b989ff88166962ea5d2b78ea032933bd6fda41fbdcc1c6d1a009
             logging.debug(f"sha256(file0): {sha256sum(files[0])}")
             logging.debug(f"sha256(file1): {sha256sum(files[1])}")
-
             self.assertTrue(comparison)