From 867829115c02f333119d9cb52464e66a64d2cc3a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 1 Oct 2026 15:04:08 -0400 Subject: [PATCH 01/44] fixup log and pbar for ctf est --- src/aspire/ctf/ctf_estimator.py | 3 ++- src/aspire/storage/starfile.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index b6983cb071..535da85a85 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -32,6 +32,7 @@ from numpy import linalg as npla from scipy.optimize import linprog from scipy.signal.windows import dpss +from aspire.utils import tqdm from aspire.basis import Coef, FFBBasis2D from aspire.image import Image @@ -761,7 +762,7 @@ def estimate_ctf( ffbbasis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) results = {} - for name in file_names: + for name in tqdm(file_names, desc="Processing files"): with mrcfile.open( os.path.join(data_folder, name), mode="r", permissive=True ) as mrc: diff --git a/src/aspire/storage/starfile.py b/src/aspire/storage/starfile.py index 98b3219607..f309830230 100644 --- a/src/aspire/storage/starfile.py +++ b/src/aspire/storage/starfile.py @@ -41,7 +41,7 @@ def __init__(self, filepath="", blocks=None): logger.error(f"Could not open {self.filepath}") raise FileNotFoundError self._initialize_blocks() - logger.info(f"Created {self}") + logger.debug(f"Created {self} at {self.filepath}") def _initialize_blocks(self): """ From fb9fe127220129f063d6813780700509b2529f4e Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 8 Oct 2026 08:46:32 -0400 Subject: [PATCH 02/44] debugging, reconcile outputs with MATLAB lots TODO --- src/aspire/ctf/ctf_estimator.py | 202 ++++++++++++++++++++++++++------ 1 file changed, 164 insertions(+), 38 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 535da85a85..19f2114a4c 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -1,6 +1,6 @@ """ Contains code supporting CTF parameter estimation. -Generally, this is a port of ASPIRE-CTF from MATLAB. +Generally, this is a Python version of ASPIRE-CTF from MATLAB. See paper: @@ -27,18 +27,25 @@ import os from collections import OrderedDict +import matplotlib.pyplot as plt import mrcfile import numpy as np from numpy import linalg as npla from scipy.optimize import linprog from scipy.signal.windows import dpss -from aspire.utils import tqdm from aspire.basis import Coef, FFBBasis2D from aspire.image import Image from aspire.numeric import fft from aspire.storage import StarFile -from aspire.utils import abs2, complex_type, grid_1d, grid_2d, voltage_to_wavelength +from aspire.utils import ( + abs2, + complex_type, + grid_1d, + grid_2d, + tqdm, + voltage_to_wavelength, +) logger = logging.getLogger(__name__) @@ -215,6 +222,89 @@ def preprocess_micrograph(self, micrograph, block_size): return self.normalize_blocks(self.micrograph_to_blocks(micrograph, block_size)) + def legacy_preprocess_micrograph(self, micrograph, block_size): + """ + Returns half-overlapping blocks extracted from the micrograph. + Equivalent to MATLAB's ctf_extract_blocks. + + Note, internally transposes micrograph to match MATLAB computations. + Returns (num_blocks, block_size, block_size) + """ + h, w = micrograph.shape + + if min(h, w) < block_size: + raise ValueError( + f"Blocks of size {block_size} x {block_size} cannot be extracted from image." + ) + + block_size = block_size - (block_size % 2) # force even + half = block_size // 2 + + # Compute crop bounds + end_val1 = (h // block_size) * block_size + ( + half if (h % block_size) > half else 0 + ) + end_val2 = (w // block_size) * block_size + ( + half if (w % block_size) > half else 0 + ) + + # Crop + cropped_micrograph = micrograph[:end_val1, :end_val2] + + # Compute number of half blocks + n_bh = end_val1 // half # number of half-blocks along rows + n_bw = end_val2 // half # number of half-blocks along cols + n_total = n_bh * n_bw + + # im2col `distinct` extracts non-overlapping half blocks + # column-major block ordering to match MATLAB (down-first) + # TODO, I don't think the order will matter here later, + # but keeping now for diagnostics. + blocks = ( + cropped_micrograph.reshape(n_bh, half, n_bw, half) + .transpose(1, 3, 0, 2) + .reshape(half * half, n_total, order="F") + ) + + # Reshape to (half, half, n_total) so each slice is one half-block + blocks = blocks.reshape(half, half, n_total, order="F") + + # TODO, compare which method is faster later, if we really need these... + # First cat + # stack each half-block with the one below it (rows overlap) + # by shifting the block indices by 1 + row_shifted = np.concatenate([blocks[:, :, 1:], blocks[:, :, :1]], axis=2) + blocks = np.concatenate( + [blocks, row_shifted], axis=0 + ) # (block_size, half, n_total) + + # Second cat + # stack each block with the one to its right (cols overlap) + # by rolling the blocks + col_shifted = np.roll(blocks, -n_bh, axis=2) + blocks = np.concatenate( + [blocks, col_shifted], axis=1 + ) # (block_size, block_size, n_total) + + # Reshape to blocksize-by-blocksize grid + n_row_blocks = int(np.floor(2 * h / block_size)) + n_col_blocks = int(np.floor(2 * w / block_size)) + blocks = blocks.reshape( + block_size, block_size, n_row_blocks, n_col_blocks, order="F" + ) + + # Drop last row and col (bounds artifacts) + blocks = blocks[:, :, :-1, :-1] + # Flatten + block = blocks.reshape(block_size, block_size, -1, order="F") + + # De-mean each block + block_mean = block.sum(axis=(0, 1), keepdims=True) / (block_size**2) + block = block - block_mean + + # Transpose output to C order (n,block_size, block_size) + return block.T + def tapers(self, N, NW, L): """ Compute data tapers (which are discrete prolate spheroidal sequences (dpss)) @@ -228,9 +318,12 @@ def tapers(self, N, NW, L): :return: NumPy array of data tapers """ - # Note the original ASPIRE implementation is negated from original scipy... - # but at time of writing subsequent code was agnostic to sign. - return dpss(M=N, NW=NW, Kmax=L, return_ratios=False).T + # Note the original ASPIRE MATLAB implementation computed these. + # This call should be close up to sign, but is not identical. + # + # In the CtfEstimator application, the signs will pass through + # the mul and fft steps, then are absolute valued. + return dpss(M=N, NW=NW, Kmax=L, return_ratios=False, sym=False, norm=2).T def estimate_psd(self, blocks, tapers_1d): """ @@ -264,13 +357,12 @@ def estimate_psd(self, blocks, tapers_1d): np.multiply(blocks[m, :, :], taper_2d, out=blocks_tapered) blocks_mt_post_fft = fft.fftn(blocks_tapered, axes=(-2, -1)) blocks_mt += abs2(blocks_mt_post_fft) + # blocks_mt += np.abs(blocks_mt_post_fft)**2 + blocks_mt /= blocks.shape[1] ** 2 blocks_mt /= blocks.shape[0] ** 2 - blocks_mt /= tapers_1d.shape[0] ** 2 - amplitude_spectrum = fft.fftshift( - blocks_mt - ) # max difference 10^-13, max relative difference 10^-14 + amplitude_spectrum = fft.fftshift(blocks_mt) return Image(amplitude_spectrum) @@ -284,8 +376,11 @@ def elliptical_average(self, ffbbasis, amplitude_spectrum, circular): :return: PSD and noise as 2-tuple of NumPy arrays. """ + # TODO breakup into two functions as the noise is unused in first call? + # RCOPT, come back and change the indices for this method coefs_s = ffbbasis.evaluate_t(amplitude_spectrum).asnumpy().copy().T + # TODO only needed in one case... coefs_n = coefs_s.copy() coefs_s[np.argwhere(ffbbasis.angular_indices == 1)] = 0 @@ -302,7 +397,7 @@ def elliptical_average(self, ffbbasis, amplitude_spectrum, circular): return psd, noise def background_subtract_1d( - self, amplitude_spectrum, linprog_method="highs", n_low_freq_cutoffs=14 + self, amplitude_spectrum, linprog_method="highs-ds", n_low_freq_cutoffs=14 ): """ Estimate and subtract the background from the power spectrum @@ -344,6 +439,9 @@ def background_subtract_1d( N = amplitude_spectrum.shape[-1] - low_freq_cutoff f = np.concatenate((np.ones(N), -1 * np.ones(N)), axis=0) + lb = np.concatenate([signal, np.full_like(signal, -np.inf)]) + ub = np.concatenate([signal, np.full_like(signal, np.inf)]) + bounds = np.column_stack([lb, ub]) superposition_condition = np.concatenate( (-1 * np.eye(N), np.eye(N)), axis=1 @@ -374,28 +472,22 @@ def background_subtract_1d( axis=0, ) - # The original code used `bounds`, - # but for many problems, linprog reports infeasable constraints. - # In practice for a micrograph from the paper, and our tutorial, - # the code seems to work better without it... - # ASPIRE #417 - - x = linprog( + res = linprog( f, A_ub=A, b_ub=np.zeros(A.shape[0]), method=linprog_method, + bounds=bounds, ) - if not x.success: + if not res.success: raise RuntimeError("Linear program did not succeed. Halting") - background = x.x[N:] + background = res.x[N:] bs_psd = signal - background final_signal[low_freq_cutoff - 1, low_freq_cutoff:] = bs_psd.T - # expected difference: 10^-7 (absolute) final_background[low_freq_cutoff - 1, low_freq_cutoff:] = background.T return final_signal, final_background @@ -428,6 +520,8 @@ def opt1d( center = N // 2 grid = grid_1d(N, normalized=True, dtype=self.dtype) + + # TODO, try using use self.r_ctf? rb = grid["r"][center:] / 2 r_ctf = rb * (10 / pixel_size) # units: inverse nm @@ -441,6 +535,7 @@ def opt1d( c = np.zeros((max_defocus - min_defocus, signal.shape[1]), dtype=self.dtype) for f in range(min_defocus, max_defocus): + # TODO try replace with ASPIRE's CTF formula? ctf_im = np.abs( np.sin( np.pi * lmbd * f * r_ctf_sq @@ -456,8 +551,8 @@ def opt1d( for m in range(0, signal.shape[1]): signal[:, m] = signal[:, m] - np.mean(signal[m + 1 :, m], axis=0) ctf_im[:, m] = ctf_im[:, m] - np.mean(ctf_im[m + 1 :, m], axis=0) - ctf_im[: m + 1, m] = np.zeros((m + 1)) - signal[: m + 1, m] = np.zeros((m + 1)) + ctf_im[: m + 1, m] = 0 # np.zeros((m + 1)) + signal[: m + 1, m] = 0 # np.zeros((m + 1)) Sx = np.sqrt(np.sum(ctf_im**2, axis=0)) Sy = np.sqrt(np.sum(signal**2, axis=0)) @@ -482,9 +577,9 @@ def background_subtract_2d(self, signal, background_p1, max_col): N = signal.shape[1] grid = grid_2d(N, normalized=False, indexing="yx", dtype=self.dtype) + radii = grid["r"] # can probably get elsewhere if signal.shape is consistent. - radii = np.sqrt((grid["x"] / 2) ** 2 + (grid["y"] / 2) ** 2) - + # okay now, but review max_col and vectorize this background = np.zeros(signal.shape, dtype=self.dtype) for r in range(max_col + 2, background_p1.shape[1]): background[:, (r < radii) & (radii <= r + 1)] = background_p1[max_col, r] @@ -716,6 +811,24 @@ def write_star(self, name, params_dict, output_dir): star = StarFile(blocks=blocks) star.write(os.path.join(output_dir, os.path.splitext(name)[0]) + ".star") + def _plot_1d_psd(self, signal_1d, background_1d): + center = self.r_ctf.shape[0] // 2 + x = self.r_ctf[center, center : center + signal_1d.shape[1]] + plt.plot(x, signal_1d[1] + background_1d[1]) + # xlabel('$\vert \mathbf{k} \vert^2 \, \, [\mathrm{nm}^{-1}]$', 'Interpreter','latex'); + # ylabel('$S_y \left( \mathbf{k} \right)$', 'Interpreter','latex'); + # title('1-d PSD') + plt.show() + + def _plot_1d_bg_sub_psd(self, signal_1d): + center = self.r_ctf.shape[0] // 2 + x = self.r_ctf[center, center : center + signal_1d.shape[1]] + plt.plot(x, signal_1d[1]) + # title('Background subtracted 1-d PSD') + # xlabel('$\vert \mathbf{k} \vert^2 \, \, [\mathrm{nm}^{-1}]$', 'Interpreter','latex'); + # ylabel('$S_y \left( \mathbf{k} \right) - S_e \left( \mathbf{k} \right)$', 'Interpreter','latex'); + plt.show() + def estimate_ctf( data_folder, @@ -761,6 +874,10 @@ def estimate_ctf( # closer to original code. ffbbasis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) + # Only need to compute tapers once + # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L + tapers_1d = ctf_object.tapers(psd_size, num_tapers, num_tapers) + results = {} for name in tqdm(file_names, desc="Processing files"): with mrcfile.open( @@ -771,23 +888,27 @@ def estimate_ctf( # Try to match dtype used in Basis instance micrograph = micrograph.astype(dtype, copy=False) - micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size) - - tapers_1d = ctf_object.tapers(psd_size, num_tapers / 2, num_tapers) + # TODO compare now that remainder of code repros. + # TODO migrate code into a main method in CtfEstimator + # micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size) + micrograph_blocks = ctf_object.legacy_preprocess_micrograph( + micrograph, psd_size + ) signal_observed = ctf_object.estimate_psd(micrograph_blocks, tapers_1d) + np.save("signal_observed.npy", signal_observed) amplitude_spectrum, _ = ctf_object.elliptical_average( ffbbasis, signal_observed, True - ) # absolute differenceL 10^-14. Relative error: 10^-7 - - # Optionally changing to: linprog_method='simplex', - # will more deterministically repro results in exchange for speed. - # linprog_method was changed from 'interior-point' to 'highs' due to - # "interior-point' being deprecated. - signal_1d, background_1d = ctf_object.background_subtract_1d( - amplitude_spectrum, linprog_method="highs" ) + np.save("amplitude_spectrum.npy", amplitude_spectrum) + + # TODO expose linprog as a init var + signal_1d, background_1d = ctf_object.background_subtract_1d(amplitude_spectrum) + + # TODO complete diagnostic plots and complete implementation in class. + # ctf_object._plot_1d_psd(signal_1d, background_1d) + # ctf_object._plot_1d_bg_sub_psd(signal_1d) avg_defocus, low_freq_skip = ctf_object.opt1d( signal_1d, @@ -798,7 +919,7 @@ def estimate_ctf( signal_observed.shape[-1], ) - low_freq_skip = 12 + # TODO check the low_freq_skip +-1, it is col ind or value signal, background_2d = ctf_object.background_subtract_2d( signal_observed, background_1d, low_freq_skip ) @@ -808,7 +929,6 @@ def estimate_ctf( signal, additional_background = ctf_object.elliptical_average( ffbbasis, signal.sqrt(), False ) - background_2d = background_2d + additional_background initial_df1 = (avg_defocus * 2) / (1 + ratio) @@ -816,10 +936,13 @@ def estimate_ctf( grid = grid_2d(psd_size, normalized=True, indexing="yx", dtype=dtype) + # TODO, try to avoid why computing r_ctf again, + # similarly theta can be computed in init r_ctf = grid["r"] / 2 * (10 / pixel_size) theta = grid["phi"] angle = -5 / 12 * np.pi # Radians (-75 degrees) + # TODO replace 6 with a variable cc_array = np.zeros((6, 4)) for a in range(0, 6): df1, df2, angle_ast, p = ctf_object.gd( @@ -843,6 +966,7 @@ def estimate_ctf( cc_array[a, 3] = p ml = np.argmax(cc_array[:, 3], -1) + # TODO what is the ultimate purpose of the result dictionary? result = { "defocus_u": cc_array[ml, 0] * 10, # Convert from nm to A "defocus_v": cc_array[ml, 1] * 10, # Convert from nm to A @@ -854,11 +978,13 @@ def estimate_ctf( } results[name] = result + # TODO consolidate all the IO our of here and into functions # we write each micrograph's ctf parameters to an individual starfile if not os.path.isdir(output_dir): os.mkdir(output_dir) ctf_object.write_star(name, result, output_dir) + # TODO compare these images with expected outputs from MATLAB if save_noise_images: with mrcfile.new( os.path.join(output_dir, os.path.splitext(name)[0] + "_noise.mrc"), From 1604b8f3dd3b00d5736fcf9c09e7386576920a30 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 8 Oct 2026 09:30:21 -0400 Subject: [PATCH 03/44] move gd iters to an input parameter --- src/aspire/ctf/ctf_estimator.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 19f2114a4c..ac102b7e09 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -63,6 +63,7 @@ def __init__( voltage, psd_size, num_tapers, + gradient_descent_iterations=6, dtype=np.float32, ): """ @@ -73,7 +74,10 @@ def __init__( :param amplitude_contrast: Amplitude contrast. :param voltage: Voltage of electron microscope. :param psd_size: Block size (in pixels) for PSD estimation. - :param num_tapers: Number of tapers to apply in PSD estimation. + :param num_tapers: Optional, number of tapers to apply in PSD estimation. + Defaults to 4. + :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. + Defaults to 6. :returns: CtfEstimator instance. """ @@ -81,9 +85,10 @@ def __init__( self.cs = cs self.amplitude_contrast = amplitude_contrast self.voltage = voltage - self.psd_size = psd_size - self.num_tapers = num_tapers + self.psd_size = int(psd_size) + self.num_tapers = int(num_tapers) self.lmbd = voltage_to_wavelength(voltage) / 10.0 # Convert angstrom to nm + self.gd_iters = int(gradient_descent_iterations) self.dtype = np.dtype(dtype) grid = grid_2d(psd_size, normalized=True, indexing="yx", dtype=self.dtype) @@ -836,10 +841,11 @@ def estimate_ctf( cs=2.0, amplitude_contrast=0.07, voltage=300, - num_tapers=2, + num_tapers=4, psd_size=512, g_min=30, g_max=5, + gradient_descent_iterations=6, output_dir="results", dtype=np.float32, save_ctf_images=False, @@ -867,7 +873,14 @@ def estimate_ctf( lmbd = voltage_to_wavelength(voltage) / 10 # Convert from angstrom to nm ctf_object = CtfEstimator( - pixel_size, cs, amplitude_contrast, voltage, psd_size, num_tapers, dtype=dtype + pixel_size, + cs, + amplitude_contrast, + voltage, + psd_size, + num_tapers=num_tapers, + gradient_descent_iterations=gradient_descent_iterations, + dtype=dtype, ) # Note for repro debugging, suggest use of doubles, @@ -942,9 +955,9 @@ def estimate_ctf( theta = grid["phi"] angle = -5 / 12 * np.pi # Radians (-75 degrees) - # TODO replace 6 with a variable - cc_array = np.zeros((6, 4)) - for a in range(0, 6): + + cc_array = np.zeros((self.gd_iters, 4)) + for a in range(0, self.gd_iters): df1, df2, angle_ast, p = ctf_object.gd( signal, initial_df1, From cf206f5279444180f607b11224424151adee5e20 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 8 Oct 2026 09:52:15 -0400 Subject: [PATCH 04/44] split up elliptical average --- src/aspire/ctf/ctf_estimator.py | 59 +++++++++++++++++++-------------- 1 file changed, 35 insertions(+), 24 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index ac102b7e09..d050d47a3e 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -371,33 +371,46 @@ def estimate_psd(self, blocks, tapers_1d): return Image(amplitude_spectrum) - def elliptical_average(self, ffbbasis, amplitude_spectrum, circular): + def radial_average(self, basis, amplitude_spectrum): """ Computes radial/elliptical average of the power spectrum - :param ffbbasis: FFBBasis instance. - :param amplitude_spectrum: Power spectrum. - :param circular: True for radial averaging and False for elliptical averaging. + :param basis: SteerableBasis instance (typically FFBBasis). + :param amplitude_spectrum: Power spectrum to average. :return: PSD and noise as 2-tuple of NumPy arrays. """ + noise = amplitude_spectrum # pass through + coefs_s = basis.evaluate_t(amplitude_spectrum).asnumpy().copy() - # TODO breakup into two functions as the noise is unused in first call? + mask_s = (basis.angular_indices == 1) | (basis.angular_indices == 2) + coefs_s[:, mask_s] = 0 - # RCOPT, come back and change the indices for this method - coefs_s = ffbbasis.evaluate_t(amplitude_spectrum).asnumpy().copy().T - # TODO only needed in one case... + psd = Coef(basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size).evaluate() + + return psd, noise + + def elliptical_average(self, basis, amplitude_spectrum): + """ + Computes elliptical average of the power spectrum + + :param basis: SteerableBasis instance (typically FFBBasis). + :param amplitude_spectrum: Power spectrum to average. + :return: PSD and noise as 2-tuple of NumPy arrays. + """ + + coefs_s = basis.evaluate_t(amplitude_spectrum).asnumpy().copy() coefs_n = coefs_s.copy() - coefs_s[np.argwhere(ffbbasis.angular_indices == 1)] = 0 - if circular: - coefs_s[np.argwhere(ffbbasis.angular_indices == 2)] = 0 - noise = amplitude_spectrum - else: - coefs_n[np.argwhere(ffbbasis.angular_indices == 0)] = 0 - coefs_n[np.argwhere(ffbbasis.angular_indices == 2)] = 0 - noise = Coef(ffbbasis, coefs_n.T).evaluate() + mask_s = basis.angular_indices == 1 + coefs_s[:, mask_s] = 0 - psd = Coef(ffbbasis, coefs_s.T).evaluate() + mask_n = (basis.angular_indices == 0) | (basis.angular_indices == 2) + coefs_n[:, mask_n] = 0 + + noise = Coef( + basis, coefs_n, pixel_size=amplitude_spectrum.pixel_size + ).evaluate() + psd = Coef(basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size).evaluate() return psd, noise @@ -885,7 +898,7 @@ def estimate_ctf( # Note for repro debugging, suggest use of doubles, # closer to original code. - ffbbasis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) + basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) # Only need to compute tapers once # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L @@ -911,9 +924,7 @@ def estimate_ctf( signal_observed = ctf_object.estimate_psd(micrograph_blocks, tapers_1d) np.save("signal_observed.npy", signal_observed) - amplitude_spectrum, _ = ctf_object.elliptical_average( - ffbbasis, signal_observed, True - ) + amplitude_spectrum, _ = ctf_object.radial_average(basis, signal_observed) np.save("amplitude_spectrum.npy", amplitude_spectrum) # TODO expose linprog as a init var @@ -940,7 +951,7 @@ def estimate_ctf( ratio = ctf_object.pca(signal_observed, pixel_size, g_min, g_max) signal, additional_background = ctf_object.elliptical_average( - ffbbasis, signal.sqrt(), False + basis, signal.sqrt() ) background_2d = background_2d + additional_background @@ -956,8 +967,8 @@ def estimate_ctf( angle = -5 / 12 * np.pi # Radians (-75 degrees) - cc_array = np.zeros((self.gd_iters, 4)) - for a in range(0, self.gd_iters): + cc_array = np.zeros((ctf_object.gd_iters, 4)) + for a in range(0, ctf_object.gd_iters): df1, df2, angle_ast, p = ctf_object.gd( signal, initial_df1, From 01d70454526d567c85f86d31913b92413f09e5cd Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 9 Oct 2026 10:52:25 -0400 Subject: [PATCH 05/44] cleanup some of the easier TODO items --- src/aspire/ctf/ctf_estimator.py | 52 ++++++++++++++++----------------- 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index d050d47a3e..e8ab7cfbdc 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -64,6 +64,7 @@ def __init__( psd_size, num_tapers, gradient_descent_iterations=6, + linprog_method="highs-ds", dtype=np.float32, ): """ @@ -78,6 +79,8 @@ def __init__( Defaults to 4. :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. Defaults to 6. + :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). + Defaults to 'highs-ds'. :returns: CtfEstimator instance. """ @@ -89,6 +92,7 @@ def __init__( self.num_tapers = int(num_tapers) self.lmbd = voltage_to_wavelength(voltage) / 10.0 # Convert angstrom to nm self.gd_iters = int(gradient_descent_iterations) + self.linprog_method = linprog_method self.dtype = np.dtype(dtype) grid = grid_2d(psd_size, normalized=True, indexing="yx", dtype=self.dtype) @@ -539,7 +543,6 @@ def opt1d( grid = grid_1d(N, normalized=True, dtype=self.dtype) - # TODO, try using use self.r_ctf? rb = grid["r"][center:] / 2 r_ctf = rb * (10 / pixel_size) # units: inverse nm @@ -670,8 +673,6 @@ def gd( df1, df2, angle_ast, - r, - theta, pixel_size, g_min, g_max, @@ -686,8 +687,6 @@ def gd( :param df1: Defocus value in the direction perpendicular to df2. :param df2: Defocus value in the direction perpendicular to df1. :param angle_ast: Angle between df1 and the x-axis, Radians. - :param r: Magnitude of spatial frequencies. - :param theta: Phase of spatial frequencies. :param pixel_size: Pixel size in \u212b (angstrom). :param g_min: Inverse of minimun resolution for PSD. :param g_max: Inverse of maximum resolution for PSD. @@ -706,9 +705,9 @@ def gd( y = (df1 - df2) * np.cos(2 * angle_ast) z = (df1 - df2) * np.sin(2 * angle_ast) - a = np.pi * lmbd * r**2 / 2 - b = np.pi * lmbd**3 * cs * 1e6 * r**4 / 2 - np.full( - shape=r.shape, fill_value=amplitude_contrast, dtype=self.dtype + a = np.pi * lmbd * self.r_ctf**2 / 2 + b = np.pi * lmbd**3 * cs * 1e6 * self.r_ctf**4 / 2 - np.full( + shape=self.r_ctf.shape, fill_value=amplitude_contrast, dtype=self.dtype ) signal = signal.asnumpy()[0].T @@ -718,15 +717,15 @@ def gd( rad_sq_min = N * pixel_size / g_min rad_sq_max = N * pixel_size / g_max - max_val = r[center, int(center - 1 + np.floor(rad_sq_max))] - min_val = r[center, int(center - 1 + np.ceil(rad_sq_min))] + max_val = self.r_ctf[center, int(center - 1 + np.floor(rad_sq_max))] + min_val = self.r_ctf[center, int(center - 1 + np.ceil(rad_sq_min))] - mask = (r <= max_val) & (r > min_val) + mask = (self.r_ctf <= max_val) & (self.r_ctf > min_val) a = a[mask] b = b[mask] signal = signal[..., mask] - r = r[mask] - theta = theta[mask] + # r = self.r_ctf[mask] + theta = self.theta[mask] sum_A = np.sum(signal**2) @@ -927,8 +926,9 @@ def estimate_ctf( amplitude_spectrum, _ = ctf_object.radial_average(basis, signal_observed) np.save("amplitude_spectrum.npy", amplitude_spectrum) - # TODO expose linprog as a init var - signal_1d, background_1d = ctf_object.background_subtract_1d(amplitude_spectrum) + signal_1d, background_1d = ctf_object.background_subtract_1d( + amplitude_spectrum, linprog_method=ctf_object.linprog_method + ) # TODO complete diagnostic plots and complete implementation in class. # ctf_object._plot_1d_psd(signal_1d, background_1d) @@ -958,13 +958,6 @@ def estimate_ctf( initial_df1 = (avg_defocus * 2) / (1 + ratio) initial_df2 = (avg_defocus * 2) - initial_df1 - grid = grid_2d(psd_size, normalized=True, indexing="yx", dtype=dtype) - - # TODO, try to avoid why computing r_ctf again, - # similarly theta can be computed in init - r_ctf = grid["r"] / 2 * (10 / pixel_size) - theta = grid["phi"] - angle = -5 / 12 * np.pi # Radians (-75 degrees) cc_array = np.zeros((ctf_object.gd_iters, 4)) @@ -974,8 +967,6 @@ def estimate_ctf( initial_df1, initial_df2, angle + a * np.pi / 6.0, # Radians, + a*30degrees - r_ctf, - theta, pixel_size, g_min, g_max, @@ -1023,12 +1014,19 @@ def estimate_ctf( ctf_object.set_angle(cc_array[ml, 2]) # Radians ctf_object.generate_ctf() df = (cc_array[ml, 0] + cc_array[ml, 1]) * np.ones( - theta.shape, theta.dtype + ctf_object.theta.shape, ctf_object.theta.dtype ) + (cc_array[ml, 0] - cc_array[ml, 1]) * np.cos( - 2 * theta - 2 * cc_array[ml, 2] * np.ones(theta.shape, theta.dtype) + 2 * ctf_object.theta + - 2 + * cc_array[ml, 2] + * np.ones(ctf_object.theta.shape, ctf_object.theta.dtype) ) ctf_im = -np.sin( - np.pi * lmbd * r_ctf**2 / 2 * (df - lmbd**2 * r_ctf**2 * cs * 1e6) + np.pi + * lmbd + * ctf_object.r_ctf**2 + / 2 + * (df - lmbd**2 * ctf_object.r_ctf**2 * cs * 1e6) + amplitude_contrast ) ctf_signal = np.zeros(ctf_im.shape, ctf_im.dtype) From 24ae4ddd5b27620611e916e18227a96c917d5841 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 12 Oct 2026 09:24:41 -0400 Subject: [PATCH 06/44] initial migration ctf est algo to estimate method --- src/aspire/ctf/ctf_estimator.py | 266 ++++++++++++++++---------------- 1 file changed, 135 insertions(+), 131 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index e8ab7cfbdc..89eae080e4 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -106,6 +106,13 @@ def __init__( self.angle = 0 # Radians self.h = 0 + # TODO, try to extend to any steerable basis + self.basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) + + # Only need to compute tapers once + # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L + self.tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) + def set_df1(self, df): """ Sets defocus. @@ -846,6 +853,88 @@ def _plot_1d_bg_sub_psd(self, signal_1d): # ylabel('$S_y \left( \mathbf{k} \right) - S_e \left( \mathbf{k} \right)$', 'Interpreter','latex'); plt.show() + def estimate(self, micrograph, g_min=30, g_max=5): + # TODO compare now that remainder of code repros. + # micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size) + micrograph_blocks = self.legacy_preprocess_micrograph(micrograph, self.psd_size) + + signal_observed = self.estimate_psd(micrograph_blocks, self.tapers_1d) + # np.save("signal_observed.npy", signal_observed) + + amplitude_spectrum, _ = self.radial_average(self.basis, signal_observed) + # np.save("amplitude_spectrum.npy", amplitude_spectrum) + + signal_1d, background_1d = self.background_subtract_1d( + amplitude_spectrum, linprog_method=self.linprog_method + ) + + # TODO complete diagnostic plots and complete implementation in class. + # self._plot_1d_psd(signal_1d, background_1d) + # self._plot_1d_bg_sub_psd(signal_1d) + + avg_defocus, low_freq_skip = self.opt1d( + signal_1d, + self.pixel_size, + self.cs, + self.lmbd, # (angstrom) + self.amplitude_contrast, + signal_observed.shape[-1], + ) + + # TODO check the low_freq_skip +-1, it is col ind or value + signal, background_2d = self.background_subtract_2d( + signal_observed, background_1d, low_freq_skip + ) + + ratio = self.pca(signal_observed, self.pixel_size, g_min, g_max) + + signal, additional_background = self.elliptical_average( + self.basis, signal.sqrt() + ) + background_2d = background_2d + additional_background + + initial_df1 = (avg_defocus * 2) / (1 + ratio) + initial_df2 = (avg_defocus * 2) - initial_df1 + + angle = -5 / 12 * np.pi # Radians (-75 degrees) + + # This array holds results from each of the `gd_iters` + cc_array = np.zeros((self.gd_iters, 4)) + for a in range(0, self.gd_iters): + df1, df2, angle_ast, p = self.gd( + signal, + initial_df1, + initial_df2, + angle + a * np.pi / 6.0, # Radians, + a*30degrees + self.pixel_size, + g_min, + g_max, + self.amplitude_contrast, + self.lmbd, # (angstrom) + self.cs, + ) + + cc_array[a, 0] = df1 + cc_array[a, 1] = df2 + cc_array[a, 2] = angle_ast # Radians + cc_array[a, 3] = p + + # Select the best result index + ml = np.argmax(cc_array[:, 3], -1) + + # TODO what is the ultimate purpose of the result dictionary? + result = { + "defocus_u": cc_array[ml, 0] * 10, # Convert from nm to A + "defocus_v": cc_array[ml, 1] * 10, # Convert from nm to A + "defocus_ang": cc_array[ml, 2], # Radians + "cs": self.cs, + "voltage": self.voltage, + "pixel_size": self.pixel_size, + "amplitude_contrast": self.amplitude_contrast, + } + + return result + def estimate_ctf( data_folder, @@ -871,19 +960,19 @@ def estimate_ctf( dtype = np.dtype(dtype) assert dtype in (np.float32, np.float64) + # TODO confirm if scans multi-dir dir_content = os.scandir(data_folder) + # TODO change to full path mrc_files = [f.name for f in dir_content if os.path.splitext(f)[1] == ".mrc"] mrcs_files = [f.name for f in dir_content if os.path.splitext(f)[1] == ".mrcs"] file_names = mrc_files + mrcs_files - amp = amplitude_contrast + # Why amplitude_contrast = np.arctan( amplitude_contrast / np.sqrt(1 - amplitude_contrast**2) ) - lmbd = voltage_to_wavelength(voltage) / 10 # Convert from angstrom to nm - ctf_object = CtfEstimator( pixel_size, cs, @@ -895,14 +984,6 @@ def estimate_ctf( dtype=dtype, ) - # Note for repro debugging, suggest use of doubles, - # closer to original code. - basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) - - # Only need to compute tapers once - # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L - tapers_1d = ctf_object.tapers(psd_size, num_tapers, num_tapers) - results = {} for name in tqdm(file_names, desc="Processing files"): with mrcfile.open( @@ -910,87 +991,10 @@ def estimate_ctf( ) as mrc: micrograph = mrc.data - # Try to match dtype used in Basis instance micrograph = micrograph.astype(dtype, copy=False) - # TODO compare now that remainder of code repros. - # TODO migrate code into a main method in CtfEstimator - # micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size) - micrograph_blocks = ctf_object.legacy_preprocess_micrograph( - micrograph, psd_size - ) - - signal_observed = ctf_object.estimate_psd(micrograph_blocks, tapers_1d) - np.save("signal_observed.npy", signal_observed) - - amplitude_spectrum, _ = ctf_object.radial_average(basis, signal_observed) - np.save("amplitude_spectrum.npy", amplitude_spectrum) - - signal_1d, background_1d = ctf_object.background_subtract_1d( - amplitude_spectrum, linprog_method=ctf_object.linprog_method - ) - - # TODO complete diagnostic plots and complete implementation in class. - # ctf_object._plot_1d_psd(signal_1d, background_1d) - # ctf_object._plot_1d_bg_sub_psd(signal_1d) - - avg_defocus, low_freq_skip = ctf_object.opt1d( - signal_1d, - pixel_size, - cs, - lmbd, # (angstrom) - amplitude_contrast, - signal_observed.shape[-1], - ) - - # TODO check the low_freq_skip +-1, it is col ind or value - signal, background_2d = ctf_object.background_subtract_2d( - signal_observed, background_1d, low_freq_skip - ) - - ratio = ctf_object.pca(signal_observed, pixel_size, g_min, g_max) - - signal, additional_background = ctf_object.elliptical_average( - basis, signal.sqrt() - ) - background_2d = background_2d + additional_background - - initial_df1 = (avg_defocus * 2) / (1 + ratio) - initial_df2 = (avg_defocus * 2) - initial_df1 - - angle = -5 / 12 * np.pi # Radians (-75 degrees) - - cc_array = np.zeros((ctf_object.gd_iters, 4)) - for a in range(0, ctf_object.gd_iters): - df1, df2, angle_ast, p = ctf_object.gd( - signal, - initial_df1, - initial_df2, - angle + a * np.pi / 6.0, # Radians, + a*30degrees - pixel_size, - g_min, - g_max, - amplitude_contrast, - lmbd, # (angstrom) - cs, - ) - - cc_array[a, 0] = df1 - cc_array[a, 1] = df2 - cc_array[a, 2] = angle_ast # Radians - cc_array[a, 3] = p - ml = np.argmax(cc_array[:, 3], -1) + result = ctf_object.estimate(micrograph, g_min=g_min, g_max=g_max) - # TODO what is the ultimate purpose of the result dictionary? - result = { - "defocus_u": cc_array[ml, 0] * 10, # Convert from nm to A - "defocus_v": cc_array[ml, 1] * 10, # Convert from nm to A - "defocus_ang": cc_array[ml, 2], # Radians - "cs": cs, - "voltage": voltage, - "pixel_size": pixel_size, - "amplitude_contrast": amp, - } results[name] = result # TODO consolidate all the IO our of here and into functions @@ -999,47 +1003,47 @@ def estimate_ctf( os.mkdir(output_dir) ctf_object.write_star(name, result, output_dir) - # TODO compare these images with expected outputs from MATLAB - if save_noise_images: - with mrcfile.new( - os.path.join(output_dir, os.path.splitext(name)[0] + "_noise.mrc"), - overwrite=True, - ) as mrc: - mrc.set_data(background_2d.asnumpy()[0].astype(np.float32)) - mrc.voxel_size = pixel_size - - if save_ctf_images: - ctf_object.set_df1(cc_array[ml, 0]) # angstrom - ctf_object.set_df2(cc_array[ml, 1]) # angstrom - ctf_object.set_angle(cc_array[ml, 2]) # Radians - ctf_object.generate_ctf() - df = (cc_array[ml, 0] + cc_array[ml, 1]) * np.ones( - ctf_object.theta.shape, ctf_object.theta.dtype - ) + (cc_array[ml, 0] - cc_array[ml, 1]) * np.cos( - 2 * ctf_object.theta - - 2 - * cc_array[ml, 2] - * np.ones(ctf_object.theta.shape, ctf_object.theta.dtype) - ) - ctf_im = -np.sin( - np.pi - * lmbd - * ctf_object.r_ctf**2 - / 2 - * (df - lmbd**2 * ctf_object.r_ctf**2 * cs * 1e6) - + amplitude_contrast - ) - ctf_signal = np.zeros(ctf_im.shape, ctf_im.dtype) - ctf_signal[: ctf_im.shape[0] // 2, :] = ctf_im[: ctf_im.shape[0] // 2, :] - ctf_signal[ctf_im.shape[0] // 2 + 1 :, :] = signal.asnumpy()[ - :, :, ctf_im.shape[0] // 2 + 1 - ] - - with mrcfile.new( - os.path.join(output_dir, os.path.splitext(name)[0] + ".ctf"), - overwrite=True, - ) as mrc: - mrc.set_data(np.float32(ctf_signal)) - mrc.voxel_size = pixel_size + # # TODO compare these images with expected outputs from MATLAB + # if save_noise_images: + # with mrcfile.new( + # os.path.join(output_dir, os.path.splitext(name)[0] + "_noise.mrc"), + # overwrite=True, + # ) as mrc: + # mrc.set_data(background_2d.asnumpy()[0].astype(np.float32)) + # mrc.voxel_size = pixel_size + + # if save_ctf_images: + # ctf_object.set_df1(cc_array[ml, 0]) # angstrom + # ctf_object.set_df2(cc_array[ml, 1]) # angstrom + # ctf_object.set_angle(cc_array[ml, 2]) # Radians + # ctf_object.generate_ctf() + # df = (cc_array[ml, 0] + cc_array[ml, 1]) * np.ones( + # ctf_object.theta.shape, ctf_object.theta.dtype + # ) + (cc_array[ml, 0] - cc_array[ml, 1]) * np.cos( + # 2 * ctf_object.theta + # - 2 + # * cc_array[ml, 2] + # * np.ones(ctf_object.theta.shape, ctf_object.theta.dtype) + # ) + # ctf_im = -np.sin( + # np.pi + # * ctf_object.lmbd + # * ctf_object.r_ctf**2 + # / 2 + # * (df - ctf_object.lmbd**2 * ctf_object.r_ctf**2 * cs * 1e6) + # + self.amplitude_contrast + # ) + # ctf_signal = np.zeros(ctf_im.shape, ctf_im.dtype) + # ctf_signal[: ctf_im.shape[0] // 2, :] = ctf_im[: ctf_im.shape[0] // 2, :] + # ctf_signal[ctf_im.shape[0] // 2 + 1 :, :] = signal.asnumpy()[ + # :, :, ctf_im.shape[0] // 2 + 1 + # ] + + # with mrcfile.new( + # os.path.join(output_dir, os.path.splitext(name)[0] + ".ctf"), + # overwrite=True, + # ) as mrc: + # mrc.set_data(np.float32(ctf_signal)) + # mrc.voxel_size = pixel_size return results From e217d1874a9400027715c83ada879e84f0756678 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 12 Oct 2026 09:43:32 -0400 Subject: [PATCH 07/44] begin cleanup image savig --- src/aspire/ctf/ctf_estimator.py | 87 +++++---------------------------- 1 file changed, 11 insertions(+), 76 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 89eae080e4..992269e852 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -101,10 +101,6 @@ def __init__( self.r_ctf = grid["r"] / 2 * (10 / pixel_size) # units: inverse nm self.theta = grid["phi"] - self.defocus1 = 0 - self.defocus2 = 0 - self.angle = 0 # Radians - self.h = 0 # TODO, try to extend to any steerable basis self.basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) @@ -113,64 +109,6 @@ def __init__( # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L self.tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) - def set_df1(self, df): - """ - Sets defocus. - - :param df: Defocus value in the direction perpendicular to df2. - """ - - self.defocus1 = df - - def set_df2(self, df): - """ - Sets defocus. - - :param df: Defocus value in the direction perpendicular to df1. - """ - - self.defocus2 = df - - def set_angle(self, angle): - """ - Sets angle. - - :param angle: Angle (in Radians) between df1 and the x-axis. - """ - - self.angle = angle - - def generate_ctf(self): - """ - Generates internal representation of the Contrast Transfer Function using parameters from this instance. - """ - - astigmatism_angle = np.full( - shape=self.theta.shape, fill_value=self.angle, dtype=self.dtype - ) - - defocus_sum = np.full( - shape=self.theta.shape, - fill_value=self.defocus1 + self.defocus2, - dtype=self.dtype, - ) - - defocus = defocus_sum + ( - (self.defocus1 - self.defocus2) - * np.cos(2 * (self.theta - astigmatism_angle)) - ) - defocus_factor = np.pi * self.lmbd * self.r_ctf * defocus / 2 - amplitude_contrast_term = self.amplitude_contrast / np.sqrt( - 1 - self.amplitude_contrast**2 - ) - - chi = ( - defocus_factor - - np.pi * self.lmbd**3 * self.cs * 1e6 * self.r_ctf**2 / 2 - + amplitude_contrast_term - ) - h = -np.sin(chi) - self.h = h def micrograph_to_blocks(self, micrograph, block_size): """ @@ -617,7 +555,7 @@ def background_subtract_2d(self, signal, background_p1, max_col): signal = signal - background signal = np.maximum(0, signal) - return Image(signal), Image(background) + return Image(signal, pixel_size=self.pixel_size), Image(background, pixel_size=self.pixel_size) def pca(self, signal, pixel_size, g_min, g_max): """ @@ -923,7 +861,7 @@ def estimate(self, micrograph, g_min=30, g_max=5): ml = np.argmax(cc_array[:, 3], -1) # TODO what is the ultimate purpose of the result dictionary? - result = { + params = { "defocus_u": cc_array[ml, 0] * 10, # Convert from nm to A "defocus_v": cc_array[ml, 1] * 10, # Convert from nm to A "defocus_ang": cc_array[ml, 2], # Radians @@ -933,7 +871,7 @@ def estimate(self, micrograph, g_min=30, g_max=5): "amplitude_contrast": self.amplitude_contrast, } - return result + return params, background_2d, signal def estimate_ctf( @@ -993,24 +931,21 @@ def estimate_ctf( micrograph = micrograph.astype(dtype, copy=False) - result = ctf_object.estimate(micrograph, g_min=g_min, g_max=g_max) + params, background_2d, signal = ctf_object.estimate(micrograph, g_min=g_min, g_max=g_max) - results[name] = result + results[name] = params # TODO consolidate all the IO our of here and into functions # we write each micrograph's ctf parameters to an individual starfile if not os.path.isdir(output_dir): os.mkdir(output_dir) - ctf_object.write_star(name, result, output_dir) + ctf_object.write_star(name, params, output_dir) + + # Save the noise + if save_noise_images: + noise_filepath = os.path.join(output_dir, os.path.splitext(name)[0] + "_noise.mrc") + background_2d.save(noise_filepath, overwrite=True) - # # TODO compare these images with expected outputs from MATLAB - # if save_noise_images: - # with mrcfile.new( - # os.path.join(output_dir, os.path.splitext(name)[0] + "_noise.mrc"), - # overwrite=True, - # ) as mrc: - # mrc.set_data(background_2d.asnumpy()[0].astype(np.float32)) - # mrc.voxel_size = pixel_size # if save_ctf_images: # ctf_object.set_df1(cc_array[ml, 0]) # angstrom From 3c82ce4e124a034620a613e41820999642d52816 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 12 Oct 2026 11:52:00 -0400 Subject: [PATCH 08/44] continue cleanup rounds --- src/aspire/ctf/ctf_estimator.py | 73 ++++++++++++++++++--------------- 1 file changed, 41 insertions(+), 32 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 992269e852..0429f8d90e 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -84,10 +84,10 @@ def __init__( :returns: CtfEstimator instance. """ - self.pixel_size = pixel_size - self.cs = cs - self.amplitude_contrast = amplitude_contrast - self.voltage = voltage + self.pixel_size = float(pixel_size) + self.cs = float(cs) + self.amplitude_contrast = float(amplitude_contrast) + self.voltage = float(voltage) self.psd_size = int(psd_size) self.num_tapers = int(num_tapers) self.lmbd = voltage_to_wavelength(voltage) / 10.0 # Convert angstrom to nm @@ -95,11 +95,10 @@ def __init__( self.linprog_method = linprog_method self.dtype = np.dtype(dtype) + # Grids grid = grid_2d(psd_size, normalized=True, indexing="yx", dtype=self.dtype) - # Note range is -half to half. self.r_ctf = grid["r"] / 2 * (10 / pixel_size) # units: inverse nm - self.theta = grid["phi"] # TODO, try to extend to any steerable basis @@ -109,7 +108,6 @@ def __init__( # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L self.tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) - def micrograph_to_blocks(self, micrograph, block_size): """ Preprocess micrograph into blocks using block_size. @@ -555,7 +553,9 @@ def background_subtract_2d(self, signal, background_p1, max_col): signal = signal - background signal = np.maximum(0, signal) - return Image(signal, pixel_size=self.pixel_size), Image(background, pixel_size=self.pixel_size) + return Image(signal, pixel_size=self.pixel_size), Image( + background, pixel_size=self.pixel_size + ) def pca(self, signal, pixel_size, g_min, g_max): """ @@ -753,17 +753,15 @@ def gd( # Note, This doesn't actually use anything from the class. # It is used in a solver loop of some sort, so it may not be correct # to just use what is avail in the obj. - def write_star(self, name, params_dict, output_dir): + def write_star(self, star_filepath, params_dict): """ Writes CTF parameters to starfile for a single micrograph. """ data_block = {} - data_block["_rlnMicrographName"] = name + data_block["_rlnMicrographName"] = star_filepath data_block["_rlnDefocusU"] = params_dict["defocus_u"] # Should already be in A data_block["_rlnDefocusV"] = params_dict["defocus_v"] # Should already be in A - data_block["_rlnDefocusAngle"] = ( - params_dict["defocus_ang"] * 180 / np.pi - ) # Convert from radian to degree + data_block["_rlnDefocusAngle"] = np.rad2deg(params_dict["defocus_ang"]) data_block["_rlnSphericalAberration"] = params_dict["cs"] data_block["_rlnAmplitudeContrast"] = params_dict["amplitude_contrast"] data_block["_rlnVoltage"] = params_dict["voltage"] @@ -771,7 +769,7 @@ def write_star(self, name, params_dict, output_dir): blocks = OrderedDict() blocks["root"] = data_block star = StarFile(blocks=blocks) - star.write(os.path.join(output_dir, os.path.splitext(name)[0]) + ".star") + star.write(star_filepath) def _plot_1d_psd(self, signal_1d, background_1d): center = self.r_ctf.shape[0] // 2 @@ -881,7 +879,7 @@ def estimate_ctf( amplitude_contrast=0.07, voltage=300, num_tapers=4, - psd_size=512, + psd_size=512, # should change to 1024? g_min=30, g_max=5, gradient_descent_iterations=6, @@ -895,9 +893,15 @@ def estimate_ctf( and returns CTF as a mrc file. """ + # Checks ------------------------------------------------------------------ dtype = np.dtype(dtype) assert dtype in (np.float32, np.float64) + # Ensure we can write outputs before we get too far + if not os.path.isdir(output_dir): + os.mkdir(output_dir) + + # File inputs ------------------------------------------------------------- # TODO confirm if scans multi-dir dir_content = os.scandir(data_folder) @@ -906,12 +910,13 @@ def estimate_ctf( mrcs_files = [f.name for f in dir_content if os.path.splitext(f)[1] == ".mrcs"] file_names = mrc_files + mrcs_files + # Setup a CTFEstimator ----------------------------------------------------- # Why amplitude_contrast = np.arctan( amplitude_contrast / np.sqrt(1 - amplitude_contrast**2) ) - ctf_object = CtfEstimator( + ctf_est = CtfEstimator( pixel_size, cs, amplitude_contrast, @@ -922,6 +927,7 @@ def estimate_ctf( dtype=dtype, ) + # Main Loop --------------------------------------------------------------- results = {} for name in tqdm(file_names, desc="Processing files"): with mrcfile.open( @@ -931,41 +937,44 @@ def estimate_ctf( micrograph = micrograph.astype(dtype, copy=False) - params, background_2d, signal = ctf_object.estimate(micrograph, g_min=g_min, g_max=g_max) + params, background_2d, signal = ctf_est.estimate( + micrograph, g_min=g_min, g_max=g_max + ) results[name] = params # TODO consolidate all the IO our of here and into functions # we write each micrograph's ctf parameters to an individual starfile - if not os.path.isdir(output_dir): - os.mkdir(output_dir) - ctf_object.write_star(name, params, output_dir) + star_filepath = os.path.join(output_dir, os.path.splitext(name)[0]) + ".star" + ctf_est.write_star(star_filepath, params) # Save the noise if save_noise_images: - noise_filepath = os.path.join(output_dir, os.path.splitext(name)[0] + "_noise.mrc") + noise_filepath = os.path.join( + output_dir, os.path.splitext(name)[0] + "_noise.mrc" + ) background_2d.save(noise_filepath, overwrite=True) - + # TODO, do after resolve CTF formula/grids issues, uses CTF formula. # if save_ctf_images: - # ctf_object.set_df1(cc_array[ml, 0]) # angstrom - # ctf_object.set_df2(cc_array[ml, 1]) # angstrom - # ctf_object.set_angle(cc_array[ml, 2]) # Radians - # ctf_object.generate_ctf() + # ctf_est.set_df1(cc_array[ml, 0]) # angstrom + # ctf_est.set_df2(cc_array[ml, 1]) # angstrom + # ctf_est.set_angle(cc_array[ml, 2]) # Radians + # ctf_est.generate_ctf() # df = (cc_array[ml, 0] + cc_array[ml, 1]) * np.ones( - # ctf_object.theta.shape, ctf_object.theta.dtype + # ctf_est.theta.shape, ctf_est.theta.dtype # ) + (cc_array[ml, 0] - cc_array[ml, 1]) * np.cos( - # 2 * ctf_object.theta + # 2 * ctf_est.theta # - 2 # * cc_array[ml, 2] - # * np.ones(ctf_object.theta.shape, ctf_object.theta.dtype) + # * np.ones(ctf_est.theta.shape, ctf_est.theta.dtype) # ) # ctf_im = -np.sin( # np.pi - # * ctf_object.lmbd - # * ctf_object.r_ctf**2 + # * ctf_est.lmbd + # * ctf_est.r_ctf**2 # / 2 - # * (df - ctf_object.lmbd**2 * ctf_object.r_ctf**2 * cs * 1e6) + # * (df - ctf_est.lmbd**2 * ctf_est.r_ctf**2 * cs * 1e6) # + self.amplitude_contrast # ) # ctf_signal = np.zeros(ctf_im.shape, ctf_im.dtype) From c57e99c355ac7c5658fffaceeee059c0c8e0efaf Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 12 Oct 2026 12:02:44 -0400 Subject: [PATCH 09/44] remove redundant taper transpose --- src/aspire/ctf/ctf_estimator.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 0429f8d90e..f2374ae506 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -275,7 +275,7 @@ def tapers(self, N, NW, L): # # In the CtfEstimator application, the signs will pass through # the mul and fft steps, then are absolute valued. - return dpss(M=N, NW=NW, Kmax=L, return_ratios=False, sym=False, norm=2).T + return dpss(M=N, NW=NW, Kmax=L, return_ratios=False, sym=False, norm=2) def estimate_psd(self, blocks, tapers_1d): """ @@ -286,7 +286,7 @@ def estimate_psd(self, blocks, tapers_1d): :return: NumPy array of estimated power spectrum. """ - num_1d_tapers = tapers_1d.shape[-1] + num_1d_tapers = tapers_1d.shape[0] tapers_1d = tapers_1d.astype(complex_type(self.dtype), copy=False) blocks_mt = np.zeros(blocks[0, :, :].shape, dtype=self.dtype) @@ -300,8 +300,8 @@ def estimate_psd(self, blocks, tapers_1d): for ax1 in range(num_1d_tapers): for ax2 in range(num_1d_tapers): np.matmul( - tapers_1d[:, ax1, np.newaxis], - tapers_1d[:, ax2, np.newaxis].T, + tapers_1d[ax1, :, np.newaxis], + tapers_1d[ax2, :, np.newaxis].T, out=taper_2d, ) From 26356d477449dc0dcd5d2dee3be189fd5c170430 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 12 Oct 2026 13:48:36 -0400 Subject: [PATCH 10/44] cleanup pca code/call --- src/aspire/ctf/ctf_estimator.py | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index f2374ae506..a75db009f6 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -101,6 +101,10 @@ def __init__( self.r_ctf = grid["r"] / 2 * (10 / pixel_size) # units: inverse nm self.theta = grid["phi"] + # Re-scale to unnormalized grids for X, Y (avoids recompute) + self.X = grid['x'] * psd_size/2 + self.Y = grid['y'] * psd_size/2 + # TODO, try to extend to any steerable basis self.basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) @@ -567,38 +571,29 @@ def pca(self, signal, pixel_size, g_min, g_max): :return: ratio. """ - # RCOPT - signal = signal.asnumpy()[0].T + signal = signal.asnumpy()[0] # Convert from Image to square array. N = signal.shape[0] center = N // 2 - grid = grid_2d(N, normalized=True, indexing="yx", dtype=self.dtype) - - r_ctf = grid["r"] / 2 * (10 / pixel_size) - - grid = grid_2d(N, normalized=False, indexing="yx", dtype=self.dtype) - X = grid["x"] - Y = grid["y"] - signal = signal - np.min(signal) rad_sq_min = N * pixel_size / g_min rad_sq_max = N * pixel_size / g_max - min_limit = r_ctf[center, (center + np.floor(rad_sq_min)).astype(int)] - signal[r_ctf < min_limit] = 0 + min_limit = self.r_ctf[center, (center + np.floor(rad_sq_min)).astype(int)] + signal[self.r_ctf < min_limit] = 0 - max_limit = r_ctf[center, (center + np.ceil(rad_sq_max)).astype(int)] - signal = np.where(r_ctf > max_limit, 0, signal) + max_limit = self.r_ctf[center, (center + np.ceil(rad_sq_max)).astype(int)] + signal[self.r_ctf > max_limit] = 0 - moment_02 = Y**2 * signal + moment_02 = self.Y**2 * signal moment_02 = np.sum(moment_02, axis=(0, 1)) - moment_11 = Y * X * signal + moment_11 = self.Y * self.X * signal moment_11 = np.sum(moment_11, axis=(0, 1)) - moment_20 = X**2 * signal + moment_20 = self.X**2 * signal moment_20 = np.sum(moment_20, axis=(0, 1)) moment_mat = np.zeros((2, 2)) From 99689adb7714ddb0cb784699b0797644eb47819a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 12 Oct 2026 13:58:15 -0400 Subject: [PATCH 11/44] temp timer line [skip ci] --- src/aspire/ctf/ctf_estimator.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index a75db009f6..2d5d86f758 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -932,9 +932,13 @@ def estimate_ctf( micrograph = micrograph.astype(dtype, copy=False) + from time import perf_counter + tic = perf_counter() params, background_2d, signal = ctf_est.estimate( micrograph, g_min=g_min, g_max=g_max ) + toc = perf_counter() + print('TIME:', toc-tic) results[name] = params From c55fea19195ed42ee32ca59ee0ba6153a5d16eac Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 12 Oct 2026 14:57:49 -0400 Subject: [PATCH 12/44] migrate the taper application to GPU [skip ci] --- src/aspire/ctf/ctf_estimator.py | 36 +++++++++++++++++---------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 2d5d86f758..8e37cf60b9 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -36,7 +36,7 @@ from aspire.basis import Coef, FFBBasis2D from aspire.image import Image -from aspire.numeric import fft +from aspire.numeric import fft, xp from aspire.storage import StarFile from aspire.utils import ( abs2, @@ -102,8 +102,8 @@ def __init__( self.theta = grid["phi"] # Re-scale to unnormalized grids for X, Y (avoids recompute) - self.X = grid['x'] * psd_size/2 - self.Y = grid['y'] * psd_size/2 + self.X = grid["x"] * psd_size / 2 + self.Y = grid["y"] * psd_size / 2 # TODO, try to extend to any steerable basis self.basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) @@ -291,26 +291,29 @@ def estimate_psd(self, blocks, tapers_1d): """ num_1d_tapers = tapers_1d.shape[0] - tapers_1d = tapers_1d.astype(complex_type(self.dtype), copy=False) + tapers_1d = xp.asarray(tapers_1d, dtype=complex_type(self.dtype)) - blocks_mt = np.zeros(blocks[0, :, :].shape, dtype=self.dtype) + blocks_mt = xp.zeros(blocks[0, :, :].shape, dtype=self.dtype) - blocks_tapered = np.zeros(blocks[0, :, :].shape, dtype=complex_type(self.dtype)) + blocks_tapered = xp.zeros(blocks[0, :, :].shape, dtype=complex_type(self.dtype)) - taper_2d = np.zeros( + taper_2d = xp.zeros( (blocks.shape[1], blocks.shape[2]), dtype=complex_type(self.dtype) ) + # Ensure blocks are on GPU + blocks = xp.asarray(blocks) + for ax1 in range(num_1d_tapers): for ax2 in range(num_1d_tapers): - np.matmul( + xp.matmul( tapers_1d[ax1, :, np.newaxis], tapers_1d[ax2, :, np.newaxis].T, out=taper_2d, ) for m in range(blocks.shape[0]): - np.multiply(blocks[m, :, :], taper_2d, out=blocks_tapered) + xp.multiply(blocks[m, :, :], taper_2d, out=blocks_tapered) blocks_mt_post_fft = fft.fftn(blocks_tapered, axes=(-2, -1)) blocks_mt += abs2(blocks_mt_post_fft) # blocks_mt += np.abs(blocks_mt_post_fft)**2 @@ -320,7 +323,7 @@ def estimate_psd(self, blocks, tapers_1d): amplitude_spectrum = fft.fftshift(blocks_mt) - return Image(amplitude_spectrum) + return Image(xp.asnumpy(amplitude_spectrum)) def radial_average(self, basis, amplitude_spectrum): """ @@ -571,7 +574,7 @@ def pca(self, signal, pixel_size, g_min, g_max): :return: ratio. """ - signal = signal.asnumpy()[0] # Convert from Image to square array. + signal = signal.asnumpy()[0] # Convert from Image to square array. N = signal.shape[0] center = N // 2 @@ -743,7 +746,7 @@ def gd( p = c1 / c2 - return df1, df2, angle_ast, p + return df1, df2, angle_ast, p, iter_no # Note, This doesn't actually use anything from the class. # It is used in a solver loop of some sort, so it may not be correct @@ -832,7 +835,7 @@ def estimate(self, micrograph, g_min=30, g_max=5): # This array holds results from each of the `gd_iters` cc_array = np.zeros((self.gd_iters, 4)) for a in range(0, self.gd_iters): - df1, df2, angle_ast, p = self.gd( + df1, df2, angle_ast, p, iters = self.gd( signal, initial_df1, initial_df2, @@ -849,6 +852,9 @@ def estimate(self, micrograph, g_min=30, g_max=5): cc_array[a, 1] = df2 cc_array[a, 2] = angle_ast # Radians cc_array[a, 3] = p + logger.debug( + f"Gradient Descent {a+1} of {self.gd_iters}:\tp={p}\titers={iters}" + ) # Select the best result index ml = np.argmax(cc_array[:, 3], -1) @@ -932,13 +938,9 @@ def estimate_ctf( micrograph = micrograph.astype(dtype, copy=False) - from time import perf_counter - tic = perf_counter() params, background_2d, signal = ctf_est.estimate( micrograph, g_min=g_min, g_max=g_max ) - toc = perf_counter() - print('TIME:', toc-tic) results[name] = params From f0fdd7ce4c323b7bf6b422734edb8353c6c4eb9d Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 13 Oct 2026 10:23:19 -0400 Subject: [PATCH 13/44] initial replace ctf_formula --- src/aspire/ctf/ctf_estimator.py | 51 ++++++++++++++------------------- 1 file changed, 22 insertions(+), 29 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 8e37cf60b9..31f5959cd5 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -37,15 +37,9 @@ from aspire.basis import Coef, FFBBasis2D from aspire.image import Image from aspire.numeric import fft, xp +from aspire.operators import CTFFilter from aspire.storage import StarFile -from aspire.utils import ( - abs2, - complex_type, - grid_1d, - grid_2d, - tqdm, - voltage_to_wavelength, -) +from aspire.utils import abs2, complex_type, grid_2d, tqdm, voltage_to_wavelength logger = logging.getLogger(__name__) @@ -100,6 +94,10 @@ def __init__( # Note range is -half to half. self.r_ctf = grid["r"] / 2 * (10 / pixel_size) # units: inverse nm self.theta = grid["phi"] + # Construct radial omega grid for ASPIRE's ctf_formula + center = self.psd_size // 2 + r = np.pi * grid["r"][center, center:] + self.radial_omega = np.stack((r, np.zeros_like(r))) # Re-scale to unnormalized grids for X, Y (avoids recompute) self.X = grid["x"] * psd_size / 2 @@ -471,7 +469,6 @@ def opt1d( cs, lmbd, w, - N, min_defocus=500, max_defocus=10000, ): @@ -483,38 +480,35 @@ def opt1d( :param cs: Spherical aberration in mm. :param lmbd: Electron wavelength \u212b (angstrom). :param w: Amplitude contrast. - :param N: Number of rows (or columns) in the estimate power spectrum. - :param min_defocus: Start of defocus loop scan. - :param max_defocus: End of defocus loop scan. + :param min_defocus: Start of defocus loop scan (xxx units). + :param max_defocus: End of defocus loop scan (xxx units). :return: 2-tuple of NumPy arrays (Estimated average of defocus and low_freq_cutoff) """ - center = N // 2 - - grid = grid_1d(N, normalized=True, dtype=self.dtype) - - rb = grid["r"][center:] / 2 - - r_ctf = rb * (10 / pixel_size) # units: inverse nm - signal = amplitude_spectrum.T signal = np.maximum(0.0, signal) signal = np.sqrt(signal) signal = signal[: 3 * signal.shape[0] // 4] - r_ctf_sq = r_ctf**2 + # XXX Note ask Ayelet about original range units, I read as nm, but that is 5000-100000A ...? c = np.zeros((max_defocus - min_defocus, signal.shape[1]), dtype=self.dtype) for f in range(min_defocus, max_defocus): - # TODO try replace with ASPIRE's CTF formula? + ctf_im = np.abs( - np.sin( - np.pi * lmbd * f * r_ctf_sq - - 0.5 * np.pi * (lmbd**3) * cs * 1e6 * r_ctf_sq**2 - + w - ) + CTFFilter.ctf_formula( + self.radial_omega, + pixel_size=self.pixel_size, + voltage=self.voltage, + defocus_u_a=10 * f, + defocus_v_a=10 * f, + defocus_ang=0, + Cs=self.cs, + alpha=self.amplitude_contrast, + B=0, + )[0] ) - ctf_im = ctf_im[: signal.shape[0]] + ctf_im = ctf_im[: signal.shape[0]] # xxx I think can remove this? ctf_im = np.reshape(ctf_im, (ctf_im.shape[0], 1)) ctf_im = np.tile(ctf_im, (1, signal.shape[1])) @@ -812,7 +806,6 @@ def estimate(self, micrograph, g_min=30, g_max=5): self.cs, self.lmbd, # (angstrom) self.amplitude_contrast, - signal_observed.shape[-1], ) # TODO check the low_freq_skip +-1, it is col ind or value From 76333427438080b4476bc9c561abd5689f26d89e Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 13 Oct 2026 11:22:38 -0400 Subject: [PATCH 14/44] continue replacing formula and fixing plots [skip ci] --- src/aspire/ctf/ctf_estimator.py | 89 +++++++++++++++++---------------- 1 file changed, 45 insertions(+), 44 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 31f5959cd5..3c0b5cca90 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -94,14 +94,16 @@ def __init__( # Note range is -half to half. self.r_ctf = grid["r"] / 2 * (10 / pixel_size) # units: inverse nm self.theta = grid["phi"] - # Construct radial omega grid for ASPIRE's ctf_formula + # Construct omega grids for use in ASPIRE's ctf_formula + self.omega = np.pi * np.vstack((grid["x"].flatten(), grid["y"].flatten())) center = self.psd_size // 2 r = np.pi * grid["r"][center, center:] - self.radial_omega = np.stack((r, np.zeros_like(r))) + self.radial_omega = np.stack((r, np.zeros_like(r))) # 1d (ray) - # Re-scale to unnormalized grids for X, Y (avoids recompute) + # Re-scale to unnormalized grids for X, Y, and R (avoids recompute) self.X = grid["x"] * psd_size / 2 self.Y = grid["y"] * psd_size / 2 + self.R = grid["r"] * psd_size / 2 # TODO, try to extend to any steerable basis self.basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) @@ -540,15 +542,11 @@ def background_subtract_2d(self, signal, background_p1, max_col): signal = signal.asnumpy() - N = signal.shape[1] - grid = grid_2d(N, normalized=False, indexing="yx", dtype=self.dtype) - radii = grid["r"] # can probably get elsewhere if signal.shape is consistent. - - # okay now, but review max_col and vectorize this + # TODO, review max_col and vectorize this background = np.zeros(signal.shape, dtype=self.dtype) for r in range(max_col + 2, background_p1.shape[1]): - background[:, (r < radii) & (radii <= r + 1)] = background_p1[max_col, r] - mask = radii <= max_col + 2 + background[:, (r < self.R) & (self.R <= r + 1)] = background_p1[max_col, r] + mask = self.R <= max_col + 2 background[:, mask] = signal[:, mask] signal = signal - background @@ -781,6 +779,16 @@ def _plot_1d_bg_sub_psd(self, signal_1d): # ylabel('$S_y \left( \mathbf{k} \right) - S_e \left( \mathbf{k} \right)$', 'Interpreter','latex'); plt.show() + def _plot_2d_psd(self, psd): + plt.imshow(psd) + # title('2-d PSD') + plt.show() + + def _plot_2d_bg_sub_psd(self, bg_sub_psd): + plt.imshow(bg_sub_psd) + # title('Background subtracted 2-d PSD') + plt.show() + def estimate(self, micrograph, g_min=30, g_max=5): # TODO compare now that remainder of code repros. # micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size) @@ -813,6 +821,9 @@ def estimate(self, micrograph, g_min=30, g_max=5): signal_observed, background_1d, low_freq_skip ) + # self._plot_2d_psd(amplitude_spectrum.asnumpy()[0]) + # self._plot_2d_bg_sub_psd(signal.asnumpy()[0]) + ratio = self.pca(signal_observed, self.pixel_size, g_min, g_max) signal, additional_background = self.elliptical_average( @@ -852,7 +863,6 @@ def estimate(self, micrograph, g_min=30, g_max=5): # Select the best result index ml = np.argmax(cc_array[:, 3], -1) - # TODO what is the ultimate purpose of the result dictionary? params = { "defocus_u": cc_array[ml, 0] * 10, # Convert from nm to A "defocus_v": cc_array[ml, 1] * 10, # Convert from nm to A @@ -950,38 +960,29 @@ def estimate_ctf( background_2d.save(noise_filepath, overwrite=True) # TODO, do after resolve CTF formula/grids issues, uses CTF formula. - # if save_ctf_images: - # ctf_est.set_df1(cc_array[ml, 0]) # angstrom - # ctf_est.set_df2(cc_array[ml, 1]) # angstrom - # ctf_est.set_angle(cc_array[ml, 2]) # Radians - # ctf_est.generate_ctf() - # df = (cc_array[ml, 0] + cc_array[ml, 1]) * np.ones( - # ctf_est.theta.shape, ctf_est.theta.dtype - # ) + (cc_array[ml, 0] - cc_array[ml, 1]) * np.cos( - # 2 * ctf_est.theta - # - 2 - # * cc_array[ml, 2] - # * np.ones(ctf_est.theta.shape, ctf_est.theta.dtype) - # ) - # ctf_im = -np.sin( - # np.pi - # * ctf_est.lmbd - # * ctf_est.r_ctf**2 - # / 2 - # * (df - ctf_est.lmbd**2 * ctf_est.r_ctf**2 * cs * 1e6) - # + self.amplitude_contrast - # ) - # ctf_signal = np.zeros(ctf_im.shape, ctf_im.dtype) - # ctf_signal[: ctf_im.shape[0] // 2, :] = ctf_im[: ctf_im.shape[0] // 2, :] - # ctf_signal[ctf_im.shape[0] // 2 + 1 :, :] = signal.asnumpy()[ - # :, :, ctf_im.shape[0] // 2 + 1 - # ] - - # with mrcfile.new( - # os.path.join(output_dir, os.path.splitext(name)[0] + ".ctf"), - # overwrite=True, - # ) as mrc: - # mrc.set_data(np.float32(ctf_signal)) - # mrc.voxel_size = pixel_size + if save_ctf_images: + ctf_im = CTFFilter.ctf_formula( + ctf_est.omega, + pixel_size=ctf_est.pixel_size, + voltage=ctf_est.voltage, + defocus_u_a=params["defocus_u"], + defocus_v_a=params["defocus_v"], + defocus_ang=params["defocus_ang"], + Cs=ctf_est.cs, + alpha=ctf_est.amplitude_contrast, + B=0, + ).reshape(psd_size, psd_size) + + center = psd_size // 2 + # Overwrite half the `ctf_im` with `signal` + ctf_im[center:, :] = signal.asnumpy()[0, center:, :] + # debug plots + # plt.imshow(ctf_im) + # plt.show() + + ctf_filepath = os.path.join(output_dir, os.path.splitext(name)[0] + ".ctf") + with mrcfile.new(ctf_filepath, overwrite=True) as mrc: + mrc.set_data(ctf_im.astype(np.float32)) # MRC only supports singles + mrc.voxel_size = ctf_est.pixel_size return results From 3d7fe312e49aa6a08773816a9b626a5fd69d2e68 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 14 Oct 2026 10:43:19 -0400 Subject: [PATCH 15/44] continue ctf est code cleanup instance vars and plots --- src/aspire/ctf/ctf_estimator.py | 192 ++++++++++++++++---------------- 1 file changed, 94 insertions(+), 98 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 3c0b5cca90..8899ea8307 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -59,6 +59,7 @@ def __init__( num_tapers, gradient_descent_iterations=6, linprog_method="highs-ds", + diagnostic_plots=False, dtype=np.float32, ): """ @@ -75,6 +76,9 @@ def __init__( Defaults to 6. :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). Defaults to 'highs-ds'. + :param diagnostic_plots: Optionally generate diagnostic_plots. + Note, these are blocking plot calls. + Defaults to False. :returns: CtfEstimator instance. """ @@ -87,6 +91,7 @@ def __init__( self.lmbd = voltage_to_wavelength(voltage) / 10.0 # Convert angstrom to nm self.gd_iters = int(gradient_descent_iterations) self.linprog_method = linprog_method + self.diagnostic_plots = diagnostic_plots self.dtype = np.dtype(dtype) # Grids @@ -110,7 +115,8 @@ def __init__( # Only need to compute tapers once # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L - self.tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) + tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) + self.tapers_1d = xp.asarray(tapers_1d, dtype=complex_type(self.dtype)) def micrograph_to_blocks(self, micrograph, block_size): """ @@ -281,17 +287,15 @@ def tapers(self, N, NW, L): # the mul and fft steps, then are absolute valued. return dpss(M=N, NW=NW, Kmax=L, return_ratios=False, sym=False, norm=2) - def estimate_psd(self, blocks, tapers_1d): + def estimate_psd(self, blocks): """ Estimate the power spectrum of the micrograph using the multi-taper method :param blocks: 3-D NumPy array containing windows extracted from the micrograph in the preprocess function. - :param tapers_1d: NumPy array of data tapers. :return: NumPy array of estimated power spectrum. """ - num_1d_tapers = tapers_1d.shape[0] - tapers_1d = xp.asarray(tapers_1d, dtype=complex_type(self.dtype)) + num_1d_tapers = self.tapers_1d.shape[0] blocks_mt = xp.zeros(blocks[0, :, :].shape, dtype=self.dtype) @@ -307,8 +311,8 @@ def estimate_psd(self, blocks, tapers_1d): for ax1 in range(num_1d_tapers): for ax2 in range(num_1d_tapers): xp.matmul( - tapers_1d[ax1, :, np.newaxis], - tapers_1d[ax2, :, np.newaxis].T, + self.tapers_1d[ax1, :, np.newaxis], + self.tapers_1d[ax2, :, np.newaxis].T, out=taper_2d, ) @@ -325,46 +329,48 @@ def estimate_psd(self, blocks, tapers_1d): return Image(xp.asnumpy(amplitude_spectrum)) - def radial_average(self, basis, amplitude_spectrum): + def radial_average(self, amplitude_spectrum): """ Computes radial/elliptical average of the power spectrum - :param basis: SteerableBasis instance (typically FFBBasis). :param amplitude_spectrum: Power spectrum to average. :return: PSD and noise as 2-tuple of NumPy arrays. """ noise = amplitude_spectrum # pass through - coefs_s = basis.evaluate_t(amplitude_spectrum).asnumpy().copy() + coefs_s = self.basis.evaluate_t(amplitude_spectrum).asnumpy().copy() - mask_s = (basis.angular_indices == 1) | (basis.angular_indices == 2) + mask_s = (self.basis.angular_indices == 1) | (self.basis.angular_indices == 2) coefs_s[:, mask_s] = 0 - psd = Coef(basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size).evaluate() + psd = Coef( + self.basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size + ).evaluate() return psd, noise - def elliptical_average(self, basis, amplitude_spectrum): + def elliptical_average(self, amplitude_spectrum): """ Computes elliptical average of the power spectrum - :param basis: SteerableBasis instance (typically FFBBasis). :param amplitude_spectrum: Power spectrum to average. :return: PSD and noise as 2-tuple of NumPy arrays. """ - coefs_s = basis.evaluate_t(amplitude_spectrum).asnumpy().copy() + coefs_s = self.basis.evaluate_t(amplitude_spectrum).asnumpy().copy() coefs_n = coefs_s.copy() - mask_s = basis.angular_indices == 1 + mask_s = self.basis.angular_indices == 1 coefs_s[:, mask_s] = 0 - mask_n = (basis.angular_indices == 0) | (basis.angular_indices == 2) + mask_n = (self.basis.angular_indices == 0) | (self.basis.angular_indices == 2) coefs_n[:, mask_n] = 0 noise = Coef( - basis, coefs_n, pixel_size=amplitude_spectrum.pixel_size + self.basis, coefs_n, pixel_size=amplitude_spectrum.pixel_size + ).evaluate() + psd = Coef( + self.basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size ).evaluate() - psd = Coef(basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size).evaluate() return psd, noise @@ -467,10 +473,6 @@ def background_subtract_1d( def opt1d( self, amplitude_spectrum, - pixel_size, - cs, - lmbd, - w, min_defocus=500, max_defocus=10000, ): @@ -478,10 +480,6 @@ def opt1d( Find optimal defocus for the radially symmetric case (where no astigmatism is present) :param amplitude_spectrum: Estimated power specrtum. - :param pixel_size: Pixel size in \u212b (angstrom). - :param cs: Spherical aberration in mm. - :param lmbd: Electron wavelength \u212b (angstrom). - :param w: Amplitude contrast. :param min_defocus: Start of defocus loop scan (xxx units). :param max_defocus: End of defocus loop scan (xxx units). :return: 2-tuple of NumPy arrays (Estimated average of defocus and low_freq_cutoff) @@ -492,7 +490,7 @@ def opt1d( signal = np.sqrt(signal) signal = signal[: 3 * signal.shape[0] // 4] - # XXX Note ask Ayelet about original range units, I read as nm, but that is 5000-100000A ...? + # Ask Ay about original range units, I read as nm, but that is 5000-100000A ...? xxx c = np.zeros((max_defocus - min_defocus, signal.shape[1]), dtype=self.dtype) for f in range(min_defocus, max_defocus): @@ -510,7 +508,9 @@ def opt1d( B=0, )[0] ) - ctf_im = ctf_im[: signal.shape[0]] # xxx I think can remove this? + # Shouldn't normally be needed within class calls, + # but if user running manually clip to signal len. + ctf_im = ctf_im[: signal.shape[0]] ctf_im = np.reshape(ctf_im, (ctf_im.shape[0], 1)) ctf_im = np.tile(ctf_im, (1, signal.shape[1])) @@ -518,8 +518,8 @@ def opt1d( for m in range(0, signal.shape[1]): signal[:, m] = signal[:, m] - np.mean(signal[m + 1 :, m], axis=0) ctf_im[:, m] = ctf_im[:, m] - np.mean(ctf_im[m + 1 :, m], axis=0) - ctf_im[: m + 1, m] = 0 # np.zeros((m + 1)) - signal[: m + 1, m] = 0 # np.zeros((m + 1)) + ctf_im[: m + 1, m] = 0 + signal[: m + 1, m] = 0 Sx = np.sqrt(np.sum(ctf_im**2, axis=0)) Sy = np.sqrt(np.sum(signal**2, axis=0)) @@ -556,11 +556,10 @@ def background_subtract_2d(self, signal, background_p1, max_col): background, pixel_size=self.pixel_size ) - def pca(self, signal, pixel_size, g_min, g_max): + def pca(self, signal, g_min, g_max): """ :param signal: Estimated power spectrum. - :param pixel_size: Pixel size in \u212b (angstrom). :param g_min: Inverse of minimun resolution for PSD. :param g_max: Inverse of maximum resolution for PSD. :return: ratio. @@ -573,8 +572,8 @@ def pca(self, signal, pixel_size, g_min, g_max): signal = signal - np.min(signal) - rad_sq_min = N * pixel_size / g_min - rad_sq_max = N * pixel_size / g_max + rad_sq_min = N * self.pixel_size / g_min + rad_sq_max = N * self.pixel_size / g_max min_limit = self.r_ctf[center, (center + np.floor(rad_sq_min)).astype(int)] signal[self.r_ctf < min_limit] = 0 @@ -608,12 +607,8 @@ def gd( df1, df2, angle_ast, - pixel_size, g_min, g_max, - amplitude_contrast, - lmbd, - cs, ): """ Runs gradient ascent to optimize defocus parameters @@ -622,12 +617,8 @@ def gd( :param df1: Defocus value in the direction perpendicular to df2. :param df2: Defocus value in the direction perpendicular to df1. :param angle_ast: Angle between df1 and the x-axis, Radians. - :param pixel_size: Pixel size in \u212b (angstrom). :param g_min: Inverse of minimun resolution for PSD. :param g_max: Inverse of maximum resolution for PSD. - :param amplitude_contrast: Amplitude contrast. - :param lmbd: Electron wavelength \u212b (angstrom). - :param cs: Spherical aberration in mm. :return: Optimal defocus parameters """ @@ -640,17 +631,21 @@ def gd( y = (df1 - df2) * np.cos(2 * angle_ast) z = (df1 - df2) * np.sin(2 * angle_ast) - a = np.pi * lmbd * self.r_ctf**2 / 2 - b = np.pi * lmbd**3 * cs * 1e6 * self.r_ctf**4 / 2 - np.full( - shape=self.r_ctf.shape, fill_value=amplitude_contrast, dtype=self.dtype + a = np.pi * self.lmbd * self.r_ctf**2 / 2 + # Ask Ay, Does this amplitude_contrast need to be the arctan'd one? + # amplitude_contrast = np.arctan( + # self.amplitude_contrast / np.sqrt(1 - self.amplitude_contrast**2) + # ) + b = np.pi * self.lmbd**3 * self.cs * 1e6 * self.r_ctf**4 / 2 - np.full( + shape=self.r_ctf.shape, fill_value=self.amplitude_contrast, dtype=self.dtype ) signal = signal.asnumpy()[0].T N = signal.shape[1] center = N // 2 - rad_sq_min = N * pixel_size / g_min - rad_sq_max = N * pixel_size / g_max + rad_sq_min = N * self.pixel_size / g_min + rad_sq_max = N * self.pixel_size / g_max max_val = self.r_ctf[center, int(center - 1 + np.floor(rad_sq_max))] min_val = self.r_ctf[center, int(center - 1 + np.ceil(rad_sq_min))] @@ -794,41 +789,35 @@ def estimate(self, micrograph, g_min=30, g_max=5): # micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size) micrograph_blocks = self.legacy_preprocess_micrograph(micrograph, self.psd_size) - signal_observed = self.estimate_psd(micrograph_blocks, self.tapers_1d) + signal_observed = self.estimate_psd(micrograph_blocks) # np.save("signal_observed.npy", signal_observed) - amplitude_spectrum, _ = self.radial_average(self.basis, signal_observed) + amplitude_spectrum, _ = self.radial_average(signal_observed) # np.save("amplitude_spectrum.npy", amplitude_spectrum) signal_1d, background_1d = self.background_subtract_1d( amplitude_spectrum, linprog_method=self.linprog_method ) - # TODO complete diagnostic plots and complete implementation in class. - # self._plot_1d_psd(signal_1d, background_1d) - # self._plot_1d_bg_sub_psd(signal_1d) + if self.diagnostic_plots: + self._plot_1d_psd(signal_1d, background_1d) + self._plot_1d_bg_sub_psd(signal_1d) - avg_defocus, low_freq_skip = self.opt1d( - signal_1d, - self.pixel_size, - self.cs, - self.lmbd, # (angstrom) - self.amplitude_contrast, - ) + avg_defocus, low_freq_skip = self.opt1d(signal_1d) # TODO check the low_freq_skip +-1, it is col ind or value signal, background_2d = self.background_subtract_2d( signal_observed, background_1d, low_freq_skip ) - # self._plot_2d_psd(amplitude_spectrum.asnumpy()[0]) - # self._plot_2d_bg_sub_psd(signal.asnumpy()[0]) + if self.diagnostic_plots: + self._plot_2d_psd(amplitude_spectrum.asnumpy()[0]) + self._plot_2d_bg_sub_psd(signal.asnumpy()[0]) - ratio = self.pca(signal_observed, self.pixel_size, g_min, g_max) + ratio = self.pca(signal_observed, g_min, g_max) + + signal, additional_background = self.elliptical_average(signal.sqrt()) - signal, additional_background = self.elliptical_average( - self.basis, signal.sqrt() - ) background_2d = background_2d + additional_background initial_df1 = (avg_defocus * 2) / (1 + ratio) @@ -838,18 +827,14 @@ def estimate(self, micrograph, g_min=30, g_max=5): # This array holds results from each of the `gd_iters` cc_array = np.zeros((self.gd_iters, 4)) - for a in range(0, self.gd_iters): + for a, angle in enumerate(np.linspace(0, np.pi, self.gd_iters)): df1, df2, angle_ast, p, iters = self.gd( signal, initial_df1, initial_df2, - angle + a * np.pi / 6.0, # Radians, + a*30degrees - self.pixel_size, + angle, g_min, g_max, - self.amplitude_contrast, - self.lmbd, # (angstrom) - self.cs, ) cc_array[a, 0] = df1 @@ -875,6 +860,35 @@ def estimate(self, micrograph, g_min=30, g_max=5): return params, background_2d, signal + def _create_ctf_image(self, signal, df1, df2, dfang): + """ + Create a diagnostic array comparing estimated CTF and signal. + + :param signal: PSD of signal + :param df1: Major defocus (Angstrom) + :param df2: Minor defocus (Angstrom) + :param dfang: Defocus angle (radians) + :return: 2d array containing image data. + """ + + ctf_im = CTFFilter.ctf_formula( + self.omega, + pixel_size=self.pixel_size, + voltage=self.voltage, + defocus_u_a=df1, + defocus_v_a=df2, + defocus_ang=dfang, + Cs=self.cs, + alpha=self.amplitude_contrast, + B=0, + ).reshape(self.psd_size, self.psd_size) + + # Overwrite half the `ctf_im` with `signal` + center = self.psd_size // 2 + ctf_im[center:, :] = signal.asnumpy()[0, center:, :] + + return ctf_im + def estimate_ctf( data_folder, @@ -915,11 +929,6 @@ def estimate_ctf( file_names = mrc_files + mrcs_files # Setup a CTFEstimator ----------------------------------------------------- - # Why - amplitude_contrast = np.arctan( - amplitude_contrast / np.sqrt(1 - amplitude_contrast**2) - ) - ctf_est = CtfEstimator( pixel_size, cs, @@ -928,6 +937,7 @@ def estimate_ctf( psd_size, num_tapers=num_tapers, gradient_descent_iterations=gradient_descent_iterations, + diagnostic_plots=False, dtype=dtype, ) @@ -947,8 +957,7 @@ def estimate_ctf( results[name] = params - # TODO consolidate all the IO our of here and into functions - # we write each micrograph's ctf parameters to an individual starfile + # Write each micrograph's ctf parameters to an individual STAR file star_filepath = os.path.join(output_dir, os.path.splitext(name)[0]) + ".star" ctf_est.write_star(star_filepath, params) @@ -959,26 +968,13 @@ def estimate_ctf( ) background_2d.save(noise_filepath, overwrite=True) - # TODO, do after resolve CTF formula/grids issues, uses CTF formula. if save_ctf_images: - ctf_im = CTFFilter.ctf_formula( - ctf_est.omega, - pixel_size=ctf_est.pixel_size, - voltage=ctf_est.voltage, - defocus_u_a=params["defocus_u"], - defocus_v_a=params["defocus_v"], - defocus_ang=params["defocus_ang"], - Cs=ctf_est.cs, - alpha=ctf_est.amplitude_contrast, - B=0, - ).reshape(psd_size, psd_size) - - center = psd_size // 2 - # Overwrite half the `ctf_im` with `signal` - ctf_im[center:, :] = signal.asnumpy()[0, center:, :] - # debug plots - # plt.imshow(ctf_im) - # plt.show() + ctf_im = ctf_est._create_ctf_image( + signal, + params["defocus_u"], + params["defocus_v"], + params["defocus_ang"], + ) ctf_filepath = os.path.join(output_dir, os.path.splitext(name)[0] + ".ctf") with mrcfile.new(ctf_filepath, overwrite=True) as mrc: From 3ff9a9d65acb7550e0ab62b07f295f1f72b94d06 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 14 Oct 2026 11:16:32 -0400 Subject: [PATCH 16/44] cleanup path searching, attempt to make recursive --- src/aspire/ctf/ctf_estimator.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 8899ea8307..dce82f658e 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -26,6 +26,7 @@ import logging import os from collections import OrderedDict +from pathlib import Path import matplotlib.pyplot as plt import mrcfile @@ -897,7 +898,7 @@ def estimate_ctf( amplitude_contrast=0.07, voltage=300, num_tapers=4, - psd_size=512, # should change to 1024? + psd_size=1024, g_min=30, g_max=5, gradient_descent_iterations=6, @@ -920,13 +921,12 @@ def estimate_ctf( os.mkdir(output_dir) # File inputs ------------------------------------------------------------- - # TODO confirm if scans multi-dir - dir_content = os.scandir(data_folder) - - # TODO change to full path - mrc_files = [f.name for f in dir_content if os.path.splitext(f)[1] == ".mrc"] - mrcs_files = [f.name for f in dir_content if os.path.splitext(f)[1] == ".mrcs"] - file_names = mrc_files + mrcs_files + data_folder = Path(data_folder) + # Handles recursive directory searching for both mrc and mrcs extensions + file_paths = list(data_folder.rglob("*.mrc[s]")) + if len(file_paths) == 0: + raise RuntimeError(f"No mrc or mrcs files found at {data_folder}.") + logger.info("Found {len(file_paths)} at {data_folder}.") # Setup a CTFEstimator ----------------------------------------------------- ctf_est = CtfEstimator( @@ -943,7 +943,7 @@ def estimate_ctf( # Main Loop --------------------------------------------------------------- results = {} - for name in tqdm(file_names, desc="Processing files"): + for name in tqdm(file_paths, desc="Processing files"): with mrcfile.open( os.path.join(data_folder, name), mode="r", permissive=True ) as mrc: From e9db82937e745d9834ebc76694cab4758e39ab8f Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 14 Oct 2026 14:10:44 -0400 Subject: [PATCH 17/44] optimize background fill, improve filepath handing hopefully --- src/aspire/ctf/ctf_estimator.py | 59 +++++++++++++++++++++------------ 1 file changed, 37 insertions(+), 22 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index dce82f658e..a42dd04864 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -24,7 +24,6 @@ """ import logging -import os from collections import OrderedDict from pathlib import Path @@ -543,10 +542,10 @@ def background_subtract_2d(self, signal, background_p1, max_col): signal = signal.asnumpy() - # TODO, review max_col and vectorize this background = np.zeros(signal.shape, dtype=self.dtype) - for r in range(max_col + 2, background_p1.shape[1]): - background[:, (r < self.R) & (self.R <= r + 1)] = background_p1[max_col, r] + idx = np.ceil(self.R).astype(int) - 1 + rings = (self.R > max_col + 2) & (self.R <= background_p1.shape[1]) + background[:, rings] = background_p1[max_col, idx[rings]] mask = self.R <= max_col + 2 background[:, mask] = signal[:, mask] @@ -906,27 +905,33 @@ def estimate_ctf( dtype=np.float32, save_ctf_images=False, save_noise_images=False, + rglob=r"*.mrc", ): """ Given paramaters estimates CTF from experimental data and returns CTF as a mrc file. + + :param rglob: File search pattern. + By default will search for **/*.mrc files. + Using "*.mrc[s]" will additonally include mrcs files, etc. """ # Checks ------------------------------------------------------------------ dtype = np.dtype(dtype) assert dtype in (np.float32, np.float64) - # Ensure we can write outputs before we get too far - if not os.path.isdir(output_dir): - os.mkdir(output_dir) - # File inputs ------------------------------------------------------------- data_folder = Path(data_folder) - # Handles recursive directory searching for both mrc and mrcs extensions - file_paths = list(data_folder.rglob("*.mrc[s]")) + # Handles recursive directory searching for both `rglob` patten + file_paths = list(data_folder.rglob(rglob)) if len(file_paths) == 0: - raise RuntimeError(f"No mrc or mrcs files found at {data_folder}.") - logger.info("Found {len(file_paths)} at {data_folder}.") + raise RuntimeError(f"No mrc or mrcs files found in {data_folder}.") + logger.info(f"Found {len(file_paths)} files in '{data_folder}'.") + + # File outputs ----------------------------------------------------------- + # Ensure we can write outputs before we get too far + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) # Setup a CTFEstimator ----------------------------------------------------- ctf_est = CtfEstimator( @@ -943,10 +948,21 @@ def estimate_ctf( # Main Loop --------------------------------------------------------------- results = {} - for name in tqdm(file_paths, desc="Processing files"): - with mrcfile.open( - os.path.join(data_folder, name), mode="r", permissive=True - ) as mrc: + for input_file in tqdm(file_paths, desc="Processing files"): + + # Handle path stuff first, in case it fails + relative_input_path = input_file.relative_to( + data_folder + ) # File path relative to data_folder + directory_stub = relative_input_path.parent # directory relative to data_folder + output_file_dir = output_dir / directory_stub # Join onto output_dir + output_file_dir.mkdir( + parents=True, exist_ok=True + ) # Ensure output_file_dir exists + output_file_prefix = output_file_dir / input_file.stem + breakpoint() + + with mrcfile.open(input_file, mode="r", permissive=True) as mrc: micrograph = mrc.data micrograph = micrograph.astype(dtype, copy=False) @@ -955,17 +971,16 @@ def estimate_ctf( micrograph, g_min=g_min, g_max=g_max ) - results[name] = params + # Store files' params into results dictionary + results[relative_input_path] = params # Write each micrograph's ctf parameters to an individual STAR file - star_filepath = os.path.join(output_dir, os.path.splitext(name)[0]) + ".star" + star_filepath = output_file_prefix.with_suffix(".star") ctf_est.write_star(star_filepath, params) # Save the noise if save_noise_images: - noise_filepath = os.path.join( - output_dir, os.path.splitext(name)[0] + "_noise.mrc" - ) + noise_filepath = Path(f"{output_file_prefix}_noise.mrc") background_2d.save(noise_filepath, overwrite=True) if save_ctf_images: @@ -976,7 +991,7 @@ def estimate_ctf( params["defocus_ang"], ) - ctf_filepath = os.path.join(output_dir, os.path.splitext(name)[0] + ".ctf") + ctf_filepath = output_file_prefix.with_suffix(".ctf") with mrcfile.new(ctf_filepath, overwrite=True) as mrc: mrc.set_data(ctf_im.astype(np.float32)) # MRC only supports singles mrc.voxel_size = ctf_est.pixel_size From bdf585dd95ebdbbab08c1d4862af0ce1b291b0f3 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 14 Oct 2026 14:20:41 -0400 Subject: [PATCH 18/44] fixup low_freq_cutoff --- src/aspire/ctf/ctf_estimator.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index a42dd04864..215ba876be 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -525,10 +525,10 @@ def opt1d( Sy = np.sqrt(np.sum(signal**2, axis=0)) c[f - min_defocus, :] = np.sum(ctf_im * signal, axis=0) / (Sx * Sy) - avg_defocus, low_freq_cutoff = np.unravel_index(np.argmax(c), c.shape)[:2] + avg_defocus, low_freq_cutoff_col = np.unravel_index(np.argmax(c), c.shape)[:2] avg_defocus += min_defocus - return avg_defocus, low_freq_cutoff + return avg_defocus, low_freq_cutoff_col def background_subtract_2d(self, signal, background_p1, max_col): """ @@ -544,9 +544,9 @@ def background_subtract_2d(self, signal, background_p1, max_col): background = np.zeros(signal.shape, dtype=self.dtype) idx = np.ceil(self.R).astype(int) - 1 - rings = (self.R > max_col + 2) & (self.R <= background_p1.shape[1]) + rings = (self.R > max_col + 1) & (self.R <= background_p1.shape[1]) background[:, rings] = background_p1[max_col, idx[rings]] - mask = self.R <= max_col + 2 + mask = self.R <= max_col + 1 background[:, mask] = signal[:, mask] signal = signal - background @@ -803,11 +803,10 @@ def estimate(self, micrograph, g_min=30, g_max=5): self._plot_1d_psd(signal_1d, background_1d) self._plot_1d_bg_sub_psd(signal_1d) - avg_defocus, low_freq_skip = self.opt1d(signal_1d) + avg_defocus, low_freq_skip_col = self.opt1d(signal_1d) - # TODO check the low_freq_skip +-1, it is col ind or value signal, background_2d = self.background_subtract_2d( - signal_observed, background_1d, low_freq_skip + signal_observed, background_1d, low_freq_skip_col ) if self.diagnostic_plots: @@ -960,7 +959,6 @@ def estimate_ctf( parents=True, exist_ok=True ) # Ensure output_file_dir exists output_file_prefix = output_file_dir / input_file.stem - breakpoint() with mrcfile.open(input_file, mode="r", permissive=True) as mrc: micrograph = mrc.data From b5ca44fa54ab73840a1d6eaec6f83a33900569d0 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 14 Oct 2026 15:14:06 -0400 Subject: [PATCH 19/44] add optimized gd --- src/aspire/ctf/ctf_estimator.py | 103 ++++++++++++++++++++++++++++++-- 1 file changed, 99 insertions(+), 4 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 215ba876be..0a805166d6 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -613,6 +613,105 @@ def gd( """ Runs gradient ascent to optimize defocus parameters + :param signal: Estimated power spectrum + :param df1: Defocus value in the direction perpendicular to df2. + :param df2: Defocus value in the direction perpendicular to df1. + :param angle_ast: Angle between df1 and the x-axis, Radians. + :param g_min: Inverse of minimun resolution for PSD. + :param g_max: Inverse of maximum resolution for PSD. + :return: Optimal defocus parameters + """ + # step sizes for (x, y, z) + alpha = np.array([1e5, 1e4, 1e4]) + + # initialization + x = df1 + df2 + y = (df1 - df2) * np.cos(2 * angle_ast) + z = (df1 - df2) * np.sin(2 * angle_ast) + + signal = signal.asnumpy()[0].T + N = signal.shape[1] + center = N // 2 + + rad_sq_min = N * self.pixel_size / g_min + rad_sq_max = N * self.pixel_size / g_max + + max_val = self.r_ctf[center, int(center - 1 + np.floor(rad_sq_max))] + min_val = self.r_ctf[center, int(center - 1 + np.ceil(rad_sq_min))] + mask = (self.r_ctf <= max_val) & (self.r_ctf > min_val) + + # everything below depends only on the ring of pixels + r2 = self.r_ctf[mask] ** 2 + a = (np.pi * self.lmbd / 2) * r2 + b = (np.pi * self.lmbd**3 * self.cs * 1e6 / 2) * r2**2 - self.amplitude_contrast + + sig = np.ascontiguousarray(signal[mask]) # repack for compute below + theta2 = 2 * self.theta[mask] + + # Referencing old implementation, + # psi = x*a + y*a*cos2t + z*a*sin2t - b == v @ A - b + A = np.ascontiguousarray(np.stack([a, a * np.cos(theta2), a * np.sin(theta2)])) + + sum_A = sig @ sig + v = np.array([x, y, z], dtype=A.dtype) + W = np.empty((sig.size, 2), dtype=A.dtype) + + def evaluate(v, need_grad=True): + psi = v @ A + psi -= b + s = np.sin(psi) + c1 = np.abs(s) @ sig + c2 = np.sqrt(sum_A * (s @ s)) + grad = None + if need_grad: + cs = np.cos(psi) + np.multiply(np.sign(s), cs, out=W[:, 0]) + W[:, 0] *= sig # sign(s) * cos * signal + np.multiply(s, cs, out=W[:, 1]) + W[:, 1] *= 2 # d(sin^2)/d(psi) + G = A @ W # compute all gradient sums + dc1 = G[:, 0] + dc2 = sum_A * G[:, 1] / (2 * c2) # derivative_sqrt = 1 / (2 * c2) + grad = (dc1 * c2 - dc2 * c1) / c2**2 + return c1, c2, grad + + dx = dy = dz = 1.0 + stop_cond = 1e-20 + iter_no = 1 + + while max(dx, dy, dz) > stop_cond: + _, _, grad = evaluate(v) + dx, dy, dz = grad + v += alpha * grad + + if iter_no < 2: + stop_cond = min(dx, dy, dz) / 1000 + if iter_no > 400: + stop_cond = max(dx, dy, dz) + 1 + iter_no += 1 + + x, y, z = v + df1 = (x + np.hypot(y, z)) / 2 + df2 = (x - np.hypot(y, z)) / 2 + angle_ast = np.arctan2(z, y) / 2 + + c1, c2, _ = evaluate(v, need_grad=False) + p = c1 / c2 + + return df1, df2, angle_ast, p, iter_no + + def _gd( + self, + signal, + df1, + df2, + angle_ast, + g_min, + g_max, + ): + """ + Runs gradient ascent to optimize defocus parameters + :param signal: Estimated power spectrum :param df1: Defocus value in the direction perpendicular to df2. :param df2: Defocus value in the direction perpendicular to df1. @@ -654,7 +753,6 @@ def gd( a = a[mask] b = b[mask] signal = signal[..., mask] - # r = self.r_ctf[mask] theta = self.theta[mask] sum_A = np.sum(signal**2) @@ -735,9 +833,6 @@ def gd( return df1, df2, angle_ast, p, iter_no - # Note, This doesn't actually use anything from the class. - # It is used in a solver loop of some sort, so it may not be correct - # to just use what is avail in the obj. def write_star(self, star_filepath, params_dict): """ Writes CTF parameters to starfile for a single micrograph. From f3436e7938ad1adf2eabaaf8c725690d51459f60 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 14 Oct 2026 15:51:03 -0400 Subject: [PATCH 20/44] update ctf cli opts --- src/aspire/commands/estimate_ctf.py | 30 +++++++++++++++++------------ 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/src/aspire/commands/estimate_ctf.py b/src/aspire/commands/estimate_ctf.py index 985b7c7571..6f09c7e69b 100644 --- a/src/aspire/commands/estimate_ctf.py +++ b/src/aspire/commands/estimate_ctf.py @@ -28,13 +28,13 @@ ) @click.option( "--num_tapers", - default=2, + default=4, type=int, help="Number of tapers to apply in PSD estimation", ) @click.option( "--psd_size", - default=512, + default=1024, type=int, help="Size of blocks for use in Power Spectrum estimation", ) @@ -57,6 +57,8 @@ default="results", help="Path to output files, defaults to './results'", ) +@click.option('--save_ctf_images', is_flag=True, default=False, help="Save ctf diagnostic images.") +@click.option('--save_noise_images', is_flag=True, default=False, help="Save noise diagnostic images.") @log_level_option def estimate_ctf( data_folder, @@ -70,6 +72,8 @@ def estimate_ctf( g_max, output_dir, dtype, + save_ctf_images, + save_noise_images, loglevel, ): """ @@ -85,14 +89,16 @@ def estimate_ctf( return ctf.estimate_ctf( data_folder, - pixel_size, - cs, - amplitude_contrast, - voltage, - num_tapers, - psd_size, - g_min, - g_max, - output_dir, - dtype, + pixel_size=pixel_size, + cs=cs, + amplitude_contrast=amplitude_contrast, + voltage=voltage, + num_tapers=num_tapers, + psd_size=psd_size, + g_min=g_min, + g_max=g_max, + output_dir=output_dir, + dtype=dtype, + save_ctf_images=save_ctf_images, + save_noise_images=save_noise_images, ) From 951ed1bf326fe50308f87e8d531728616b1c6352 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 Oct 2026 09:38:08 -0400 Subject: [PATCH 21/44] continue arg cleanup g_min g_max stop_cond iterations --- src/aspire/ctf/ctf_estimator.py | 69 +++++++++++++++++---------------- 1 file changed, 35 insertions(+), 34 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 0a805166d6..809795720e 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -56,7 +56,9 @@ def __init__( amplitude_contrast, voltage, psd_size, - num_tapers, + num_tapers=4, + g_min=30, + g_max=5, gradient_descent_iterations=6, linprog_method="highs-ds", diagnostic_plots=False, @@ -72,6 +74,8 @@ def __init__( :param psd_size: Block size (in pixels) for PSD estimation. :param num_tapers: Optional, number of tapers to apply in PSD estimation. Defaults to 4. + :param g_min: Inverse of minimun resolution for PSD. Default 30. + :param g_max: Inverse of maximum resolution for PSD. Default 5. :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. Defaults to 6. :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). @@ -89,6 +93,8 @@ def __init__( self.psd_size = int(psd_size) self.num_tapers = int(num_tapers) self.lmbd = voltage_to_wavelength(voltage) / 10.0 # Convert angstrom to nm + self.g_min=float(g_min) + self.g_max=float(g_max) self.gd_iters = int(gradient_descent_iterations) self.linprog_method = linprog_method self.diagnostic_plots = diagnostic_plots @@ -556,12 +562,10 @@ def background_subtract_2d(self, signal, background_p1, max_col): background, pixel_size=self.pixel_size ) - def pca(self, signal, g_min, g_max): + def pca(self, signal): """ :param signal: Estimated power spectrum. - :param g_min: Inverse of minimun resolution for PSD. - :param g_max: Inverse of maximum resolution for PSD. :return: ratio. """ @@ -572,8 +576,8 @@ def pca(self, signal, g_min, g_max): signal = signal - np.min(signal) - rad_sq_min = N * self.pixel_size / g_min - rad_sq_max = N * self.pixel_size / g_max + rad_sq_min = N * self.pixel_size / self.g_min + rad_sq_max = N * self.pixel_size / self.g_max min_limit = self.r_ctf[center, (center + np.floor(rad_sq_min)).astype(int)] signal[self.r_ctf < min_limit] = 0 @@ -607,8 +611,8 @@ def gd( df1, df2, angle_ast, - g_min, - g_max, + stop_cond=None, + max_gd_iterations=400, ): """ Runs gradient ascent to optimize defocus parameters @@ -617,10 +621,16 @@ def gd( :param df1: Defocus value in the direction perpendicular to df2. :param df2: Defocus value in the direction perpendicular to df1. :param angle_ast: Angle between df1 and the x-axis, Radians. - :param g_min: Inverse of minimun resolution for PSD. - :param g_max: Inverse of maximum resolution for PSD. + :param stop_cond: Optionally select value to stop iterations early. + Default `None` is 1e-10 for float32 `dtype` and 1e20 otherwise. + :param max_gd_iterations: 400 :return: Optimal defocus parameters """ + + if stop_cond is None: + stop_cond = 1e-10 if self.dtype==np.float32 else 1e-20 + stop_cond = float(stop_cond) + # step sizes for (x, y, z) alpha = np.array([1e5, 1e4, 1e4]) @@ -633,8 +643,8 @@ def gd( N = signal.shape[1] center = N // 2 - rad_sq_min = N * self.pixel_size / g_min - rad_sq_max = N * self.pixel_size / g_max + rad_sq_min = N * self.pixel_size / self.g_min + rad_sq_max = N * self.pixel_size / self.g_max max_val = self.r_ctf[center, int(center - 1 + np.floor(rad_sq_max))] min_val = self.r_ctf[center, int(center - 1 + np.ceil(rad_sq_min))] @@ -665,9 +675,9 @@ def evaluate(v, need_grad=True): grad = None if need_grad: cs = np.cos(psi) - np.multiply(np.sign(s), cs, out=W[:, 0]) + W[:,0] = np.sign(s) * cs W[:, 0] *= sig # sign(s) * cos * signal - np.multiply(s, cs, out=W[:, 1]) + W[:, 1] = s * cs W[:, 1] *= 2 # d(sin^2)/d(psi) G = A @ W # compute all gradient sums dc1 = G[:, 0] @@ -676,7 +686,6 @@ def evaluate(v, need_grad=True): return c1, c2, grad dx = dy = dz = 1.0 - stop_cond = 1e-20 iter_no = 1 while max(dx, dy, dz) > stop_cond: @@ -686,7 +695,7 @@ def evaluate(v, need_grad=True): if iter_no < 2: stop_cond = min(dx, dy, dz) / 1000 - if iter_no > 400: + if iter_no > max_gd_iterations: stop_cond = max(dx, dy, dz) + 1 iter_no += 1 @@ -706,8 +715,6 @@ def _gd( df1, df2, angle_ast, - g_min, - g_max, ): """ Runs gradient ascent to optimize defocus parameters @@ -716,8 +723,6 @@ def _gd( :param df1: Defocus value in the direction perpendicular to df2. :param df2: Defocus value in the direction perpendicular to df1. :param angle_ast: Angle between df1 and the x-axis, Radians. - :param g_min: Inverse of minimun resolution for PSD. - :param g_max: Inverse of maximum resolution for PSD. :return: Optimal defocus parameters """ @@ -743,8 +748,8 @@ def _gd( N = signal.shape[1] center = N // 2 - rad_sq_min = N * self.pixel_size / g_min - rad_sq_max = N * self.pixel_size / g_max + rad_sq_min = N * self.pixel_size / self.g_min + rad_sq_max = N * self.pixel_size / self.g_max max_val = self.r_ctf[center, int(center - 1 + np.floor(rad_sq_max))] min_val = self.r_ctf[center, int(center - 1 + np.ceil(rad_sq_min))] @@ -879,16 +884,14 @@ def _plot_2d_bg_sub_psd(self, bg_sub_psd): # title('Background subtracted 2-d PSD') plt.show() - def estimate(self, micrograph, g_min=30, g_max=5): + def estimate(self, micrograph): # TODO compare now that remainder of code repros. - # micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size) - micrograph_blocks = self.legacy_preprocess_micrograph(micrograph, self.psd_size) + micrograph_blocks = self.preprocess_micrograph(micrograph, self.psd_size) + #micrograph_blocks = self.legacy_preprocess_micrograph(micrograph, self.psd_size) signal_observed = self.estimate_psd(micrograph_blocks) - # np.save("signal_observed.npy", signal_observed) amplitude_spectrum, _ = self.radial_average(signal_observed) - # np.save("amplitude_spectrum.npy", amplitude_spectrum) signal_1d, background_1d = self.background_subtract_1d( amplitude_spectrum, linprog_method=self.linprog_method @@ -908,7 +911,7 @@ def estimate(self, micrograph, g_min=30, g_max=5): self._plot_2d_psd(amplitude_spectrum.asnumpy()[0]) self._plot_2d_bg_sub_psd(signal.asnumpy()[0]) - ratio = self.pca(signal_observed, g_min, g_max) + ratio = self.pca(signal_observed) signal, additional_background = self.elliptical_average(signal.sqrt()) @@ -921,14 +924,12 @@ def estimate(self, micrograph, g_min=30, g_max=5): # This array holds results from each of the `gd_iters` cc_array = np.zeros((self.gd_iters, 4)) - for a, angle in enumerate(np.linspace(0, np.pi, self.gd_iters)): + for a, angle in enumerate(np.linspace(0, np.pi, self.gd_iters, endpoint=False)): df1, df2, angle_ast, p, iters = self.gd( signal, initial_df1, initial_df2, angle, - g_min, - g_max, ) cc_array[a, 0] = df1 @@ -1035,6 +1036,8 @@ def estimate_ctf( voltage, psd_size, num_tapers=num_tapers, + g_min=g_min, + g_max=g_max, gradient_descent_iterations=gradient_descent_iterations, diagnostic_plots=False, dtype=dtype, @@ -1060,9 +1063,7 @@ def estimate_ctf( micrograph = micrograph.astype(dtype, copy=False) - params, background_2d, signal = ctf_est.estimate( - micrograph, g_min=g_min, g_max=g_max - ) + params, background_2d, signal = ctf_est.estimate(micrograph) # Store files' params into results dictionary results[relative_input_path] = params From 547194c293228a415b64f54d58fc338ec6a54c87 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 Oct 2026 10:18:34 -0400 Subject: [PATCH 22/44] rm stray todos in legacy code --- src/aspire/ctf/ctf_estimator.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 809795720e..a10f81825a 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -226,7 +226,6 @@ def legacy_preprocess_micrograph(self, micrograph, block_size): # im2col `distinct` extracts non-overlapping half blocks # column-major block ordering to match MATLAB (down-first) - # TODO, I don't think the order will matter here later, # but keeping now for diagnostics. blocks = ( cropped_micrograph.reshape(n_bh, half, n_bw, half) @@ -237,7 +236,6 @@ def legacy_preprocess_micrograph(self, micrograph, block_size): # Reshape to (half, half, n_total) so each slice is one half-block blocks = blocks.reshape(half, half, n_total, order="F") - # TODO, compare which method is faster later, if we really need these... # First cat # stack each half-block with the one below it (rows overlap) # by shifting the block indices by 1 From cc59c87f36336b74127b2ae8e89567bfb56bb0a2 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 Oct 2026 13:56:13 -0400 Subject: [PATCH 23/44] extend to additional basis and more arg work --- src/aspire/commands/estimate_ctf.py | 11 +++++-- src/aspire/ctf/ctf_estimator.py | 51 ++++++++++++++++++++++------- tests/test_CtfEstimate.py | 29 ++++++++++++++-- tox.ini | 1 + 4 files changed, 76 insertions(+), 16 deletions(-) diff --git a/src/aspire/commands/estimate_ctf.py b/src/aspire/commands/estimate_ctf.py index 6f09c7e69b..c74290d5e2 100644 --- a/src/aspire/commands/estimate_ctf.py +++ b/src/aspire/commands/estimate_ctf.py @@ -57,8 +57,15 @@ default="results", help="Path to output files, defaults to './results'", ) -@click.option('--save_ctf_images', is_flag=True, default=False, help="Save ctf diagnostic images.") -@click.option('--save_noise_images', is_flag=True, default=False, help="Save noise diagnostic images.") +@click.option( + "--save_ctf_images", is_flag=True, default=False, help="Save ctf diagnostic images." +) +@click.option( + "--save_noise_images", + is_flag=True, + default=False, + help="Save noise diagnostic images.", +) @log_level_option def estimate_ctf( data_folder, diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index a10f81825a..d3903c810a 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -34,7 +34,7 @@ from scipy.optimize import linprog from scipy.signal.windows import dpss -from aspire.basis import Coef, FFBBasis2D +from aspire.basis import Coef, FFBBasis2D, SteerableBasis2D from aspire.image import Image from aspire.numeric import fft, xp from aspire.operators import CTFFilter @@ -60,6 +60,7 @@ def __init__( g_min=30, g_max=5, gradient_descent_iterations=6, + basis=None, linprog_method="highs-ds", diagnostic_plots=False, dtype=np.float32, @@ -78,6 +79,7 @@ def __init__( :param g_max: Inverse of maximum resolution for PSD. Default 5. :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. Defaults to 6. + :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). Defaults to 'highs-ds'. :param diagnostic_plots: Optionally generate diagnostic_plots. @@ -93,8 +95,8 @@ def __init__( self.psd_size = int(psd_size) self.num_tapers = int(num_tapers) self.lmbd = voltage_to_wavelength(voltage) / 10.0 # Convert angstrom to nm - self.g_min=float(g_min) - self.g_max=float(g_max) + self.g_min = float(g_min) + self.g_max = float(g_max) self.gd_iters = int(gradient_descent_iterations) self.linprog_method = linprog_method self.diagnostic_plots = diagnostic_plots @@ -116,8 +118,12 @@ def __init__( self.Y = grid["y"] * psd_size / 2 self.R = grid["r"] * psd_size / 2 - # TODO, try to extend to any steerable basis - self.basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) + # Setup Basis + if basis is None: + basis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) + if not isinstance(basis, SteerableBasis2D): + raise RuntimeError("Basis must be subclass of SteerableBasis2D.") + self.basis = basis # Only need to compute tapers once # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L @@ -626,7 +632,7 @@ def gd( """ if stop_cond is None: - stop_cond = 1e-10 if self.dtype==np.float32 else 1e-20 + stop_cond = 1e-10 if self.dtype == np.float32 else 1e-20 stop_cond = float(stop_cond) # step sizes for (x, y, z) @@ -673,7 +679,7 @@ def evaluate(v, need_grad=True): grad = None if need_grad: cs = np.cos(psi) - W[:,0] = np.sign(s) * cs + W[:, 0] = np.sign(s) * cs W[:, 0] *= sig # sign(s) * cos * signal W[:, 1] = s * cs W[:, 1] *= 2 # d(sin^2)/d(psi) @@ -885,7 +891,7 @@ def _plot_2d_bg_sub_psd(self, bg_sub_psd): def estimate(self, micrograph): # TODO compare now that remainder of code repros. micrograph_blocks = self.preprocess_micrograph(micrograph, self.psd_size) - #micrograph_blocks = self.legacy_preprocess_micrograph(micrograph, self.psd_size) + # micrograph_blocks = self.legacy_preprocess_micrograph(micrograph, self.psd_size) signal_observed = self.estimate_psd(micrograph_blocks) @@ -994,6 +1000,7 @@ def estimate_ctf( g_min=30, g_max=5, gradient_descent_iterations=6, + basis=None, output_dir="results", dtype=np.float32, save_ctf_images=False, @@ -1001,12 +1008,33 @@ def estimate_ctf( rglob=r"*.mrc", ): """ - Given paramaters estimates CTF from experimental data - and returns CTF as a mrc file. - + Given paramaters estimates CTF from experimental data in `data_folder + and returns dictionary of estimated CTF parameters. + + Optionally saves diagnostic CTF and noise image (mrc) files. + + :param pixel_size: Size of the pixel in \u212b (angstrom). + :param cs: Spherical aberration in mm. + :param amplitude_contrast: Amplitude contrast. + :param voltage: Voltage of electron microscope. + :param num_tapers: Optional, number of tapers to apply in PSD estimation. + Defaults to 4. + :param psd_size: Block size (in pixels) for PSD estimation. + :param g_min: Inverse of minimun resolution for PSD. Default 30. + :param g_max: Inverse of maximum resolution for PSD. Default 5. + :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. + Defaults to 6. + :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. + :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). + Defaults to 'highs-ds'. + :param diagnostic_plots: Optionally generate diagnostic_plots. + Note, these are blocking plot calls. + Defaults to False. :param rglob: File search pattern. By default will search for **/*.mrc files. Using "*.mrc[s]" will additonally include mrcs files, etc. + :param save_ctf_images: Save diagnostic ctf images. + :param save_noise_images: Save diagnostic noise images. """ # Checks ------------------------------------------------------------------ @@ -1037,6 +1065,7 @@ def estimate_ctf( g_min=g_min, g_max=g_max, gradient_descent_iterations=gradient_descent_iterations, + basis=basis, diagnostic_plots=False, dtype=dtype, ) diff --git a/tests/test_CtfEstimate.py b/tests/test_CtfEstimate.py index 26cc168d9d..3bc7a7cd24 100644 --- a/tests/test_CtfEstimate.py +++ b/tests/test_CtfEstimate.py @@ -7,12 +7,15 @@ import numpy as np import pytest +from aspire.basis import FBBasis2D, FLEBasis2D, FPSWFBasis2D, PSWFBasis2D from aspire.ctf import estimate_ctf logger = logging.getLogger(__name__) DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data") +DTYPE = np.float64 +PSD_SIZE = 512 TEST_INPUT_FN = "sample.mrc" # These values are from CTFFIND4 TEST_OUTPUT = { @@ -25,8 +28,27 @@ "amplitude_contrast": 0.07, } +BASES = [ + None, # resolves to FFBBasis2D, + pytest.param(FLEBasis2D, marks=pytest.mark.expensive), + # The following are very long and may require a higher memory machine + pytest.param(FBBasis2D, marks=pytest.mark.ctf), + pytest.param(PSWFBasis2D, marks=pytest.mark.ctf), + pytest.param(FPSWFBasis2D, marks=pytest.mark.ctf), +] -def test_estimate_CTF(): + +@pytest.fixture(params=BASES, ids=lambda x: f"basis={x}", scope="module") +def basis(request): + cls = request.param + if cls is None: + return None + # Setup a Basis + basis = cls(PSD_SIZE, dtype=DTYPE) + return basis + + +def test_estimate_CTF(basis): with tempfile.TemporaryDirectory() as tmp_input_dir: # Copy input file copyfile( @@ -43,11 +65,12 @@ def test_estimate_CTF(): amplitude_contrast=0.07, voltage=300.0, num_tapers=2, - psd_size=512, + psd_size=PSD_SIZE, g_min=30.0, g_max=5.0, output_dir=tmp_output_dir, - dtype=np.float64, + basis=basis, + dtype=DTYPE, save_ctf_images=True, save_noise_images=True, ) diff --git a/tox.ini b/tox.ini index 37b5bcfcbe..f6c4bd41c0 100644 --- a/tox.ini +++ b/tox.ini @@ -93,6 +93,7 @@ testpaths = tests markers = expensive: mark a test as a long running test. covar: extended (long) tests for covar components + ctf: extended (long/hi-mem) tests for ctf components scheduled: tests that should only run in the scheduled workflow addopts = -m "not expensive and not scheduled and not covar" From 5951e6c7971985aac28399d68d822426254bf1c0 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 16 Oct 2026 08:11:24 -0400 Subject: [PATCH 24/44] initial add bfgs --- src/aspire/ctf/ctf_estimator.py | 248 +++++++++++++++----------------- 1 file changed, 118 insertions(+), 130 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index d3903c810a..0c5492ef2c 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -31,7 +31,7 @@ import mrcfile import numpy as np from numpy import linalg as npla -from scipy.optimize import linprog +from scipy.optimize import linprog, minimize from scipy.signal.windows import dpss from aspire.basis import Coef, FFBBasis2D, SteerableBasis2D @@ -61,6 +61,7 @@ def __init__( g_max=5, gradient_descent_iterations=6, basis=None, + solver_method="bfgs", linprog_method="highs-ds", diagnostic_plots=False, dtype=np.float32, @@ -98,6 +99,12 @@ def __init__( self.g_min = float(g_min) self.g_max = float(g_max) self.gd_iters = int(gradient_descent_iterations) + _solver_methods = ("gd", "bfgs") + if solver_method not in _solver_methods: + raise RuntimeError( + f"Provided solver_method {solver_method} must be one of {_solver_methods}" + ) + self.solver_method = solver_method self.linprog_method = linprog_method self.diagnostic_plots = diagnostic_plots self.dtype = np.dtype(dtype) @@ -490,8 +497,8 @@ def opt1d( Find optimal defocus for the radially symmetric case (where no astigmatism is present) :param amplitude_spectrum: Estimated power specrtum. - :param min_defocus: Start of defocus loop scan (xxx units). - :param max_defocus: End of defocus loop scan (xxx units). + :param min_defocus: Start of defocus loop scan + :param max_defocus: End of defocus loop scan :return: 2-tuple of NumPy arrays (Estimated average of defocus and low_freq_cutoff) """ @@ -500,7 +507,6 @@ def opt1d( signal = np.sqrt(signal) signal = signal[: 3 * signal.shape[0] // 4] - # Ask Ay about original range units, I read as nm, but that is 5000-100000A ...? xxx c = np.zeros((max_defocus - min_defocus, signal.shape[1]), dtype=self.dtype) for f in range(min_defocus, max_defocus): @@ -609,7 +615,7 @@ def pca(self, signal): return ratio - def gd( + def _gd( self, signal, df1, @@ -619,7 +625,7 @@ def gd( max_gd_iterations=400, ): """ - Runs gradient ascent to optimize defocus parameters + Runs inner gradient ascent to optimize defocus parameters. :param signal: Estimated power spectrum :param df1: Defocus value in the direction perpendicular to df2. @@ -713,134 +719,135 @@ def evaluate(v, need_grad=True): return df1, df2, angle_ast, p, iter_no - def _gd( - self, - signal, - df1, - df2, - angle_ast, - ): + def gd(self, signal, initial_df1, initial_df2): """ - Runs gradient ascent to optimize defocus parameters + Calls gradient descent loop `gd_iters` times and + returns best optimum found. - :param signal: Estimated power spectrum - :param df1: Defocus value in the direction perpendicular to df2. - :param df2: Defocus value in the direction perpendicular to df1. - :param angle_ast: Angle between df1 and the x-axis, Radians. - :return: Optimal defocus parameters + This implementation was adapted and optimized from the original MATLAB code. + + :param signal: Signal PSD array. + :param initial_df1: Initial major defocus (nm) + :param initial_df2: Initial minor defocus (nm) + :return: defocus_u_nm, defocus_v_nm, defocus_ang_radian """ + # This array holds results from each of the `gd_iters` + cc_array = np.zeros((self.gd_iters, 4)) + for a, angle in enumerate(np.linspace(0, np.pi, self.gd_iters, endpoint=False)): + df1, df2, angle_ast, p, iters = self._gd( + signal, + initial_df1, + initial_df2, + angle, + ) - # step size - alpha1 = 1e5 - alpha2 = 1e4 + cc_array[a, 0] = df1 + cc_array[a, 1] = df2 + cc_array[a, 2] = angle_ast # Radians + cc_array[a, 3] = p + logger.debug( + f"Gradient Descent {a+1} of {self.gd_iters}:\tp={p}\titers={iters}" + ) - # initialization - x = df1 + df2 - y = (df1 - df2) * np.cos(2 * angle_ast) - z = (df1 - df2) * np.sin(2 * angle_ast) + # Select the best result index + ml = np.argmax(cc_array[:, 3], -1) + return cc_array[ml, 0:3] - a = np.pi * self.lmbd * self.r_ctf**2 / 2 - # Ask Ay, Does this amplitude_contrast need to be the arctan'd one? - # amplitude_contrast = np.arctan( - # self.amplitude_contrast / np.sqrt(1 - self.amplitude_contrast**2) - # ) - b = np.pi * self.lmbd**3 * self.cs * 1e6 * self.r_ctf**4 / 2 - np.full( - shape=self.r_ctf.shape, fill_value=self.amplitude_contrast, dtype=self.dtype + def bfgs( + self, signal, initial_df1, initial_df2, stop_cond=1e-10, max_gd_iterations=200 + ): + """ + Calls quasi-newton method BFGS and returns found optimal values. + + :param signal: Signal PSD array. + :param initial_df1: Initial major defocus (nm) + :param initial_df2: Initial minor defocus (nm) + :param stop_cond: Optionally select value to stop iterations. + Default is 1e-10. + :param max_gd_iterations: Optionally provide a max number of iterations for BFGS. + Defaults to 200. + :return: defocus_u_nm, defocus_v_nm, defocus_ang_radian + """ + + # Initial v = [x, y, z] + angle_ast = np.deg2rad(-75) # legacy value, shouldn't matter in bfgs + v0 = np.array( + [ + initial_df1 + initial_df2, + (initial_df1 - initial_df2) * np.cos(2 * angle_ast), + (initial_df1 - initial_df2) * np.sin(2 * angle_ast), + ] ) signal = signal.asnumpy()[0].T N = signal.shape[1] center = N // 2 - rad_sq_min = N * self.pixel_size / self.g_min rad_sq_max = N * self.pixel_size / self.g_max max_val = self.r_ctf[center, int(center - 1 + np.floor(rad_sq_max))] min_val = self.r_ctf[center, int(center - 1 + np.ceil(rad_sq_min))] - mask = (self.r_ctf <= max_val) & (self.r_ctf > min_val) - a = a[mask] - b = b[mask] - signal = signal[..., mask] - theta = self.theta[mask] - - sum_A = np.sum(signal**2) - - dx = 1 - dy = 1 - dz = 1 - - stop_cond = 1e-20 - iter_no = 1 - while np.maximum(np.maximum(dx, dy), dz) > stop_cond: - inner_cosine = y * np.cos(2 * theta) + z * np.sin(2 * theta) - psi = a * x + a * inner_cosine - b - outer_sine = np.sin(psi) - outer_cosine = np.cos(psi) - - sine_x_term = a - sine_y_term = a * np.cos(2 * theta) - sine_z_term = a * np.sin(2 * theta) - - c1 = np.sum(np.abs(outer_sine) * signal) - c2 = np.sqrt(sum_A * np.sum(outer_sine**2)) - - # gradients of numerator - dx_c1 = np.sum(np.sign(outer_sine) * outer_cosine * a * signal) - - dy_c1 = np.sum( - np.sign(outer_sine) * outer_cosine * a * np.cos(2 * theta) * signal - ) - dz_c1 = np.sum( - np.sign(outer_sine) * outer_cosine * a * np.sin(2 * theta) * signal - ) - - derivative_sqrt = 1 / (2 * np.sqrt(sum_A * np.sum(outer_sine**2))) - - derivative_sine2 = 2 * outer_sine * outer_cosine - - # gradients of denomenator - dx_c2 = derivative_sqrt * sum_A * np.sum(derivative_sine2 * sine_x_term) - dy_c2 = derivative_sqrt * sum_A * np.sum(derivative_sine2 * sine_y_term) - dz_c2 = derivative_sqrt * sum_A * np.sum(derivative_sine2 * sine_z_term) + # Everything that does not depend on (x, y, z) + # psi = x*a + y*a*cos(2t) + z*a*sin(2t) - b == v @ A - b + r2 = self.r_ctf[mask] ** 2 + a = (np.pi * self.lmbd / 2) * r2 + b = (np.pi * self.lmbd**3 * self.cs * 1e6 / 2) * r2**2 - self.amplitude_contrast + sig = np.ascontiguousarray(signal[mask]) + theta2 = 2 * self.theta[mask] + A = np.ascontiguousarray(np.stack([a, a * np.cos(theta2), a * np.sin(theta2)])) + sum_A = sig @ sig - # gradients - dx = (dx_c1 * c2 - dx_c2 * c1) / c2**2 - dy = (dy_c1 * c2 - dy_c2 * c1) / c2**2 - dz = (dz_c1 * c2 - dz_c2 * c1) / c2**2 + # Variable scaling + scale = np.array([1e2, 1e1, 1e1]) - # update - x = x + alpha1 * dx - y = y + alpha2 * dy - z = z + alpha2 * dz + def neg_p_and_grad(u): + v = u * scale + psi = v @ A + psi -= b + s = np.sin(psi) + c = np.cos(psi) - if iter_no < 2: - stop_cond = np.minimum(np.minimum(dx, dy), dz) / 1000 + c1 = np.abs(s) @ sig + c2 = np.sqrt(sum_A * (s @ s)) - if iter_no > 400: - stop_cond = np.maximum(np.maximum(dx, dy), dz) + 1 + dc1 = A @ (np.sign(s) * c * sig) + dc2 = (sum_A / (2 * c2)) * (A @ (2 * s * c)) + grad = (dc1 * c2 - dc2 * c1) / c2**2 - iter_no = iter_no + 1 + return -c1 / c2, -grad * scale - df1 = (x + np.abs(y + z * 1j)) / 2 - df2 = (x - np.abs(y + z * 1j)) / 2 - angle_ast = np.angle(y + z * 1j) / 2 # Radians + res = minimize( + neg_p_and_grad, + v0 / scale, + jac=True, + method="BFGS", + options={"gtol": stop_cond, "maxiter": max_gd_iterations}, + ) - inner_cosine = y * np.cos(2 * theta) + z * np.sin(2 * theta) - outer_sine = np.sin(a * x + a * inner_cosine - b) - outer_cosine = np.cos(a * x + a * inner_cosine - b) + x, y, z = res.x * scale + df1 = (x + np.hypot(y, z)) / 2 # nm + df2 = (x - np.hypot(y, z)) / 2 # nm + angle_ast = np.arctan2(z, y) / 2 # Radians - sine_x_term = a - sine_y_term = a * np.cos(2 * theta) - sine_z_term = a * np.sin(2 * theta) + logger.debug(f"BFGS :\tp={-res.fun}\titers={res.nit}") - c1 = np.sum(np.abs(outer_sine) * signal) - c2 = np.sqrt(sum_A * np.sum(outer_sine**2)) + return df1, df2, angle_ast - p = c1 / c2 + def maximize_correlation(self, signal, initial_df1, initial_df2): + """ + Selects which solver to use for optimizing CTF estimates. - return df1, df2, angle_ast, p, iter_no + :param signal: Signal PSD array + :param initial_df1: Initial major defocus (nm) + :param initial_df2: Initial minor defocus (nm) + :return: defocus_u_nm, defocus_v_nm, defocus_ang_rad + """ + if self.solver_method == "gd": + return self.gd(signal, initial_df1, initial_df2) + elif self.solver_method == "bfgs": + return self.bfgs(signal, initial_df1, initial_df2) def write_star(self, star_filepath, params_dict): """ @@ -924,33 +931,14 @@ def estimate(self, micrograph): initial_df1 = (avg_defocus * 2) / (1 + ratio) initial_df2 = (avg_defocus * 2) - initial_df1 - angle = -5 / 12 * np.pi # Radians (-75 degrees) - - # This array holds results from each of the `gd_iters` - cc_array = np.zeros((self.gd_iters, 4)) - for a, angle in enumerate(np.linspace(0, np.pi, self.gd_iters, endpoint=False)): - df1, df2, angle_ast, p, iters = self.gd( - signal, - initial_df1, - initial_df2, - angle, - ) - - cc_array[a, 0] = df1 - cc_array[a, 1] = df2 - cc_array[a, 2] = angle_ast # Radians - cc_array[a, 3] = p - logger.debug( - f"Gradient Descent {a+1} of {self.gd_iters}:\tp={p}\titers={iters}" - ) - - # Select the best result index - ml = np.argmax(cc_array[:, 3], -1) + defocus_u, defocus_v, defocus_ang = self.maximize_correlation( + signal, initial_df1, initial_df2 + ) params = { - "defocus_u": cc_array[ml, 0] * 10, # Convert from nm to A - "defocus_v": cc_array[ml, 1] * 10, # Convert from nm to A - "defocus_ang": cc_array[ml, 2], # Radians + "defocus_u": defocus_u * 10, # Convert from nm to A + "defocus_v": defocus_v * 10, # Convert from nm to A + "defocus_ang": defocus_ang, # radians "cs": self.cs, "voltage": self.voltage, "pixel_size": self.pixel_size, From a82470a239b1acb835401fddc98784192f80b9e4 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 16 Oct 2026 11:21:31 -0400 Subject: [PATCH 25/44] continue cleanup --- src/aspire/ctf/ctf_estimator.py | 126 ++++++++++++++++++++------------ 1 file changed, 80 insertions(+), 46 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 0c5492ef2c..1c0560d8d9 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -1,6 +1,8 @@ """ Contains code supporting CTF parameter estimation. -Generally, this is a Python version of ASPIRE-CTF from MATLAB. +This implements a version of ASPIRE-CTF originally written in MATLAB, + +This Python version supports multiple basis and solvers. See paper: @@ -20,7 +22,9 @@ Created on Sep 10, 2019 @author: Ayelet Heimowitz, Amit Moscovich -Integrated into ASPIRE-Python by Garrett Wright Feb 2021. +Originally integrated into ASPIRE-Python Feb 2021, +optimized and refactored by Garrett Wright Oct 2026. + """ import logging @@ -79,8 +83,11 @@ def __init__( :param g_min: Inverse of minimun resolution for PSD. Default 30. :param g_max: Inverse of maximum resolution for PSD. Default 5. :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. - Defaults to 6. + Defaults to 6. Unused for `bfgs`. :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. + :param solver_method: Solver used to optimize defocus params. `gd` or `bfgs`. + Default uses single call to `bfgs`. + Original implementation used repeated gradient descent and is provided for reproducibility. :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). Defaults to 'highs-ds'. :param diagnostic_plots: Optionally generate diagnostic_plots. @@ -137,7 +144,7 @@ def __init__( tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) self.tapers_1d = xp.asarray(tapers_1d, dtype=complex_type(self.dtype)) - def micrograph_to_blocks(self, micrograph, block_size): + def micrograph_to_blocks(self, micrograph): """ Preprocess micrograph into blocks using block_size. @@ -146,8 +153,8 @@ def micrograph_to_blocks(self, micrograph, block_size): :return: NumPy array of blocks extracted from the micrograph. """ - # verify block_size is even - assert block_size % 2 == 0 + # verify self.psd_size is even + assert self.psd_size % 2 == 0 if micrograph.ndim == 3: assert ( @@ -158,7 +165,7 @@ def micrograph_to_blocks(self, micrograph, block_size): size_x = micrograph.shape[-1] size_y = micrograph.shape[-2] - step_size = block_size // 2 + step_size = self.psd_size // 2 range_y = size_y // step_size - 1 range_x = size_x // step_size - 1 @@ -192,33 +199,35 @@ def normalize_blocks(self, blocks): return blocks - def preprocess_micrograph(self, micrograph, block_size): + def preprocess_micrograph(self, micrograph): """ Preprocess micrograph into normalized blocks using block_size. - :param micrograph: Micrograph as NumPy array. #NOTE looks like F order - :param blocksize: Size of the square blocks to partition micrograph. + :param micrograph: Micrograph as NumPy array. :return: NumPy array of normalized blocks extracted from the micrograph. """ - return self.normalize_blocks(self.micrograph_to_blocks(micrograph, block_size)) + return self.normalize_blocks(self.micrograph_to_blocks(micrograph)) - def legacy_preprocess_micrograph(self, micrograph, block_size): + def legacy_preprocess_micrograph(self, micrograph): """ Returns half-overlapping blocks extracted from the micrograph. Equivalent to MATLAB's ctf_extract_blocks. Note, internally transposes micrograph to match MATLAB computations. Returns (num_blocks, block_size, block_size) + + :param micrograph: Micrograph as NumPy array. + :return: NumPy array of normalized blocks extracted from the micrograph. """ h, w = micrograph.shape - if min(h, w) < block_size: + if min(h, w) < self.psd_size: raise ValueError( - f"Blocks of size {block_size} x {block_size} cannot be extracted from image." + f"Blocks of size {self.psd_size} x {self.psd_size} cannot be extracted from image." ) - block_size = block_size - (block_size % 2) # force even + block_size = self.psd_size - (self.psd_size % 2) # force even half = block_size // 2 # Compute crop bounds @@ -281,7 +290,7 @@ def legacy_preprocess_micrograph(self, micrograph, block_size): block_mean = block.sum(axis=(0, 1), keepdims=True) / (block_size**2) block = block - block_mean - # Transpose output to C order (n,block_size, block_size) + # Transpose output to (n,block_size, block_size) return block.T def tapers(self, N, NW, L): @@ -312,35 +321,28 @@ def estimate_psd(self, blocks): :return: NumPy array of estimated power spectrum. """ + # Ensure blocks are on GPU + blocks = xp.asarray(blocks) + num_1d_tapers = self.tapers_1d.shape[0] blocks_mt = xp.zeros(blocks[0, :, :].shape, dtype=self.dtype) - - blocks_tapered = xp.zeros(blocks[0, :, :].shape, dtype=complex_type(self.dtype)) - - taper_2d = xp.zeros( - (blocks.shape[1], blocks.shape[2]), dtype=complex_type(self.dtype) - ) - - # Ensure blocks are on GPU - blocks = xp.asarray(blocks) + blocks_tapered = xp.empty(blocks[0, :, :].shape, dtype=complex_type(self.dtype)) + taper_2d = xp.empty(blocks[0, :, :].shape, dtype=complex_type(self.dtype)) for ax1 in range(num_1d_tapers): for ax2 in range(num_1d_tapers): - xp.matmul( - self.tapers_1d[ax1, :, np.newaxis], - self.tapers_1d[ax2, :, np.newaxis].T, - out=taper_2d, + taper_2d[:, :] = ( + self.tapers_1d[ax1, :, np.newaxis] + @ self.tapers_1d[ax2, :, np.newaxis].T ) - for m in range(blocks.shape[0]): - xp.multiply(blocks[m, :, :], taper_2d, out=blocks_tapered) + blocks_tapered[:, :] = blocks[m] * taper_2d blocks_mt_post_fft = fft.fftn(blocks_tapered, axes=(-2, -1)) blocks_mt += abs2(blocks_mt_post_fft) - # blocks_mt += np.abs(blocks_mt_post_fft)**2 - blocks_mt /= blocks.shape[1] ** 2 - blocks_mt /= blocks.shape[0] ** 2 + # Normalize + blocks_mt /= blocks.shape[1] ** 2 * blocks.shape[0] ** 2 amplitude_spectrum = fft.fftshift(blocks_mt) @@ -574,6 +576,7 @@ def background_subtract_2d(self, signal, background_p1, max_col): def pca(self, signal): """ + Compute PCA to yield ratio for estimating defocus values. :param signal: Estimated power spectrum. :return: ratio. @@ -757,7 +760,7 @@ def bfgs( self, signal, initial_df1, initial_df2, stop_cond=1e-10, max_gd_iterations=200 ): """ - Calls quasi-newton method BFGS and returns found optimal values. + Calls quasi-newton method BFGS and returns found optimized values. :param signal: Signal PSD array. :param initial_df1: Initial major defocus (nm) @@ -868,37 +871,66 @@ def write_star(self, star_filepath, params_dict): star.write(star_filepath) def _plot_1d_psd(self, signal_1d, background_1d): + """ + Helper function plots 1D PSD. + + :param signal_1d: Background subtracted radial PSD + :param background_1d: Background radial PSD. + """ center = self.r_ctf.shape[0] // 2 x = self.r_ctf[center, center : center + signal_1d.shape[1]] plt.plot(x, signal_1d[1] + background_1d[1]) - # xlabel('$\vert \mathbf{k} \vert^2 \, \, [\mathrm{nm}^{-1}]$', 'Interpreter','latex'); - # ylabel('$S_y \left( \mathbf{k} \right)$', 'Interpreter','latex'); - # title('1-d PSD') + plt.title("1-d PSD") + plt.xlabel(r"$\vert \mathbf{k} \vert^2 \, \, [\mathrm{nm}^{-1}]$") + plt.ylabel(r"$S_y \left( \mathbf{k} \right)$") plt.show() def _plot_1d_bg_sub_psd(self, signal_1d): + """ + Helper function plots 1D PSD. + + :param signal_1d: Background subtracted radial PSD + """ center = self.r_ctf.shape[0] // 2 x = self.r_ctf[center, center : center + signal_1d.shape[1]] plt.plot(x, signal_1d[1]) - # title('Background subtracted 1-d PSD') - # xlabel('$\vert \mathbf{k} \vert^2 \, \, [\mathrm{nm}^{-1}]$', 'Interpreter','latex'); - # ylabel('$S_y \left( \mathbf{k} \right) - S_e \left( \mathbf{k} \right)$', 'Interpreter','latex'); + plt.title("Background subtracted 1-d PSD") + plt.xlabel(r"$\vert \mathbf{k} \vert^2 \, \, [\mathrm{nm}^{-1}]$") + plt.ylabel(r"$S_y \left( \mathbf{k} \right) - S_e \left( \mathbf{k} \right)$") plt.show() def _plot_2d_psd(self, psd): + """ + Helper function plots 2D PSD. + + :param psd: 2D PSD array + """ plt.imshow(psd) - # title('2-d PSD') + plt.title("2-d PSD") plt.show() def _plot_2d_bg_sub_psd(self, bg_sub_psd): + """ + Helper function plots 2D background subtracted PSD. + + :param psd: 2D background subtracted PSD + """ plt.imshow(bg_sub_psd) - # title('Background subtracted 2-d PSD') + plt.title("Background subtracted 2-d PSD") plt.show() def estimate(self, micrograph): - # TODO compare now that remainder of code repros. - micrograph_blocks = self.preprocess_micrograph(micrograph, self.psd_size) - # micrograph_blocks = self.legacy_preprocess_micrograph(micrograph, self.psd_size) + """ + Estimate the CTF parameters of `micrograph`. + + This is the main method implementing the algorithm at a high level. + + :param micrograph: 2D Array of micrograph data. + :return: CTF params, estimated background PSD, estimated signal PSD + """ + # TODO Ask Ayelet about legacy version intent and rm if not needed. + micrograph_blocks = self.preprocess_micrograph(micrograph) + # micrograph_blocks = self.legacy_preprocess_micrograph(micrograph) signal_observed = self.estimate_psd(micrograph_blocks) @@ -999,6 +1031,8 @@ def estimate_ctf( Given paramaters estimates CTF from experimental data in `data_folder and returns dictionary of estimated CTF parameters. + Saves STAR files to `output_dir`, which can be loaded into ASPIRE etc. + Optionally saves diagnostic CTF and noise image (mrc) files. :param pixel_size: Size of the pixel in \u212b (angstrom). From 132e7074c9e8007544aecd30954122b5a425232f Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 16 Oct 2026 14:45:23 -0400 Subject: [PATCH 26/44] skip long/hi-mem ctf tests --- tox.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tox.ini b/tox.ini index f6c4bd41c0..c6513345a1 100644 --- a/tox.ini +++ b/tox.ini @@ -95,7 +95,7 @@ markers = covar: extended (long) tests for covar components ctf: extended (long/hi-mem) tests for ctf components scheduled: tests that should only run in the scheduled workflow -addopts = -m "not expensive and not scheduled and not covar" +addopts = -m "not expensive and not scheduled and not covar and not ctf" [gh-actions] python = From fe82ea1ce653d070fadd78cdfa8f5168899997d4 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 16 Oct 2026 14:58:33 -0400 Subject: [PATCH 27/44] fixup scalar,0d,1d arg handling --- src/aspire/ctf/ctf_estimator.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 1c0560d8d9..101ffc6358 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -96,15 +96,17 @@ def __init__( :returns: CtfEstimator instance. """ - self.pixel_size = float(pixel_size) - self.cs = float(cs) - self.amplitude_contrast = float(amplitude_contrast) - self.voltage = float(voltage) + self.pixel_size = np.asarray(pixel_size).item() + self.cs = np.asarray(cs).item() + self.amplitude_contrast = np.asarray(amplitude_contrast).item() + self.voltage = np.asarray(voltage).item() self.psd_size = int(psd_size) self.num_tapers = int(num_tapers) - self.lmbd = voltage_to_wavelength(voltage) / 10.0 # Convert angstrom to nm - self.g_min = float(g_min) - self.g_max = float(g_max) + self.lmbd = ( + voltage_to_wavelength(np.asarray(voltage).item()) / 10.0 + ) # Convert angstrom to nm + self.g_min = np.asarray(g_min).item() + self.g_max = np.asarray(g_max).item() self.gd_iters = int(gradient_descent_iterations) _solver_methods = ("gd", "bfgs") if solver_method not in _solver_methods: From 999cc5aec867126aef213382a44ba8449fb80371 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 19 Oct 2026 08:48:16 -0400 Subject: [PATCH 28/44] fixup path key for gallery --- src/aspire/ctf/ctf_estimator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 101ffc6358..e2d1913acc 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -1117,7 +1117,7 @@ def estimate_ctf( params, background_2d, signal = ctf_est.estimate(micrograph) # Store files' params into results dictionary - results[relative_input_path] = params + results[str(relative_input_path)] = params # Write each micrograph's ctf parameters to an individual STAR file star_filepath = output_file_prefix.with_suffix(".star") From d79b4d666a9dd2f642d6e81cebdd0110280262b4 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 19 Oct 2026 09:31:12 -0400 Subject: [PATCH 29/44] more flexible arg handling --- src/aspire/ctf/ctf_estimator.py | 83 ++++++++++++++++++++++----------- 1 file changed, 57 insertions(+), 26 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index e2d1913acc..bb77fe6b4f 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -53,6 +53,21 @@ class CtfEstimator: CtfEstimator Class ... """ + # Allow configuration for advanced users and future solvers. + default_solver_opts = { + "gd": { + "gradient_descent_iterations": 6, + "max_iterations": 400, + "stop_cond": None, # Base on dtype + }, + "bfgs": { + "maxiter": 200, + "gtol": 1e-10, + }, + } + + _solver_methods = default_solver_opts.keys() + def __init__( self, pixel_size, @@ -63,9 +78,9 @@ def __init__( num_tapers=4, g_min=30, g_max=5, - gradient_descent_iterations=6, basis=None, solver_method="bfgs", + solver_opts=None, linprog_method="highs-ds", diagnostic_plots=False, dtype=np.float32, @@ -82,12 +97,13 @@ def __init__( Defaults to 4. :param g_min: Inverse of minimun resolution for PSD. Default 30. :param g_max: Inverse of maximum resolution for PSD. Default 5. - :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. - Defaults to 6. Unused for `bfgs`. :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. :param solver_method: Solver used to optimize defocus params. `gd` or `bfgs`. Default uses single call to `bfgs`. Original implementation used repeated gradient descent and is provided for reproducibility. + Each solver may have unique configuration options, optionally passed via a dictionary. + :param solver_opts: Optional dictionary of solver options. + Populated with `default_solver_opts[solver]` by default. :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). Defaults to 'highs-ds'. :param diagnostic_plots: Optionally generate diagnostic_plots. @@ -105,16 +121,21 @@ def __init__( self.lmbd = ( voltage_to_wavelength(np.asarray(voltage).item()) / 10.0 ) # Convert angstrom to nm + self.g_min = np.asarray(g_min).item() self.g_max = np.asarray(g_max).item() - self.gd_iters = int(gradient_descent_iterations) - _solver_methods = ("gd", "bfgs") - if solver_method not in _solver_methods: + + # Setup solvers and respective solver options + if solver_method not in self._solver_methods: raise RuntimeError( - f"Provided solver_method {solver_method} must be one of {_solver_methods}" + f"Provided solver_method {solver_method} must be one of {self._solver_methods}" ) self.solver_method = solver_method + if solver_opts is None: + solver_opts = self.default_solver_opts[solver_method] + self.solver_opts = solver_opts self.linprog_method = linprog_method + self.diagnostic_plots = diagnostic_plots self.dtype = np.dtype(dtype) @@ -627,7 +648,7 @@ def _gd( df2, angle_ast, stop_cond=None, - max_gd_iterations=400, + max_iterations=400, ): """ Runs inner gradient ascent to optimize defocus parameters. @@ -638,7 +659,8 @@ def _gd( :param angle_ast: Angle between df1 and the x-axis, Radians. :param stop_cond: Optionally select value to stop iterations early. Default `None` is 1e-10 for float32 `dtype` and 1e20 otherwise. - :param max_gd_iterations: 400 + :param max_iterations: Maximum number of inner gd steps. + Default 400. :return: Optimal defocus parameters """ @@ -710,7 +732,7 @@ def evaluate(v, need_grad=True): if iter_no < 2: stop_cond = min(dx, dy, dz) / 1000 - if iter_no > max_gd_iterations: + if iter_no > max_iterations: stop_cond = max(dx, dy, dz) + 1 iter_no += 1 @@ -724,9 +746,11 @@ def evaluate(v, need_grad=True): return df1, df2, angle_ast, p, iter_no - def gd(self, signal, initial_df1, initial_df2): + def gd( + self, signal, initial_df1, initial_df2, gradient_descent_iterations=6, **kwargs + ): """ - Calls gradient descent loop `gd_iters` times and + Calls gradient descent loop multiple times and returns best optimum found. This implementation was adapted and optimized from the original MATLAB code. @@ -734,16 +758,22 @@ def gd(self, signal, initial_df1, initial_df2): :param signal: Signal PSD array. :param initial_df1: Initial major defocus (nm) :param initial_df2: Initial minor defocus (nm) + :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. + Defaults to 6. :return: defocus_u_nm, defocus_v_nm, defocus_ang_radian """ - # This array holds results from each of the `gd_iters` - cc_array = np.zeros((self.gd_iters, 4)) - for a, angle in enumerate(np.linspace(0, np.pi, self.gd_iters, endpoint=False)): + + # This array holds results from each of the `gradient_descent_iterations` + cc_array = np.zeros((gradient_descent_iterations, 4)) + for a, angle in enumerate( + np.linspace(0, np.pi, gradient_descent_iterations, endpoint=False) + ): df1, df2, angle_ast, p, iters = self._gd( signal, initial_df1, initial_df2, angle, + **kwargs, ) cc_array[a, 0] = df1 @@ -751,7 +781,7 @@ def gd(self, signal, initial_df1, initial_df2): cc_array[a, 2] = angle_ast # Radians cc_array[a, 3] = p logger.debug( - f"Gradient Descent {a+1} of {self.gd_iters}:\tp={p}\titers={iters}" + f"Gradient Descent {a+1} of {gradient_descent_iterations}:\tp={p}\titers={iters}" ) # Select the best result index @@ -759,7 +789,12 @@ def gd(self, signal, initial_df1, initial_df2): return cc_array[ml, 0:3] def bfgs( - self, signal, initial_df1, initial_df2, stop_cond=1e-10, max_gd_iterations=200 + self, + signal, + initial_df1, + initial_df2, + gtol=1e-10, + maxiter=200, ): """ Calls quasi-newton method BFGS and returns found optimized values. @@ -767,9 +802,9 @@ def bfgs( :param signal: Signal PSD array. :param initial_df1: Initial major defocus (nm) :param initial_df2: Initial minor defocus (nm) - :param stop_cond: Optionally select value to stop iterations. + :param gtol: Optionally select value to stop iterations. Default is 1e-10. - :param max_gd_iterations: Optionally provide a max number of iterations for BFGS. + :param maxiter: Optionally provide a max number of iterations for BFGS. Defaults to 200. :return: defocus_u_nm, defocus_v_nm, defocus_ang_radian """ @@ -828,7 +863,7 @@ def neg_p_and_grad(u): v0 / scale, jac=True, method="BFGS", - options={"gtol": stop_cond, "maxiter": max_gd_iterations}, + options={"gtol": gtol, "maxiter": maxiter}, ) x, y, z = res.x * scale @@ -850,9 +885,9 @@ def maximize_correlation(self, signal, initial_df1, initial_df2): :return: defocus_u_nm, defocus_v_nm, defocus_ang_rad """ if self.solver_method == "gd": - return self.gd(signal, initial_df1, initial_df2) + return self.gd(signal, initial_df1, initial_df2, **self.solver_opts) elif self.solver_method == "bfgs": - return self.bfgs(signal, initial_df1, initial_df2) + return self.bfgs(signal, initial_df1, initial_df2, **self.solver_opts) def write_star(self, star_filepath, params_dict): """ @@ -1021,7 +1056,6 @@ def estimate_ctf( psd_size=1024, g_min=30, g_max=5, - gradient_descent_iterations=6, basis=None, output_dir="results", dtype=np.float32, @@ -1046,8 +1080,6 @@ def estimate_ctf( :param psd_size: Block size (in pixels) for PSD estimation. :param g_min: Inverse of minimun resolution for PSD. Default 30. :param g_max: Inverse of maximum resolution for PSD. Default 5. - :param gradient_descent_iterations: Optional, number of Gradient Descent attempts. - Defaults to 6. :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). Defaults to 'highs-ds'. @@ -1088,7 +1120,6 @@ def estimate_ctf( num_tapers=num_tapers, g_min=g_min, g_max=g_max, - gradient_descent_iterations=gradient_descent_iterations, basis=basis, diagnostic_plots=False, dtype=dtype, From a1a3a2a200c294286fc9acb8b26c7027d71551bf Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 19 Oct 2026 10:31:49 -0400 Subject: [PATCH 30/44] add option handling --- src/aspire/commands/estimate_ctf.py | 6 +++++- src/aspire/ctf/ctf_estimator.py | 11 +++++++++++ tests/test_CtfEstimate.py | 23 ++++++++++++++++++++--- 3 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/aspire/commands/estimate_ctf.py b/src/aspire/commands/estimate_ctf.py index c74290d5e2..12711a9d11 100644 --- a/src/aspire/commands/estimate_ctf.py +++ b/src/aspire/commands/estimate_ctf.py @@ -88,7 +88,11 @@ def estimate_ctf( and returns CTF as a mrc file. This is a Click command line interface wrapper for - the aspire.ctf module. + the aspire.ctf.CtfEstimator. + + Note not every option to `CtfEstimator` is exposed via CLI. Users + may want to consider calling `estimate_ctf` directly for advanced + applications. The CLI is provided convenient general use. """ # Set desired logging option for the command line diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index bb77fe6b4f..73f73bf599 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -1062,6 +1062,9 @@ def estimate_ctf( save_ctf_images=False, save_noise_images=False, rglob=r"*.mrc", + solver_method="gd", + solver_opts=None, + linprog_method="highs-ds", ): """ Given paramaters estimates CTF from experimental data in `data_folder @@ -1091,6 +1094,14 @@ def estimate_ctf( Using "*.mrc[s]" will additonally include mrcs files, etc. :param save_ctf_images: Save diagnostic ctf images. :param save_noise_images: Save diagnostic noise images. + :param solver_method: Solver used to optimize defocus params. `gd` or `bfgs`. + Default uses single call to `bfgs`. + Original implementation used repeated gradient descent and is provided for reproducibility. + Each solver may have unique configuration options, optionally passed via a dictionary. + :param solver_opts: Optional dictionary of solver options. + Populated with `CtfEstimator.default_solver_opts[solver]` by default. + :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). + Defaults to 'highs-ds'. """ # Checks ------------------------------------------------------------------ diff --git a/tests/test_CtfEstimate.py b/tests/test_CtfEstimate.py index 3bc7a7cd24..7046bf381e 100644 --- a/tests/test_CtfEstimate.py +++ b/tests/test_CtfEstimate.py @@ -7,7 +7,7 @@ import numpy as np import pytest -from aspire.basis import FBBasis2D, FLEBasis2D, FPSWFBasis2D, PSWFBasis2D +from aspire.basis import FBBasis2D, FFBBasis2D, FLEBasis2D, FPSWFBasis2D, PSWFBasis2D from aspire.ctf import estimate_ctf logger = logging.getLogger(__name__) @@ -48,7 +48,8 @@ def basis(request): return basis -def test_estimate_CTF(basis): +def test_estimate_CTF(basis, **customized_args): + _dtype = customized_args.pop("dtype", DTYPE) # avoid duplicate dtype arg with tempfile.TemporaryDirectory() as tmp_input_dir: # Copy input file copyfile( @@ -70,9 +71,10 @@ def test_estimate_CTF(basis): g_max=5.0, output_dir=tmp_output_dir, basis=basis, - dtype=DTYPE, + dtype=_dtype, save_ctf_images=True, save_noise_images=True, + **customized_args, ) logger.debug(f"results: {results}") @@ -114,6 +116,21 @@ def test_estimate_CTF(basis): np.testing.assert_allclose(result[param], TEST_OUTPUT[param]) +def test_customized_estimate_CTF(): + basis = FFBBasis2D(PSD_SIZE, dtype=np.float32) + customized = dict( + dtype=basis.dtype, + solver_method="gd", + solver_opts={ + "gradient_descent_iterations": 4, + "max_iterations": 300, + "stop_cond": 1e-8, + }, + linprog_method="highs", + ) + return test_estimate_CTF(basis, **customized) + + # we are chopping the micrograph into a vertical and a horizontal rectangle # as small as possible to save testing duration @pytest.mark.parametrize( From 08afcd4066cb2f2af28ac5e71707f0501e0fcc0b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 19 Oct 2026 14:13:14 -0400 Subject: [PATCH 31/44] optimize estimated psd half psd, remove large tmps, redux kernel --- src/aspire/ctf/ctf_estimator.py | 144 +++++++++++++++++++++----------- 1 file changed, 95 insertions(+), 49 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 73f73bf599..62ae1d58f4 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -34,6 +34,7 @@ import matplotlib.pyplot as plt import mrcfile import numpy as np +from line_profiler import profile from numpy import linalg as npla from scipy.optimize import linprog, minimize from scipy.signal.windows import dpss @@ -43,7 +44,7 @@ from aspire.numeric import fft, xp from aspire.operators import CTFFilter from aspire.storage import StarFile -from aspire.utils import abs2, complex_type, grid_2d, tqdm, voltage_to_wavelength +from aspire.utils import complex_type, grid_2d, tqdm, voltage_to_wavelength logger = logging.getLogger(__name__) @@ -165,7 +166,7 @@ def __init__( # Only need to compute tapers once # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) - self.tapers_1d = xp.asarray(tapers_1d, dtype=complex_type(self.dtype)) + self.tapers_1d = xp.asarray(tapers_1d, dtype=self.dtype) def micrograph_to_blocks(self, micrograph): """ @@ -344,30 +345,47 @@ def estimate_psd(self, blocks): :return: NumPy array of estimated power spectrum. """ - # Ensure blocks are on GPU - blocks = xp.asarray(blocks) - + # Shapes + num_blocks, N, _ = blocks.shape + Nhalf = N // 2 + 1 # rfft size num_1d_tapers = self.tapers_1d.shape[0] - blocks_mt = xp.zeros(blocks[0, :, :].shape, dtype=self.dtype) - blocks_tapered = xp.empty(blocks[0, :, :].shape, dtype=complex_type(self.dtype)) - taper_2d = xp.empty(blocks[0, :, :].shape, dtype=complex_type(self.dtype)) - - for ax1 in range(num_1d_tapers): - for ax2 in range(num_1d_tapers): - taper_2d[:, :] = ( - self.tapers_1d[ax1, :, np.newaxis] - @ self.tapers_1d[ax2, :, np.newaxis].T - ) - for m in range(blocks.shape[0]): - blocks_tapered[:, :] = blocks[m] * taper_2d - blocks_mt_post_fft = fft.fftn(blocks_tapered, axes=(-2, -1)) - blocks_mt += abs2(blocks_mt_post_fft) + # Ensure blocks are on GPU and consistent dtype + blocks = xp.asarray(blocks, dtype=self.dtype) + + # Array to accumulate PSD + # Use hermitian symmetry to cut PSD compute in half + psd_half = xp.zeros((N, Nhalf), dtype=self.dtype) + accum = xp.empty((N, Nhalf), dtype=self.dtype) + + for ax2 in range(num_1d_tapers): + # taper along columns + A = fft.rfft( + blocks * self.tapers_1d[ax2][None, None, :], axis=-1 + ) # (num_1d_tapers, N, Nhalf) + + for ax1 in range(num_1d_tapers): + # taper along rows + Bf = fft.fft( + A * self.tapers_1d[ax1][None, :, None], axis=-2 + ) # (num_1d_tapers, N, Nhalf) + # Accumulate abs2 summed over all blocks into temp array `accum` + _abs2_sum0(Bf, accum) + psd_half += accum + + # Rebuild the full spectrum from psd_half + amplitude_spectrum = xp.empty((N, N), dtype=self.dtype) + amplitude_spectrum[:, :Nhalf] = psd_half + # reflect all the positive freq columns (1:N - Nhalf + 1), roll the single zero column. + amplitude_spectrum[:, Nhalf:] = xp.roll( + psd_half[:, 1 : N - Nhalf + 1][::-1, ::-1], 1, axis=0 + ) # Normalize - blocks_mt /= blocks.shape[1] ** 2 * blocks.shape[0] ** 2 + # Ask Ayelet about the num_blocks factor. Needed to match MATLAB intermediate. + amplitude_spectrum /= N**2 * num_blocks**2 - amplitude_spectrum = fft.fftshift(blocks_mt) + amplitude_spectrum = fft.fftshift(amplitude_spectrum) return Image(xp.asnumpy(amplitude_spectrum)) @@ -512,6 +530,7 @@ def background_subtract_1d( return final_signal, final_background + @profile def opt1d( self, amplitude_spectrum, @@ -527,46 +546,51 @@ def opt1d( :return: 2-tuple of NumPy arrays (Estimated average of defocus and low_freq_cutoff) """ - signal = amplitude_spectrum.T - signal = np.maximum(0.0, signal) - signal = np.sqrt(signal) + signal = xp.asarray(amplitude_spectrum.T) + signal = xp.maximum(0.0, signal) + signal = xp.sqrt(signal) signal = signal[: 3 * signal.shape[0] // 4] - c = np.zeros((max_defocus - min_defocus, signal.shape[1]), dtype=self.dtype) - - for f in range(min_defocus, max_defocus): - - ctf_im = np.abs( - CTFFilter.ctf_formula( - self.radial_omega, - pixel_size=self.pixel_size, - voltage=self.voltage, - defocus_u_a=10 * f, - defocus_v_a=10 * f, - defocus_ang=0, - Cs=self.cs, - alpha=self.amplitude_contrast, - B=0, - )[0] + c = xp.zeros((max_defocus - min_defocus, signal.shape[1]), dtype=self.dtype) + + defoci_angstrom = xp.arange(min_defocus, max_defocus, dtype=int) * 10 + ctf_ims = np.abs( + CTFFilter.ctf_formula( + xp.asarray(self.radial_omega), + pixel_size=self.pixel_size, + voltage=self.voltage, + defocus_u_a=defoci_angstrom, + defocus_v_a=defoci_angstrom, + defocus_ang=0, + Cs=self.cs, + alpha=self.amplitude_contrast, + B=0, ) + ) + defoci_angstrom = xp.asnumpy(defoci_angstrom // 10) # nm + + # Think we can optimize this way better... + for f, ctf_im in zip(defoci_angstrom, ctf_ims): # Shouldn't normally be needed within class calls, - # but if user running manually clip to signal len. + # but if user running manually, ensure clip to signal len. ctf_im = ctf_im[: signal.shape[0]] - ctf_im = np.reshape(ctf_im, (ctf_im.shape[0], 1)) - ctf_im = np.tile(ctf_im, (1, signal.shape[1])) + ctf_im = ctf_im.reshape((ctf_im.shape[0], 1)) + ctf_im = xp.tile(ctf_im, (1, signal.shape[1])) for m in range(0, signal.shape[1]): - signal[:, m] = signal[:, m] - np.mean(signal[m + 1 :, m], axis=0) - ctf_im[:, m] = ctf_im[:, m] - np.mean(ctf_im[m + 1 :, m], axis=0) + signal[:, m] = signal[:, m] - xp.mean(signal[m + 1 :, m], axis=0) + ctf_im[:, m] = ctf_im[:, m] - xp.mean(ctf_im[m + 1 :, m], axis=0) ctf_im[: m + 1, m] = 0 signal[: m + 1, m] = 0 - Sx = np.sqrt(np.sum(ctf_im**2, axis=0)) - Sy = np.sqrt(np.sum(signal**2, axis=0)) - c[f - min_defocus, :] = np.sum(ctf_im * signal, axis=0) / (Sx * Sy) + Sx = xp.sqrt(xp.sum(ctf_im**2, axis=0)) + Sy = xp.sqrt(xp.sum(signal**2, axis=0)) + c[f - min_defocus, :] = xp.sum(ctf_im * signal, axis=0) / (Sx * Sy) - avg_defocus, low_freq_cutoff_col = np.unravel_index(np.argmax(c), c.shape)[:2] + avg_defocus, low_freq_cutoff_col = np.unravel_index( + xp.argmax(c).item(), c.shape + )[:2] avg_defocus += min_defocus return avg_defocus, low_freq_cutoff_col @@ -1184,3 +1208,25 @@ def estimate_ctf( mrc.voxel_size = ctf_est.pixel_size return results + + +# Helper function for sum(abs^2(x)) +# numpy expression fallback +def _abs2_sum0(x, out): + out[:] = (x.real**2 + x.imag**2).sum(axis=0) + + +# CuPy +if hasattr(xp, "ReductionKernel"): + _abs2_sum0_kernel = xp.ReductionKernel( + "T x", + "Y y", + "x.real() * x.real() + x.imag() * x.imag()", + "a + b", + "y = a", + "0", + name="abs2_sum0", + ) + + def _abs2_sum0(x, out): + _abs2_sum0_kernel(x, axis=0, out=out) From 31a36638135f9d5af2e0659d422339b157cbd157 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 19 Oct 2026 15:08:49 -0400 Subject: [PATCH 32/44] optimized opt1d stats/correlation compute --- src/aspire/ctf/ctf_estimator.py | 62 +++++++++++++++++++-------------- 1 file changed, 35 insertions(+), 27 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 62ae1d58f4..f9ffc0c964 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -44,7 +44,7 @@ from aspire.numeric import fft, xp from aspire.operators import CTFFilter from aspire.storage import StarFile -from aspire.utils import complex_type, grid_2d, tqdm, voltage_to_wavelength +from aspire.utils import grid_2d, tqdm, voltage_to_wavelength logger = logging.getLogger(__name__) @@ -550,11 +550,11 @@ def opt1d( signal = xp.maximum(0.0, signal) signal = xp.sqrt(signal) signal = signal[: 3 * signal.shape[0] // 4] + N, M = signal.shape - c = xp.zeros((max_defocus - min_defocus, signal.shape[1]), dtype=self.dtype) - + # Compute the (radial) CTF filter stack in bulk defoci_angstrom = xp.arange(min_defocus, max_defocus, dtype=int) * 10 - ctf_ims = np.abs( + ctf_ims = xp.abs( CTFFilter.ctf_formula( xp.asarray(self.radial_omega), pixel_size=self.pixel_size, @@ -569,24 +569,25 @@ def opt1d( ) defoci_angstrom = xp.asnumpy(defoci_angstrom // 10) # nm - # Think we can optimize this way better... - for f, ctf_im in zip(defoci_angstrom, ctf_ims): - # Shouldn't normally be needed within class calls, - # but if user running manually, ensure clip to signal len. - ctf_im = ctf_im[: signal.shape[0]] - - ctf_im = ctf_im.reshape((ctf_im.shape[0], 1)) - ctf_im = xp.tile(ctf_im, (1, signal.shape[1])) - - for m in range(0, signal.shape[1]): - signal[:, m] = signal[:, m] - xp.mean(signal[m + 1 :, m], axis=0) - ctf_im[:, m] = ctf_im[:, m] - xp.mean(ctf_im[m + 1 :, m], axis=0) - ctf_im[: m + 1, m] = 0 - signal[: m + 1, m] = 0 - - Sx = xp.sqrt(xp.sum(ctf_im**2, axis=0)) - Sy = xp.sqrt(xp.sum(signal**2, axis=0)) - c[f - min_defocus, :] = xp.sum(ctf_im * signal, axis=0) / (Sx * Sy) + # Compute `signal` statistics, independent of defocus + LT = xp.tril(signal, k=-1) # keep lower triangle + sx = LT.sum(axis=0) # sum of row values + sxx = (LT * LT).sum(axis=0) # sum of squared row values + nrows = (N - 1 - xp.arange(M)).astype( + self.dtype + ) # rows below diagonal for each column + Sxx = sxx - sx**2 / nrows # centered + + # CTF statistics + ctf_ims = xp.asarray(ctf_ims, dtype=self.dtype)[:, :N] + sy = tail_sum(ctf_ims, M) + syy = tail_sum(ctf_ims * ctf_ims, M) + Syy = syy - sy**2 / nrows + + # Compute the cross corr using one GEMM + c = ctf_ims @ LT + c -= sy * sx[None, :] / nrows + c = c / xp.sqrt(Sxx[None, :] * Syy) avg_defocus, low_freq_cutoff_col = np.unravel_index( xp.argmax(c).item(), c.shape @@ -1211,11 +1212,6 @@ def estimate_ctf( # Helper function for sum(abs^2(x)) -# numpy expression fallback -def _abs2_sum0(x, out): - out[:] = (x.real**2 + x.imag**2).sum(axis=0) - - # CuPy if hasattr(xp, "ReductionKernel"): _abs2_sum0_kernel = xp.ReductionKernel( @@ -1230,3 +1226,15 @@ def _abs2_sum0(x, out): def _abs2_sum0(x, out): _abs2_sum0_kernel(x, axis=0, out=out) + +else: + # numpy expression fallback + def _abs2_sum0(x, out): + out[:] = (x.real**2 + x.imag**2).sum(axis=0) + + +# Helper function +def tail_sum(x, M): + """sum over i > m (exclusive reverse cumsum)""" + cs = xp.cumsum(x, axis=1) + return cs[:, -1:] - cs[:, :M] From d478da8b7aa62289f8c7f78ad3173a397823da26 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 20 Oct 2026 08:56:54 -0400 Subject: [PATCH 33/44] expose opt1d loop to class and cli --- src/aspire/commands/estimate_ctf.py | 21 ++++++++++++++++++ src/aspire/ctf/ctf_estimator.py | 34 ++++++++++++++++++++--------- 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/src/aspire/commands/estimate_ctf.py b/src/aspire/commands/estimate_ctf.py index 12711a9d11..a6ba656b09 100644 --- a/src/aspire/commands/estimate_ctf.py +++ b/src/aspire/commands/estimate_ctf.py @@ -47,6 +47,21 @@ @click.option( "--g_max", default=5, type=float, help="Inverse of maximum resolution for PSD" ) +@click.option( + "--defocus_min", + default=5000, + type=int, + help="Start of defocus loop scan (Angstrom).", +) +@click.option( + "--defocus_max", + default=100000, + type=int, + help="End of defocus loop scan (Angstrom).", +) +@click.option( + "--defocus_step", default=10, type=int, help="Step of defocus loop scan (Angstrom)." +) @click.option( "--dtype", default="float32", @@ -77,6 +92,9 @@ def estimate_ctf( psd_size, g_min, g_max, + min_defocus, + max_defocus, + step_defocus, output_dir, dtype, save_ctf_images, @@ -108,6 +126,9 @@ def estimate_ctf( psd_size=psd_size, g_min=g_min, g_max=g_max, + min_defocus=min_defocus, + max_defocus=max_defocus, + step_defocus=step_defocus, output_dir=output_dir, dtype=dtype, save_ctf_images=save_ctf_images, diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index f9ffc0c964..112e6dad43 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -79,6 +79,9 @@ def __init__( num_tapers=4, g_min=30, g_max=5, + min_defocus=5000, + max_defocus=100000, + step_defocus=10, basis=None, solver_method="bfgs", solver_opts=None, @@ -98,6 +101,9 @@ def __init__( Defaults to 4. :param g_min: Inverse of minimun resolution for PSD. Default 30. :param g_max: Inverse of maximum resolution for PSD. Default 5. + :param min_defocus: Start of defocus loop scan (Angstrom). Default 5000. + :param max_defocus: End of defocus loop scan (Angstrom). Default 100000. + :param step_defocus: Step size of defocus loop scan (Angstrom). Default 10. :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. :param solver_method: Solver used to optimize defocus params. `gd` or `bfgs`. Default uses single call to `bfgs`. @@ -125,6 +131,9 @@ def __init__( self.g_min = np.asarray(g_min).item() self.g_max = np.asarray(g_max).item() + self.min_defocus = np.asarray(min_defocus, dtype=int).item() + self.max_defocus = np.asarray(max_defocus, dtype=int).item() + self.step_defocus = np.asarray(step_defocus, dtype=int).item() # Setup solvers and respective solver options if solver_method not in self._solver_methods: @@ -534,15 +543,11 @@ def background_subtract_1d( def opt1d( self, amplitude_spectrum, - min_defocus=500, - max_defocus=10000, ): """ Find optimal defocus for the radially symmetric case (where no astigmatism is present) :param amplitude_spectrum: Estimated power specrtum. - :param min_defocus: Start of defocus loop scan - :param max_defocus: End of defocus loop scan :return: 2-tuple of NumPy arrays (Estimated average of defocus and low_freq_cutoff) """ @@ -553,7 +558,9 @@ def opt1d( N, M = signal.shape # Compute the (radial) CTF filter stack in bulk - defoci_angstrom = xp.arange(min_defocus, max_defocus, dtype=int) * 10 + defoci_angstrom = xp.arange( + self.min_defocus, self.max_defocus, self.step_defocus, dtype=int + ) ctf_ims = xp.abs( CTFFilter.ctf_formula( xp.asarray(self.radial_omega), @@ -567,7 +574,6 @@ def opt1d( B=0, ) ) - defoci_angstrom = xp.asnumpy(defoci_angstrom // 10) # nm # Compute `signal` statistics, independent of defocus LT = xp.tril(signal, k=-1) # keep lower triangle @@ -589,10 +595,9 @@ def opt1d( c -= sy * sx[None, :] / nrows c = c / xp.sqrt(Sxx[None, :] * Syy) - avg_defocus, low_freq_cutoff_col = np.unravel_index( - xp.argmax(c).item(), c.shape - )[:2] - avg_defocus += min_defocus + max_c = xp.argmax(c).item() + avg_defocus, low_freq_cutoff_col = np.unravel_index(max_c, c.shape)[:2] + avg_defocus += self.min_defocus / 10 # (nm) return avg_defocus, low_freq_cutoff_col @@ -1081,6 +1086,9 @@ def estimate_ctf( psd_size=1024, g_min=30, g_max=5, + min_defocus=5000, + max_defocus=100000, + step_defocus=10, basis=None, output_dir="results", dtype=np.float32, @@ -1108,6 +1116,9 @@ def estimate_ctf( :param psd_size: Block size (in pixels) for PSD estimation. :param g_min: Inverse of minimun resolution for PSD. Default 30. :param g_max: Inverse of maximum resolution for PSD. Default 5. + :param min_defocus: Start of defocus loop scan (Angstrom). Default 5000. + :param max_defocus: End of defocus loop scan (Angstrom). Default 100000. + :param step_defocus: Step size of defocus loop scan (Angstrom). Default 10. :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). Defaults to 'highs-ds'. @@ -1156,6 +1167,9 @@ def estimate_ctf( num_tapers=num_tapers, g_min=g_min, g_max=g_max, + min_defocus=min_defocus, + max_defocus=max_defocus, + step_defocus=step_defocus, basis=basis, diagnostic_plots=False, dtype=dtype, From 3721f3f85df34bea728fab03279b7abc93cd51da Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 20 Oct 2026 09:01:54 -0400 Subject: [PATCH 34/44] remove profiler lines --- src/aspire/ctf/ctf_estimator.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 112e6dad43..44dcef1391 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -34,7 +34,6 @@ import matplotlib.pyplot as plt import mrcfile import numpy as np -from line_profiler import profile from numpy import linalg as npla from scipy.optimize import linprog, minimize from scipy.signal.windows import dpss @@ -539,7 +538,6 @@ def background_subtract_1d( return final_signal, final_background - @profile def opt1d( self, amplitude_spectrum, From d66b41cfd05d0c0cc73522d9aa99f2b62a2a0fa4 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 20 Oct 2026 09:02:51 -0400 Subject: [PATCH 35/44] more stable with 4 tapers --- gallery/tutorials/tutorials/ctf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gallery/tutorials/tutorials/ctf.py b/gallery/tutorials/tutorials/ctf.py index a269e3c35b..53a7c63af2 100644 --- a/gallery/tutorials/tutorials/ctf.py +++ b/gallery/tutorials/tutorials/ctf.py @@ -260,7 +260,7 @@ def generate_example_image(L, noise_variance=0.1): amplitude_contrast=radial_ctf_filter.alpha, voltage=radial_ctf_filter.voltage, psd_size=512, - num_tapers=2, + num_tapers=4, dtype=np.float64, ) From 01f53db0b20637b11c79c35dd73f7df55eb45fd6 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 20 Oct 2026 10:00:03 -0400 Subject: [PATCH 36/44] fix argname typo in cli --- src/aspire/commands/estimate_ctf.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aspire/commands/estimate_ctf.py b/src/aspire/commands/estimate_ctf.py index a6ba656b09..03ecb11541 100644 --- a/src/aspire/commands/estimate_ctf.py +++ b/src/aspire/commands/estimate_ctf.py @@ -48,19 +48,19 @@ "--g_max", default=5, type=float, help="Inverse of maximum resolution for PSD" ) @click.option( - "--defocus_min", + "--min_defocus", default=5000, type=int, help="Start of defocus loop scan (Angstrom).", ) @click.option( - "--defocus_max", + "--max_defocus", default=100000, type=int, help="End of defocus loop scan (Angstrom).", ) @click.option( - "--defocus_step", default=10, type=int, help="Step of defocus loop scan (Angstrom)." + "--step_defocus", default=10, type=int, help="Step of defocus loop scan (Angstrom)." ) @click.option( "--dtype", From 65b4cee7d09064d79e0b09a48964ebc14150e37b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 20 Oct 2026 12:54:49 -0400 Subject: [PATCH 37/44] handle failure to find soln and improve linprog call: --- src/aspire/ctf/ctf_estimator.py | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 44dcef1391..077d6f6200 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -518,6 +518,13 @@ def background_subtract_1d( axis=0, ) + # Rescale to avoid small values conflicting with internal bounds checks. + # It is possible to disable with `options={'presolve': False}`, + # but this may impact the solution. + # Instead the input and bounds are rescaled, then the solution is rescaled. + _scale = 1e4 + f, bounds = f / _scale, bounds * _scale + res = linprog( f, A_ub=A, @@ -527,9 +534,11 @@ def background_subtract_1d( ) if not res.success: - raise RuntimeError("Linear program did not succeed. Halting") + raise RuntimeError( + f"Linear program did not succeed after {res.nit} iterations with condition {res.status}: '{res.message}'.\tHalting." + ) - background = res.x[N:] + background = res.x[N:] / _scale bs_psd = signal - background @@ -1175,6 +1184,7 @@ def estimate_ctf( # Main Loop --------------------------------------------------------------- results = {} + failed = {} for input_file in tqdm(file_paths, desc="Processing files"): # Handle path stuff first, in case it fails @@ -1193,7 +1203,12 @@ def estimate_ctf( micrograph = micrograph.astype(dtype, copy=False) - params, background_2d, signal = ctf_est.estimate(micrograph) + try: + params, background_2d, signal = ctf_est.estimate(micrograph) + except Exception as e: + failed[str(relative_input_path)] = e + logger.warning(f"Estimation failed for {str(relative_input_path)}") + continue # skip results for this failure # Store files' params into results dictionary results[str(relative_input_path)] = params @@ -1220,6 +1235,15 @@ def estimate_ctf( mrc.set_data(ctf_im.astype(np.float32)) # MRC only supports singles mrc.voxel_size = ctf_est.pixel_size + if len(failed) > 0: + failure_fname = "failures.txt" + logger.info( + f"{len(failed)} files failed to estimate. Writing to '{failure_fname}'" + ) + with open(failure_fname, "w") as f: + for k, v in failed.items(): + f.write(f"{k}\t{v}\n") + return results From af316212bdb7a8095d41818f0839c00cd7ebf7d6 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 20 Oct 2026 14:03:24 -0400 Subject: [PATCH 38/44] default for psd_size --- src/aspire/ctf/ctf_estimator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 077d6f6200..a2b35ce667 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -74,7 +74,7 @@ def __init__( cs, amplitude_contrast, voltage, - psd_size, + psd_size=1024, num_tapers=4, g_min=30, g_max=5, @@ -95,7 +95,7 @@ def __init__( :param cs: Spherical aberration in mm. :param amplitude_contrast: Amplitude contrast. :param voltage: Voltage of electron microscope. - :param psd_size: Block size (in pixels) for PSD estimation. + :param psd_size: Block size (in pixels) for PSD estimation. Default 1024. :param num_tapers: Optional, number of tapers to apply in PSD estimation. Defaults to 4. :param g_min: Inverse of minimun resolution for PSD. Default 30. @@ -1170,7 +1170,7 @@ def estimate_ctf( cs, amplitude_contrast, voltage, - psd_size, + psd_size=psd_size, num_tapers=num_tapers, g_min=g_min, g_max=g_max, From 1f6524e38451a1d62d700d41442b8ac2eeadc20a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 21 Oct 2026 10:53:54 -0400 Subject: [PATCH 39/44] improve numerical stability by fixing basis artifacts --- src/aspire/ctf/ctf_estimator.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index a2b35ce667..129c4a310d 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -395,7 +395,7 @@ def estimate_psd(self, blocks): amplitude_spectrum = fft.fftshift(amplitude_spectrum) - return Image(xp.asnumpy(amplitude_spectrum)) + return Image(xp.asnumpy(amplitude_spectrum), pixel_size=self.pixel_size) def radial_average(self, amplitude_spectrum): """ @@ -414,6 +414,12 @@ def radial_average(self, amplitude_spectrum): self.basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size ).evaluate() + if (_min := np.min(psd)) < 0: + logger.debug( + f"PSD should be strictly greater than 0, found {_min}, clipping to 0." + ) + psd = Image(np.maximum(psd, 0), pixel_size=self.pixel_size) + return psd, noise def elliptical_average(self, amplitude_spectrum): @@ -440,6 +446,17 @@ def elliptical_average(self, amplitude_spectrum): self.basis, coefs_s, pixel_size=amplitude_spectrum.pixel_size ).evaluate() + if (_min := np.min(psd)) < 0: + logger.debug( + f"PSD should be strictly greater than 0, found {_min}, clipping to 0." + ) + psd = Image(np.maximum(psd, 0), pixel_size=self.pixel_size) + if (_min := np.min(noise)) < 0: + logger.debug( + f"noise should be strictly greater than 0, found {_min}, clipping to 0." + ) + noise = Image(np.maximum(noise, 0), pixel_size=self.pixel_size) + return psd, noise def background_subtract_1d( From 7a9739587345be64726d45f0749015ef04d7148a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 21 Oct 2026 13:36:37 -0400 Subject: [PATCH 40/44] less redundant name in starfiles --- src/aspire/ctf/ctf_estimator.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 129c4a310d..906992f3a1 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -943,12 +943,13 @@ def maximize_correlation(self, signal, initial_df1, initial_df2): elif self.solver_method == "bfgs": return self.bfgs(signal, initial_df1, initial_df2, **self.solver_opts) - def write_star(self, star_filepath, params_dict): + def write_star(self, star_filepath, micrograph_name, params_dict): """ Writes CTF parameters to starfile for a single micrograph. """ + data_block = {} - data_block["_rlnMicrographName"] = star_filepath + data_block["_rlnMicrographName"] = micrograph_name data_block["_rlnDefocusU"] = params_dict["defocus_u"] # Should already be in A data_block["_rlnDefocusV"] = params_dict["defocus_v"] # Should already be in A data_block["_rlnDefocusAngle"] = np.rad2deg(params_dict["defocus_ang"]) @@ -1232,7 +1233,7 @@ def estimate_ctf( # Write each micrograph's ctf parameters to an individual STAR file star_filepath = output_file_prefix.with_suffix(".star") - ctf_est.write_star(star_filepath, params) + ctf_est.write_star(star_filepath, str(relative_input_path), params) # Save the noise if save_noise_images: From 43a7017a03012128435d03d67e1d34ab6728dc48 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 22 Oct 2026 08:34:28 -0400 Subject: [PATCH 41/44] fix ast angle tranpose issue (blocks.T) and cleanup some strings --- src/aspire/ctf/ctf_estimator.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 906992f3a1..6ab9e4047c 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -220,12 +220,12 @@ def normalize_blocks(self, blocks): :return: NumPy array of normalized blocks. """ - # Take block size from blocks + # Take block size from blocks, check still square block_size = blocks.shape[1] assert block_size == blocks.shape[2] - # Create a sum and reshape so it may be broadcast with `block`. - blocks_sum = np.sum(blocks, axis=(-1, -2))[:, np.newaxis, np.newaxis] + # Sum per block + blocks_sum = blocks.sum(axis=(-1, -2), keepdims=True) blocks -= blocks_sum / (block_size**2) @@ -238,8 +238,10 @@ def preprocess_micrograph(self, micrograph): :param micrograph: Micrograph as NumPy array. :return: NumPy array of normalized blocks extracted from the micrograph. """ - - return self.normalize_blocks(self.micrograph_to_blocks(micrograph)) + blocks = self.normalize_blocks(self.micrograph_to_blocks(micrograph)) + # Cheaper to transpose block (window) than whole micrograph in beginning. + # Without this tranpose the astig axis will be reflected wrt other tools' convention. + return blocks.transpose(0, 2, 1) def legacy_preprocess_micrograph(self, micrograph): """ @@ -350,7 +352,7 @@ def estimate_psd(self, blocks): Estimate the power spectrum of the micrograph using the multi-taper method :param blocks: 3-D NumPy array containing windows extracted from the micrograph in the preprocess function. - :return: NumPy array of estimated power spectrum. + :return: Image of estimated power spectrum. """ # Shapes @@ -402,7 +404,7 @@ def radial_average(self, amplitude_spectrum): Computes radial/elliptical average of the power spectrum :param amplitude_spectrum: Power spectrum to average. - :return: PSD and noise as 2-tuple of NumPy arrays. + :return: PSD and noise as 2-tuple of Image instances """ noise = amplitude_spectrum # pass through coefs_s = self.basis.evaluate_t(amplitude_spectrum).asnumpy().copy() @@ -427,7 +429,7 @@ def elliptical_average(self, amplitude_spectrum): Computes elliptical average of the power spectrum :param amplitude_spectrum: Power spectrum to average. - :return: PSD and noise as 2-tuple of NumPy arrays. + :return: PSD and noise as 2-tuple of Image instances. """ coefs_s = self.basis.evaluate_t(amplitude_spectrum).asnumpy().copy() @@ -1225,7 +1227,9 @@ def estimate_ctf( params, background_2d, signal = ctf_est.estimate(micrograph) except Exception as e: failed[str(relative_input_path)] = e - logger.warning(f"Estimation failed for {str(relative_input_path)}") + logger.warning( + f"Estimation failed for {str(relative_input_path)}: {str(e)}" + ) continue # skip results for this failure # Store files' params into results dictionary From 5729178de1a0b34913c4a30c4cb537fd1bddb46b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 22 Oct 2026 09:10:31 -0400 Subject: [PATCH 42/44] optimize preproc micrograph and rm legacy implementation --- src/aspire/ctf/ctf_estimator.py | 185 +++++++++----------------------- 1 file changed, 52 insertions(+), 133 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index 6ab9e4047c..db4ee95ad1 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -176,17 +176,16 @@ def __init__( tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) self.tapers_1d = xp.asarray(tapers_1d, dtype=self.dtype) - def micrograph_to_blocks(self, micrograph): + def preprocess_micrograph(self, micrograph): """ - Preprocess micrograph into blocks using block_size. + Preprocess micrograph into normalized half overlapping blocks of `psd_size` by `psd_size`. - :param micrograph: Micrograph as NumPy array. #NOTE looks like F order - :param blocksize: Size of the square blocks to partition micrograph. - :return: NumPy array of blocks extracted from the micrograph. + :param micrograph: Micrograph as NumPy array. + :return: Array of normalized blocks extracted from the micrograph (..., psd_size, psd_size). + Note, may be a CuPy array depending on `xp` mode. """ - # verify self.psd_size is even - assert self.psd_size % 2 == 0 + assert self.psd_size % 2 == 0, "psd_size must be even" if micrograph.ndim == 3: assert ( @@ -194,138 +193,59 @@ def micrograph_to_blocks(self, micrograph): ), f"micrograph should be 2D or stack of 1 2D image: {micrograph.shape}" micrograph = micrograph[0] - size_x = micrograph.shape[-1] - size_y = micrograph.shape[-2] - - step_size = self.psd_size // 2 - range_y = size_y // step_size - 1 - range_x = size_x // step_size - 1 - - block_list = [ - micrograph[ - j * step_size : (j + 2) * step_size, i * step_size : (i + 2) * step_size - ] - for j in range(range_y) - for i in range(range_x) - ] - blocks = np.asarray(block_list, dtype=micrograph.dtype) - - return blocks - - def normalize_blocks(self, blocks): - """ - Preprocess CTF of micrograph using block_size. - - :param blocks: NumPy array of blocks extracted from the micrograph. - :return: NumPy array of normalized blocks. - """ - - # Take block size from blocks, check still square - block_size = blocks.shape[1] - assert block_size == blocks.shape[2] - - # Sum per block - blocks_sum = blocks.sum(axis=(-1, -2), keepdims=True) - - blocks -= blocks_sum / (block_size**2) + assert ( + min(micrograph.shape) >= self.psd_size + ), f"micrograph {micrograph.shape} must be larger than psd_size {self.psd_size}" + step = self.psd_size // 2 - return blocks + # Perform the micrograph processing on the gpu when in xp mode + micrograph = xp.asarray(micrograph) - def preprocess_micrograph(self, micrograph): - """ - Preprocess micrograph into normalized blocks using block_size. - - :param micrograph: Micrograph as NumPy array. - :return: NumPy array of normalized blocks extracted from the micrograph. - """ - blocks = self.normalize_blocks(self.micrograph_to_blocks(micrograph)) - # Cheaper to transpose block (window) than whole micrograph in beginning. + # Cheaper to transpose whole micrograph than all the blocks later # Without this tranpose the astig axis will be reflected wrt other tools' convention. - return blocks.transpose(0, 2, 1) - - def legacy_preprocess_micrograph(self, micrograph): - """ - Returns half-overlapping blocks extracted from the micrograph. - Equivalent to MATLAB's ctf_extract_blocks. - - Note, internally transposes micrograph to match MATLAB computations. - Returns (num_blocks, block_size, block_size) - - :param micrograph: Micrograph as NumPy array. - :return: NumPy array of normalized blocks extracted from the micrograph. - """ - h, w = micrograph.shape - - if min(h, w) < self.psd_size: - raise ValueError( - f"Blocks of size {self.psd_size} x {self.psd_size} cannot be extracted from image." - ) - - block_size = self.psd_size - (self.psd_size % 2) # force even - half = block_size // 2 - - # Compute crop bounds - end_val1 = (h // block_size) * block_size + ( - half if (h % block_size) > half else 0 - ) - end_val2 = (w // block_size) * block_size + ( - half if (w % block_size) > half else 0 - ) + micrograph = micrograph.T - # Crop - cropped_micrograph = micrograph[:end_val1, :end_val2] - - # Compute number of half blocks - n_bh = end_val1 // half # number of half-blocks along rows - n_bw = end_val2 // half # number of half-blocks along cols - n_total = n_bh * n_bw - - # im2col `distinct` extracts non-overlapping half blocks - # column-major block ordering to match MATLAB (down-first) - # but keeping now for diagnostics. - blocks = ( - cropped_micrograph.reshape(n_bh, half, n_bw, half) - .transpose(1, 3, 0, 2) - .reshape(half * half, n_total, order="F") - ) - - # Reshape to (half, half, n_total) so each slice is one half-block - blocks = blocks.reshape(half, half, n_total, order="F") - - # First cat - # stack each half-block with the one below it (rows overlap) - # by shifting the block indices by 1 - row_shifted = np.concatenate([blocks[:, :, 1:], blocks[:, :, :1]], axis=2) - blocks = np.concatenate( - [blocks, row_shifted], axis=0 - ) # (block_size, half, n_total) - - # Second cat - # stack each block with the one to its right (cols overlap) - # by rolling the blocks - col_shifted = np.roll(blocks, -n_bh, axis=2) - blocks = np.concatenate( - [blocks, col_shifted], axis=1 - ) # (block_size, block_size, n_total) - - # Reshape to blocksize-by-blocksize grid - n_row_blocks = int(np.floor(2 * h / block_size)) - n_col_blocks = int(np.floor(2 * w / block_size)) - blocks = blocks.reshape( - block_size, block_size, n_row_blocks, n_col_blocks, order="F" + # Strided view of psd_size windows, using stride `step` + # Avoids repeatedly building the entire blocklist + # Note, modifying this view can modify underlying array. + # + # (ny, nx, psd_size, psd_size), no copy + windows = xp.lib.stride_tricks.sliding_window_view( + micrograph, (self.psd_size, self.psd_size) + )[::step, ::step] + + # The following tiling code is a little complicated, + # but yields good comprimise of performance on both CPU and GPU. + # It is equivalent to the following simpler code: + # windows_sum = windows.sum(axis=(-1, -2), keepdims=True, dtype=np.float64) + # normalized_blocks = windows - windows_sum / (self.psd_size**2) + + # Instead of looping over windows (which repeat calculations on sub-tiles), + # each window is subdivided into tiles, which are read only once each. + # Use doubles for intermediates. + ny, nx = windows.shape[:2] + tiles = micrograph[: (ny + 1) * step, : (nx + 1) * step] + tile_sums = tiles.reshape(ny + 1, step, nx + 1, step).sum( + axis=(1, 3), dtype=np.float64 ) - # Drop last row and col (bounds artifacts) - blocks = blocks[:, :, :-1, :-1] - # Flatten - block = blocks.reshape(block_size, block_size, -1, order="F") + # Now for each window, gather the respective tile_sums + window_means = ( + tile_sums[:-1, :-1] # Upper Left + + tile_sums[1:, :-1] # Lower Left + + tile_sums[:-1, 1:] # Upper Right + + tile_sums[1:, 1:] # Lower Right + ) / ( + 4 * step**2 + ) # divide sum to yield mean - # De-mean each block - block_mean = block.sum(axis=(0, 1), keepdims=True) / (block_size**2) - block = block - block_mean + normalized_blocks = xp.empty(windows.shape, dtype=self.dtype) + normalized_blocks[:, :] = windows - window_means[:, :, None, None] + # This should be a free reshape + normalized_blocks = normalized_blocks.reshape(-1, self.psd_size, self.psd_size) - # Transpose output to (n,block_size, block_size) - return block.T + # Since this might be large-ish, keep on GPU when in xp mode (ie skip the usual asarray()) + return normalized_blocks def tapers(self, N, NW, L): """ @@ -1022,9 +942,8 @@ def estimate(self, micrograph): :param micrograph: 2D Array of micrograph data. :return: CTF params, estimated background PSD, estimated signal PSD """ - # TODO Ask Ayelet about legacy version intent and rm if not needed. + micrograph_blocks = self.preprocess_micrograph(micrograph) - # micrograph_blocks = self.legacy_preprocess_micrograph(micrograph) signal_observed = self.estimate_psd(micrograph_blocks) From 6311a5012c953d9b5dd03e55ae5fff0dd97882bf Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 26 Oct 2026 11:15:32 -0400 Subject: [PATCH 43/44] review string and arg cleanup --- src/aspire/ctf/ctf_estimator.py | 37 ++++++++++++++++----------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index db4ee95ad1..a3f1307d5a 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -2,7 +2,7 @@ Contains code supporting CTF parameter estimation. This implements a version of ASPIRE-CTF originally written in MATLAB, -This Python version supports multiple basis and solvers. +This Python version supports multiple bases and solvers. See paper: @@ -98,7 +98,7 @@ def __init__( :param psd_size: Block size (in pixels) for PSD estimation. Default 1024. :param num_tapers: Optional, number of tapers to apply in PSD estimation. Defaults to 4. - :param g_min: Inverse of minimun resolution for PSD. Default 30. + :param g_min: Inverse of minimum resolution for PSD. Default 30. :param g_max: Inverse of maximum resolution for PSD. Default 5. :param min_defocus: Start of defocus loop scan (Angstrom). Default 5000. :param max_defocus: End of defocus loop scan (Angstrom). Default 100000. @@ -151,7 +151,7 @@ def __init__( # Grids grid = grid_2d(psd_size, normalized=True, indexing="yx", dtype=self.dtype) # Note range is -half to half. - self.r_ctf = grid["r"] / 2 * (10 / pixel_size) # units: inverse nm + self.r_ctf = grid["r"] / 2 * (10 / self.pixel_size) # units: inverse nm self.theta = grid["phi"] # Construct omega grids for use in ASPIRE's ctf_formula self.omega = np.pi * np.vstack((grid["x"].flatten(), grid["y"].flatten())) @@ -175,6 +175,7 @@ def __init__( # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) self.tapers_1d = xp.asarray(tapers_1d, dtype=self.dtype) + breakpoint() def preprocess_micrograph(self, micrograph): """ @@ -202,7 +203,7 @@ def preprocess_micrograph(self, micrograph): micrograph = xp.asarray(micrograph) # Cheaper to transpose whole micrograph than all the blocks later - # Without this tranpose the astig axis will be reflected wrt other tools' convention. + # Without this transpose the astig axis will be reflected wrt other tools' convention. micrograph = micrograph.T # Strided view of psd_size windows, using stride `step` @@ -215,7 +216,7 @@ def preprocess_micrograph(self, micrograph): )[::step, ::step] # The following tiling code is a little complicated, - # but yields good comprimise of performance on both CPU and GPU. + # but yields good compromise of performance on both CPU and GPU. # It is equivalent to the following simpler code: # windows_sum = windows.sum(axis=(-1, -2), keepdims=True, dtype=np.float64) # normalized_blocks = windows - windows_sum / (self.psd_size**2) @@ -321,7 +322,7 @@ def estimate_psd(self, blocks): def radial_average(self, amplitude_spectrum): """ - Computes radial/elliptical average of the power spectrum + Computes radial average of the power spectrum. :param amplitude_spectrum: Power spectrum to average. :return: PSD and noise as 2-tuple of Image instances @@ -346,7 +347,7 @@ def radial_average(self, amplitude_spectrum): def elliptical_average(self, amplitude_spectrum): """ - Computes elliptical average of the power spectrum + Computes elliptical average of the power spectrum. :param amplitude_spectrum: Power spectrum to average. :return: PSD and noise as 2-tuple of Image instances. @@ -634,7 +635,7 @@ def _gd( :param df2: Defocus value in the direction perpendicular to df1. :param angle_ast: Angle between df1 and the x-axis, Radians. :param stop_cond: Optionally select value to stop iterations early. - Default `None` is 1e-10 for float32 `dtype` and 1e20 otherwise. + Default `None` is 1e-10 for float32 `dtype` and 1e-20 otherwise. :param max_iterations: Maximum number of inner gd steps. Default 400. :return: Optimal defocus parameters @@ -1041,7 +1042,7 @@ def estimate_ctf( save_ctf_images=False, save_noise_images=False, rglob=r"*.mrc", - solver_method="gd", + solver_method="bfgs", solver_opts=None, linprog_method="highs-ds", ): @@ -1060,17 +1061,12 @@ def estimate_ctf( :param num_tapers: Optional, number of tapers to apply in PSD estimation. Defaults to 4. :param psd_size: Block size (in pixels) for PSD estimation. - :param g_min: Inverse of minimun resolution for PSD. Default 30. + :param g_min: Inverse of minimum resolution for PSD. Default 30. :param g_max: Inverse of maximum resolution for PSD. Default 5. :param min_defocus: Start of defocus loop scan (Angstrom). Default 5000. :param max_defocus: End of defocus loop scan (Angstrom). Default 100000. :param step_defocus: Step size of defocus loop scan (Angstrom). Default 10. :param basis: Optional user provided SteerableBasis2D. Defaults to `FFBBasis2D`. - :param linprog_method: Method passed to linear program solver (scipy.optimize.linprog). - Defaults to 'highs-ds'. - :param diagnostic_plots: Optionally generate diagnostic_plots. - Note, these are blocking plot calls. - Defaults to False. :param rglob: File search pattern. By default will search for **/*.mrc files. Using "*.mrc[s]" will additonally include mrcs files, etc. @@ -1105,10 +1101,10 @@ def estimate_ctf( # Setup a CTFEstimator ----------------------------------------------------- ctf_est = CtfEstimator( - pixel_size, - cs, - amplitude_contrast, - voltage, + pixel_size=pixel_size, + cs=cs, + amplitude_contrast=amplitude_contrast, + voltage=voltage, psd_size=psd_size, num_tapers=num_tapers, g_min=g_min, @@ -1118,6 +1114,9 @@ def estimate_ctf( step_defocus=step_defocus, basis=basis, diagnostic_plots=False, + solver_method=solver_method, + solver_opts=solver_opts, + linprog_method=linprog_method, dtype=dtype, ) From 3da767877fd32d2340af30ae1f36cf411da9da55 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 26 Oct 2026 11:27:38 -0400 Subject: [PATCH 44/44] fix the implicitly hardcoded 10 step size --- src/aspire/ctf/ctf_estimator.py | 18 ++++++++++-------- tests/test_CtfEstimate.py | 1 + 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py index a3f1307d5a..b23d37efbb 100644 --- a/src/aspire/ctf/ctf_estimator.py +++ b/src/aspire/ctf/ctf_estimator.py @@ -175,7 +175,6 @@ def __init__( # Relating to MATLAB code, NW:= R * N / 2 = (2*L)/N * N / 2 = L tapers_1d = self.tapers(psd_size, num_tapers, num_tapers) self.tapers_1d = xp.asarray(tapers_1d, dtype=self.dtype) - breakpoint() def preprocess_micrograph(self, micrograph): """ @@ -495,7 +494,7 @@ def opt1d( Find optimal defocus for the radially symmetric case (where no astigmatism is present) :param amplitude_spectrum: Estimated power specrtum. - :return: 2-tuple of NumPy arrays (Estimated average of defocus and low_freq_cutoff) + :return: 2-tuple of NumPy arrays (Estimated average of defocus (nm), low_freq_cutoff) """ signal = xp.asarray(amplitude_spectrum.T) @@ -543,10 +542,13 @@ def opt1d( c = c / xp.sqrt(Sxx[None, :] * Syy) max_c = xp.argmax(c).item() - avg_defocus, low_freq_cutoff_col = np.unravel_index(max_c, c.shape)[:2] - avg_defocus += self.min_defocus / 10 # (nm) + avg_defocus_steps, low_freq_cutoff_col = np.unravel_index(max_c, c.shape)[:2] + # Convert from number of steps to nm defocus units + avg_defocus_nm = ( + self.min_defocus + avg_defocus_steps * self.step_defocus + ) / 10 # (nm) - return avg_defocus, low_freq_cutoff_col + return avg_defocus_nm, low_freq_cutoff_col def background_subtract_2d(self, signal, background_p1, max_col): """ @@ -958,7 +960,7 @@ def estimate(self, micrograph): self._plot_1d_psd(signal_1d, background_1d) self._plot_1d_bg_sub_psd(signal_1d) - avg_defocus, low_freq_skip_col = self.opt1d(signal_1d) + avg_defocus_nm, low_freq_skip_col = self.opt1d(signal_1d) signal, background_2d = self.background_subtract_2d( signal_observed, background_1d, low_freq_skip_col @@ -974,8 +976,8 @@ def estimate(self, micrograph): background_2d = background_2d + additional_background - initial_df1 = (avg_defocus * 2) / (1 + ratio) - initial_df2 = (avg_defocus * 2) - initial_df1 + initial_df1 = (avg_defocus_nm * 2) / (1 + ratio) + initial_df2 = (avg_defocus_nm * 2) - initial_df1 defocus_u, defocus_v, defocus_ang = self.maximize_correlation( signal, initial_df1, initial_df2 diff --git a/tests/test_CtfEstimate.py b/tests/test_CtfEstimate.py index 7046bf381e..dabd36735f 100644 --- a/tests/test_CtfEstimate.py +++ b/tests/test_CtfEstimate.py @@ -127,6 +127,7 @@ def test_customized_estimate_CTF(): "stop_cond": 1e-8, }, linprog_method="highs", + step_defocus=25, ) return test_estimate_CTF(basis, **customized)