Estimate Shifts w/ Symmetry

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
@@ -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.
+        # 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
-        n_equations = self._estimate_num_shift_equations(n_img)
+        # 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
 
         # Allocate local variables for estimating 2D shifts based on the estimated number
         # of equations. The shift equations are represented using a sparse matrix,
@@ -290,86 +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
-        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).
+        # Generate base [i, j] image pairs before symmetry expansion.
+        idx_i, idx_j = self._generate_index_pairs(n_pair_equations)
+
+        # 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
+
+            # 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
+            )
+
+            # 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
-            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
+            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 of common ray in image i
+            # Angle(s) of common ray(s) in image i
             shift_alpha = c_ij * d_theta
-            # Angle of common ray in image j.
+            # Angle(s) of common ray(s) 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),
-                    ]
+            # 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]),
                 )
-            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(
@@ -458,6 +465,29 @@ def _get_cl_indices(self, rotations, i, j, n_theta):
 
         return c_ij, c_ji
 
+    def _get_cl_indices_from_rot_pairs(self, Ri, Rjs, n_theta):
+        """
+        Get common-line indices for one rotation Ri and multiple rotations Rjs.
+        """
+        ell = 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
+
     def _apply_filter_and_norm(self, subscripts, pf, r_max, h):
         """
         Apply common line filter and normalize each ray

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)