[WIP] C3/C4 missing components

src/aspire/abinitio/commonline_c3_c4.py

@@ -5,6 +5,7 @@
 
 from aspire.abinitio import CLOrient3D, JSync
 from aspire.abinitio.sync_voting import _syncmatrix_ij_vote_3n
+from aspire.numeric import fft
 from aspire.operators import PolarFT
 from aspire.utils import J_conjugate, Rotation, all_pairs, anorm, tqdm, trange
 
@@ -279,6 +280,33 @@ def _self_clmatrix_c3_c4(self):
             shift, scl1, scl2 = np.unravel_index(np.argmax(np.real(corrs)), corrs.shape)
             sclmatrix[i] = [scl1, scl2]
 
+        # Correct C4 equator self-common-lines using the shared symmetry axis.
+        if order == 4:
+            # Detect equator images.
+            equator_indices = self._detect_equator_images()
+
+            if len(equator_indices) >= 2:
+                # Equator self-common-lines are 180 degrees apart and cannot
+                # be distinguished from conjugate symmetry. All equator views
+                # share the symmetry axis, including an image and its
+                # symmetry-related copies. Estimate this axis from each
+                # image's median common-line with the other equator views.
+                clmatrix_eq = self.clmatrix[equator_indices][:, equator_indices]
+                n_theta_half = n_theta // 2
+
+                for i, image_index in enumerate(equator_indices):
+                    # Exclude diagonal entries and missing common-lines.
+                    lines = clmatrix_eq[i]
+                    lines = lines[lines >= 0]  # diagonal vals are -1
+
+                    # Map common-lines into [0, 180 degrees) before
+                    # taking the median.
+                    line = np.median(lines % n_theta_half)
+                    sclmatrix[image_index] = [
+                        line,
+                        (line + n_theta_half) % n_theta,
+                    ]
+
         return sclmatrix
 
     def _estimate_all_Riis_c3_c4(self, sclmatrix):
@@ -437,3 +465,88 @@ def _local_J_sync_c3_c4(self, Rijs, Riis):
         pbar.close()
 
         return vijs, viis
+
+    def _detect_equator_images(self, max_shift=5, res_factor=10, fraction=0.1):
+        """
+        Detect C4 equator images using angular reflection symmetry.
+
+        Follow MATLAB's integer translation search, angular correlation,
+        binning, and selection, using ASPIRE's polar frequency conventions.
+
+        :param max_shift: Integer translation search radius in pixels.
+        :param res_factor: Number of angular samples per score bin.
+        :param fraction: Fraction of images selected as equator images.
+
+        :return: Selected image indices in descending score order.
+        """
+        pf = self.pf
+
+        # Build the polar frequency grid used to generate translation phases.
+        polar_ft = PolarFT(
+            (self.n_res, self.n_res),
+            nrad=self.n_rad,
+            ntheta=self.n_theta,
+            dtype=self.dtype,
+        )
+
+        # Search a square grid of integer (x, y) translations, as MATLAB
+        # does. The default radius of 5 gives 121 candidate translations.
+        # This search is independent of the common-line shift search.
+        shifts = np.array(
+            [
+                (x, y)
+                for x in range(-max_shift, max_shift + 1)
+                for y in range(-max_shift, max_shift + 1)
+            ],
+            dtype=self.dtype,
+        )
+
+        # Reshape the frequency grid into (coordinate, angle, radius),
+        # then remove DC to match the radial samples stored in self.pf.
+        freqs = polar_ft.freqs.reshape(2, self.n_theta // 2, self.n_rad)[:, :, 1:]
+
+        # Precompute the Fourier phase multiplier for every translation.
+        # Take the dot product of each translation with each frequency
+        # coordinate, producing shape (n_shifts, n_theta // 2, n_r - 1).
+        # Reversing (x, y) and using a positive exponential matches
+        # PolarFT.shift. These phases are reused for every image.
+        shift_phases = np.exp(1j * np.einsum("sa,atr->str", shifts[:, ::-1], freqs))
+
+        # Store each image's strongest reflection-correlation score.
+        scores = np.empty(self.n_img, dtype=np.float64)
+
+        for i in trange(self.n_img, desc="Detecting C4 equator images"):
+            # Apply shifts and reconstruct the full circle using conjugate symmetry.
+            # The result has shape (n_shifts, n_theta, n_r - 1).
+            shifted = PolarFT.half_to_full(shift_phases * pf[i][None])
+
+            # Normalize over angles.
+            shifted /= norm(shifted, axis=1, keepdims=True)
+
+            # Reverse angular order to form a reflected version.
+            reflected = shifted[:, ::-1, :]
+
+            # Compute circular cross-correlation along the angular axis:
+            # FFT(image) * conjugate(FFT(reflection)), followed by IFFT.
+            cross_spectrum = (
+                fft.fft(shifted, axis=1) * fft.fft(reflected, axis=1).conj()
+            )
+            correlations = fft.ifft(cross_spectrum, axis=1).real
+
+            # Sum correlation contributions over radial frequencies. Group
+            # neighboring angular offsets into fixed bins of res_factor samples.
+            # Keep the best score across all translations and angular bins.
+            angular_scores = correlations.sum(axis=2)
+            binned_scores = angular_scores.reshape(
+                len(shifts), self.n_theta // res_factor, res_factor
+            ).sum(axis=2)
+            scores[i] = binned_scores.max()
+
+        # Match MATLAB's selection policy: retain the top fraction by score.
+        # Note: this imposes no absolute quality threshold and selects this
+        # many images even if the dataset contains no true equator views.
+        count = int(np.floor(fraction * self.n_img))
+        self._equator_scores = scores
+        self._equator_indices = np.argsort(-scores, kind="stable")[:count]
+
+        return self._equator_indices.copy()

tests/test_orient_symmetric.py

@@ -479,6 +479,54 @@ def test_dtype_pass_through(dtype):
     assert src.dtype == cl_symm.dtype
 
 
+def test_detect_equator_images():
+    """
+    Check that _detect_equator_images detects most ground truth equator views.
+
+    The method is heuristic and selects the highest-scoring 10% of images
+    by default, so some missed detections and false positives are expected.
+    """
+    n_img = 200
+    L = 44
+    dtype = np.float32
+
+    vol = CnSymmetricVolume(
+        L=L,
+        C=1,
+        K=100,
+        order=4,
+        dtype=dtype,
+    ).generate()
+
+    src = Simulation(
+        vols=vol,
+        n=n_img,
+        C=1,
+        offsets=0,
+        amplitudes=1,
+        dtype=dtype,
+    )
+
+    # Find viewing directions within 5 degrees of the equator.
+    equator_indices_gt = np.flatnonzero(
+        np.abs(src.rotations[:, 2, 2]) < np.sin(np.deg2rad(5))
+    )
+    if len(equator_indices_gt) == 0:
+        pytest.skip("Simulation contains no equator views.")
+
+    # Detect equator images and check at least 80% are correctly identified.
+    cl_symm = CLSymmetryC3C4(src, symmetry="C4")
+
+    equator_indices = cl_symm._detect_equator_images()
+
+    recall = np.isin(equator_indices_gt, equator_indices).mean()
+    assert recall >= 0.8, (
+        f"Equator recall: {recall}; "
+        f"ground truth: {equator_indices_gt}; "
+        f"detected: {equator_indices}"
+    )
+
+
 def build_self_commonlines_matrix(n_theta, rots, order):
     # Construct rotatation matrices associated with cyclic order.
     rots_symm = cyclic_rotations(order, rots.dtype).matrices