replace scipy calls with interop wrapper

Author: garrettwrong

src/aspire/basis/basis.py

@@ -1,9 +1,10 @@
 import logging
 
 import numpy as np
-from scipy.sparse.linalg import LinearOperator, cg
+from scipy.sparse.linalg import LinearOperator
 
 from aspire.image import Image
+from aspire.numeric.scipy import cg
 from aspire.utils import mdim_mat_fun_conj
 from aspire.volume import Volume
 
@@ -580,7 +581,7 @@ def expand(self, x, tol=None, atol=0):
         for isample in range(0, n_data):
             b = self.evaluate_t(self._cls(x[isample])).asnumpy().T
             # TODO: need check the initial condition x0 can improve the results or not.
-            v[isample], info = cg(operator, b, tol=tol, atol=atol)
+            v[isample], info = cg(operator, b, rtol=tol, atol=atol)
             if info != 0:
                 raise RuntimeError(f"Unable to converge! cg info={info}")
 

src/aspire/covariance/covar.py

@@ -2,12 +2,12 @@
 from functools import partial
 
 import numpy as np
-import scipy.sparse.linalg
 from scipy.linalg import norm
 from scipy.sparse.linalg import LinearOperator
 
 from aspire.nufft import anufft
 from aspire.numeric import fft
+from aspire.numeric.scipy import cg
 from aspire.operators import evaluate_src_filters_on_grid
 from aspire.reconstruction import Estimator, FourierKernel, MeanEstimator
 from aspire.utils import (
@@ -127,9 +127,7 @@ def cb(xk):
                 f"Delta {norm(b_coef - self.apply_kernel(xk, packed=True))} (target {target_residual})"
             )
 
-        x, info = scipy.sparse.linalg.cg(
-            operator, b_coef, M=M, callback=cb, tol=tol, atol=0
-        )
+        x, info = cg(operator, b_coef, M=M, callback=cb, rtol=tol, atol=0)
 
         if info != 0:
             raise RuntimeError("Unable to converge!")

src/aspire/numeric/scipy.py

@@ -2,15 +2,16 @@
 Utility wrappers for scipy methods.
 """
 
+import scipy
+from packaging.version import Version
 
-from scipy.sparse.linalg import cg
 
-def cg(*args,**kwargs):
+def cg(*args, **kwargs):
     """
     Supports scipy cg before and after 1.14.0.
     """
 
     # older scipy cg interface uses `tol` instead of `rtol`
     if Version(scipy.__version__) < Version("1.14.0"):
         kwargs["tol"] = kwargs.pop("rtol", None)
-    return cg(*args,**kwargs)
+    return scipy.sparse.linalg.cg(*args, **kwargs)

src/aspire/reconstruction/mean.py

@@ -3,12 +3,13 @@
 
 import numpy as np
 from scipy.linalg import norm
-from scipy.sparse.linalg import LinearOperator, cg
+from scipy.sparse.linalg import LinearOperator
 
 from aspire import config
 from aspire.basis import Coef
 from aspire.nufft import anufft
 from aspire.numeric import fft
+from aspire.numeric.scipy import cg
 from aspire.operators import evaluate_src_filters_on_grid
 from aspire.reconstruction import Estimator, FourierKernel, FourierKernelMatrix
 from aspire.volume import Volume, rotated_grids
@@ -251,7 +252,7 @@ def cb(xk):
             x0=x0,
             M=M,
             callback=cb,
-            tol=tol,
+            rtol=tol,
             atol=0,
             maxiter=self.maxiter,
         )