tunable proximal refinement iterations

Author: j-c-c

src/aspire/abinitio/commonline_nug.py

@@ -38,7 +38,7 @@ def __init__(
         mult=1.5,
         S2_grid=441,
         Nstep_yI=10,
-        perform_pr=False,
+        pr_iters=None,
         verbose=True,
         **kwargs,
     ):
@@ -62,7 +62,8 @@ def __init__(
         :param mult: Step-size multiplier for the ADMM primal update.
         :param S2_grid: Number of sphere samples used to discretize SO(3).
         :param Nstep_yI: Number of inequality-multiplier updates per ADMM iteration.
-        :param perform_pr: Whether to apply proximal refinement after ADMM.
+        :param pr_iters: Number of proximal refinement iterations. Default of None
+            does not perform proximal refinement. Recommended value is 4.
         :param verbose: Whether to log ADMM progress.
         """
 
@@ -85,7 +86,6 @@ def __init__(
         self.mult = mult
         self.S2_grid = S2_grid
         self.Nstep_yI = Nstep_yI
-        self.perform_pr = perform_pr
         self.verbose = verbose
 
         # Handle symmetry
@@ -110,6 +110,13 @@ def __init__(
         self.sym_euler = self.sym_grp.rotations.angles
         self.n_sym = len(self.sym_euler)
 
+        # Set up proximal refinement terms
+        if pr_iters is not None:
+            self.pr_weights = 1 / (1 + np.arange(self.Lmax))
+            self.pr_penalty = [1] * pr_iters
+            self.pr_rank = list(range(pr_iters - 1, -1, -1))  # [pr_iters - 1,..., 0]
+        self.pr_iters = pr_iters
+
         self._build_full_pft()
 
     def _build_full_pft(self):
@@ -276,15 +283,12 @@ def perform_admm(self):
         """
         X_est = self.admm_sym_J(self.C, self.verbose)
 
-        if self.perform_pr:
-            weight = 1 / (1 + np.arange(self.Lmax))
-            Penalty = [1, 1, 1, 1]
-            r = [3, 2, 1, 0]
+        if self.pr_iters is not None:
             X_est = self.proximal_refine(
                 X_est,
-                weight,
-                Penalty,
-                r,
+                self.pr_weights,
+                self.pr_penalty,
+                self.pr_rank,
             )
         self.X_est = X_est
 
@@ -980,11 +984,10 @@ def low_rank_proj(X, r_step):
                 Xproj.append(tmp)
             return Xproj
 
-        Niter = len(r)
         CC = [None] * self.Lmax
         current = [np.copy(Xk) for Xk in X_admm]
 
-        for step in range(Niter):
+        for step in range(self.pr_iters):
             X_proj = low_rank_proj(current, r[step])
 
             for k in range(self.Lmax):
@@ -999,7 +1002,7 @@ def low_rank_proj(X, r_step):
                 logger.info(
                     "Proximal refine step %d/%d: relative update %.3e",
                     step + 1,
-                    Niter,
+                    self.pr_iters,
                     rel_change(X_next, current),
                 )
 

tests/test_nug.py

@@ -11,7 +11,7 @@
 N_IMG = [15]
 OFFSETS = [0, None]
 ORDER = [3, 4]
-PR = [False]
+PR = [None]
 SEED = 1980
 VOLUME = [
     CnSymmetricVolume,
@@ -89,7 +89,8 @@ def orient_est(source, proximal_refine):
         source,
         max_shift=max_shift,
         shift_step=shift_step,
-        perform_pr=proximal_refine,
+        max_iter=201,
+        pr_iters=proximal_refine,
         verbose=False,
     )
     orient_est.estimate_rotations()
@@ -101,7 +102,7 @@ def orient_est(source, proximal_refine):
 #########
 
 
-def test_smoke_nug(dtype, Volume):
+def test_smoke(dtype, Volume):
     """
     Perform quick smoke test since other tests are long running.
     """
@@ -122,7 +123,6 @@ def test_smoke_nug(dtype, Volume):
         max_iter=10,
         S2_grid=50,
         max_shift=0,
-        mask=False,
         verbose=False,
     )
 
@@ -151,7 +151,11 @@ def test_dtypes(orient_est):
 
 @pytest.mark.expensive
 def test_estimate_rotations_pairwise(orient_est):
-    """ """
+    """
+    Check mean squared error between estimates and ground truth
+    pairwise rotations, Rij. This serves as a reference to the error
+    metric used by the researcher in the related publication.
+    """
     MSE = compare_rots_sym(
         orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp
     )
@@ -160,13 +164,21 @@ def test_estimate_rotations_pairwise(orient_est):
 
 @pytest.mark.expensive
 def test_estimate_rotations(orient_est):
+    """
+    Check that the mean angular distance between estimates and ground
+    truth, after symmetry synchronization and global alignment, are
+    within 10 degrees.
+    """
     gt_rots_synced = g_sync(
         orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp
     )
-    mean_aligned_angular_distance(orient_est.rotations, gt_rots_synced, 8.0)
+    mean_aligned_angular_distance(orient_est.rotations, gt_rots_synced, 10.0)
 
 
 def test_unspupported_symmetry_raises(dtype):
+    """
+    Check that we raise for symmetries other than Cn/Dn.
+    """
     vol = TSymmetricVolume(L=16, C=1, K=10, dtype=dtype).generate()
     src = Simulation(n=3, vols=vol)