Merge pull request #583 from ComputationalCryoEM/deprecate_ensure

Deprecate ensure

Author: garrettwrong

src/aspire/abinitio/commonline_sync.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from aspire.abinitio import CLOrient3D
-from aspire.utils import Rotation, ensure
+from aspire.utils import Rotation
 from aspire.utils.matlab_compat import stable_eigsh
 
 logger = logging.getLogger(__name__)
@@ -53,8 +53,8 @@ def estimate_rotations(self):
 
         S = self.syncmatrix
         sz = S.shape
-        ensure(sz[0] == sz[1], "syncmatrix must be a square matrix.")
-        ensure(sz[0] % 2 == 0, "syncmatrix must be a square matrix of size 2Kx2K.")
+        assert sz[0] == sz[1], "syncmatrix must be a square matrix."
+        assert sz[0] % 2 == 0, "syncmatrix must be a square matrix of size 2Kx2K."
 
         n_img = sz[0] // 2
 
@@ -153,7 +153,7 @@ def syncmatrix_vote(self):
         sz = clmatrix.shape
         n_theta = self.n_theta
 
-        ensure(sz[0] == sz[1], "clmatrix must be a square matrix.")
+        assert sz[0] == sz[1], "clmatrix must be a square matrix."
 
         n_img = sz[0]
         S = np.eye(2 * n_img, dtype=self.dtype).reshape(n_img, 2, n_img, 2)

src/aspire/basis/basis.py

@@ -5,7 +5,7 @@
 
 from aspire.basis.basis_utils import num_besselj_zeros
 from aspire.image import Image
-from aspire.utils import ensure, mdim_mat_fun_conj
+from aspire.utils import mdim_mat_fun_conj
 from aspire.utils.matlab_compat import m_reshape
 from aspire.volume import Volume
 
@@ -185,10 +185,9 @@ def expand(self, x):
 
         x = x.reshape((-1, *self.sz))
 
