Logical rather than hardcoded downsample tests (#567)

* downsample tests

* cleanup and deleting test downsample from preprocess_pipeline

* remove hardcoded ds / whiten test

* flake8

* 3d volume downsample test

* remove redundant test from tests/test_preprocess.py

* remove unused import

* remove numbered tests for uniformity

* typo

* skip one of the ds tests

* use gaussian_3d to generate test volume

* isort

* checkCenterPoint

* max_resolution -> L_ds

* parametrize ds tests

Author: chris-langfield

tests/test_downsample.py

@@ -0,0 +1,90 @@
+import unittest
+from unittest import TestCase
+
+import numpy as np
+
+from aspire.source import Simulation
+from aspire.utils import utest_tolerance
+from aspire.utils.matrix import anorm
+from aspire.utils.misc import gaussian_3d
+from aspire.volume import Volume
+
+
+class DownsampleTestCase(TestCase):
+    def setUp(self):
+        self.n = 128
+        self.dtype = np.float32
+
+    def tearDown(self):
+        pass
+
+    def _testDownsample2DCase(self, L, L_ds):
+        # downsampling from size L to L_ds
+        imgs_org, imgs_ds = self.createImages(L, L_ds)
+        # check resolution is correct
+        self.assertEqual((self.n, L_ds, L_ds), imgs_ds.shape)
+        # check center points for all images
+        self.assertTrue(self.checkCenterPoint(imgs_org, imgs_ds))
+        # check signal energy is conserved
+        self.assertTrue(self.checkSignalEnergy(imgs_org, imgs_ds))
+
+    def testDownsample2D_EvenEven(self):
+        # source resolution: 64
+        # target resolution: 32
+        self._testDownsample2DCase(64, 32)
+
+    @unittest.skip(
+        "Signal energy test fails for this case in current DS implementation"
+    )
+    def testDownsample2D_EvenOdd(self):
+        # source resolution: 64
+        # target resolution: 33
+        self._testDownsample2DCase(64, 33)
+
+    def testDownsample2D_OddOdd(self):
+        # source resolution: 65
+        # target resolution: 33
+        self._testDownsample2DCase(65, 33)
+
+    def testDownsample2D_OddEven(self):
+        # source resolution: 65
+        # target resolution: 32
+        self._testDownsample2DCase(65, 32)
+
+    def checkCenterPoint(self, imgs_org, imgs_ds):
+        # Check that center point is the same after ds
+        L = imgs_org.res
+        max_resolution = imgs_ds.res
+        return np.allclose(
+            imgs_org[:, L // 2, L // 2],
+            imgs_ds[:, max_resolution // 2, max_resolution // 2],
+            atol=utest_tolerance(self.dtype),
+        )
+
+    def checkSignalEnergy(self, imgs_org, imgs_ds):
+        # check conservation of energy after downsample
+        L = imgs_org.res
+        max_resolution = imgs_ds.res
+        return np.allclose(
+            anorm(imgs_org.asnumpy(), axes=(1, 2)) / L,
+            anorm(imgs_ds.asnumpy(), axes=(1, 2)) / max_resolution,
+            atol=utest_tolerance(self.dtype),
+        )
+
+    def createImages(self, L, L_ds):
+        # generate a 3D Gaussian volume
+        sigma = 0.1
+        vol = gaussian_3d(L, sigma=(L * sigma,) * 3, dtype=self.dtype)
+        # initialize a Simulation object to generate projections of the volume
+        sim = Simulation(
+            L, self.n, vols=Volume(vol), offsets=0.0, amplitudes=1.0, dtype=self.dtype
+        )
+
+        # get images before downsample
+        imgs_org = sim.images(start=0, num=self.n)
+
+        # get images after downsample
+        sim.downsample(L_ds)
+        imgs_ds = sim.images(start=0, num=self.n)
+
+        return imgs_org, imgs_ds

tests/test_preprocess.py

@@ -4,7 +4,7 @@
 import numpy as np
 from scipy.fftpack import fftn, fftshift
 
-from aspire.image import crop_pad, downsample, fuzzy_mask
+from aspire.image import crop_pad, fuzzy_mask
 from aspire.volume import Volume
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data")
@@ -17,30 +17,22 @@ def setUp(self):
     def tearDown(self):
         pass
 
-    def test01CropPad(self):
+    def testCropPad(self):
         results = np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol_crop8.npy"))
         vols = np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy"))
         vols = vols[..., np.newaxis]
         vols_f = crop_pad(fftshift(fftn(vols[:, :, :, 0])), 8)
         self.assertTrue(np.allclose(results, vols_f, atol=1e-7))
 
-    def test02Downsample(self):
-        results = np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol_down8.npy"))
-        results = results[np.newaxis, ...]
-        vols = np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy"))
-        vols = vols[np.newaxis, ...]
-        vols = downsample(vols, (8, 8, 8))
-        self.assertTrue(np.allclose(results, vols, atol=1e-7))
-
-    def test03Vol2img(self):
+    def testVol2img(self):
         results = np.load(os.path.join(DATA_DIR, "clean70SRibosome_down8_imgs32.npy"))
         vols = Volume(np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol_down8.npy")))
         rots = np.load(os.path.join(DATA_DIR, "rand_rot_matrices32.npy"))
         rots = np.moveaxis(rots, 2, 0)
         imgs_clean = vols.project(0, rots).asnumpy()
         self.assertTrue(np.allclose(results, imgs_clean, atol=1e-7))
 
-    def test04FuzzyMask(self):
+    def testFuzzyMask(self):
         results = np.array(
             [
                 [

tests/test_preprocess_pipeline.py

@@ -4,12 +4,11 @@
 import numpy as np
 
 from aspire.noise import AnisotropicNoiseEstimator
-from aspire.operators.filters import FunctionFilter, RadialCTFFilter, ScalarFilter
+from aspire.operators.filters import FunctionFilter, RadialCTFFilter
 from aspire.source import ArrayImageSource
 from aspire.source.simulation import Simulation
-from aspire.utils import grid_2d, grid_3d, utest_tolerance
+from aspire.utils import grid_2d
 from aspire.utils.matrix import anorm
-from aspire.volume import Volume
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data")
 
@@ -45,58 +44,6 @@ def testPhaseFlip(self):
             )
         )
 
-    def testDownsample(self):
-        # generate a 3D map with density decays as Gaussian function
-        g3d = grid_3d(self.L, indexing="zyx", dtype=self.dtype)
-        coords = np.array([g3d["x"].flatten(), g3d["y"].flatten(), g3d["z"].flatten()])
-        sigma = 0.2
-        vol = np.exp(-0.5 * np.sum(np.abs(coords / sigma) ** 2, axis=0)).astype(
-            self.dtype
-        )
-        vol = np.reshape(vol, g3d["x"].shape)
-        vols = Volume(vol)
-
-        # set noise to zero and CFT filters to unity for simulation object
-        noise_var = 0
-        noise_filter = ScalarFilter(dim=2, value=noise_var)
-        sim = Simulation(
-            L=self.L,
-            n=self.n,
-            vols=vols,
-            offsets=0.0,
-            amplitudes=1.0,
-            unique_filters=[
-                ScalarFilter(dim=2, value=1) for d in np.linspace(1.5e4, 2.5e4, 7)
-            ],
-            noise_filter=noise_filter,
-            dtype=self.dtype,
-        )
-        # get images before downsample
-        imgs_org = sim.images(start=0, num=self.n)
-        # get images after downsample
-        max_resolution = 32
-        sim.downsample(max_resolution)
-        imgs_ds = sim.images(start=0, num=self.n)
-
-        # Check individual grid points
-        self.assertTrue(
-            np.allclose(
-                imgs_org[:, 32, 32],
-                imgs_ds[:, 16, 16],
-                atol=utest_tolerance(self.dtype),
-            )
-        )
-        # check resolution
-        self.assertTrue(np.allclose(max_resolution, imgs_ds.shape[1]))
-        # check energy conservation after downsample
-        self.assertTrue(
-            np.allclose(
-                anorm(imgs_org.asnumpy(), axes=(1, 2)) / self.L,
-                anorm(imgs_ds.asnumpy(), axes=(1, 2)) / max_resolution,
-                atol=utest_tolerance(self.dtype),
-            )
-        )
-
     def testNormBackground(self):
         bg_radius = 1.0
         grid = grid_2d(self.L, indexing="yx")

tests/test_starfile_stack.py

@@ -8,7 +8,6 @@
 
 import tests.saved_test_data
 from aspire.image import Image
-from aspire.operators import ScalarFilter
 from aspire.source.relion import RelionSource
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data")
@@ -61,20 +60,6 @@ def testImageDownsample(self):
         first_image = self.src.images(0, 1)[0]
         self.assertEqual(first_image.shape, (16, 16))
 
-    def testImageDownsampleAndWhiten(self):
-        self.src.downsample(16)
-        self.src.whiten(noise_filter=ScalarFilter(dim=2, value=0.02450909546680349))
-        first_whitened_image = self.src.images(0, 1)[0]
-        self.assertTrue(
-            np.allclose(
-                first_whitened_image,
-                np.load(
-                    os.path.join(DATA_DIR, "starfile_image_0_whitened.npy")
-                ).T,  # RCOPT
-                atol=1e-6,
-            )
-        )
-
 
 class StarFileSingleImage(StarFileTestCase):
     def setUp(self):

tests/test_volume.py

@@ -9,6 +9,7 @@
 from pytest import raises
 
 from aspire.utils import Rotation, grid_3d, powerset
+from aspire.utils.matrix import anorm
 from aspire.utils.types import utest_tolerance
 from aspire.volume import Volume, gaussian_blob_vols
 
@@ -313,11 +314,25 @@ def testFlip(self):
             self.assertTrue(isinstance(result, Volume))
 
     def testDownsample(self):
-        # Data files re-used from test_preprocess
         vols = Volume(np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")))
-
-        resv = Volume(np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol_down8.npy")))
-
         result = vols.downsample((8, 8, 8))
-        self.assertTrue(np.allclose(result, resv))
-        self.assertTrue(isinstance(result, Volume))
+        res = vols.resolution
+        ds_res = result.resolution
+
+        # check signal energy
+        self.assertTrue(
+            np.allclose(
+                anorm(vols.asnumpy(), axes=(1, 2, 3)) / res,
+                anorm(result.asnumpy(), axes=(1, 2, 3)) / ds_res,
+                atol=1e-3,
+            )
+        )
+
+        # check gridpoints
+        self.assertTrue(
+            np.allclose(
+                vols[:, res // 2, res // 2, res // 2],
+                result[:, ds_res // 2, ds_res // 2, ds_res // 2],
+                atol=1e-4,
+            )
+        )