diff --git a/src/aspire/abinitio/commonline_c3_c4.py b/src/aspire/abinitio/commonline_c3_c4.py index 7ef9a22c8d..2b7e45d5a6 100644 --- a/src/aspire/abinitio/commonline_c3_c4.py +++ b/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() diff --git a/tests/test_orient_symmetric.py b/tests/test_orient_symmetric.py index 71e3eb743b..6a1f5c6457 100644 --- a/tests/test_orient_symmetric.py +++ b/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