-        ensure(
-            x.shape[-self.ndim :] == self.sz,
-            f"Last {self.ndim} dimensions of x must match {self.sz}.",
-        )
+        assert (
+            x.shape[-self.ndim :] == self.sz
+        ), f"Last {self.ndim} dimensions of x must match {self.sz}."
 
         operator = LinearOperator(
             shape=(self.count, self.count),

src/aspire/basis/basis_utils.py

@@ -10,7 +10,6 @@
 from numpy.polynomial.legendre import leggauss
 from scipy.special import jn, jv, sph_harm
 
-from aspire.utils import ensure
 from aspire.utils.coor_trans import grid_2d, grid_3d
 
 logger = logging.getLogger(__name__)
@@ -158,8 +157,8 @@ def real_sph_harmonic(j, m, theta, phi):
 
 
 def besselj_zeros(nu, k):
-    ensure(k >= 3, "k must be >= 3")
-    ensure(0 <= nu <= 1e7, "nu must be between 0 and 1e7")
+    assert k >= 3, "k must be >= 3"
+    assert 0 <= nu <= 1e7, "nu must be between 0 and 1e7"
 
     z = np.zeros(k)
 
@@ -198,10 +197,9 @@ def besselj_zeros(nu, k):
         z[n : n + j] = besselj_newton(nu, z0)
 
         # Check to see that the sequence of zeros makes sense
-        ensure(
-            check_besselj_zeros(nu, z[n - 2 : n + j]),
-            "Unable to properly estimate Bessel function zeros.",
-        )
+        assert check_besselj_zeros(
+            nu, z[n - 2 : n + j]
+        ), "Unable to properly estimate Bessel function zeros."
 
         # Check how far off we are
         err = (z[n : n + j] - z0) / np.diff(z[n - 1 : n + j])
@@ -236,10 +234,11 @@ def unique_coords_nd(N, ndim, shifted=False, normalized=True, dtype=np.float32):
     :param normalized: normalize the grid or not.
     :return: The unique polar coordinates in 2D or 3D
     """
-    ensure(
-        ndim in (2, 3), "Only two- or three-dimensional basis functions are supported."
-    )
-    ensure(N > 0, "Number of grid points should be greater than 0.")
+    assert ndim in (
+        2,
+        3,
+    ), "Only two- or three-dimensional basis functions are supported."
+    assert N > 0, "Number of grid points should be greater than 0."
 
     if ndim == 2:
         grid = grid_2d(

src/aspire/basis/fb_2d.py

@@ -6,7 +6,7 @@
 from aspire.basis import SteerableBasis2D
 from aspire.basis.basis_utils import unique_coords_nd
 from aspire.image import Image
-from aspire.utils import complex_type, ensure, real_type, roll_dim, unroll_dim
+from aspire.utils import complex_type, real_type, roll_dim, unroll_dim
 from aspire.utils.matlab_compat import m_flatten, m_reshape
 
 logger = logging.getLogger(__name__)
@@ -37,8 +37,8 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
         """
 
         ndim = len(size)
-        ensure(ndim == 2, "Only two-dimensional basis functions are supported.")
-        ensure(len(set(size)) == 1, "Only square domains are supported.")
+        assert ndim == 2, "Only two-dimensional basis functions are supported."
+        assert len(set(size)) == 1, "Only square domains are supported."
         super().__init__(size, ell_max, dtype=dtype)
 
     def _build(self):

src/aspire/basis/fb_3d.py

@@ -4,7 +4,7 @@
 
 from aspire.basis import Basis
 from aspire.basis.basis_utils import real_sph_harmonic, sph_bessel, unique_coords_nd
-from aspire.utils import ensure, roll_dim, unroll_dim
+from aspire.utils import roll_dim, unroll_dim
 from aspire.utils.matlab_compat import m_flatten, m_reshape
 
 logger = logging.getLogger(__name__)
@@ -30,8 +30,8 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
             below the Nyquist frequency (default Inf).
         """
         ndim = len(size)
-        ensure(ndim == 3, "Only three-dimensional basis functions are supported.")
-        ensure(len(set(size)) == 1, "Only cubic domains are supported.")
+        assert ndim == 3, "Only three-dimensional basis functions are supported."
+        assert len(set(size)) == 1, "Only cubic domains are supported."
 
         super().__init__(size, ell_max, dtype=dtype)
 

src/aspire/basis/polar_2d.py

@@ -5,7 +5,7 @@
 from aspire.basis import Basis
 from aspire.image import Image
 from aspire.nufft import anufft, nufft
-from aspire.utils import ensure, real_type
+from aspire.utils import real_type
 
 logger = logging.getLogger(__name__)
 
@@ -26,8 +26,8 @@ def __init__(self, size, nrad=None, ntheta=None, dtype=np.float32):
         """
 
         ndim = len(size)
-        ensure(ndim == 2, "Only two-dimensional grids are supported.")
-        ensure(len(set(size)) == 1, "Only square domains are supported.")
+        assert ndim == 2, "Only two-dimensional grids are supported."
+        assert len(set(size)) == 1, "Only square domains are supported."
 
         self.nrad = nrad
         self.ntheta = ntheta

src/aspire/covariance/covar.py

@@ -14,7 +14,6 @@
 from aspire.operators import evaluate_src_filters_on_grid
 from aspire.reconstruction import Estimator, FourierKernel, MeanEstimator
 from aspire.utils import (
-    ensure,
     make_symmat,
     symmat_to_vec_iso,
     vec_to_symmat_iso,
@@ -197,10 +196,12 @@ def _shrink(self, covar_b_coeff, noise_variance, method=None):
         :param method: One of None/'frobenius_norm'/'operator_norm'/'soft_threshold'
         :return: Shrunk covariance matrix
         """
-        ensure(
-            method in (None, "frobenius_norm", "operator_norm", "soft_threshold"),
-            "Unsupported shrink method",
-        )
+        assert method in (
+            None,
+            "frobenius_norm",
+            "operator_norm",
+            "soft_threshold",
+        ), "Unsupported shrink method"
 
         An = self.basis.mat_evaluate_t(self.mean_kernel.toeplitz())
         if method is None:

src/aspire/covariance/covar2d.py

@@ -6,7 +6,7 @@
 
 from aspire.operators import BlkDiagMatrix, RadialCTFFilter
 from aspire.optimization import conj_grad, fill_struct
-from aspire.utils import ensure, make_symmat
+from aspire.utils import make_symmat
 from aspire.utils.matlab_compat import m_reshape
 
 logger = logging.getLogger(__name__)
@@ -22,10 +22,11 @@ def shrink_covar(covar, noise_var, gamma, shrinker="frobenius_norm"):
     :return: The shrinked covariance matrix
     """
 
-    ensure(
-        shrinker in ("frobenius_norm", "operator_norm", "soft_threshold"),
-        "Unsupported shrink method",
-    )
+    assert shrinker in (
+        "frobenius_norm",
+        "operator_norm",
+        "soft_threshold",
+    ), "Unsupported shrink method"
 
     lambs, eig_vec = eig(make_symmat(covar))
 
@@ -96,7 +97,7 @@ def __init__(self, basis):
         """
         self.basis = basis
         self.dtype = self.basis.dtype
-        ensure(basis.ndim == 2, "Only two-dimensional basis functions are needed.")
+        assert basis.ndim == 2, "Only two-dimensional basis functions are needed."
 
     def _get_mean(self, coeffs):
         """
@@ -327,7 +328,7 @@ def identity(x):
 
         def precond_fun(S, x):
             p = np.size(S, 0)
-            ensure(np.size(x) == p * p, "The sizes of S and x are not consistent.")
+            assert np.size(x) == p * p, "The sizes of S and x are not consistent."
             x = m_reshape(x, (p, p))
             y = S @ x @ S
             y = m_reshape(y, (p**2,))
@@ -632,7 +633,7 @@ def _solve_covar(self, A_covar, b_covar, M, covar_est_opt):
 
         def precond_fun(S, x):
             p = np.size(S, 0)
-            ensure(np.size(x) == p * p, "The sizes of S and x are not consistent.")
+            assert np.size(x) == p * p, "The sizes of S and x are not consistent."
             x = m_reshape(x, (p, p))
             y = S @ x @ S
             y = m_reshape(y, (p**2,))

src/aspire/image/image.py

@@ -9,7 +9,6 @@
 import aspire.volume
 from aspire.nufft import anufft
 from aspire.numeric import fft, xp
-from aspire.utils import ensure
 from aspire.utils.coor_trans import grid_2d
 from aspire.utils.matrix import anorm
 
@@ -144,7 +143,7 @@ def __init__(self, data, dtype=None):
         self.n_images = self.shape[0]
         self.res = self.shape[1]
 
-        ensure(data.shape[1] == data.shape[2], "Only square ndarrays are supported.")
+        assert data.shape[1] == data.shape[2], "Only square ndarrays are supported."
 
     def __getitem__(self, item):
         return self.data[item]
@@ -279,12 +278,11 @@ def _im_translate(self, shifts):
             shifts = shifts[np.newaxis, :]
         n_shifts = shifts.shape[0]
 
-        ensure(shifts.shape[-1] == 2, "shifts must be nx2")
+        assert shifts.shape[-1] == 2, "shifts must be nx2"
 
-        ensure(
-            n_shifts == 1 or n_shifts == self.n_images,
-            "number of shifts must be 1 or match the number of images",
-        )
+        assert (
+            n_shifts == 1 or n_shifts == self.n_images
+        ), "number of shifts must be 1 or match the number of images"
         # Cast shifts to this instance's internal dtype
         shifts = shifts.astype(self.dtype)
 
@@ -330,10 +328,9 @@ def backproject(self, rot_matrices):
 
         L = self.res
 
-        ensure(
-            self.n_images == rot_matrices.shape[0],
-            "Number of rotation matrices must match the number of images",
-        )
+        assert (
+            self.n_images == rot_matrices.shape[0]
+        ), "Number of rotation matrices must match the number of images"
 
         # TODO: rotated_grids might as well give us correctly shaped array in the first place
         pts_rot = aspire.volume.rotated_grids(L, rot_matrices)

src/aspire/image/preprocess.py

@@ -5,7 +5,6 @@
 from scipy.special import erf
 
 from aspire.numeric import fft, xp
-from aspire.utils import ensure
 
 logger = logging.getLogger(__name__)
 
@@ -129,10 +128,9 @@ def downsample(insamples, szout, mask=None):
     :return: An array consists of the blurred and downsampled objects.
     """
 
-    ensure(
-        insamples.ndim - 1 == np.size(szout),
-        "The number of downsampling dimensions is not the same as that of objects.",
-    )
+    assert insamples.ndim - 1 == np.size(
+        szout
+    ), "The number of downsampling dimensions is not the same as that of objects."
 
     L_in = insamples.shape[1]
     L_out = szout[0]