From 1f303a2f38afbd0d33d6be8bc58ff4af84543c44 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 7 Jul 2026 10:41:00 -0400 Subject: [PATCH 1/5] Initial implementation of adding symmetry to estimate_shifts. It's slow. Need to vectorize. --- src/aspire/abinitio/commonline_base.py | 184 ++++++++++++++----------- 1 file changed, 106 insertions(+), 78 deletions(-) diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index 94c41d0e8f..69263559a2 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -259,14 +259,19 @@ def _get_shift_equations_approx(self): n_theta_half = self.n_theta // 2 n_img = self.n_img + pf = self.pf.copy() # `estimate_shifts()` requires that rotations have already been estimated. - rotations = Rotation(self.rotations) + rotations = self.rotations - pf = self.pf.copy() + # Apply symmetry group to rotations + sym_rots = self.src.symmetry_group.matrices.astype(self.dtype, copy=False) + n_sym = len(sym_rots) - # Estimate number of equations that will be used to calculate the shifts - n_equations = self._estimate_num_shift_equations(n_img) + # Estimate number of equations that will be used to calculate the shifts, + # taking into account particle symmetry. + n_pair_equations = self._estimate_num_shift_equations(n_img) + n_equations = n_pair_equations * n_sym # Allocate local variables for estimating 2D shifts based on the estimated number # of equations. The shift equations are represented using a sparse matrix, @@ -291,86 +296,94 @@ def _get_shift_equations_approx(self): d_theta = np.pi / n_theta_half # Generate two index lists for [i, j] pairs of images - idx_i, idx_j = self._generate_index_pairs(n_equations) + idx_i, idx_j = self._generate_index_pairs(n_pair_equations) # Go through all shift equations in the size of n_equations # Iterate over the common lines pairs and for each pair find the 1D # relative shift between the two Fourier lines in the pair. - for shift_eq_idx in range(n_equations): - i = idx_i[shift_eq_idx] - j = idx_j[shift_eq_idx] - # get the common line indices based on the rotations from i and j images - c_ij, c_ji = self._get_cl_indices(rotations, i, j, n_theta_half) - - # Extract the Fourier rays that correspond to the common line - pf_i = pf[i, c_ij] - - # Check whether need to flip or not Fourier ray of j image - # Is the common line in image j in the positive - # direction of the ray (is_pf_j_flipped=False) or in the - # negative direction (is_pf_j_flipped=True). - is_pf_j_flipped = c_ji >= n_theta_half - if not is_pf_j_flipped: - pf_j = pf[j, c_ji] - else: - pf_j = pf[j, c_ji - n_theta_half] - - # Use ray from opposite side of origin. - # Correpsonds to `freqs` convention in PFT, - # where the legacy code used a negated frequency grid. - pf_i, pf_j = np.conj(pf_i), np.conj(pf_j) - - # perform bandpass filter, normalize each ray of each image, - pf_i = self._apply_filter_and_norm("i, i -> i", pf_i, r_max, h) - pf_j = self._apply_filter_and_norm("i, i -> i", pf_j, r_max, h) - - # apply the shifts to images - pf_i_flipped = np.conj(pf_i) - pf_i_stack = pf_i[:, None] * shift_phases.T - pf_i_flipped_stack = pf_i_flipped[:, None] * shift_phases.T - - c1 = 2 * np.dot(pf_i_stack.T.conj(), pf_j).real - c2 = 2 * np.dot(pf_i_flipped_stack.T.conj(), pf_j).real - - # find the indices for the maximum values - # and apply corresponding shifts - sidx1 = np.argmax(c1) - sidx2 = np.argmax(c2) - sidx = sidx1 if c1[sidx1] > c2[sidx2] else sidx2 - dx = -self.offsets_max_shift + sidx * self.offsets_shift_step - - # angle of common ray in image i - shift_alpha = c_ij * d_theta - # Angle of common ray in image j. - shift_beta = c_ji * d_theta - # Row index to construct the sparse equations - shift_i[shift_eq_idx] = shift_eq_idx - # Columns of the shift variables that correspond to the current pair [i, j] - shift_j[shift_eq_idx] = [2 * i, 2 * i + 1, 2 * j, 2 * j + 1] - # Right hand side of the current equation - shift_b[shift_eq_idx] = dx - - # Compute the coefficients of the current equation - if not is_pf_j_flipped: - shift_eq[shift_eq_idx] = np.array( - [ - np.sin(shift_alpha), - np.cos(shift_alpha), - -np.sin(shift_beta), - -np.cos(shift_beta), - ] - ) - else: - shift_beta = shift_beta - np.pi - shift_eq[shift_eq_idx] = np.array( - [ - -np.sin(shift_alpha), - -np.cos(shift_alpha), - -np.sin(shift_beta), - -np.cos(shift_beta), - ] + for pair_eq_idx in range(n_pair_equations): + i = idx_i[pair_eq_idx] + j = idx_j[pair_eq_idx] + + for sym_idx, g in enumerate(sym_rots): + shift_eq_idx = pair_eq_idx + sym_idx * n_pair_equations + + # get the common line indices based on the rotations from i and j images + c_ij, c_ji = self._get_cl_indices_from_rot_pair( + rotations[i], + g @ rotations[j], + n_theta_half, ) + # Extract the Fourier rays that correspond to the common line + pf_i = pf[i, c_ij] + + # Check whether need to flip or not Fourier ray of j image + # Is the common line in image j in the positive + # direction of the ray (is_pf_j_flipped=False) or in the + # negative direction (is_pf_j_flipped=True). + is_pf_j_flipped = c_ji >= n_theta_half + if not is_pf_j_flipped: + pf_j = pf[j, c_ji] + else: + pf_j = pf[j, c_ji - n_theta_half] + + # Use ray from opposite side of origin. + # Correpsonds to `freqs` convention in PFT, + # where the legacy code used a negated frequency grid. + pf_i, pf_j = np.conj(pf_i), np.conj(pf_j) + + # perform bandpass filter, normalize each ray of each image, + pf_i = self._apply_filter_and_norm("i, i -> i", pf_i, r_max, h) + pf_j = self._apply_filter_and_norm("i, i -> i", pf_j, r_max, h) + + # apply the shifts to images + pf_i_flipped = np.conj(pf_i) + pf_i_stack = pf_i[:, None] * shift_phases.T + pf_i_flipped_stack = pf_i_flipped[:, None] * shift_phases.T + + c1 = 2 * np.dot(pf_i_stack.T.conj(), pf_j).real + c2 = 2 * np.dot(pf_i_flipped_stack.T.conj(), pf_j).real + + # find the indices for the maximum values + # and apply corresponding shifts + sidx1 = np.argmax(c1) + sidx2 = np.argmax(c2) + sidx = sidx1 if c1[sidx1] > c2[sidx2] else sidx2 + dx = -self.offsets_max_shift + sidx * self.offsets_shift_step + + # angle of common ray in image i + shift_alpha = c_ij * d_theta + # Angle of common ray in image j. + shift_beta = c_ji * d_theta + # Row index to construct the sparse equations + shift_i[shift_eq_idx] = shift_eq_idx + # Columns of the shift variables that correspond to the current pair [i, j] + shift_j[shift_eq_idx] = [2 * i, 2 * i + 1, 2 * j, 2 * j + 1] + # Right hand side of the current equation + shift_b[shift_eq_idx] = dx + + # Compute the coefficients of the current equation + if not is_pf_j_flipped: + shift_eq[shift_eq_idx] = np.array( + [ + np.sin(shift_alpha), + np.cos(shift_alpha), + -np.sin(shift_beta), + -np.cos(shift_beta), + ] + ) + else: + shift_beta = shift_beta - np.pi + shift_eq[shift_eq_idx] = np.array( + [ + -np.sin(shift_alpha), + -np.cos(shift_alpha), + -np.sin(shift_beta), + -np.cos(shift_beta), + ] + ) + # create sparse matrix object only containing non-zero elements shift_equations = sparse.csr_matrix( (shift_eq.flatten(), (shift_i.flatten(), shift_j.flatten())), @@ -458,6 +471,21 @@ def _get_cl_indices(self, rotations, i, j, n_theta): return c_ij, c_ji + def _get_cl_indices_from_rot_pair(self, Ri, Rj, n_theta): + """ + Get common-line indices for an explicit pair of rotation matrices. + """ + rotations = Rotation(np.stack((Ri, Rj))).invert() + c_ij, c_ji = rotations.common_lines(0, 1, 2 * n_theta) + + if c_ij >= n_theta: + c_ij -= n_theta + c_ji -= n_theta + if c_ji < 0: + c_ji += 2 * n_theta + + return c_ij, c_ji + def _apply_filter_and_norm(self, subscripts, pf, r_max, h): """ Apply common line filter and normalize each ray From c54bbd8ae9dcafda407ec7938442e21af903f61b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 8 Jul 2026 10:21:15 -0400 Subject: [PATCH 2/5] Vectorize estimate_shifts to remove symmetry loop --- src/aspire/abinitio/commonline_base.py | 184 +++++++++++++------------ 1 file changed, 93 insertions(+), 91 deletions(-) diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index 69263559a2..416aef775d 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -264,12 +264,12 @@ def _get_shift_equations_approx(self): # `estimate_shifts()` requires that rotations have already been estimated. rotations = self.rotations - # Apply symmetry group to rotations + # Apply symmetry group to rotations. + # Symmetry copies add more equations for the same per-image shift unknowns. sym_rots = self.src.symmetry_group.matrices.astype(self.dtype, copy=False) n_sym = len(sym_rots) - # Estimate number of equations that will be used to calculate the shifts, - # taking into account particle symmetry. + # Estimate base image-pair equations, then expand each pair by symmetry. n_pair_equations = self._estimate_num_shift_equations(n_img) n_equations = n_pair_equations * n_sym @@ -295,94 +295,88 @@ def _get_shift_equations_approx(self): d_theta = np.pi / n_theta_half - # Generate two index lists for [i, j] pairs of images + # Generate base [i, j] image pairs before symmetry expansion. idx_i, idx_j = self._generate_index_pairs(n_pair_equations) - # Go through all shift equations in the size of n_equations - # Iterate over the common lines pairs and for each pair find the 1D - # relative shift between the two Fourier lines in the pair. + # Filter, normalize, and conjugate all rays once instead of once per equation. + # Conjugation uses ray from opposite side of origin. + # Correpsonds to `freqs` convention in PFT, + # where the legacy code used a negated frequency grid. + pf = np.conj(self._apply_filter_and_norm("ijk, k -> ijk", pf, r_max, h)) + + # Iterate over image pairs; each iteration fills one block of n_sym + # symmetry-induced common-line shift equations. for pair_eq_idx in range(n_pair_equations): i = idx_i[pair_eq_idx] j = idx_j[pair_eq_idx] + rows = pair_eq_idx + np.arange(n_sym) * n_pair_equations - for sym_idx, g in enumerate(sym_rots): - shift_eq_idx = pair_eq_idx + sym_idx * n_pair_equations + # Common lines for Ri against all symmetry copies g @ Rj. + Rjs = sym_rots @ rotations[j] + c_ij, c_ji = self._get_cl_indices_from_rot_pairs( + rotations[i], Rjs, n_theta_half + ) - # get the common line indices based on the rotations from i and j images - c_ij, c_ji = self._get_cl_indices_from_rot_pair( - rotations[i], - g @ rotations[j], - n_theta_half, + # Extract the Fourier rays that correspond to the common lines + pf_i = pf[i, c_ij] # shape (n_sym, n_rad) + + # Track which symmetry-induced rays in image j use the opposite ray direction. + is_pf_j_flipped = c_ji >= n_theta_half + pf_j = pf[j, c_ji % n_theta_half] + + # Apply candidate 1D shifts to all symmetry-induced rays for this image pair. + pf_i_stack = pf_i[:, :, None] * shift_phases.T[None, :, :] + pf_i_flipped_stack = np.conj(pf_i)[:, :, None] * shift_phases.T[None, :, :] + + c1 = 2 * np.sum(pf_i_stack.conj() * pf_j[:, :, None], axis=1).real + c2 = 2 * np.sum(pf_i_flipped_stack.conj() * pf_j[:, :, None], axis=1).real + + # Pick the best candidate shift for each symmetry-induced ray pair. + sidx1 = np.argmax(c1, axis=1) + sidx2 = np.argmax(c2, axis=1) + + score1 = c1[np.arange(n_sym), sidx1] + score2 = c2[np.arange(n_sym), sidx2] + sidx = np.where(score1 > score2, sidx1, sidx2) + dx = -self.offsets_max_shift + sidx * self.offsets_shift_step + + # Angle(s) of common ray(s) in image i + shift_alpha = c_ij * d_theta + # Angle(s) of common ray(s) in image j + shift_beta = c_ji * d_theta + # Row indices to construct the sparse equations + shift_i[rows] = rows[:, None] + # All symmetry rows for this pair use the same set of image shift unknowns. + shift_j[rows] = [2 * i, 2 * i + 1, 2 * j, 2 * j + 1] + # Right hand side of the current equation(s) + shift_b[rows] = dx + + # Initialize shift equation block. + # One four-coefficient equation row per symmetry-induced common line. + eq = np.empty((n_sym, 4), dtype=self.dtype) + + # Compute the coefficients of the current block of equations. + not_flipped = ~is_pf_j_flipped + eq[not_flipped] = np.column_stack( + ( + np.sin(shift_alpha[not_flipped]), + np.cos(shift_alpha[not_flipped]), + -np.sin(shift_beta[not_flipped]), + -np.cos(shift_beta[not_flipped]), ) + ) - # Extract the Fourier rays that correspond to the common line - pf_i = pf[i, c_ij] - - # Check whether need to flip or not Fourier ray of j image - # Is the common line in image j in the positive - # direction of the ray (is_pf_j_flipped=False) or in the - # negative direction (is_pf_j_flipped=True). - is_pf_j_flipped = c_ji >= n_theta_half - if not is_pf_j_flipped: - pf_j = pf[j, c_ji] - else: - pf_j = pf[j, c_ji - n_theta_half] - - # Use ray from opposite side of origin. - # Correpsonds to `freqs` convention in PFT, - # where the legacy code used a negated frequency grid. - pf_i, pf_j = np.conj(pf_i), np.conj(pf_j) - - # perform bandpass filter, normalize each ray of each image, - pf_i = self._apply_filter_and_norm("i, i -> i", pf_i, r_max, h) - pf_j = self._apply_filter_and_norm("i, i -> i", pf_j, r_max, h) - - # apply the shifts to images - pf_i_flipped = np.conj(pf_i) - pf_i_stack = pf_i[:, None] * shift_phases.T - pf_i_flipped_stack = pf_i_flipped[:, None] * shift_phases.T - - c1 = 2 * np.dot(pf_i_stack.T.conj(), pf_j).real - c2 = 2 * np.dot(pf_i_flipped_stack.T.conj(), pf_j).real - - # find the indices for the maximum values - # and apply corresponding shifts - sidx1 = np.argmax(c1) - sidx2 = np.argmax(c2) - sidx = sidx1 if c1[sidx1] > c2[sidx2] else sidx2 - dx = -self.offsets_max_shift + sidx * self.offsets_shift_step - - # angle of common ray in image i - shift_alpha = c_ij * d_theta - # Angle of common ray in image j. - shift_beta = c_ji * d_theta - # Row index to construct the sparse equations - shift_i[shift_eq_idx] = shift_eq_idx - # Columns of the shift variables that correspond to the current pair [i, j] - shift_j[shift_eq_idx] = [2 * i, 2 * i + 1, 2 * j, 2 * j + 1] - # Right hand side of the current equation - shift_b[shift_eq_idx] = dx - - # Compute the coefficients of the current equation - if not is_pf_j_flipped: - shift_eq[shift_eq_idx] = np.array( - [ - np.sin(shift_alpha), - np.cos(shift_alpha), - -np.sin(shift_beta), - -np.cos(shift_beta), - ] - ) - else: - shift_beta = shift_beta - np.pi - shift_eq[shift_eq_idx] = np.array( - [ - -np.sin(shift_alpha), - -np.cos(shift_alpha), - -np.sin(shift_beta), - -np.cos(shift_beta), - ] - ) + beta_flipped = shift_beta[is_pf_j_flipped] - np.pi + eq[is_pf_j_flipped] = np.column_stack( + ( + -np.sin(shift_alpha[is_pf_j_flipped]), + -np.cos(shift_alpha[is_pf_j_flipped]), + -np.sin(beta_flipped), + -np.cos(beta_flipped), + ) + ) + + shift_eq[rows] = eq # create sparse matrix object only containing non-zero elements shift_equations = sparse.csr_matrix( @@ -471,18 +465,26 @@ def _get_cl_indices(self, rotations, i, j, n_theta): return c_ij, c_ji - def _get_cl_indices_from_rot_pair(self, Ri, Rj, n_theta): + def _get_cl_indices_from_rot_pairs(self, Ri, Rjs, n_theta): """ - Get common-line indices for an explicit pair of rotation matrices. + Get common-line indices for one rotation Ri and multiple rotations Rjs. """ - rotations = Rotation(np.stack((Ri, Rj))).invert() - c_ij, c_ji = rotations.common_lines(0, 1, 2 * n_theta) + ell = 2 * n_theta - if c_ij >= n_theta: - c_ij -= n_theta - c_ji -= n_theta - if c_ji < 0: - c_ji += 2 * n_theta + # Match _get_cl_indices, which calls + # Rotation(np.stack((Ri, Rj))).invert().common_lines(i, j, ...). + ut = np.swapaxes(Rjs, -1, -2) @ Ri + + alpha_ij = np.arctan2(ut[:, 2, 0], -ut[:, 2, 1]) + np.pi + alpha_ji = np.arctan2(-ut[:, 0, 2], ut[:, 1, 2]) + np.pi + + c_ij = np.mod(np.round(alpha_ij * ell / (2 * np.pi)), ell).astype(int) + c_ji = np.mod(np.round(alpha_ji * ell / (2 * np.pi)), ell).astype(int) + + mask = c_ij >= n_theta + c_ij[mask] -= n_theta + c_ji[mask] -= n_theta + c_ji[c_ji < 0] += ell return c_ij, c_ji From f68a4984968a251ce55ae25ad83ddea04c92362f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 8 Jul 2026 12:01:02 -0400 Subject: [PATCH 3/5] tox --- src/aspire/abinitio/commonline_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index 416aef775d..2ad9529132 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -6,7 +6,7 @@ from aspire.image import Image from aspire.operators import PolarFT -from aspire.utils import Rotation, fuzzy_mask +from aspire.utils import fuzzy_mask from aspire.utils.random import choice from .commonline_utils import _generate_shift_phase_and_filter From 671e91dd9e942391d049bb85612a375b9290b67a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 9 Jul 2026 16:09:58 -0400 Subject: [PATCH 4/5] Add estimate_shifts test parametrized over symmetry groups --- tests/test_estimate_shifts.py | 162 ++++++++++++++++++++++++++++++++++ 1 file changed, 162 insertions(+) create mode 100644 tests/test_estimate_shifts.py diff --git a/tests/test_estimate_shifts.py b/tests/test_estimate_shifts.py new file mode 100644 index 0000000000..4c80b57039 --- /dev/null +++ b/tests/test_estimate_shifts.py @@ -0,0 +1,162 @@ +import numpy as np +import pytest +from scipy import sparse + +from aspire.abinitio import Orient3D +from aspire.source import Simulation +from aspire.volume import ( + AsymmetricVolume, + CnSymmetricVolume, + DnSymmetricVolume, + OSymmetricVolume, + TSymmetricVolume, +) + +DTYPES = [np.float64] +RES = [89] +SYMMETRIES = [ + None, + "C2", + "C3", + "C4", + "C5", + "C6", + "D2", + "D3", + "D4", + "D5", + "D6", + "D7", + "T", + "O", +] +N_IMGS = 100 +SEED = 1980 + + +@pytest.fixture(params=RES, ids=lambda x: f"resolution={x}", scope="module") +def resolution(request): + return request.param + + +@pytest.fixture(params=SYMMETRIES, ids=lambda x: f"symmetry={x}", scope="module") +def symmetry(request): + return request.param + + +@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module") +def dtype(request): + return request.param + + +@pytest.fixture(scope="module") +def volume(resolution, symmetry, dtype): + if symmetry is None: + return AsymmetricVolume( + L=resolution, C=1, K=25, dtype=dtype, seed=SEED + ).generate() + + if symmetry.startswith("C"): + order = int(symmetry[1:]) + return CnSymmetricVolume( + L=resolution, C=1, order=order, K=25, dtype=dtype, seed=SEED + ).generate() + + if symmetry.startswith("D"): + order = int(symmetry[1:]) + return DnSymmetricVolume( + L=resolution, C=1, order=order, K=25, dtype=dtype, seed=SEED + ).generate() + + if symmetry == "T": + return TSymmetricVolume( + L=resolution, C=1, K=25, dtype=dtype, seed=SEED + ).generate() + + if symmetry == "O": + return OSymmetricVolume( + L=resolution, C=1, K=25, dtype=dtype, seed=SEED + ).generate() + + +@pytest.fixture(scope="module") +def estimator(volume): + """ + Build a simulated source and use ground-truth rotations so this test isolates + shift-equation construction and shift recovery from orientation estimation error. + """ + offset_scale = 1.5 # standard deviation of shifts + offsets = np.random.normal(scale=offset_scale, size=(N_IMGS, 2)) + + src = Simulation( + n=N_IMGS, + vols=volume, + amplitudes=1, + offsets=offsets, + seed=SEED, + ).cache() + + orient_est = Orient3D(src) + orient_est.rotations = src.rotations + + return orient_est + + +def test_estimate_shifts(estimator): + """ + Compare estimated shifts to ground truth after removing the nullspace of + the shift equation matrix. See the following publication for more info on + measuring shift estimation error: + + Y. Shkolnisky and A. Singer, + Center of Mass Operators for Cryo-EM - Theory and Implementation, + Modeling Nanoscale Imaging in Electron Microscopy, + T. Vogt, W. Dahmen, and P. Binev (Eds.) + Nanostructure Science and Technology Series, + Springer, 2012, pp. 147–177 + """ + # Build the sparse common-line shift system Ax = b and solve it directly, + # matching the solver used by estimate_shifts(). + A, b = estimator._get_shift_equations_approx() + lsqr_result = sparse.linalg.lsqr(A, b, atol=1e-8, btol=1e-8, iter_lim=100) + x_est = lsqr_result[0] + + # Convert Simulation offsets to the internal LSQR convention: + # estimate_shifts returns -x_est.reshape(n, 2)[:, ::-1]. + x_ref_internal = (-estimator.src.offsets[:, ::-1]).reshape(-1) + + # Use the SVD to separate the constrained directions from the nullspace, + # which corresponds to global 3D translation ambiguity. + _, s, Vt = np.linalg.svd(A.toarray(), full_matrices=False) + + # Estimate the effective rank of A and keep the constrained directions. + sv_tol = 1e-2 + rank = int(np.sum(s > sv_tol * s[0])) + V_nonnull = Vt[:rank].T + + # Compute relative error after projecting out the nullspace. + num = np.linalg.norm(V_nonnull.T @ (x_ref_internal - x_est)) + den = np.linalg.norm(V_nonnull.T @ x_ref_internal) + projected_rel_err = num / den + + # Check the shift error is within 15% of the reference shift norm. + np.testing.assert_array_less(projected_rel_err, 0.15) + + # The projected relative error follows the legacy diagnostic, but it is not + # a pixel-scale quantity. Below we check the same solution after aligning away + # the nullspace component so the error is easier to interpret. + V_null = Vt[rank:].T + + # Add the nullspace component to the estimate before comparing directly + # against the reference shifts. + x_err = x_ref_internal - x_est + x_est_aligned = x_est + V_null @ (V_null.T @ x_err) + + # Convert back to ASPIRE shift convention and compute per-image Euclidean + # shift error in pixels. + est_shifts_aligned = -x_est_aligned.reshape(estimator.src.n, 2)[:, ::-1] + per_img_err = np.linalg.norm(estimator.src.offsets - est_shifts_aligned, axis=1) + mean_aligned_px = per_img_err.mean() + + # Check that aligned estimate errors are within 0.25 pixels on average. + np.testing.assert_array_less(mean_aligned_px, 0.25) From dc9a73c10cbc431db71285b3a98320d072983bad Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 28 Jul 2026 13:52:59 -0400 Subject: [PATCH 5/5] Leave legacy estimate_shifts. Separate code path for symmetric shift estimation. --- src/aspire/abinitio/commonline_base.py | 166 ++++++++++++++++++++++++- tests/test_estimate_shifts.py | 2 +- 2 files changed, 161 insertions(+), 7 deletions(-) diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index 2ad9529132..a97ad7ddde 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -6,7 +6,7 @@ from aspire.image import Image from aspire.operators import PolarFT -from aspire.utils import fuzzy_mask +from aspire.utils import Rotation, fuzzy_mask from aspire.utils.random import choice from .commonline_utils import _generate_shift_phase_and_filter @@ -212,7 +212,7 @@ def estimate_shifts(self): """ # Generate approximated shift equations from estimated rotations - shift_equations, shift_b = self._get_shift_equations_approx() + shift_equations, shift_b = self._get_shift_equations() # Solve the linear equation, optionally printing numerical debug details. show = False @@ -241,6 +241,22 @@ def estimate(self, **kwargs): return self.rotations, self.shifts + def _get_shift_equations(self): + """ + Generate shift equations from the estimated rotations. + + Dispatches to the legacy asymmetric shift-equation construction for C1 + sources, and to the symmetry-expanded construction for sources with + nontrivial symmetry. This keeps the C1 code path unchanged while allowing + symmetric molecules to contribute multiple common-line equations per image + pair without introducing additional shift unknowns. + + :return: The sparse shift-equation matrix and right-hand side vector. + """ + if str(self.src.symmetry_group) == "C1": + return self._get_shift_equations_approx() + return self._get_shift_equations_approx_symmetric() + def _get_shift_equations_approx(self): """ Generate approximated shift equations from estimated rotations @@ -257,6 +273,142 @@ def _get_shift_equations_approx(self): :return; The left and right-hand side of shift equations """ + n_theta_half = self.n_theta // 2 + n_img = self.n_img + + # `estimate_shifts()` requires that rotations have already been estimated. + rotations = Rotation(self.rotations) + + pf = self.pf.copy() + + # Estimate number of equations that will be used to calculate the shifts + n_equations = self._estimate_num_shift_equations(n_img) + + # Allocate local variables for estimating 2D shifts based on the estimated number + # of equations. The shift equations are represented using a sparse matrix, + # since each row in the system contains four non-zeros (as it involves + # exactly four unknowns). The variables below are used to construct + # this sparse system. The k'th non-zero element of the equations matrix + # is stored at index (shift_i(k),shift_j(k)). + shift_i = np.zeros((n_equations, 4), dtype=self.dtype) + shift_j = np.zeros((n_equations, 4), dtype=self.dtype) + shift_eq = np.zeros((n_equations, 4), dtype=self.dtype) + shift_b = np.zeros(n_equations, dtype=self.dtype) + + # Prepare the shift phases to try and generate filter for common-line detection + # The shift phases are pre-defined in a range of max_shift that can be + # applied to maximize the common line calculation. The common-line filter + # is also applied to the radial direction for easier detection. + r_max = pf.shape[2] + _, shift_phases, h = _generate_shift_phase_and_filter( + r_max, self.offsets_max_shift, self.offsets_shift_step, self.dtype + ) + + d_theta = np.pi / n_theta_half + + # Generate two index lists for [i, j] pairs of images + idx_i, idx_j = self._generate_index_pairs(n_equations) + + # Go through all shift equations in the size of n_equations + # Iterate over the common lines pairs and for each pair find the 1D + # relative shift between the two Fourier lines in the pair. + for shift_eq_idx in range(n_equations): + i = idx_i[shift_eq_idx] + j = idx_j[shift_eq_idx] + # get the common line indices based on the rotations from i and j images + c_ij, c_ji = self._get_cl_indices(rotations, i, j, n_theta_half) + + # Extract the Fourier rays that correspond to the common line + pf_i = pf[i, c_ij] + + # Check whether need to flip or not Fourier ray of j image + # Is the common line in image j in the positive + # direction of the ray (is_pf_j_flipped=False) or in the + # negative direction (is_pf_j_flipped=True). + is_pf_j_flipped = c_ji >= n_theta_half + if not is_pf_j_flipped: + pf_j = pf[j, c_ji] + else: + pf_j = pf[j, c_ji - n_theta_half] + + # Use ray from opposite side of origin. + # Correpsonds to `freqs` convention in PFT, + # where the legacy code used a negated frequency grid. + pf_i, pf_j = np.conj(pf_i), np.conj(pf_j) + + # perform bandpass filter, normalize each ray of each image, + pf_i = self._apply_filter_and_norm("i, i -> i", pf_i, r_max, h) + pf_j = self._apply_filter_and_norm("i, i -> i", pf_j, r_max, h) + + # apply the shifts to images + pf_i_flipped = np.conj(pf_i) + pf_i_stack = pf_i[:, None] * shift_phases.T + pf_i_flipped_stack = pf_i_flipped[:, None] * shift_phases.T + + c1 = 2 * np.dot(pf_i_stack.T.conj(), pf_j).real + c2 = 2 * np.dot(pf_i_flipped_stack.T.conj(), pf_j).real + + # find the indices for the maximum values + # and apply corresponding shifts + sidx1 = np.argmax(c1) + sidx2 = np.argmax(c2) + sidx = sidx1 if c1[sidx1] > c2[sidx2] else sidx2 + dx = -self.offsets_max_shift + sidx * self.offsets_shift_step + + # angle of common ray in image i + shift_alpha = c_ij * d_theta + # Angle of common ray in image j. + shift_beta = c_ji * d_theta + # Row index to construct the sparse equations + shift_i[shift_eq_idx] = shift_eq_idx + # Columns of the shift variables that correspond to the current pair [i, j] + shift_j[shift_eq_idx] = [2 * i, 2 * i + 1, 2 * j, 2 * j + 1] + # Right hand side of the current equation + shift_b[shift_eq_idx] = dx + + # Compute the coefficients of the current equation + if not is_pf_j_flipped: + shift_eq[shift_eq_idx] = np.array( + [ + np.sin(shift_alpha), + np.cos(shift_alpha), + -np.sin(shift_beta), + -np.cos(shift_beta), + ] + ) + else: + shift_beta = shift_beta - np.pi + shift_eq[shift_eq_idx] = np.array( + [ + -np.sin(shift_alpha), + -np.cos(shift_alpha), + -np.sin(shift_beta), + -np.cos(shift_beta), + ] + ) + + # create sparse matrix object only containing non-zero elements + shift_equations = sparse.csr_matrix( + (shift_eq.flatten(), (shift_i.flatten(), shift_j.flatten())), + shape=(n_equations, 2 * n_img), + dtype=self.dtype, + ) + + return shift_equations, shift_b + + def _get_shift_equations_approx_symmetric(self): + """ + Generate symmetry-expanded approximate shift equations from estimated rotations. + + For each sampled image pair, this method computes the common lines induced by + the first image rotation and every symmetry-transformed copy of the second + image rotation. Each symmetry copy contributes one shift equation involving + the same two 2D image-shift unknowns, adding constraints without duplicating + images or introducing independent shift variables for symmetry copies. + + :return: The sparse shift-equation matrix and right-hand side vector. + """ + n_theta_half = self.n_theta // 2 n_img = self.n_img pf = self.pf.copy() @@ -270,7 +422,7 @@ def _get_shift_equations_approx(self): n_sym = len(sym_rots) # Estimate base image-pair equations, then expand each pair by symmetry. - n_pair_equations = self._estimate_num_shift_equations(n_img) + n_pair_equations = self._estimate_num_shift_equations(n_img, n_sym=n_sym) n_equations = n_pair_equations * n_sym # Allocate local variables for estimating 2D shifts based on the estimated number @@ -387,7 +539,7 @@ def _get_shift_equations_approx(self): return shift_equations, shift_b - def _estimate_num_shift_equations(self, n_img): + def _estimate_num_shift_equations(self, n_img, n_sym=1): """ Estimate total number of shift equations in images @@ -395,7 +547,9 @@ def _estimate_num_shift_equations(self, n_img): number of images and preselected memory factor. :param n_img: The total number of input images - :return: Estimated number of shift equations + :param n_sym: Number of symmetry-expanded rows generated per sampled image pair. + Defaults to 1 for the legacy asymmetric path. + :return: Number of base image-pair equations to sample before any symmetry expansion. """ # Number of equations that will be used to estimation the shifts n_equations_total = int(np.ceil(n_img * (self.n_check - 1) / 2)) @@ -404,7 +558,7 @@ def _estimate_num_shift_equations(self, n_img): # This ignores the sparsity of the system, since backslash seems to # ignore it. memory_total = self.offsets_equations_factor * ( - n_equations_total * 2 * n_img * self.dtype.itemsize + n_equations_total * n_sym * 2 * n_img * self.dtype.itemsize ) if memory_total < (self.offsets_max_memory * 10**6): diff --git a/tests/test_estimate_shifts.py b/tests/test_estimate_shifts.py index 4c80b57039..9630d111f3 100644 --- a/tests/test_estimate_shifts.py +++ b/tests/test_estimate_shifts.py @@ -117,7 +117,7 @@ def test_estimate_shifts(estimator): """ # Build the sparse common-line shift system Ax = b and solve it directly, # matching the solver used by estimate_shifts(). - A, b = estimator._get_shift_equations_approx() + A, b = estimator._get_shift_equations() lsqr_result = sparse.linalg.lsqr(A, b, atol=1e-8, btol=1e-8, iter_lim=100) x_est = lsqr_result[0]