From 1d9698787a0ed860d13a9b25abbce0b9e8e1454c Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 27 Nov 2023 14:41:48 -0500 Subject: [PATCH 1/9] Add simple rotation test across bases --- tests/test_rotation.py | 330 ++++++++++++++++++++++++++--------------- 1 file changed, 210 insertions(+), 120 deletions(-) diff --git a/tests/test_rotation.py b/tests/test_rotation.py index caffb17f0c..10e32e7487 100644 --- a/tests/test_rotation.py +++ b/tests/test_rotation.py @@ -1,131 +1,191 @@ import logging -from unittest import TestCase import numpy as np +import PIL.Image as PILImage import pytest +import scipy as sp from scipy.spatial.transform import Rotation as sp_rot -from aspire.utils import Rotation, utest_tolerance +from aspire.basis import FBBasis2D, FFBBasis2D, FLEBasis2D, FPSWFBasis2D, PSWFBasis2D +from aspire.image import Image +from aspire.utils import Rotation, grid_2d, utest_tolerance logger = logging.getLogger(__name__) +# Parameters -class UtilsTestCase(TestCase): - def setUp(self): - self.dtype = np.float32 - self.num_rots = 32 - self.rot_obj = Rotation.generate_random_rotations( - self.num_rots, seed=0, dtype=self.dtype - ) - self.angles = self.rot_obj.angles - self.matrices = self.rot_obj.matrices - - def testRotMatrices(self): - rot_ref = sp_rot.from_matrix(self.matrices) - matrices = rot_ref.as_matrix().astype(self.dtype) - self.assertTrue( - np.allclose(self.matrices, matrices, atol=utest_tolerance(self.dtype)) +NUM_ROTS = 32 +SEED = 0 + +DTYPES = [ + np.float64, + np.float32, +] + +BASES = [ + FFBBasis2D, + FBBasis2D, + FLEBasis2D, + PSWFBasis2D, + FPSWFBasis2D, +] + +IMG_SIZES = [ + 32, + pytest.param(31, marks=pytest.mark.expensive), +] + +# Fixtures + + +@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module") +def dtype(request): + return request.param + + +@pytest.fixture(scope="module") +def rot_obj(dtype): + return Rotation.generate_random_rotations(NUM_ROTS, seed=SEED, dtype=dtype) + + +@pytest.fixture(params=IMG_SIZES, ids=lambda x: f"img_size={x}", scope="module") +def img_size(request): + return request.param + + +@pytest.fixture(params=BASES, ids=lambda x: f"basis={x}", scope="module") +def basis(request, img_size, dtype): + cls = request.param + # Setup a Basis + basis = cls(img_size, dtype=dtype) + return basis + + +# Rotation Class Tests + + +def test_matrices(rot_obj): + rot_ref = sp_rot.from_matrix(rot_obj.matrices) + matrices = rot_ref.as_matrix() + np.testing.assert_allclose( + rot_obj.matrices, matrices, atol=utest_tolerance(rot_obj.dtype) + ) + + +def test_as_angles(rot_obj): + rot_ref = sp_rot.from_euler("ZYZ", rot_obj.angles, degrees=False) + angles = rot_ref.as_euler("ZYZ", degrees=False) + np.testing.assert_allclose(rot_obj.angles, angles) + + +def test_from_matrix(rot_obj): + rot_ref = sp_rot.from_matrix(rot_obj.matrices) + angles = rot_ref.as_euler("ZYZ", degrees=False) + rot = Rotation.from_matrix(rot_obj.matrices) + np.testing.assert_allclose(rot.angles, angles) + + +def test_from_euler(rot_obj): + rot_ref = sp_rot.from_euler("ZYZ", rot_obj.angles, degrees=False) + matrices = rot_ref.as_matrix() + rot = Rotation.from_euler(rot_obj.angles, dtype=rot_obj.dtype) + np.testing.assert_allclose(rot._matrices, matrices) + + +def test_invert(rot_obj): + rot_mat = rot_obj.matrices + rot_mat_t = rot_obj.invert() + np.testing.assert_allclose(rot_mat_t, np.transpose(rot_mat, (0, 2, 1))) + + +def test_multiplication(rot_obj): + result = (rot_obj * rot_obj.invert()).matrices + for i in range(len(rot_obj)): + np.testing.assert_allclose( + np.eye(3), result[i], atol=utest_tolerance(rot_obj.dtype) ) - def testRotAngles(self): - rot_ref = sp_rot.from_euler("ZYZ", self.angles, degrees=False) - angles = rot_ref.as_euler("ZYZ", degrees=False).astype(self.dtype) - self.assertTrue(np.allclose(self.angles, angles)) - - def testFromMatrix(self): - rot_ref = sp_rot.from_matrix(self.matrices) - angles = rot_ref.as_euler("ZYZ", degrees=False).astype(self.dtype) - rot = Rotation.from_matrix(self.matrices, dtype=self.dtype) - self.assertTrue(np.allclose(rot.angles, angles)) - - def testFromEuler(self): - rot_ref = sp_rot.from_euler("ZYZ", self.angles, degrees=False) - matrices = rot_ref.as_matrix().astype(self.dtype) - rot = Rotation.from_euler(self.angles, dtype=self.dtype) - self.assertTrue(np.allclose(rot._matrices, matrices)) - - def testInvert(self): - rot_mat = self.rot_obj.matrices - rot_mat_t = self.rot_obj.invert() - self.assertTrue(np.allclose(rot_mat_t, np.transpose(rot_mat, (0, 2, 1)))) - - def testMultiplication(self): - result = (self.rot_obj * self.rot_obj.invert()).matrices - for i in range(len(self.rot_obj)): - self.assertTrue( - np.allclose(np.eye(3), result[i], atol=utest_tolerance(self.dtype)) - ) - - def testRegisterRots(self): - q_mat = Rotation.generate_random_rotations(1, dtype=self.dtype)[0] - for flag in [0, 1]: - regrots_ref = self.rot_obj.apply_registration(q_mat, flag) - q_mat_est, flag_est = self.rot_obj.find_registration(regrots_ref) - self.assertTrue( - np.allclose(flag_est, flag) - and np.allclose(q_mat_est, q_mat, atol=utest_tolerance(self.dtype)) - ) - - def testRegister(self): - # These will yield two more distinct sets of random rotations wrt self.rot_obj - set1 = Rotation.generate_random_rotations(self.num_rots, dtype=self.dtype) - set2 = Rotation.generate_random_rotations( - self.num_rots, dtype=self.dtype, seed=7 + +def test_register_rots(rot_obj): + q_mat = Rotation.generate_random_rotations(1, dtype=rot_obj.dtype)[0] + for flag in [0, 1]: + regrots_ref = rot_obj.apply_registration(q_mat, flag) + q_mat_est, flag_est = rot_obj.find_registration(regrots_ref) + np.testing.assert_allclose(flag_est, flag) + np.testing.assert_allclose( + q_mat_est, q_mat, atol=utest_tolerance(rot_obj.dtype) ) - # Align both sets of random rotations to rot_obj - aligned_rots1 = self.rot_obj.register(set1) - aligned_rots2 = self.rot_obj.register(set2) - self.assertTrue(aligned_rots1.mse(aligned_rots2) < utest_tolerance(self.dtype)) - self.assertTrue(aligned_rots2.mse(aligned_rots1) < utest_tolerance(self.dtype)) - - def testMSE(self): - q_ang = [np.random.random(3)] - q_mat = sp_rot.from_euler("ZYZ", q_ang, degrees=False).as_matrix()[0] - for flag in [0, 1]: - regrots_ref = self.rot_obj.apply_registration(q_mat, flag) - mse = self.rot_obj.mse(regrots_ref) - self.assertTrue(mse < utest_tolerance(self.dtype)) - - def testCommonLines(self): - ell_ij, ell_ji = self.rot_obj.common_lines(8, 11, 360) - self.assertTrue(ell_ij == 235 and ell_ji == 284) - - def testString(self): - logger.debug(str(self.rot_obj)) - - def testRepr(self): - logger.debug(repr(self.rot_obj)) - - def testLen(self): - self.assertTrue(len(self.rot_obj) == self.num_rots) - - def testSetterGetter(self): - # Excute set - tmp = np.arange(9).reshape((3, 3)) - self.rot_obj[13] = tmp - # Execute get - self.assertTrue(np.all(self.rot_obj[13] == tmp)) - - def testDtype(self): - self.assertTrue(self.dtype == self.rot_obj.dtype) - - def testFromRotvec(self): - # Build random rotation vectors. - axis = np.array([1, 0, 0], dtype=self.dtype) - angles = np.random.uniform(0, 2 * np.pi, 10) - rot_vecs = np.array([angle * axis for angle in angles], dtype=self.dtype) - - # Build rotations using from_rotvec and about_axis (as reference). - rotations = Rotation.from_rotvec(rot_vecs, dtype=self.dtype) - ref_rots = Rotation.about_axis("x", angles, dtype=self.dtype) - - self.assertTrue(isinstance(rotations, Rotation)) - self.assertTrue(rotations.matrices.dtype == self.dtype) - self.assertTrue(np.allclose(rotations.matrices, ref_rots.matrices)) - - -def test_angle_dist(): - dtype = np.float32 + + +def test_register(rot_obj): + # These will yield two more distinct sets of random rotations wrt rot_obj + set1 = Rotation.generate_random_rotations(NUM_ROTS, dtype=rot_obj.dtype) + set2 = Rotation.generate_random_rotations(NUM_ROTS, dtype=rot_obj.dtype, seed=7) + # Align both sets of random rotations to rot_obj + aligned_rots1 = rot_obj.register(set1) + aligned_rots2 = rot_obj.register(set2) + tol = utest_tolerance(rot_obj.dtype) + np.testing.assert_array_less(aligned_rots1.mse(aligned_rots2), tol) + np.testing.assert_array_less(aligned_rots2.mse(aligned_rots1), tol) + + +def test_mse(rot_obj): + q_ang = [np.random.random(3)] + q_mat = sp_rot.from_euler("ZYZ", q_ang, degrees=False).as_matrix()[0] + for flag in [0, 1]: + regrots_ref = rot_obj.apply_registration(q_mat, flag) + mse = rot_obj.mse(regrots_ref) + np.testing.assert_array_less(mse, utest_tolerance(rot_obj.dtype)) + + +def test_common_lines(rot_obj): + ell_ij, ell_ji = rot_obj.common_lines(8, 11, 360) + np.testing.assert_equal([ell_ij, ell_ji], [235, 284]) + + +def test_string(rot_obj): + logger.debug(str(rot_obj)) + + +def test_repr(rot_obj): + logger.debug(repr(rot_obj)) + + +def test_len(rot_obj): + assert len(rot_obj) == NUM_ROTS + + +def test_setter_getter(rot_obj): + # Excute set + tmp = np.arange(9).reshape((3, 3)) + rot_obj[13] = tmp + # Execute get + np.testing.assert_equal(rot_obj[13], tmp) + + +def test_dtype(dtype, rot_obj): + assert dtype == rot_obj.dtype + + +def test_from_rotvec(rot_obj): + # Build random rotation vectors. + axis = np.array([1, 0, 0], dtype=rot_obj.dtype) + angles = np.random.uniform(0, 2 * np.pi, 10) + rot_vecs = np.array([angle * axis for angle in angles], dtype=rot_obj.dtype) + + # Build rotations using from_rotvec and about_axis (as reference). + rotations = Rotation.from_rotvec(rot_vecs, dtype=rot_obj.dtype) + ref_rots = Rotation.about_axis("x", angles, dtype=rot_obj.dtype) + + assert isinstance(rotations, Rotation) + assert rotations.matrices.dtype == rot_obj.dtype + np.testing.assert_allclose(rotations.matrices, ref_rots.matrices) + + +# Angular Distance Tests + + +def test_angle_dist(dtype): angles = np.array([i * np.pi / 360 for i in range(360)], dtype=dtype) rots = Rotation.about_axis("x", angles, dtype=dtype) @@ -140,12 +200,42 @@ def test_angle_dist(): _ = Rotation.angle_dist(rots[:3], rots[:5]) -def test_mean_angular_distance(): - rots_z = Rotation.about_axis( - "z", [0, np.pi / 4, np.pi / 2], dtype=np.float32 - ).matrices - rots_id = Rotation.about_axis("z", [0, 0, 0], dtype=np.float32).matrices +def test_mean_angular_distance(dtype): + rots_z = Rotation.about_axis("z", [0, np.pi / 4, np.pi / 2], dtype=dtype).matrices + rots_id = Rotation.about_axis("z", [0, 0, 0], dtype=dtype).matrices mean_ang_dist = Rotation.mean_angular_distance(rots_z, rots_id) assert np.allclose(mean_ang_dist, np.pi / 4) + + +# Basis Rotations + + +def test_basis_rotation_2d(basis): + """ + Test steerable basis rotation performs similar operation to PIL real space image rotation. + + Checks both orientation and rough values. + """ + # Set a rotation amount + rot_radians = np.pi / 6 + + # Create an empty image + L = basis.nres + img = np.zeros((L, L), dtype=basis.dtype) + # Set one pixel (between I and IV quadrants) + img[L // 2, 3 * L // 4] = 1 + # Roundtrip using the basis. Smooths out the discontinuous pixel. + img = basis.expand(Image(img)).evaluate() + + # Rotate with ASPIRE Steerable Basis, returning to real space. + rot_img = basis.expand(img).rotate(rot_radians).evaluate() + + # Rotate image with PIL, returning to Numpy array. + pil_rot_img = np.asarray( + PILImage.fromarray(img.asnumpy()[0]).rotate(rot_radians * 180 / np.pi) + ) + + # Rough compare arrays. + np.testing.assert_allclose(rot_img.asnumpy()[0], pil_rot_img, atol=0.25) From b774f92d48e7c4691cf6fb8c89fe6316bbbd350d Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 27 Nov 2023 14:57:46 -0500 Subject: [PATCH 2/9] Convert basis rotation test to use Gaussian blob --- tests/test_rotation.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/test_rotation.py b/tests/test_rotation.py index 10e32e7487..4f4756667e 100644 --- a/tests/test_rotation.py +++ b/tests/test_rotation.py @@ -3,12 +3,11 @@ import numpy as np import PIL.Image as PILImage import pytest -import scipy as sp from scipy.spatial.transform import Rotation as sp_rot from aspire.basis import FBBasis2D, FFBBasis2D, FLEBasis2D, FPSWFBasis2D, PSWFBasis2D from aspire.image import Image -from aspire.utils import Rotation, grid_2d, utest_tolerance +from aspire.utils import Rotation, gaussian_2d, utest_tolerance logger = logging.getLogger(__name__) @@ -221,20 +220,18 @@ def test_basis_rotation_2d(basis): # Set a rotation amount rot_radians = np.pi / 6 - # Create an empty image + # Create an Image containing a smooth blob. L = basis.nres - img = np.zeros((L, L), dtype=basis.dtype) - # Set one pixel (between I and IV quadrants) - img[L // 2, 3 * L // 4] = 1 - # Roundtrip using the basis. Smooths out the discontinuous pixel. - img = basis.expand(Image(img)).evaluate() + img = Image(gaussian_2d(L, mu=(L // 4, 0), dtype=basis.dtype)) # Rotate with ASPIRE Steerable Basis, returning to real space. rot_img = basis.expand(img).rotate(rot_radians).evaluate() # Rotate image with PIL, returning to Numpy array. pil_rot_img = np.asarray( - PILImage.fromarray(img.asnumpy()[0]).rotate(rot_radians * 180 / np.pi) + PILImage.fromarray(img.asnumpy()[0]).rotate( + rot_radians * 180 / np.pi, resample=PILImage.BILINEAR + ) ) # Rough compare arrays. From 9cf2a6ede201a3a3334a37a084a2e0f959874eef Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 28 Nov 2023 09:15:01 -0500 Subject: [PATCH 3/9] Update rot test params to run doubles in long running reduces regular ut time --- tests/test_rotation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_rotation.py b/tests/test_rotation.py index 4f4756667e..46fdbdf403 100644 --- a/tests/test_rotation.py +++ b/tests/test_rotation.py @@ -11,14 +11,15 @@ logger = logging.getLogger(__name__) + # Parameters NUM_ROTS = 32 SEED = 0 DTYPES = [ - np.float64, np.float32, + pytest.param(np.float64, marks=pytest.mark.expensive), ] BASES = [ From d91552ba8501e23b3e3b694cc27b0631d9f2eb4d Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 28 Nov 2023 09:22:14 -0500 Subject: [PATCH 4/9] Remove old FFB and FLE specific rotation tests --- tests/test_FFBbasis2D.py | 90 ---------------------------------------- tests/test_FLEbasis2D.py | 72 -------------------------------- 2 files changed, 162 deletions(-) diff --git a/tests/test_FFBbasis2D.py b/tests/test_FFBbasis2D.py index 04b7a9b4c7..73bb8667e8 100644 --- a/tests/test_FFBbasis2D.py +++ b/tests/test_FFBbasis2D.py @@ -81,96 +81,6 @@ def testElements(self, basis): for ell, k, sgn in zip(ells, ks, sgns): self._testElement(basis, ell, k, sgn) - def testRotate(self, basis): - # Now low res (8x8) had problems; - # better with odd (7x7), but still not good. - # We'll use a higher res test image. - # fh = np.load(os.path.join(DATA_DIR, 'ffbbasis2d_xcoef_in_8_8.npy'))[:7,:7] - # Use a real data volume to generate a clean test image. - v = Volume( - np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype( - basis.dtype - ) - ) - src = Simulation(L=v.resolution, n=1, vols=v, dtype=v.dtype) - # Extract, this is the original image to transform. - x1 = src.images[0] - - # Rotate 90 degrees CCW in cartesian coordinates. - x2 = Image(np.rot90(x1.asnumpy(), axes=(1, 2))) - - # Express in an FB basis - basis = FFBBasis2D(x1.resolution, dtype=x1.dtype) - v1 = basis.evaluate_t(x1) - v2 = basis.evaluate_t(x2) - - # Reflect in the FB basis space - v4 = basis.rotate(v1, 0, refl=[True]) - - # Rotate in the FB basis space - v3 = basis.rotate(v1, 2 * np.pi) - v1 = basis.rotate(v1, np.pi / 2) - - # Evaluate back into cartesian - y1 = basis.evaluate(v1).asnumpy() - y2 = basis.evaluate(v2).asnumpy() - y3 = basis.evaluate(v3).asnumpy() - y4 = basis.evaluate(v4).asnumpy() - - # Rotate 90 - assert np.allclose(y1[0], y2[0], atol=1e-5) - - # 2*pi Identity - assert np.allclose(x1[0], y3[0], atol=1e-5) - - # Refl (flipped using flipud) - assert np.allclose(np.flipud(x1.asnumpy()[0]), y4[0], atol=1e-5) - - def testRotateComplex(self, basis): - # Now low res (8x8) had problems; - # better with odd (7x7), but still not good. - # We'll use a higher res test image. - # fh = np.load(os.path.join(DATA_DIR, 'ffbbasis2d_xcoef_in_8_8.npy'))[:7,:7] - # Use a real data volume to generate a clean test image. - v = Volume( - np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype( - basis.dtype - ) - ) - src = Simulation(L=v.resolution, n=1, vols=v, dtype=v.dtype) - # Extract, this is the original image to transform. - x1 = src.images[0] - - # Rotate 90 degrees CCW in cartesian coordinates. - x2 = Image(np.rot90(x1.asnumpy(), axes=(1, 2))) - - # Express in an FB basis - basis = FFBBasis2D(x1.resolution, dtype=x1.dtype) - v1 = basis.evaluate_t(x1) - v2 = basis.evaluate_t(x2) - - # Reflect in the FB basis space - v4 = basis.to_real(basis.complex_rotate(basis.to_complex(v1), 0, refl=[True])) - - # Complex Rotate in the FB basis space - v3 = basis.to_real(basis.complex_rotate(basis.to_complex(v1), 2 * np.pi)) - v1 = basis.to_real(basis.complex_rotate(basis.to_complex(v1), np.pi / 2)) - - # Evaluate back into cartesian - y1 = basis.evaluate(v1).asnumpy() - y2 = basis.evaluate(v2).asnumpy() - y3 = basis.evaluate(v3).asnumpy() - y4 = basis.evaluate(v4).asnumpy() - - # Rotate 90 - assert np.allclose(y1[0], y2[0], atol=1e-5) - - # 2*pi Identity - assert np.allclose(x1[0].asnumpy(), y3[0], atol=1e-5) - - # Refl (flipped using flipud) - assert np.allclose(np.flipud(x1.asnumpy()[0]), y4[0], atol=1e-5) - def testShift(self, basis): """ Compare shifting using Image with shifting provided by the Basis. diff --git a/tests/test_FLEbasis2D.py b/tests/test_FLEbasis2D.py index 9fb5504202..f25885684e 100644 --- a/tests/test_FLEbasis2D.py +++ b/tests/test_FLEbasis2D.py @@ -228,78 +228,6 @@ def testLowPass(): _ = basis.lowpass(coefs[0, :], L) -def testRotate(): - # test ability to accurately rotate images via - # FLE coefficients - - L = 128 - basis = FLEBasis2D(L, match_fb=False) - - # sample image - ims = create_images(L, 1) - # rotate 90 degrees in cartesian coordinates - ims_90 = Image(np.rot90(ims.asnumpy(), axes=(1, 2))) - - # get FLE coefficients - coefs = basis.evaluate_t(ims) - coefs_cart_rot = basis.evaluate_t(ims_90) - - # rotate original image in FLE space using Steerable rotate method - coefs_fle_rot = basis.rotate(coefs, np.pi / 2) - - # back to cartesian - ims_cart_rot = basis.evaluate(coefs_cart_rot) - ims_fle_rot = basis.evaluate(coefs_fle_rot) - - # test rot90 close - assert np.allclose(ims_cart_rot[0], ims_fle_rot[0], atol=1e-4) - - # 2Pi identity in FLE space (rotate by 2Pi) - coefs_fle_2pi = basis.rotate(coefs, 2 * np.pi) - ims_fle_2pi = basis.evaluate(coefs_fle_2pi) - - # test 2Pi identity - assert np.allclose(ims[0], ims_fle_2pi[0], atol=utest_tolerance(basis.dtype)) - - # Reflect in FLE space (rotate by Pi) - coefs_fle_pi = basis.rotate(coefs, np.pi) - ims_fle_pi = basis.evaluate(coefs_fle_pi) - - # test reflection - assert np.allclose(np.fliplr(np.flipud(ims.asnumpy()[0])), ims_fle_pi[0], atol=1e-4) - - # make sure you can pass in a 1-D array if you want - _ = basis.lowpass(Coef(basis, np.zeros((basis.count,))), np.pi) - - -def testRotate45(): - # test ability to accurately rotate images via - # FLE coefficients - dtype = np.float64 - - L = 128 - fb_basis = FFBBasis2D(L, dtype=dtype) - basis = FLEBasis2D(L, match_fb=True, dtype=dtype) - - # sample image - ims = create_images(L, 1) - - # get FLE coefficients - fb_coefs = fb_basis.evaluate_t(ims) - coefs = basis.evaluate_t(ims) - - # rotate original image in FLE space using Steerable rotate method - fb_coefs_rot = fb_basis.rotate(fb_coefs, np.pi / 4) - coefs_rot = basis.rotate(coefs, np.pi / 4) - - # back to cartesian - fb_ims_rot = fb_basis.evaluate(fb_coefs_rot) - ims_rot = basis.evaluate(coefs_rot) - - # test close - assert np.allclose(ims_rot[0], fb_ims_rot[0], atol=1e-4) - - def testRadialConvolution(): # test ability to accurately convolve with a radial # (e.g. CTF) function via FLE coefficients From f0ae0a5ddfad8e12149a8aa94199ccc75fc2ac94 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 28 Nov 2023 09:47:09 -0500 Subject: [PATCH 5/9] move new steerable rotate tests out of test_rotation --- tests/test_FFBbasis2D.py | 1 - tests/test_FLEbasis2D.py | 3 +- tests/test_rotation.py | 62 +----------------------- tests/test_steerable_bases_2d.py | 82 ++++++++++++++++++++++++++++++++ 4 files changed, 85 insertions(+), 63 deletions(-) create mode 100644 tests/test_steerable_bases_2d.py diff --git a/tests/test_FFBbasis2D.py b/tests/test_FFBbasis2D.py index 73bb8667e8..8acf7201d1 100644 --- a/tests/test_FFBbasis2D.py +++ b/tests/test_FFBbasis2D.py @@ -6,7 +6,6 @@ from scipy.special import jv from aspire.basis import Coef, FFBBasis2D -from aspire.image import Image from aspire.source import Simulation from aspire.utils.misc import grid_2d from aspire.volume import Volume diff --git a/tests/test_FLEbasis2D.py b/tests/test_FLEbasis2D.py index f25885684e..7d6b3f5c47 100644 --- a/tests/test_FLEbasis2D.py +++ b/tests/test_FLEbasis2D.py @@ -4,12 +4,11 @@ import numpy as np import pytest -from aspire.basis import Coef, FBBasis2D, FFBBasis2D, FLEBasis2D +from aspire.basis import Coef, FBBasis2D, FLEBasis2D from aspire.image import Image from aspire.nufft import backend_available from aspire.numeric import fft from aspire.source import Simulation -from aspire.utils import utest_tolerance from aspire.volume import Volume from ._basis_util import UniversalBasisMixin diff --git a/tests/test_rotation.py b/tests/test_rotation.py index 46fdbdf403..aaa1262695 100644 --- a/tests/test_rotation.py +++ b/tests/test_rotation.py @@ -1,13 +1,10 @@ import logging import numpy as np -import PIL.Image as PILImage import pytest from scipy.spatial.transform import Rotation as sp_rot -from aspire.basis import FBBasis2D, FFBBasis2D, FLEBasis2D, FPSWFBasis2D, PSWFBasis2D -from aspire.image import Image -from aspire.utils import Rotation, gaussian_2d, utest_tolerance +from aspire.utils import Rotation, utest_tolerance logger = logging.getLogger(__name__) @@ -19,21 +16,9 @@ DTYPES = [ np.float32, - pytest.param(np.float64, marks=pytest.mark.expensive), + np.float64, ] -BASES = [ - FFBBasis2D, - FBBasis2D, - FLEBasis2D, - PSWFBasis2D, - FPSWFBasis2D, -] - -IMG_SIZES = [ - 32, - pytest.param(31, marks=pytest.mark.expensive), -] # Fixtures @@ -48,19 +33,6 @@ def rot_obj(dtype): return Rotation.generate_random_rotations(NUM_ROTS, seed=SEED, dtype=dtype) -@pytest.fixture(params=IMG_SIZES, ids=lambda x: f"img_size={x}", scope="module") -def img_size(request): - return request.param - - -@pytest.fixture(params=BASES, ids=lambda x: f"basis={x}", scope="module") -def basis(request, img_size, dtype): - cls = request.param - # Setup a Basis - basis = cls(img_size, dtype=dtype) - return basis - - # Rotation Class Tests @@ -207,33 +179,3 @@ def test_mean_angular_distance(dtype): mean_ang_dist = Rotation.mean_angular_distance(rots_z, rots_id) assert np.allclose(mean_ang_dist, np.pi / 4) - - -# Basis Rotations - - -def test_basis_rotation_2d(basis): - """ - Test steerable basis rotation performs similar operation to PIL real space image rotation. - - Checks both orientation and rough values. - """ - # Set a rotation amount - rot_radians = np.pi / 6 - - # Create an Image containing a smooth blob. - L = basis.nres - img = Image(gaussian_2d(L, mu=(L // 4, 0), dtype=basis.dtype)) - - # Rotate with ASPIRE Steerable Basis, returning to real space. - rot_img = basis.expand(img).rotate(rot_radians).evaluate() - - # Rotate image with PIL, returning to Numpy array. - pil_rot_img = np.asarray( - PILImage.fromarray(img.asnumpy()[0]).rotate( - rot_radians * 180 / np.pi, resample=PILImage.BILINEAR - ) - ) - - # Rough compare arrays. - np.testing.assert_allclose(rot_img.asnumpy()[0], pil_rot_img, atol=0.25) diff --git a/tests/test_steerable_bases_2d.py b/tests/test_steerable_bases_2d.py new file mode 100644 index 0000000000..863ad9c829 --- /dev/null +++ b/tests/test_steerable_bases_2d.py @@ -0,0 +1,82 @@ +import logging + +import numpy as np +import PIL.Image as PILImage +import pytest + +from aspire.basis import FBBasis2D, FFBBasis2D, FLEBasis2D, FPSWFBasis2D, PSWFBasis2D +from aspire.image import Image +from aspire.utils import gaussian_2d + +logger = logging.getLogger(__name__) + + +# Parameters + +DTYPES = [ + np.float32, + pytest.param(np.float64, marks=pytest.mark.expensive), +] + +BASES = [ + FFBBasis2D, + FBBasis2D, + FLEBasis2D, + PSWFBasis2D, + FPSWFBasis2D, +] + +IMG_SIZES = [ + 32, + pytest.param(31, marks=pytest.mark.expensive), +] + +# Fixtures + + +@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module") +def dtype(request): + return request.param + + +@pytest.fixture(params=IMG_SIZES, ids=lambda x: f"img_size={x}", scope="module") +def img_size(request): + return request.param + + +@pytest.fixture(params=BASES, ids=lambda x: f"basis={x}", scope="module") +def basis(request, img_size, dtype): + cls = request.param + # Setup a Basis + basis = cls(img_size, dtype=dtype) + return basis + + +# Basis Rotations + + +def test_basis_rotation_2d(basis): + """ + Test steerable basis rotation performs similar operation to PIL real space image rotation. + + Checks both orientation and rough values. + """ + # Set a rotation amount + rot_radians = np.pi / 6 + + # Create an Image containing a smooth blob. + L = basis.nres + img = Image(gaussian_2d(L, mu=(L // 4, 0), dtype=basis.dtype)) + + # Rotate with ASPIRE Steerable Basis, returning to real space. + rot_img = basis.expand(img).rotate(rot_radians).evaluate() + + # Rotate image with PIL, returning to Numpy array. + pil_rot_img = np.asarray( + PILImage.fromarray(img.asnumpy()[0]).rotate( + rot_radians * 180 / np.pi, resample=PILImage.BILINEAR + ) + ) + + # Rough compare arrays. + np.testing.assert_allclose(rot_img.asnumpy()[0], pil_rot_img, atol=0.25) From bb4212279b5d88b920073b07aff79f298f5c99e2 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 28 Nov 2023 11:50:10 -0500 Subject: [PATCH 6/9] Add a shared steerable reflection test. --- tests/test_steerable_bases_2d.py | 35 ++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/test_steerable_bases_2d.py b/tests/test_steerable_bases_2d.py index 863ad9c829..9a9d181b0d 100644 --- a/tests/test_steerable_bases_2d.py +++ b/tests/test_steerable_bases_2d.py @@ -80,3 +80,38 @@ def test_basis_rotation_2d(basis): # Rough compare arrays. np.testing.assert_allclose(rot_img.asnumpy()[0], pil_rot_img, atol=0.25) + + +def test_basis_reflection_2d(basis): + """ + Test steerable basis reflection performs similar operation to Numpy flips. + + Checks both orientation and rough values. + """ + + # Create an Image containing a smooth blob. + L = basis.nres + img = Image(gaussian_2d(L, mu=(L // 4, L // 5), dtype=basis.dtype)) + + # Reflect with ASPIRE Steerable Basis, returning to real space. + refl_img = basis.expand(img).rotate(0, refl=True).evaluate() + + # Reflect image with Numpy. + # Note for odd images we can simply use Numpy, + # but evens have the expected offset issue + # when compared to a row/col based flip. + flip = np.flipud + if isinstance(basis, PSWFBasis2D): + # TODO, reconcile PSWF reflection axis + flip = np.fliplr + + refl_img_np = flip(img.asnumpy()[0]) + + # Rough compare arrays. + atol = 0.01 + if L % 2 == 0: + # Even images test is crude, + # but is enough ensure flipping without complicating test. + atol = 0.5 + + np.testing.assert_allclose(refl_img.asnumpy()[0], refl_img_np, atol=atol) From 0b2a5306a171a15b2359af645a3e1b0da0df17c2 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 28 Nov 2023 13:14:03 -0500 Subject: [PATCH 7/9] Touch up tests --- tests/test_rotation.py | 4 +++- tests/test_steerable_bases_2d.py | 12 ++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/test_rotation.py b/tests/test_rotation.py index aaa1262695..9e0dba4ec6 100644 --- a/tests/test_rotation.py +++ b/tests/test_rotation.py @@ -92,7 +92,9 @@ def test_register_rots(rot_obj): def test_register(rot_obj): # These will yield two more distinct sets of random rotations wrt rot_obj set1 = Rotation.generate_random_rotations(NUM_ROTS, dtype=rot_obj.dtype) - set2 = Rotation.generate_random_rotations(NUM_ROTS, dtype=rot_obj.dtype, seed=7) + set2 = Rotation.generate_random_rotations( + NUM_ROTS, dtype=rot_obj.dtype, seed=SEED + 7 + ) # Align both sets of random rotations to rot_obj aligned_rots1 = rot_obj.register(set1) aligned_rots2 = rot_obj.register(set2) diff --git a/tests/test_steerable_bases_2d.py b/tests/test_steerable_bases_2d.py index 9a9d181b0d..88a90ac15d 100644 --- a/tests/test_steerable_bases_2d.py +++ b/tests/test_steerable_bases_2d.py @@ -27,8 +27,8 @@ ] IMG_SIZES = [ - 32, - pytest.param(31, marks=pytest.mark.expensive), + 31, + pytest.param(32, marks=pytest.mark.expensive), ] # Fixtures @@ -68,7 +68,7 @@ def test_basis_rotation_2d(basis): L = basis.nres img = Image(gaussian_2d(L, mu=(L // 4, 0), dtype=basis.dtype)) - # Rotate with ASPIRE Steerable Basis, returning to real space. + # Rotate with an ASPIRE steerable basis, returning to real space. rot_img = basis.expand(img).rotate(rot_radians).evaluate() # Rotate image with PIL, returning to Numpy array. @@ -93,13 +93,13 @@ def test_basis_reflection_2d(basis): L = basis.nres img = Image(gaussian_2d(L, mu=(L // 4, L // 5), dtype=basis.dtype)) - # Reflect with ASPIRE Steerable Basis, returning to real space. + # Reflect with an ASPIRE steerable basis, returning to real space. refl_img = basis.expand(img).rotate(0, refl=True).evaluate() # Reflect image with Numpy. - # Note for odd images we can simply use Numpy, + # Note for odd images we can accurately use Numpy, # but evens have the expected offset issue - # when compared to a row/col based flip. + # when compared to a purely row/col based flip. flip = np.flipud if isinstance(basis, PSWFBasis2D): # TODO, reconcile PSWF reflection axis From f32f023a8294aadad9dcbbe1e07bcb7c79748015 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 29 Nov 2023 14:20:09 -0500 Subject: [PATCH 8/9] normalize the images in steerable tests --- tests/test_steerable_bases_2d.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_steerable_bases_2d.py b/tests/test_steerable_bases_2d.py index 88a90ac15d..d2cf3cbb4b 100644 --- a/tests/test_steerable_bases_2d.py +++ b/tests/test_steerable_bases_2d.py @@ -66,7 +66,8 @@ def test_basis_rotation_2d(basis): # Create an Image containing a smooth blob. L = basis.nres - img = Image(gaussian_2d(L, mu=(L // 4, 0), dtype=basis.dtype)) + img = gaussian_2d(L, mu=(L // 4, 0), dtype=basis.dtype) + img = Image(img / np.linalg.norm(img)) # Normalize # Rotate with an ASPIRE steerable basis, returning to real space. rot_img = basis.expand(img).rotate(rot_radians).evaluate() @@ -91,7 +92,8 @@ def test_basis_reflection_2d(basis): # Create an Image containing a smooth blob. L = basis.nres - img = Image(gaussian_2d(L, mu=(L // 4, L // 5), dtype=basis.dtype)) + img = gaussian_2d(L, mu=(L // 4, L // 5), dtype=basis.dtype) + img = Image(img / np.linalg.norm(img)) # Normalize # Reflect with an ASPIRE steerable basis, returning to real space. refl_img = basis.expand(img).rotate(0, refl=True).evaluate() From aff113148e324ed106942df449cf5fa90e8d78a1 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 30 Nov 2023 08:41:23 -0500 Subject: [PATCH 9/9] Tighten up steerable rotate tol after normalization --- tests/test_steerable_bases_2d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_steerable_bases_2d.py b/tests/test_steerable_bases_2d.py index d2cf3cbb4b..9777c3d2ed 100644 --- a/tests/test_steerable_bases_2d.py +++ b/tests/test_steerable_bases_2d.py @@ -75,12 +75,12 @@ def test_basis_rotation_2d(basis): # Rotate image with PIL, returning to Numpy array. pil_rot_img = np.asarray( PILImage.fromarray(img.asnumpy()[0]).rotate( - rot_radians * 180 / np.pi, resample=PILImage.BILINEAR + rot_radians * 180 / np.pi, resample=PILImage.BICUBIC ) ) # Rough compare arrays. - np.testing.assert_allclose(rot_img.asnumpy()[0], pil_rot_img, atol=0.25) + np.testing.assert_allclose(rot_img.asnumpy()[0], pil_rot_img, atol=0.15) def test_basis_reflection_2d(basis):