improve test coverage

Author: garrettwrong

src/aspire/covariance/covar2d.py

@@ -736,7 +736,7 @@ def _noise_correct_covar_rhs(self, b_covar, b_noise, noise_var, shrinker):
 
     def _solve_covar(self, A_covar, b_covar, M, covar_est_opt):
         method = self._solve_covar_cg
-        if all(isinstance(a, DiagMatrix) for a in A_covar):
+        if all(isinstance(a, DiagMatrix) or a is None for a in A_covar):
             method = self._solve_covar_direct
 
         return method(A_covar, b_covar, M, covar_est_opt)

src/aspire/numeric/cupy.py

@@ -17,6 +17,6 @@ def atleast_1d(x):
         Returns same type as input.
         """
         _fn = np.atleast_1d
-        if cp and isinstance(x, cp.ndarray):
+        if isinstance(x, cp.ndarray):
             _fn = cp.atleast_1d
         return _fn(x)

tests/test_FFBbasis2D.py

@@ -128,6 +128,12 @@ def testShift(self, basis):
         # Check pixel_size passthrough
         np.testing.assert_array_equal(f_imgs.pixel_size, f_shifted_imgs.pixel_size)
 
+    def testBadFilterLength(self, basis):
+        filt = RadialCTFFilter(defocus=np.linspace(1000, 10000, 3))  # len 3
+        with pytest.raises(RuntimeError, match=r".*Unexpected filter length.*"):
+            # filter_to_basis_mat should enforce len 1
+            _ = basis.filter_to_basis_mat(filt)
+
 
 params = [pytest.param(512, np.float32, marks=pytest.mark.expensive)]
 

tests/test_FLEbasis2D.py

@@ -9,6 +9,7 @@
 from aspire.image import Image
 from aspire.nufft import backend_available
 from aspire.numeric import fft
+from aspire.operators import RadialCTFFilter
 from aspire.source import Simulation
 from aspire.volume import Volume
 
@@ -122,6 +123,12 @@ def testEvaluateExpand(self, basis):
 
         assert relerr(expand.asnumpy(), evaluate_t.asnumpy()) < basis.epsilon
 
+    def testBadFilterLength(self, basis):
+        filt = RadialCTFFilter(defocus=np.linspace(1000, 10000, 3))  # len 3
+        with pytest.raises(RuntimeError, match=r".*Unexpected filter length.*"):
+            # filter_to_basis_mat should enforce len 1
+            _ = basis.filter_to_basis_mat(filt)
+
 
 @pytest.mark.parametrize("basis", test_bases_match_fb, ids=show_fle_params)
 def testMatchFBEvaluate(basis):

tests/test_large_covar.py

@@ -15,7 +15,8 @@
 
 from aspire.basis import FFBBasis2D, FLEBasis2D
 from aspire.covariance import BatchedRotCov2D
-from aspire.source import RelionSource
+from aspire.operators import RadialCTFFilter
+from aspire.source import RelionSource, Simulation
 
 DTYPES = [
     np.float32,
@@ -139,3 +140,30 @@ def test_covar2d(preprocessed_src, basis, force_radial):
     cov2d = BatchedRotCov2D(preprocessed_src, basis, expand_method=expand_method)
     # smoke test
     _ = cov2d.get_covar()
+
+
+def test_covar2d_many_ctf():
+    """
+    Smoke test for many CTF case using optimized radial expansion code path.
+    """
+    # N must be >= 2048 to enable auto GPU filter eval branch
+    #   in covar2d _radial_filter_stack_to_basis_mats
+    N = 2500
+    L = 33
+    dt = np.float32
+    src = Simulation(
+        C=1,
+        n=N,
+        L=L,
+        filter_stack=RadialCTFFilter(defocus=np.linspace(10000, 20000, N)),
+        offsets=0,
+        amplitudes=1,
+        dtype=dt,
+    ).cache()
+
+    basis = FLEBasis2D(L, dtype=dt)
+
+    cov2d = BatchedRotCov2D(src, basis, expand_method="radial")
+
+    # smoke test
+    _ = cov2d.get_covar()

tests/test_preprocess_pipeline.py

@@ -192,6 +192,28 @@ def testWhiten2(dtype):
     np.testing.assert_allclose(np.eye(2), corr_coef, atol=2e-1)
 
 
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def testWhitenNoCTF(dtype):
+    """
+    Test operation with `filter_stack=None`.
+    """
+    # Note this atol holds only for L even. Odd tested in testWhiten2.
+    L = 64
+    sim = get_sim_object(L, dtype)
+    sim.filter_stack = None
+    noise_estimator = AnisotropicNoiseEstimator(sim)
+    sim = sim.whiten(noise_estimator)
+    imgs_wt = sim.images[:num_images].asnumpy()
+
+    # calculate correlation between two neighboring pixels from background
+    corr_coef = np.corrcoef(imgs_wt[:, L - 1, L - 1], imgs_wt[:, L - 2, L - 1])
+
+    # correlation matrix should be close to identity
+    np.testing.assert_allclose(np.eye(2), corr_coef, atol=1e-1)
+    # dtype of returned images should be the same
+    assert dtype == imgs_wt.dtype
+
+
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
 def test_whiten_safeguard(dtype):
     """Test that whitening safeguard works as expected."""

tests/test_utils.py

@@ -13,6 +13,7 @@
 
 import aspire
 from aspire import __version__
+from aspire.numeric import xp
 from aspire.utils import (
     LogFilterByCount,
     all_pairs,
@@ -498,3 +499,24 @@ def wrapper(*args, **kwargs):
             return func(*args, **kwargs)
 
     return wrapper
+
+
+def test_atleast1d():
+    """
+    Test xp.atleast_1d is agnostic to inputs.
+    """
+    x = 1.0
+    y = xp.asnumpy(xp.atleast_1d(x))
+    ref = np.atleast_1d(x)
+    np.testing.assert_allclose(y, ref)
+
+    x = np.arange(9)
+    # x is host input
+    y = xp.atleast_1d(x)
+    # cp input when in GPU mode
+    #   xp functions should pass through as np otherwise
+    y2 = xp.atleast_1d(xp.asarray(x))
+
+    ref = np.atleast_1d(x)
+    np.testing.assert_allclose(xp.asnumpy(y), ref)
+    np.testing.assert_allclose(xp.asnumpy(y2), ref)