From 2cb54494414e105a9dd18251c154212670c7f412 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 4 Aug 2025 10:40:04 -0400 Subject: [PATCH 01/14] stashing, need port radialnorm methods [skip ci] --- src/aspire/operators/filters.py | 39 +++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index b20247861b..a6c3390b83 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -511,6 +511,45 @@ def __init__( ) +class m_CTFFilter(Filter): + """ + Reproduce MATLAB's cryo_CTF_relion CTF (Contrast Transfer Function) Filter + + Note if comparing to legacy MATLAB cryo_CTF_Relion, + take care regarding defocus unit conversion to nm. + """ + + def _evaluate(self, omega): + # disregard omega, we'll be making our own grids. + n = omega.shape[-1] + + # Wavelength in nm. + lamb = 1.22639 / np.sqrt(self.voltage * 1000 + 0.97845 * self.voltage**2) + + # Divide by 10 to make pixel size in nm. BW is the + # bandwidth of the signal corresponding to the given pixel size. + BW = 1 / (self.pixel_size / 10) + + # RadiusNorm returns radii such that when multiplied by the + # bandwidth of the signal, we get the correct radial frequencies + # corresponding to each pixel in our nxn grid. + s, theta = self._RadiusNorm(n, fctr(n)) + + s = s * BW + DFavg = self._defocus_mean_nm * 10 # (DefocusU+DefocusV)/2 + DFdiff = self._defocus_diff_nm * 10 # (DefocusU-DefocusV) + df = DFavg + DFdiff * cos(2 * (theta - self.defocus_ang)) / 2 + + k2 = np.pi * lamb * df + # 10*6 converts Cs from mm to nm. + k4 = np.pi / 2 * 10**6 * self.Cs * lamb**3 + chi = k4 * s**4 - k2 * s**2 + + h = np.sqrt(1 - self.alpha ^ 2) * sin(chi) - self.alpha * cos(chi) + + return h + + class BlueFilter(Filter): """ Filter where power increases with frequency. From cda9a14ed8d40e9777041756b662eb9f16e1f255 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 4 Aug 2025 12:07:07 -0400 Subject: [PATCH 02/14] quick and dirty implement MATLAB's ctf filter [skip ci] --- src/aspire/operators/__init__.py | 1 + src/aspire/operators/filters.py | 22 +++++++++++++--------- src/aspire/source/relion.py | 4 ++-- 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/aspire/operators/__init__.py b/src/aspire/operators/__init__.py index 97e47a4d1d..61f3dc5e79 100644 --- a/src/aspire/operators/__init__.py +++ b/src/aspire/operators/__init__.py @@ -17,6 +17,7 @@ ScaledFilter, ZeroFilter, evaluate_src_filters_on_grid, + m_CTFFilter, ) from .polar_ft import PolarFT from .wemd import wemd_embed, wemd_norm diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index a6c3390b83..53d504abd3 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -511,7 +511,7 @@ def __init__( ) -class m_CTFFilter(Filter): +class m_CTFFilter(CTFFilter): """ Reproduce MATLAB's cryo_CTF_relion CTF (Contrast Transfer Function) Filter @@ -520,8 +520,8 @@ class m_CTFFilter(Filter): """ def _evaluate(self, omega): - # disregard omega, we'll be making our own grids. - n = omega.shape[-1] + # disregard omega, we'll be making our own grids for now + L = int(np.sqrt(omega.shape[-1])) # Wavelength in nm. lamb = 1.22639 / np.sqrt(self.voltage * 1000 + 0.97845 * self.voltage**2) @@ -530,22 +530,26 @@ def _evaluate(self, omega): # bandwidth of the signal corresponding to the given pixel size. BW = 1 / (self.pixel_size / 10) - # RadiusNorm returns radii such that when multiplied by the + # Returns radii such that when multiplied by the # bandwidth of the signal, we get the correct radial frequencies # corresponding to each pixel in our nxn grid. - s, theta = self._RadiusNorm(n, fctr(n)) + # + # s, theta should match MATLAB's RadiusNorm + g = grid_2d(L, normalized=True, indexing="xy", dtype=np.float64) + s, theta = g["r"] / 2, g["phi"] s = s * BW - DFavg = self._defocus_mean_nm * 10 # (DefocusU+DefocusV)/2 - DFdiff = self._defocus_diff_nm * 10 # (DefocusU-DefocusV) - df = DFavg + DFdiff * cos(2 * (theta - self.defocus_ang)) / 2 + DFavg = self._defocus_mean_nm # (DefocusU+DefocusV)/2 + DFdiff = self._defocus_diff_nm # (DefocusU-DefocusV) + # Note the missing / 2 is already in _defocus_diff_nm + df = DFavg + DFdiff * np.cos(2 * (theta - self.defocus_ang)) k2 = np.pi * lamb * df # 10*6 converts Cs from mm to nm. k4 = np.pi / 2 * 10**6 * self.Cs * lamb**3 chi = k4 * s**4 - k2 * s**2 - h = np.sqrt(1 - self.alpha ^ 2) * sin(chi) - self.alpha * cos(chi) + h = np.sqrt(1 - self.alpha**2) * np.sin(chi) - self.alpha * np.cos(chi) return h diff --git a/src/aspire/source/relion.py b/src/aspire/source/relion.py index 09bf77b8ea..00938801b9 100644 --- a/src/aspire/source/relion.py +++ b/src/aspire/source/relion.py @@ -7,7 +7,7 @@ import numpy as np from aspire.image import Image -from aspire.operators import CTFFilter, IdentityFilter +from aspire.operators import IdentityFilter, m_CTFFilter from aspire.source import ImageSource from aspire.utils import RelionStarFile @@ -152,7 +152,7 @@ def __init__( # for each unique CTF configuration, create a CTFFilter object for row in filter_params: filters.append( - CTFFilter( + m_CTFFilter( pixel_size=self.pixel_size, voltage=row[0], defocus_u=row[1], From 6c80f0714f867c60360bcef67a9ca5b2b9d4aaac Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 8 Aug 2025 12:01:24 -0400 Subject: [PATCH 03/14] match matlab pf only temp file --- src/aspire/image/image.py | 36 ++++++++++++++++++++++++++++++++++++ src/aspire/image/xform.py | 2 +- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 185b9f798c..7559748ede 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -607,6 +607,42 @@ def filter(self, filter): original_stack_shape ) + def filter64(self, filter): + """ + Apply a `Filter` object to the Image and returns a new Image. + + :param filter: An object of type `Filter`. + :return: A new filtered `Image` object. + """ + original_stack_shape = self.stack_shape + + im = self.stack_reshape(-1) + + # Note image and filter data is intentionally migrated via + # `xp.asarray` because all of the subsequent calls until + # `asnumpy` are GPU when xp and fft in `cupy` mode. + # + # Second note, filter dtype may not match image dtype. + filter_values = xp.asarray( + filter.evaluate_grid(self.resolution), dtype=np.float64 + ) + + #breakpoint() + # Convolve + _im = xp.asarray(im._data, dtype=np.float64).transpose(0,2,1) + im_f = fft.fftshift(fft.fft2(_im), axes=(-1,-2)) + im_f = filter_values * im_f + im = fft.ifft2(fft.ifftshift(im_f, axes=(-1,-2))) + + im = xp.asnumpy(im.real.transpose(0,2,1)).astype( + self.dtype, copy=False + ) # restore to original dtype + # this matches for first image up to transposing, what now + + return self.__class__(im, pixel_size=self.pixel_size).stack_reshape( + original_stack_shape + ) + def rotate(self): raise NotImplementedError diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py index 316d8ca3c8..30aab2faa1 100644 --- a/src/aspire/image/xform.py +++ b/src/aspire/image/xform.py @@ -300,7 +300,7 @@ def __init__(self, filter): self.filter = filter def _forward(self, im, indices): - return im.filter(self.filter) + return im.filter64(self.filter) def __str__(self): return f"FilterXform ({self.filter})" From 5168b7a916c966795a718de13ed03b0b4298954b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 8 Aug 2025 12:05:41 -0400 Subject: [PATCH 04/14] match single precision removes redundant tranposes --- src/aspire/image/image.py | 10 ++++------ src/aspire/operators/filters.py | 4 ++-- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 7559748ede..a74d91b234 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -627,17 +627,15 @@ def filter64(self, filter): filter.evaluate_grid(self.resolution), dtype=np.float64 ) - #breakpoint() # Convolve - _im = xp.asarray(im._data, dtype=np.float64).transpose(0,2,1) - im_f = fft.fftshift(fft.fft2(_im), axes=(-1,-2)) + _im = xp.asarray(im._data, dtype=np.float64) + im_f = fft.fftshift(fft.fft2(_im), axes=(-1, -2)) im_f = filter_values * im_f - im = fft.ifft2(fft.ifftshift(im_f, axes=(-1,-2))) + im = fft.ifft2(fft.ifftshift(im_f, axes=(-1, -2))) - im = xp.asnumpy(im.real.transpose(0,2,1)).astype( + im = xp.asnumpy(im.real).astype( self.dtype, copy=False ) # restore to original dtype - # this matches for first image up to transposing, what now return self.__class__(im, pixel_size=self.pixel_size).stack_reshape( original_stack_shape diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index 53d504abd3..7308516a85 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -534,8 +534,8 @@ def _evaluate(self, omega): # bandwidth of the signal, we get the correct radial frequencies # corresponding to each pixel in our nxn grid. # - # s, theta should match MATLAB's RadiusNorm - g = grid_2d(L, normalized=True, indexing="xy", dtype=np.float64) + # s, theta should match MATLAB's RadiusNorm up to a transpose + g = grid_2d(L, normalized=True, indexing="yx", dtype=np.float64) s, theta = g["r"] / 2, g["phi"] s = s * BW From 87528c937de221580923eac66caeea914f2aee2e Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 25 Aug 2025 08:21:47 -0400 Subject: [PATCH 05/14] replace CTFFilter with m_CTFFilter code --- src/aspire/operators/__init__.py | 1 - src/aspire/operators/filters.py | 109 +++++++++++-------------------- src/aspire/source/relion.py | 4 +- 3 files changed, 39 insertions(+), 75 deletions(-) diff --git a/src/aspire/operators/__init__.py b/src/aspire/operators/__init__.py index 61f3dc5e79..97e47a4d1d 100644 --- a/src/aspire/operators/__init__.py +++ b/src/aspire/operators/__init__.py @@ -17,7 +17,6 @@ ScaledFilter, ZeroFilter, evaluate_src_filters_on_grid, - m_CTFFilter, ) from .polar_ft import PolarFT from .wemd import wemd_embed, wemd_norm diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index 7308516a85..df687e62e3 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -36,6 +36,7 @@ def evaluate_src_filters_on_grid(src, indices=None): idx_k = np.where(src.filter_indices[indices] == i)[0] if len(idx_k) > 0: filter_values = filt.evaluate(omega) + breakpoint() h[:, idx_k] = np.column_stack((filter_values,) * len(idx_k)) h = np.reshape(h, grid2d["x"].shape + (len(indices),)) @@ -406,6 +407,13 @@ def __init__(self, dim=None): class CTFFilter(Filter): + """ + Reproduce MATLAB's cryo_CTF_relion CTF (Contrast Transfer Function) Filter + + Note if comparing to legacy MATLAB cryo_CTF_Relion, + take care regarding defocus unit conversion to/from nm. + """ + def __init__( self, pixel_size=1, @@ -432,7 +440,8 @@ def __init__( :param alpha: Amplitude contrast phase in radians :param B: Envelope decay in inverse square angstrom (default 0) """ - super().__init__(dim=2, radial=defocus_u == defocus_v) + #super().__init__(dim=2, radial=defocus_u == defocus_v) + super().__init__(dim=2) self.pixel_size = float(pixel_size) self.voltage = voltage self.wavelength = voltage_to_wavelength(self.voltage) @@ -448,39 +457,38 @@ def __init__( self._defocus_diff_nm = 0.05 * (self.defocus_u - self.defocus_v) def _evaluate(self, omega): - # Note the grid is wrt nm. - om_y, om_x = np.vsplit(omega / (2 * np.pi * self.pixel_size / 10), 2) - - eps = np.finfo(np.pi).eps - ind_nz = (np.abs(om_x) > eps) | (np.abs(om_y) > eps) - angles_nz = np.arctan2(om_y[ind_nz], om_x[ind_nz]) - angles_nz -= self.defocus_ang - - defocus = np.zeros_like(om_x) - # Note the division by 2 for _defocus_diff_nm is in `__init__`. - defocus[ind_nz] = self._defocus_mean_nm + self._defocus_diff_nm * np.cos( - 2 * angles_nz - ) + # disregard omega, we'll be making our own grids for now + L = int(np.sqrt(omega.shape[-1])) - # Note lambda must be in nm, and `Cs` must be converted from mm to nm. - lambda_nm = self.wavelength / 10 - c2 = -np.pi * lambda_nm * defocus - c4 = 0.5 * np.pi * (self.Cs * 1e6) * lambda_nm**3 + # Wavelength in nm. + lamb = 1.22639 / np.sqrt(self.voltage * 1000 + 0.97845 * self.voltage**2) - r2 = om_x**2 + om_y**2 - r4 = r2**2 - gamma = c2 * r2 + c4 * r4 - h = np.sqrt(1 - self.alpha**2) * np.sin(gamma) - self.alpha * np.cos(gamma) + # Divide by 10 to make pixel size in nm. BW is the + # bandwidth of the signal corresponding to the given pixel size. + BW = 1 / (self.pixel_size / 10) - # For historical reference, below is a translated formula from the legacy MATLAB code. - # The two implementations seem to agree for odd images, but the original MATLAB code - # behaves differently for even image sizes. - # h = np.sin(c2*r2 + c4*r2*r2 - self.alpha) + # Returns radii such that when multiplied by the + # bandwidth of the signal, we get the correct radial frequencies + # corresponding to each pixel in our nxn grid. + # + # s, theta should match MATLAB's RadiusNorm up to a transpose + g = grid_2d(L, normalized=True, indexing="yx", dtype=np.float64) + s, theta = g["r"] / 2, g["phi"] + + s = s * BW + DFavg = self._defocus_mean_nm # (DefocusU+DefocusV)/2 + DFdiff = self._defocus_diff_nm # (DefocusU-DefocusV) + # Note the missing / 2 is already in _defocus_diff_nm + df = DFavg + DFdiff * np.cos(2 * (theta - self.defocus_ang)) + + k2 = np.pi * lamb * df + # 10*6 converts Cs from mm to nm. + k4 = np.pi / 2 * 10**6 * self.Cs * lamb**3 + chi = k4 * s**4 - k2 * s**2 - if self.B: - h *= np.exp(-self.B * r2) + h = np.sqrt(1 - self.alpha**2) * np.sin(chi) - self.alpha * np.cos(chi) - return h.squeeze() + return h def scale(self, c=1): return CTFFilter( @@ -511,49 +519,6 @@ def __init__( ) -class m_CTFFilter(CTFFilter): - """ - Reproduce MATLAB's cryo_CTF_relion CTF (Contrast Transfer Function) Filter - - Note if comparing to legacy MATLAB cryo_CTF_Relion, - take care regarding defocus unit conversion to nm. - """ - - def _evaluate(self, omega): - # disregard omega, we'll be making our own grids for now - L = int(np.sqrt(omega.shape[-1])) - - # Wavelength in nm. - lamb = 1.22639 / np.sqrt(self.voltage * 1000 + 0.97845 * self.voltage**2) - - # Divide by 10 to make pixel size in nm. BW is the - # bandwidth of the signal corresponding to the given pixel size. - BW = 1 / (self.pixel_size / 10) - - # Returns radii such that when multiplied by the - # bandwidth of the signal, we get the correct radial frequencies - # corresponding to each pixel in our nxn grid. - # - # s, theta should match MATLAB's RadiusNorm up to a transpose - g = grid_2d(L, normalized=True, indexing="yx", dtype=np.float64) - s, theta = g["r"] / 2, g["phi"] - - s = s * BW - DFavg = self._defocus_mean_nm # (DefocusU+DefocusV)/2 - DFdiff = self._defocus_diff_nm # (DefocusU-DefocusV) - # Note the missing / 2 is already in _defocus_diff_nm - df = DFavg + DFdiff * np.cos(2 * (theta - self.defocus_ang)) - - k2 = np.pi * lamb * df - # 10*6 converts Cs from mm to nm. - k4 = np.pi / 2 * 10**6 * self.Cs * lamb**3 - chi = k4 * s**4 - k2 * s**2 - - h = np.sqrt(1 - self.alpha**2) * np.sin(chi) - self.alpha * np.cos(chi) - - return h - - class BlueFilter(Filter): """ Filter where power increases with frequency. diff --git a/src/aspire/source/relion.py b/src/aspire/source/relion.py index 00938801b9..09bf77b8ea 100644 --- a/src/aspire/source/relion.py +++ b/src/aspire/source/relion.py @@ -7,7 +7,7 @@ import numpy as np from aspire.image import Image -from aspire.operators import IdentityFilter, m_CTFFilter +from aspire.operators import CTFFilter, IdentityFilter from aspire.source import ImageSource from aspire.utils import RelionStarFile @@ -152,7 +152,7 @@ def __init__( # for each unique CTF configuration, create a CTFFilter object for row in filter_params: filters.append( - m_CTFFilter( + CTFFilter( pixel_size=self.pixel_size, voltage=row[0], defocus_u=row[1], From 0069faabce492ad22b64bf6991172b5d80755028 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 25 Aug 2025 09:05:11 -0400 Subject: [PATCH 06/14] rm dbg breakpoint --- src/aspire/operators/filters.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index df687e62e3..4d4ef66345 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -36,7 +36,6 @@ def evaluate_src_filters_on_grid(src, indices=None): idx_k = np.where(src.filter_indices[indices] == i)[0] if len(idx_k) > 0: filter_values = filt.evaluate(omega) - breakpoint() h[:, idx_k] = np.column_stack((filter_values,) * len(idx_k)) h = np.reshape(h, grid2d["x"].shape + (len(indices),)) @@ -440,7 +439,7 @@ def __init__( :param alpha: Amplitude contrast phase in radians :param B: Envelope decay in inverse square angstrom (default 0) """ - #super().__init__(dim=2, radial=defocus_u == defocus_v) + # super().__init__(dim=2, radial=defocus_u == defocus_v) super().__init__(dim=2) self.pixel_size = float(pixel_size) self.voltage = voltage From bc12ce116725a52b78da12112769c592978a0a33 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 25 Aug 2025 11:37:32 -0400 Subject: [PATCH 07/14] update filter unit tests --- src/aspire/operators/filters.py | 17 ++- tests/test_filters.py | 237 ++------------------------------ 2 files changed, 24 insertions(+), 230 deletions(-) diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index 4d4ef66345..80ce90fd64 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -439,8 +439,7 @@ def __init__( :param alpha: Amplitude contrast phase in radians :param B: Envelope decay in inverse square angstrom (default 0) """ - # super().__init__(dim=2, radial=defocus_u == defocus_v) - super().__init__(dim=2) + super().__init__(dim=2, radial=defocus_u == defocus_v) self.pixel_size = float(pixel_size) self.voltage = voltage self.wavelength = voltage_to_wavelength(self.voltage) @@ -455,9 +454,7 @@ def __init__( self._defocus_mean_nm = 0.05 * (self.defocus_u + self.defocus_v) self._defocus_diff_nm = 0.05 * (self.defocus_u - self.defocus_v) - def _evaluate(self, omega): - # disregard omega, we'll be making our own grids for now - L = int(np.sqrt(omega.shape[-1])) + def _evaluate(self, L): # Wavelength in nm. lamb = 1.22639 / np.sqrt(self.voltage * 1000 + 0.97845 * self.voltage**2) @@ -501,6 +498,16 @@ def scale(self, c=1): B=self.B, ) + def evaluate(self, omega): + # disregard omega, CTF will be making its own grids for now + L = int(np.sqrt(omega.shape[-1])) + h = self._evaluate(L) + return h.flatten() + + @lru_cache(maxsize=config["cache"]["filter_cache_size"].get()) # noqa: B019 + def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs): + return self._evaluate(L) + class RadialCTFFilter(CTFFilter): def __init__( diff --git a/tests/test_filters.py b/tests/test_filters.py index bf0bbbc2cb..bfd3908d71 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -26,6 +26,7 @@ def setUp(self): self.dtype = np.float32 # A 2 x 256 ndarray of spatial frequencies self.omega = np.load(os.path.join(DATA_DIR, "omega_2_256.npy")) + self.test_filter = ArrayFilter(np.random.randn(8, 8)) def tearDown(self): pass @@ -93,246 +94,28 @@ def testCTFFilter(self): self.assertEqual(result.shape, (256,)) def testScaledFilter(self): - filt1 = CTFFilter(defocus_u=1.5e4, defocus_v=1.5e4) scale_value = 2.5 - result1 = filt1.evaluate(self.omega) + result1 = self.test_filter.evaluate(self.omega) # ScaledFilter scales the pixel size which cancels out # a corresponding scaling in omega - filt2 = ScaledFilter(filt1, scale_value) + filt2 = ScaledFilter(self.test_filter, scale_value) result2 = filt2.evaluate(self.omega * scale_value) self.assertTrue(np.allclose(result1, result2, atol=utest_tolerance(self.dtype))) - def testCTFScale(self): - filt = CTFFilter(defocus_u=1.5e4, defocus_v=1.5e4) - result1 = filt.evaluate(self.omega) - scale_value = 2.5 - filt = filt.scale(scale_value) - # scaling a CTFFilter scales the pixel size which cancels out - # a corresponding scaling in omega - result2 = filt.evaluate(self.omega * scale_value) - self.assertTrue(np.allclose(result1, result2, atol=utest_tolerance(self.dtype))) - def testRadialCTFFilter(self): filter = RadialCTFFilter(defocus=2.5e4) result = filter.evaluate(self.omega) self.assertEqual(result.shape, (256,)) - def testRadialCTFFilterGrid(self): - # Set legacy pixel size - filter = RadialCTFFilter(pixel_size=10, defocus=2.5e4) - result = filter.evaluate_grid(8, dtype=self.dtype) - - self.assertEqual(result.shape, (8, 8)) - - # Setting tolerence to 1e-4. - # After precision was improved on `voltage_to_wavelength` method this reference array - # is no longer within utest_tolerance: np.max(abs(result - reference)) = 5.2729227306036464e-05 - self.assertTrue( - np.allclose( - result, - np.array( - [ - [ - 0.461755701877834, - -0.995184514498978, - 0.063120922443392, - 0.833250206225063, - 0.961464660252150, - 0.833250206225063, - 0.063120922443392, - -0.995184514498978, - ], - [ - -0.995184514498978, - 0.626977423649552, - 0.799934516166400, - 0.004814348317439, - -0.298096205735759, - 0.004814348317439, - 0.799934516166400, - 0.626977423649552, - ], - [ - 0.063120922443392, - 0.799934516166400, - -0.573061561512667, - -0.999286510416273, - -0.963805291282899, - -0.999286510416273, - -0.573061561512667, - 0.799934516166400, - ], - [ - 0.833250206225063, - 0.004814348317439, - -0.999286510416273, - -0.633095739808868, - -0.368890743119366, - -0.633095739808868, - -0.999286510416273, - 0.004814348317439, - ], - [ - 0.961464660252150, - -0.298096205735759, - -0.963805291282899, - -0.368890743119366, - -0.070000000000000, - -0.368890743119366, - -0.963805291282899, - -0.298096205735759, - ], - [ - 0.833250206225063, - 0.004814348317439, - -0.999286510416273, - -0.633095739808868, - -0.368890743119366, - -0.633095739808868, - -0.999286510416273, - 0.004814348317439, - ], - [ - 0.063120922443392, - 0.799934516166400, - -0.573061561512667, - -0.999286510416273, - -0.963805291282899, - -0.999286510416273, - -0.573061561512667, - 0.799934516166400, - ], - [ - -0.995184514498978, - 0.626977423649552, - 0.799934516166400, - 0.004814348317439, - -0.298096205735759, - 0.004814348317439, - 0.799934516166400, - 0.626977423649552, - ], - ] - ), - atol=1e-4, - ) - ) - - def testRadialCTFFilterMultiplierGrid(self): - # Set legacy pixel size - filter = RadialCTFFilter(pixel_size=10, defocus=2.5e4) * RadialCTFFilter( - pixel_size=10, defocus=2.5e4 - ) - result = filter.evaluate_grid(8, dtype=self.dtype) - - self.assertEqual(result.shape, (8, 8)) - - # Setting tolerence to 1e-4. - # After precision was improved on `voltage_to_wavelength` method this reference array - # is no longer within utest_tolerance: np.max(abs(result - reference)) = 4.869387449749074e-05 - self.assertTrue( - np.allclose( - result, - np.array( - [ - [ - 0.461755701877834, - -0.995184514498978, - 0.063120922443392, - 0.833250206225063, - 0.961464660252150, - 0.833250206225063, - 0.063120922443392, - -0.995184514498978, - ], - [ - -0.995184514498978, - 0.626977423649552, - 0.799934516166400, - 0.004814348317439, - -0.298096205735759, - 0.004814348317439, - 0.799934516166400, - 0.626977423649552, - ], - [ - 0.063120922443392, - 0.799934516166400, - -0.573061561512667, - -0.999286510416273, - -0.963805291282899, - -0.999286510416273, - -0.573061561512667, - 0.799934516166400, - ], - [ - 0.833250206225063, - 0.004814348317439, - -0.999286510416273, - -0.633095739808868, - -0.368890743119366, - -0.633095739808868, - -0.999286510416273, - 0.004814348317439, - ], - [ - 0.961464660252150, - -0.298096205735759, - -0.963805291282899, - -0.368890743119366, - -0.070000000000000, - -0.368890743119366, - -0.963805291282899, - -0.298096205735759, - ], - [ - 0.833250206225063, - 0.004814348317439, - -0.999286510416273, - -0.633095739808868, - -0.368890743119366, - -0.633095739808868, - -0.999286510416273, - 0.004814348317439, - ], - [ - 0.063120922443392, - 0.799934516166400, - -0.573061561512667, - -0.999286510416273, - -0.963805291282899, - -0.999286510416273, - -0.573061561512667, - 0.799934516166400, - ], - [ - -0.995184514498978, - 0.626977423649552, - 0.799934516166400, - 0.004814348317439, - -0.298096205735759, - 0.004814348317439, - 0.799934516166400, - 0.626977423649552, - ], - ] - ) - ** 2, - atol=1e-4, - ) - ) - def testDualFilter(self): - ctf_filter = CTFFilter(defocus_u=1.5e4, defocus_v=1.5e4) - result = ctf_filter.evaluate(-self.omega) - dual_filter = ctf_filter.dual() + result = self.test_filter.evaluate(-self.omega) + dual_filter = self.test_filter.dual() dual_result = dual_filter.evaluate(self.omega) self.assertTrue(np.allclose(result, dual_result)) def testFilterSigns(self): - ctf_filter = CTFFilter(defocus_u=1.5e4, defocus_v=1.5e4) - signs = np.sign(ctf_filter.evaluate(self.omega)) - sign_filter = ctf_filter.sign + signs = np.sign(self.test_filter.evaluate(self.omega)) + sign_filter = self.test_filter.sign self.assertTrue(np.allclose(sign_filter.evaluate(self.omega), signs)) @@ -419,6 +202,10 @@ def test_ctf_reference(): # Compare with MATLAB. Note DF converted to nm # >> n=5; V=200; DF1=1000; DF2=1500; theta=1.23; Cs=2.0; A=0.1; pxA=4.56; # >> ref_h=cryo_CTF_Relion(n,V,DF1,DF2,theta,Cs,pxA,A) + # + # Note we transpose the reference array. + # Python keeps the filter C order because the images we will convove with are C order. + # MATLAB is F and F respectively. ref_h = np.array( [ [-0.6152, 0.0299, -0.5638, 0.9327, 0.9736], @@ -427,7 +214,7 @@ def test_ctf_reference(): [0.1733, 0.9383, -0.7543, 0.2598, -0.9865], [0.9736, 0.9327, -0.5638, 0.0299, -0.6152], ] - ) + ).T # Test we're within 1%. # There are minor differences in the formulas for wavelength and grids. From 0d83ff7baa34ae310f9ee81331f610871c64ae64 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 25 Aug 2025 13:28:59 -0400 Subject: [PATCH 08/14] continue to update CTF and filter unit tests --- src/aspire/operators/filters.py | 33 +++++++++++++-------------------- src/aspire/utils/__init__.py | 1 + tests/test_covar2d.py | 3 ++- 3 files changed, 16 insertions(+), 21 deletions(-) diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index 80ce90fd64..00af22aa57 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -6,7 +6,7 @@ from scipy.interpolate import RegularGridInterpolator from aspire import config -from aspire.utils import grid_2d, voltage_to_wavelength +from aspire.utils import cart2pol, grid_2d, voltage_to_wavelength logger = logging.getLogger(__name__) @@ -454,7 +454,18 @@ def __init__( self._defocus_mean_nm = 0.05 * (self.defocus_u + self.defocus_v) self._defocus_diff_nm = 0.05 * (self.defocus_u - self.defocus_v) - def _evaluate(self, L): + def _evaluate(self, omega): + # s, theta should match MATLAB's RadiusNorm up to a transpose + # To accomplish this given ASPIRE-Python's defaul `omega`, + # we unpack and remove the pi scaling here, + # and further rescale the radii `s` by half below. + # Additionally we upcast so downstream computations remain in doubles. + x, y = omega.astype(np.float64, copy=False) / np.pi + # Returns radii such that when multiplied by the + # bandwidth of the signal, we get the correct radial frequencies + # corresponding to each pixel in our nxn grid. + theta, s = cart2pol(x, y) + s = s / 2 # Wavelength in nm. lamb = 1.22639 / np.sqrt(self.voltage * 1000 + 0.97845 * self.voltage**2) @@ -463,14 +474,6 @@ def _evaluate(self, L): # bandwidth of the signal corresponding to the given pixel size. BW = 1 / (self.pixel_size / 10) - # Returns radii such that when multiplied by the - # bandwidth of the signal, we get the correct radial frequencies - # corresponding to each pixel in our nxn grid. - # - # s, theta should match MATLAB's RadiusNorm up to a transpose - g = grid_2d(L, normalized=True, indexing="yx", dtype=np.float64) - s, theta = g["r"] / 2, g["phi"] - s = s * BW DFavg = self._defocus_mean_nm # (DefocusU+DefocusV)/2 DFdiff = self._defocus_diff_nm # (DefocusU-DefocusV) @@ -498,16 +501,6 @@ def scale(self, c=1): B=self.B, ) - def evaluate(self, omega): - # disregard omega, CTF will be making its own grids for now - L = int(np.sqrt(omega.shape[-1])) - h = self._evaluate(L) - return h.flatten() - - @lru_cache(maxsize=config["cache"]["filter_cache_size"].get()) # noqa: B019 - def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs): - return self._evaluate(L) - class RadialCTFFilter(CTFFilter): def __init__( diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index b7924f0fb6..ac6c4d2721 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -1,6 +1,7 @@ from .types import complex_type, real_type, utest_tolerance # isort:skip from .coor_trans import ( # isort:skip mean_aligned_angular_distance, + cart2pol, crop_pad_2d, crop_pad_3d, grid_1d, diff --git a/tests/test_covar2d.py b/tests/test_covar2d.py index 6ec5bf0b14..f27bc61855 100644 --- a/tests/test_covar2d.py +++ b/tests/test_covar2d.py @@ -292,7 +292,8 @@ def test_get_covar_ctf(cov2d_fixture, ctf_enabled): covar_coef_ctf = cov2d.get_covar(coef, h_ctf_fb, h_idx, noise_var=NOISE_VAR) for im, mat in enumerate(results.tolist()): - np.testing.assert_allclose(mat, covar_coef_ctf[im], rtol=1e-05, atol=1e-08) + # These tolerances were adjusted slightly (1e-8 to 3e-8) to accomodate MATLAB CTF repro changes + np.testing.assert_allclose(mat, covar_coef_ctf[im], rtol=3e-05, atol=3e-08) def test_get_covar_ctf_shrink(cov2d_fixture, ctf_enabled): From 700ba37062569747cba66e8a8cc4b4b47e4c7682 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 25 Aug 2025 14:37:19 -0400 Subject: [PATCH 09/14] add back B factor to CTF --- src/aspire/operators/filters.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index 00af22aa57..824885d33e 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -487,6 +487,9 @@ def _evaluate(self, omega): h = np.sqrt(1 - self.alpha**2) * np.sin(chi) - self.alpha * np.cos(chi) + if self.B: + h *= np.exp(-self.B * s**2) + return h def scale(self, c=1): From ee67a347e051068e917e00b113b9de2c067f9890 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 25 Aug 2025 14:45:27 -0400 Subject: [PATCH 10/14] default all Image filters to upcast the convolution --- src/aspire/image/image.py | 31 ------------------------------- src/aspire/image/xform.py | 2 +- 2 files changed, 1 insertion(+), 32 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index a74d91b234..38317fd5da 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -587,37 +587,6 @@ def filter(self, filter): im = self.stack_reshape(-1) - # Note image and filter data is intentionally migrated via - # `xp.asarray` because all of the subsequent calls until - # `asnumpy` are GPU when xp and fft in `cupy` mode. - # - # Second note, filter dtype may not match image dtype. - filter_values = xp.asarray( - filter.evaluate_grid(self.resolution), dtype=self.dtype - ) - - # Convolve - im_f = fft.centered_fft2(xp.asarray(im._data)) - im_f = filter_values * im_f - im = fft.centered_ifft2(im_f) - - im = xp.asnumpy(im.real) - - return self.__class__(im, pixel_size=self.pixel_size).stack_reshape( - original_stack_shape - ) - - def filter64(self, filter): - """ - Apply a `Filter` object to the Image and returns a new Image. - - :param filter: An object of type `Filter`. - :return: A new filtered `Image` object. - """ - original_stack_shape = self.stack_shape - - im = self.stack_reshape(-1) - # Note image and filter data is intentionally migrated via # `xp.asarray` because all of the subsequent calls until # `asnumpy` are GPU when xp and fft in `cupy` mode. diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py index 30aab2faa1..316d8ca3c8 100644 --- a/src/aspire/image/xform.py +++ b/src/aspire/image/xform.py @@ -300,7 +300,7 @@ def __init__(self, filter): self.filter = filter def _forward(self, im, indices): - return im.filter64(self.filter) + return im.filter(self.filter) def __str__(self): return f"FilterXform ({self.filter})" From e1fa0d58b12fd0fc16a05fe2de8f99debeb0ad13 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 26 Aug 2025 08:32:07 -0400 Subject: [PATCH 11/14] cleanup comments etc --- src/aspire/operators/filters.py | 12 +++++++++--- tests/test_filters.py | 9 +++++---- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py index 824885d33e..d3788d0d1f 100644 --- a/src/aspire/operators/filters.py +++ b/src/aspire/operators/filters.py @@ -455,12 +455,18 @@ def __init__( self._defocus_diff_nm = 0.05 * (self.defocus_u - self.defocus_v) def _evaluate(self, omega): + # Reference MATLAB code, includes reference to paper + # Mindell, J. A.; Grigorieff, N. (2003). + # https://github.com/PrincetonUniversity/aspire/blob/760a43b35453e55ff2d9354339e9ffa109a25371/projections/cryo_CTF_Relion.m#L34 + # # s, theta should match MATLAB's RadiusNorm up to a transpose - # To accomplish this given ASPIRE-Python's defaul `omega`, - # we unpack and remove the pi scaling here, + # To accomplish this given ASPIRE-Python's default `omega` grid, + # we unpack and remove the pi scaling, # and further rescale the radii `s` by half below. + # # Additionally we upcast so downstream computations remain in doubles. x, y = omega.astype(np.float64, copy=False) / np.pi + # Returns radii such that when multiplied by the # bandwidth of the signal, we get the correct radial frequencies # corresponding to each pixel in our nxn grid. @@ -477,7 +483,7 @@ def _evaluate(self, omega): s = s * BW DFavg = self._defocus_mean_nm # (DefocusU+DefocusV)/2 DFdiff = self._defocus_diff_nm # (DefocusU-DefocusV) - # Note the missing / 2 is already in _defocus_diff_nm + # Note division by 2 is pre-computed in _defocus_diff_nm df = DFavg + DFdiff * np.cos(2 * (theta - self.defocus_ang)) k2 = np.pi * lamb * df diff --git a/tests/test_filters.py b/tests/test_filters.py index bfd3908d71..8678480eef 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -204,8 +204,10 @@ def test_ctf_reference(): # >> ref_h=cryo_CTF_Relion(n,V,DF1,DF2,theta,Cs,pxA,A) # # Note we transpose the reference array. - # Python keeps the filter C order because the images we will convove with are C order. + # Python keeps the filter C order because the images we will convolve with are C order. # MATLAB is F and F respectively. + # + # The floating point values were truncated to four decimal digits. ref_h = np.array( [ [-0.6152, 0.0299, -0.5638, 0.9327, 0.9736], @@ -216,6 +218,5 @@ def test_ctf_reference(): ] ).T - # Test we're within 1%. - # There are minor differences in the formulas for wavelength and grids. - np.testing.assert_allclose(h, ref_h, rtol=0.01) + # Test match all significant digits above + np.testing.assert_allclose(h, ref_h, atol=5e-5) From f437da2425d30ee782b5f1c124f4c2a72d7b9e39 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 28 Aug 2025 15:14:03 -0400 Subject: [PATCH 12/14] extend coverage to both ArrayFilter and CTFFilter --- tests/test_filters.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_filters.py b/tests/test_filters.py index 8678480eef..6b84e29be4 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -22,11 +22,12 @@ class SimTestCase(TestCase): + test_filter = ArrayFilter(np.random.randn(8, 8)) + def setUp(self): self.dtype = np.float32 # A 2 x 256 ndarray of spatial frequencies self.omega = np.load(os.path.join(DATA_DIR, "omega_2_256.npy")) - self.test_filter = ArrayFilter(np.random.randn(8, 8)) def tearDown(self): pass @@ -119,6 +120,14 @@ def testFilterSigns(self): self.assertTrue(np.allclose(sign_filter.evaluate(self.omega), signs)) +class SimTestCaseCTFFilter(SimTestCase): + """ + Covers same tests as SimTestCase, but use CTFFilter in place of ArrayFilter. + """ + + test_filter = CTFFilter() + + DTYPES = [np.float32, np.float64] EPS = [None, 0.01] From 98bb5afb455d72696ec642964fae53d6417613e5 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 5 Sep 2025 08:35:28 -0400 Subject: [PATCH 13/14] fix dtype regression introduced in b1a0 --- src/aspire/image/image.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 38317fd5da..adc7e8ba13 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -591,9 +591,10 @@ def filter(self, filter): # `xp.asarray` because all of the subsequent calls until # `asnumpy` are GPU when xp and fft in `cupy` mode. # - # Second note, filter dtype may not match image dtype. + # Second note, filter and grid dtype may not match image dtype, + # upcast both here for most accurate convolution. filter_values = xp.asarray( - filter.evaluate_grid(self.resolution), dtype=np.float64 + filter.evaluate_grid(self.resolution, dtype=np.float64), dtype=np.float64 ) # Convolve From d6e2c1201ff2d89fca58fe2bee0b4fce9bff38f1 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 5 Sep 2025 09:22:59 -0400 Subject: [PATCH 14/14] revert to our centered wrappers --- src/aspire/image/image.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index adc7e8ba13..63f6a8b45e 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -599,9 +599,9 @@ def filter(self, filter): # Convolve _im = xp.asarray(im._data, dtype=np.float64) - im_f = fft.fftshift(fft.fft2(_im), axes=(-1, -2)) + im_f = fft.centered_fft2(_im) im_f = filter_values * im_f - im = fft.ifft2(fft.ifftshift(im_f, axes=(-1, -2))) + im = fft.centered_ifft2(im_f) im = xp.asnumpy(im.real).astype( self.dtype, copy=False