diff --git a/tests/_basis_util.py b/tests/_basis_util.py new file mode 100644 index 0000000000..cf9321427c --- /dev/null +++ b/tests/_basis_util.py @@ -0,0 +1,132 @@ +from unittest.case import SkipTest + +import numpy as np + +from aspire.image import Image +from aspire.utils import gaussian_2d, utest_tolerance +from aspire.utils.coor_trans import grid_2d +from aspire.utils.random import randn + + +class Steerable2DMixin: + def testIndices(self): + ell_max = self.basis.ell_max + k_max = self.basis.k_max + + indices = self.basis.indices() + + i = 0 + + for ell in range(ell_max + 1): + if ell == 0: + sgns = [1] + else: + sgns = [1, -1] + + for sgn in sgns: + for k in range(k_max[ell]): + self.assertTrue(indices["ells"][i] == ell) + self.assertTrue(indices["sgns"][i] == sgn) + self.assertTrue(indices["ks"][i] == k) + + i += 1 + + def testGaussianExpand(self): + # Offset slightly + x0 = 0.50 + y0 = 0.75 + + # Want sigma to be as large as possible without the Gaussian + # spilling too much outside the central disk. + sigma = self.L / 8 + im1 = gaussian_2d( + self.L, x0=x0, y0=y0, sigma_x=sigma, sigma_y=sigma, dtype=self.dtype + ) + + coef = self.basis.expand(im1) + im2 = self.basis.evaluate(coef) + + if isinstance(im2, Image): + im2 = im2.asnumpy() + im2 = im2[0] + + # For small L there's too much clipping at high freqs to get 1e-3 + # accuracy. + if self.L < 32: + atol = 1e-2 + else: + atol = 1e-3 + + self.assertTrue(im1.shape == im2.shape) + self.assertTrue(np.allclose(im1, im2, atol=atol)) + + def testIsotropic(self): + sigma = self.L / 8 + im = gaussian_2d(self.L, sigma_x=sigma, sigma_y=sigma, dtype=self.dtype) + + coef = self.basis.expand(im) + + ells = self.basis.indices()["ells"] + + energy_outside = np.sum(np.abs(coef[ells != 0]) ** 2) + energy_total = np.sum(np.abs(coef) ** 2) + + energy_ratio = energy_outside / energy_total + + self.assertTrue(energy_ratio < 0.01) + + def testModulated(self): + if self.L < 32: + raise SkipTest + + ell = 1 + + sigma = self.L / 8 + im = gaussian_2d(self.L, sigma_x=sigma, sigma_y=sigma, dtype=self.dtype) + + g2d = grid_2d(self.L) + + for trig_fun in (np.sin, np.cos): + im1 = im * trig_fun(ell * g2d["phi"]) + + coef = self.basis.expand(im1) + + ells = self.basis.indices()["ells"] + + energy_outside = np.sum(np.abs(coef[ells != ell]) ** 2) + energy_total = np.sum(np.abs(coef) ** 2) + + energy_ratio = energy_outside / energy_total + + self.assertTrue(energy_ratio < 0.10) + + def testEvaluateExpand(self): + coef1 = randn(self.basis.count, seed=self.seed) + coef1 = coef1.astype(self.dtype) + + im = self.basis.evaluate(coef1) + if isinstance(im, Image): + im = im.asnumpy() + coef2 = self.basis.expand(im)[0] + + self.assertTrue(coef1.shape == coef2.shape) + self.assertTrue(np.allclose(coef1, coef2, atol=utest_tolerance(self.dtype))) + + def testAdjoint(self): + u = randn(self.basis.count, seed=self.seed) + u = u.astype(self.dtype) + + Au = self.basis.evaluate(u) + if isinstance(Au, Image): + Au = Au.asnumpy() + + x = randn(*self.basis.sz, seed=self.seed) + x = x.astype(self.dtype) + + ATx = self.basis.evaluate_t(x) + + Au_dot_x = np.sum(Au * x) + u_dot_ATx = np.sum(u * ATx) + + self.assertTrue(Au_dot_x.shape == u_dot_ATx.shape) + self.assertTrue(np.isclose(Au_dot_x, u_dot_ATx)) diff --git a/tests/test_FBbasis2D.py b/tests/test_FBbasis2D.py index 2387374368..30eda0de73 100644 --- a/tests/test_FBbasis2D.py +++ b/tests/test_FBbasis2D.py @@ -2,330 +2,87 @@ from unittest import TestCase import numpy as np +from parameterized import parameterized_class from pytest import raises +from scipy.special import jv from aspire.basis import FBBasis2D -from aspire.image import Image -from aspire.utils import complex_type, real_type, utest_tolerance +from aspire.utils import complex_type, real_type +from aspire.utils.coor_trans import grid_2d +from aspire.utils.random import randn + +from ._basis_util import Steerable2DMixin DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data") -class FBBasis2DTestCase(TestCase): +# NOTE: Class with default values is already present, so don't list it below. +@parameterized_class( + ("L", "dtype"), + [ + (8, np.float64), + (16, np.float32), + (16, np.float64), + (32, np.float32), + (32, np.float64), + ], +) +class FBBasis2DTestCase(TestCase, Steerable2DMixin): + L = 8 + dtype = np.float32 + def setUp(self): - self.dtype = np.float32 - self.basis = FBBasis2D((8, 8), dtype=self.dtype) + self.basis = FBBasis2D((self.L, self.L), dtype=self.dtype) + self.seed = 9161341 def tearDown(self): pass - def testFBBasis2DIndices(self): + def _testElement(self, ell, k, sgn): + # This is covered by the isotropic test. + assert ell > 0 + indices = self.basis.indices() + ells = indices["ells"] + sgns = indices["sgns"] + ks = indices["ks"] + + g2d = grid_2d(self.L, dtype=self.dtype) + mask = g2d["r"] < 1 + + r0 = self.basis.r0[k, ell] + + im = np.zeros((self.L, self.L), dtype=self.dtype) + im[mask] = jv(ell, g2d["r"][mask] * r0) + im *= np.sqrt(2**2 / self.L**2) + im *= 1 / (np.sqrt(np.pi) * np.abs(jv(ell + 1, r0))) + + if sgn == 1: + im *= np.sqrt(2) * np.cos(ell * g2d["phi"]) + else: + im *= np.sqrt(2) * np.sin(ell * g2d["phi"]) + + coef_ref = np.zeros(self.basis.count, dtype=self.dtype) + coef_ref[(ells == ell) & (sgns == sgn) & (ks == k)] = 1 + + im_ref = self.basis.evaluate(coef_ref) + + coef = self.basis.expand(im) + + # TODO: These tolerances should be tighter. + self.assertTrue(np.allclose(im, im_ref, atol=1e-4)) + self.assertTrue(np.allclose(coef, coef_ref, atol=1e-4)) + + def testElements(self): + ells = [1, 1, 1, 1] + ks = [1, 2, 1, 2] + sgns = [-1, -1, 1, 1] - self.assertTrue( - np.allclose( - indices["ells"], - [ - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 2.0, - 2.0, - 2.0, - 2.0, - 2.0, - 2.0, - 3.0, - 3.0, - 3.0, - 3.0, - 4.0, - 4.0, - 4.0, - 4.0, - 5.0, - 5.0, - 5.0, - 5.0, - 6.0, - 6.0, - 7.0, - 7.0, - 8.0, - 8.0, - ], - ) - ) - - self.assertTrue( - np.allclose( - indices["ks"], - [ - 0.0, - 1.0, - 2.0, - 3.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - ], - ) - ) - - self.assertTrue( - np.allclose( - indices["sgns"], - [ - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - -1.0, - -1.0, - -1.0, - 1.0, - 1.0, - 1.0, - -1.0, - -1.0, - -1.0, - 1.0, - 1.0, - -1.0, - -1.0, - 1.0, - 1.0, - -1.0, - -1.0, - 1.0, - 1.0, - -1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - ], - ) - ) - - def testFBBasis2DNorms(self): - radial_norms, angular_norms = self.basis.norms() - self.assertTrue( - np.allclose( - radial_norms * angular_norms, - [ - 3.68065992303471, - 2.41241466684800, - 1.92454669738088, - 1.64809729313301, - 2.01913617828263, - 1.50455726188833, - 1.25183461029289, - 1.70284654929000, - 1.36051054373844, - 1.16529703804363, - 1.49532071137207, - 1.25039038364830, - 1.34537533748304, - 1.16245357319190, - 1.23042467443861, - 1.09002083501080, - 1.13867113286781, - 1.06324777330476, - 0.999841586390824, - ], - ) - ) - - def testFBBasis2DEvaluate(self): - coeffs = np.array( - [ - 1.07338590e-01, - 1.23690941e-01, - 6.44482039e-03, - -5.40484306e-02, - -4.85304586e-02, - 1.09852144e-02, - 3.87838396e-02, - 3.43796455e-02, - -6.43284705e-03, - -2.86677145e-02, - -1.42313328e-02, - -2.25684091e-03, - -3.31840727e-02, - -2.59706174e-03, - -5.91919887e-04, - -9.97433028e-03, - 9.19123928e-04, - 1.19891589e-03, - 7.49154982e-03, - 6.18865229e-03, - -8.13265715e-04, - -1.30715655e-02, - -1.44160603e-02, - 2.90379956e-03, - 2.37066082e-02, - 4.88805735e-03, - 1.47870707e-03, - 7.63376018e-03, - -5.60619559e-03, - 1.05165081e-02, - 3.30510143e-03, - -3.48652120e-03, - -4.23228797e-04, - 1.40484061e-02, - ], - dtype=self.dtype, - ) - result = self.basis.evaluate(coeffs) - - self.assertTrue( - np.allclose( - result, - np.load( - os.path.join(DATA_DIR, "fbbasis_evaluation_8_8.npy") - ).T, # RCOPT - atol=utest_tolerance(self.dtype), - ) - ) - - def testFBBasis2DEvaluate_t(self): - v = np.load(os.path.join(DATA_DIR, "fbbasis_coefficients_8_8.npy")).T # RCOPT - # While FB can accept arrays, prefable to pass FB2D and FFB2D Image instances. - img = Image(v.astype(self.dtype)) - result = self.basis.evaluate_t(img) - self.assertTrue( - np.allclose( - result, - [ - 0.10761825, - 0.12291151, - 0.00836345, - -0.0619454, - -0.0483326, - 0.01053718, - 0.03977641, - 0.03420101, - -0.0060131, - -0.02970658, - -0.0151334, - -0.00017575, - -0.03987446, - -0.00257069, - -0.0006621, - -0.00975174, - 0.00108047, - 0.00072022, - 0.00753342, - 0.00604493, - 0.00024362, - -0.01711248, - -0.01387371, - 0.00112805, - 0.02407385, - 0.00376325, - 0.00081128, - 0.00951368, - -0.00557536, - 0.01087579, - 0.00255393, - -0.00525156, - -0.00839695, - 0.00802198, - ], - atol=utest_tolerance(self.dtype), - ) - ) - - def testFBBasis2DExpand(self): - v = np.load(os.path.join(DATA_DIR, "fbbasis_coefficients_8_8.npy")).T # RCOPT - result = self.basis.expand(v.astype(self.dtype)) - self.assertTrue( - np.allclose( - result, - [ - 0.10733859, - 0.12369094, - 0.00644482, - -0.05404843, - -0.04853046, - 0.01098521, - 0.03878384, - 0.03437965, - -0.00643285, - -0.02866771, - -0.01423133, - -0.00225684, - -0.03318407, - -0.00259706, - -0.00059192, - -0.00997433, - 0.00091912, - 0.00119892, - 0.00749155, - 0.00618865, - -0.00081327, - -0.01307157, - -0.01441606, - 0.00290380, - 0.02370661, - 0.00488806, - 0.00147871, - 0.00763376, - -0.00560620, - 0.01051651, - 0.00330510, - -0.00348652, - -0.00042323, - 0.01404841, - ], - atol=utest_tolerance(self.dtype), - ) - ) + for ell, k, sgn in zip(ells, ks, sgns): + self._testElement(ell, k, sgn) def testComplexCoversion(self): - # Load a reasonable input - x = np.load(os.path.join(DATA_DIR, "fbbasis_coefficients_8_8.npy")) + x = randn(*self.basis.sz, seed=self.seed) # Express in an FB basis v1 = self.basis.expand(x.astype(self.dtype)) @@ -339,8 +96,7 @@ def testComplexCoversion(self): self.assertTrue(np.allclose(v1, v2)) def testComplexCoversionErrorsToComplex(self): - # Load a reasonable input - x = np.load(os.path.join(DATA_DIR, "fbbasis_coefficients_8_8.npy")) + x = randn(*self.basis.sz, seed=self.seed) # Express in an FB basis v1 = self.basis.expand(x.astype(self.dtype)) @@ -365,8 +121,7 @@ def testComplexCoversionErrorsToComplex(self): _ = self.basis.to_complex(v1.reshape(-1)) def testComplexCoversionErrorsToReal(self): - # Load a reasonable input - x = np.load(os.path.join(DATA_DIR, "fbbasis_coefficients_8_8.npy")) + x = randn(*self.basis.sz, seed=self.seed) # Express in an FB basis cv1 = self.basis.to_complex(self.basis.expand(x.astype(self.dtype))) diff --git a/tests/test_FFBbasis2D.py b/tests/test_FFBbasis2D.py index db76d447ee..cd37606e0d 100644 --- a/tests/test_FFBbasis2D.py +++ b/tests/test_FFBbasis2D.py @@ -1,8 +1,11 @@ import logging import os.path from unittest import TestCase +from unittest.case import SkipTest import numpy as np +from parameterized import parameterized_class +from scipy.special import jv from aspire.basis import FFBBasis2D from aspire.image import Image @@ -11,256 +14,92 @@ from aspire.utils.misc import grid_2d from aspire.volume import Volume +from ._basis_util import Steerable2DMixin + logger = logging.getLogger(__name__) DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data") -class FFBBasis2DTestCase(TestCase): +# NOTE: Class with default values is already present, so don't list it below. +@parameterized_class( + ("L", "dtype"), + [ + (8, np.float64), + (16, np.float32), + (16, np.float64), + (32, np.float32), + (32, np.float64), + ], +) +class FFBBasis2DTestCase(TestCase, Steerable2DMixin): + L = 8 + dtype = np.float32 + def setUp(self): - self.dtype = np.float32 # Required for convergence of this test - self.L = 8 self.basis = FFBBasis2D((self.L, self.L), dtype=self.dtype) + self.seed = 9161341 def tearDown(self): pass - def testFFBBasis2DIndices(self): + def _testElement(self, ell, k, sgn): indices = self.basis.indices() - - self.assertTrue( - np.allclose( - indices["ells"], - [ - 0.0, - 0.0, - 0.0, - 0.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 2.0, - 2.0, - 2.0, - 2.0, - 2.0, - 2.0, - 3.0, - 3.0, - 3.0, - 3.0, - 4.0, - 4.0, - 4.0, - 4.0, - 5.0, - 5.0, - 5.0, - 5.0, - 6.0, - 6.0, - 7.0, - 7.0, - 8.0, - 8.0, - ], - ) + ells = indices["ells"] + sgns = indices["sgns"] + ks = indices["ks"] + + g2d = grid_2d(self.L, dtype=self.dtype) + mask = g2d["r"] < 1 + + r0 = self.basis.r0[k, ell] + + # TODO: Figure out where these factors of 1 / 2 are coming from. + # Intuitively, the grid should go from -L / 2 to L / 2, not -L / 2 to + # L / 4. Furthermore, there's an extra factor of 1 / 2 in the + # definition of `im` below that may be related. + r = g2d["r"] * self.L / 4 + + im = np.zeros((self.L, self.L), dtype=self.dtype) + im[mask] = ( + (-1) ** k + * np.sqrt(np.pi) + * r0 + * jv(ell, 2 * np.pi * r[mask]) + / ((2 * np.pi * r[mask]) ** 2 - r0**2) ) - self.assertTrue( - np.allclose( - indices["ks"], - [ - 0.0, - 1.0, - 2.0, - 3.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 2.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 1.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - 0.0, - ], - ) - ) + if sgn == 1: + im *= np.sqrt(2) * np.cos(ell * g2d["phi"]) + else: + im *= np.sqrt(2) * np.sin(ell * g2d["phi"]) - self.assertTrue( - np.allclose( - indices["sgns"], - [ - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - -1.0, - -1.0, - -1.0, - 1.0, - 1.0, - 1.0, - -1.0, - -1.0, - -1.0, - 1.0, - 1.0, - -1.0, - -1.0, - 1.0, - 1.0, - -1.0, - -1.0, - 1.0, - 1.0, - -1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - ], - ) - ) + coef_ref = np.zeros(self.basis.count, dtype=self.dtype) + coef_ref[(ells == ell) & (sgns == sgn) & (ks == k)] = 1 - def testFFBBasis2DNorms(self): - radial_norms, angular_norms = self.basis.norms() - self.assertTrue( - np.allclose( - radial_norms * angular_norms, - [ - 3.68065992303471, - 2.41241466684800, - 1.92454669738088, - 1.64809729313301, - 2.01913617828263, - 1.50455726188833, - 1.25183461029289, - 1.70284654929000, - 1.36051054373844, - 1.16529703804363, - 1.49532071137207, - 1.25039038364830, - 1.34537533748304, - 1.16245357319190, - 1.23042467443861, - 1.09002083501080, - 1.13867113286781, - 1.06324777330476, - 0.999841586390824, - ], - ) - ) + im_ref = self.basis.evaluate(coef_ref).asnumpy()[0] - def testFFBBasis2DEvaluate(self): - v = np.array( - [ - 1.07338590e-01, - 1.23690941e-01, - 6.44482039e-03, - -5.40484306e-02, - -4.85304586e-02, - 1.09852144e-02, - 3.87838396e-02, - 3.43796455e-02, - -6.43284705e-03, - -2.86677145e-02, - -1.42313328e-02, - -2.25684091e-03, - -3.31840727e-02, - -2.59706174e-03, - -5.91919887e-04, - -9.97433028e-03, - 9.19123928e-04, - 1.19891589e-03, - 7.49154982e-03, - 6.18865229e-03, - -8.13265715e-04, - -1.30715655e-02, - -1.44160603e-02, - 2.90379956e-03, - 2.37066082e-02, - 4.88805735e-03, - 1.47870707e-03, - 7.63376018e-03, - -5.60619559e-03, - 1.05165081e-02, - 3.30510143e-03, - -3.48652120e-03, - -4.23228797e-04, - 1.40484061e-02, - ], - dtype=self.dtype, - ) - result = self.basis.evaluate(v) - - self.assertTrue( - np.allclose( - result.asnumpy(), # Result of evaluate is an Image after RCOPT - np.load( - os.path.join(DATA_DIR, "ffbbasis2d_xcoeff_out_8_8.npy") - ).T, # RCOPT - atol=utest_tolerance(self.dtype), - ) - ) + coef = self.basis.expand(im) - def testFFBBasis2DEvaluate_t(self): - x = np.load(os.path.join(DATA_DIR, "ffbbasis2d_xcoeff_in_8_8.npy")).T # RCOPT - result = self.basis.evaluate_t(x.astype(self.dtype)) + # NOTE: These tolerances are expected to be rather loose since the + # above expression for `im` is derived from the analytical formulation + # (eq. 6 in Zhao and Singer, 2013) and does not take into account + # discretization and other approximations. + self.assertTrue(np.allclose(im, im_ref, atol=1e-1)) + self.assertTrue(np.allclose(coef, coef_ref, atol=1e-1)) - self.assertTrue( - np.allclose( - result, - np.load(os.path.join(DATA_DIR, "ffbbasis2d_vcoeff_out_8_8.npy"))[ - ..., 0 - ], - ) - ) + def testElements(self): + ells = [1, 1, 1, 1] + ks = [1, 2, 1, 2] + sgns = [-1, -1, 1, 1] - def testFFBBasis2DExpand(self): - x = np.load(os.path.join(DATA_DIR, "ffbbasis2d_xcoeff_in_8_8.npy")).T # RCOPT - result = self.basis.expand(x.astype(self.dtype)) - self.assertTrue( - np.allclose( - result, - np.load(os.path.join(DATA_DIR, "ffbbasis2d_vcoeff_out_exp_8_8.npy"))[ - ..., 0 - ], - atol=utest_tolerance(self.dtype), - ) - ) + for ell, k, sgn in zip(ells, ks, sgns): + self._testElement(ell, k, sgn) def testRotate(self): + # Convergence issues for double precision. + if np.dtype(self.dtype) is np.dtype(np.float64): + raise SkipTest + # Now low res (8x8) had problems; # better with odd (7x7), but still not good. # We'll use a higher res test image.