Non-unique Games w/ Symmetry

gallery/experiments/experimental_nug_pipeline_10081.py

@@ -0,0 +1,112 @@
+"""
+Symmetric NUG Pipeline
+======================
+
+This notebook demonstrates reproducing simliar results to the
+experiment found in:
+
+.. admonition:: Publication
+
+   | Shi, Y., Singer, A., Wang, J., & Yang, R.
+   | Orientation estimation of cryo-EM images for molecules with
+   | arbitrary symmetry by semidefinite programming
+
+This pipeline begins with 17 class averages that were generated (as described
+in the publication) from the EMPIAR 10081 picked particle data, available here:
+
+https://www.ebi.ac.uk/empiar/EMPIAR-10081
+
+First, orientation estimation is performed using the Non-Unique Games
+algorithm adapted for handling symmetric particles. A volume is then
+reconstructed from the estimated orientations and saved to an output
+file for inspection. The reconstructed volume can be compared with the
+EMPIAR-10081 associated reconstruction found here:
+
+https://www.ebi.ac.uk/emdb/EMD-8511
+"""
+
+# %%
+# Imports
+# -------
+# First import the necessary utilities and ASPIRE
+# classes used throughout the notebook.
+import logging
+
+from aspire.abinitio import CommonlineNUG
+from aspire.downloader import nug_10081
+from aspire.reconstruction import MeanEstimator
+from aspire.source import ArrayImageSource, OrientedSource
+
+logger = logging.getLogger(__name__)
+
+
+# %%
+# Setup
+# -----
+
+# Note: Use of GPU is recommended for the NUG algorithm.
+# Below we setup some output files to save off orientation
+# and reconstruction results.
+oriented_fn = "nug_10081_oriented.star"
+volume_output_fn = "nug_10081_abinitio.mrc"
+
+
+# %%
+# Load the Dataset
+# ----------------
+#
+# The 17 class averages can be accessed as an ASPIRE ``Image`` object
+# via the built-in downloader utility.
+
+logger.info("Load Precomputed Class Averages")
+ims = nug_10081()
+avgs = ArrayImageSource(ims, pixel_size=1.3, symmetry_group="C4")
+
+
+# %%
+# Orientation Estimation
+# ----------------------
+#
+# Create an orientation estimation object for the ``avgs``.
+# The ``CommonlineNUG`` algorithm will detect the symmetry from
+# the ``avgs`` metadata and use the approriate NUG ADMM-solver to estimate
+# orientations. To replicate the original experiment as closely as possible,
+# we use the 15 iterations of proximal refinement and disable searching over
+# the commonline shift space.
+
+logger.info("Begin Orientation Estimation")
+orient_est = CommonlineNUG(avgs, max_shift=0, pr_iters=15)
+
+# Create an ``OrientedSource`` class instance that performs orientation
+# estimation in a lazy fashion upon request of images or rotations.
+oriented_src = OrientedSource(avgs, orient_est)
+
+# Estimate orientations and override the estimated shifts.
+# In the case of insufficient number of images, estimating shifts fails.
+oriented_src.rotations
+oriented_src = oriented_src.update(offsets=0)
+
+# Confirm offsets have been set to zero.
+logger.info(f"Offsets: {oriented_src.offsets}")
+
+# Save oriented source.
+oriented_src.save(oriented_fn)
+
+# %%
+# Volume Reconstruction
+# ----------------------
+#
+# Using the oriented source, attempt to reconstruct a volume.  Since
+# this is a C4 symmetric molecule the ``symmetry_group`` source attribute
+# will flow through the pipeline to ``oriented_src``. Then each image will be
+# repeated for each of the 3 additional rotations during
+# back-projection. In general, this boosts the effective number of images used in
+# the reconstruction from ``n_imgs`` to ``4 * n_imgs``, or from 17 to 68 in this case.
+
+# Setup an estimator to perform the back projection.
+logger.info("Begin Volume reconstruction")
+estimator = MeanEstimator(oriented_src)
+
+# Perform the estimation and save the reconstructed volume.
+est_vol = estimator.estimate()
+est_vol.save(volume_output_fn)

gallery/experiments/experimental_nug_pipeline_12036.py

@@ -0,0 +1,111 @@
+"""
+Dn-Symmetric NUG Pipeline
+=========================
+
+This notebook demonstrates reproducing simliar results to the
+experiment found in:
+
+.. admonition:: Publication
+
+   | Shi, Y., Singer, A., Wang, J., & Yang, R.
+   | Orientation estimation of cryo-EM images for molecules with
+   | arbitrary symmetry by semidefinite programming
+
+This pipeline begins with 93 class averages that were generated (as described
+in the publication) from the EMPIAR 12036 picked particle data, available here:
+
+https://www.ebi.ac.uk/empiar/EMPIAR-12036
+
+First, orientation estimation is performed using the Non-Unique Games
+algorithm adapted for handling symmetric particles. A volume is then
+reconstructed from the estimated orientations and saved to an output
+file for inspection. The reconstructed volume can be compared with the
+EMPIAR-12036 associated reconstruction found here:
+
+https://www.ebi.ac.uk/emdb/EMD-14803
+"""
+
+# %%
+# Imports
+# -------
+# First import the necessary utilities and ASPIRE
+# classes used throughout the notebook.
+import logging
+
+from aspire.abinitio import CommonlineNUG
+from aspire.downloader import nug_12036
+from aspire.reconstruction import MeanEstimator
+from aspire.source import ArrayImageSource, OrientedSource
+
+logger = logging.getLogger(__name__)
+
+
+# %%
+# Setup
+# -----
+
+# Note: Use of GPU is recommended for the NUG algorithm.
+# Below we setup some output files to save off orientation
+# and reconstruction results.
+oriented_fn = "nug_12036_oriented.star"
+volume_output_fn = "nug_12036_abinitio.mrc"
+
+
+# %%
+# Load the Dataset
+# ----------------
+#
+# The 93 class averages can be accessed as an ASPIRE ``Image`` object
+# via the built-in downloader utility.
+logger.info("Load Precomputed Class Averages")
+ims = nug_12036()
+avgs = ArrayImageSource(ims, symmetry_group="D3")
+
+
+# %%
+# Orientation Estimation
+# ----------------------
+#
+# Create an orientation estimation object for the ``avgs``.
+# The ``CommonlineNUG`` algorithm will detect the symmetry from
+# the ``avgs`` metadata and use the approriate NUG ADMM-solver to estimate
+# orientations. To replicate the original experiment as closely as possible,
+# we use the 5 iterations of proximal refinement and disable searching over
+# the commonline shift space.
+
+logger.info("Begin Orientation Estimation")
+orient_est = CommonlineNUG(avgs, max_shift=0, pr_iters=5)
+
+# Create an ``OrientedSource`` class instance that performs orientation
+# estimation in a lazy fashion upon request of images or rotations.
+oriented_src = OrientedSource(avgs, orient_est)
+
+# Estimate orientations and override the estimated shifts.
+# In the case of insufficient number of images, estimating shifts fails.
+oriented_src.rotations
+oriented_src = oriented_src.update(offsets=0)
+
+# Confirm offsets have been set to zero.
+logger.info(f"Offsets: {oriented_src.offsets}")
+
+# Save oriented source.
+oriented_src.save(oriented_fn)
+
+# %%
+# Volume Reconstruction
+# ----------------------
+#
+# Using the oriented source, attempt to reconstruct a volume.  Since
+# this is a D3 symmetric molecule the ``symmetry_group`` source attribute
+# will flow through the pipeline to ``oriented_src``. Then each image will be
+# repeated for each of the 5 additional symmetry induced rotations during
+# back-projection. In general, this boosts the effective number of images used in
+# the reconstruction from ``n_imgs`` to ``6 * n_imgs``, or from 93 to 558 in this case.
+
+# Setup an estimator to perform the back projection.
+logger.info("Begin Volume reconstruction")
+estimator = MeanEstimator(oriented_src)
+
+# Perform the estimation and save the reconstructed volume.
+est_vol = estimator.estimate()
+est_vol.save(volume_output_fn)

gallery/experiments/experimental_nug_simulated_pipeline_10028.py

@@ -0,0 +1,120 @@
+"""
+C1 NUG Rotation Accuracy on Simulated EMDB-2660 Images
+======================================================
+
+Estimate orientations of noisy, simulated projections of EMDB-2660 using
+the asymmetric (C1) NUG method. Compare the estimated rotations with the
+known simulation rotations at several signal-to-noise ratios (SNRs). A
+volume is reconstructed and saved off for comparison against the generating
+volume.
+
+EMDB-2660 is the volume entry reconstructed from the experimental dataset
+EMPIAR 10028, both available here:
+
+https://www.ebi.ac.uk/emdb/EMD-2660
+https://www.ebi.ac.uk/empiar/EMPIAR-10028
+
+This experiment is based on:
+
+.. admonition:: Reference
+
+   A. S. Bandeira, Y. Chen, R. R. Lederman, and A. Singer,
+   "Non-unique games over compact groups and orientation estimation
+   in cryo-EM," Inverse Problems 36(6), 064002 (2020).
+   https://doi.org/10.1088/1361-6420/ab7d2c
+
+The ASPIRE-Python implementation of NUG differs from the original MATLAB
+implementation, so this example is intended to reproduce the experimental
+workflow rather than its exact numerical results.
+"""
+
+# %%
+# Imports
+# -------
+import logging
+
+import numpy as np
+
+from aspire.abinitio import CommonlineNUG
+from aspire.downloader import emdb_2660
+from aspire.noise import WhiteNoiseAdder
+from aspire.reconstruction import MeanEstimator
+from aspire.source import OrientedSource, Simulation
+from aspire.utils import Rotation
+
+logger = logging.getLogger(__name__)
+
+
+# %%
+# Experiment settings
+# -------------------
+# Start with 100 images. The paper also reports results for 500 images and
+# SNR ranging from 1 to 1/128, which requires substantially more computation.
+SEED = 1980
+RESOLUTION = 129
+N_IMAGES = (100,)
+SNR_VALUES = (1, 1 / 2, 1 / 4, 1 / 8, 1 / 16)
+
+# Download volume once; every SNR uses the same underlying volume.
+volume = emdb_2660().astype(np.float64)
+
+# %%
+# Estimate rotations at each SNR
+# ------------------------------
+results = []
+
+for n_images in N_IMAGES:
+    for snr in SNR_VALUES:
+        logger.info(f"Estimating rotations: n={n_images}, SNR={snr}")
+        # Keeping the simulation seed fixed gives corresponding ground-truth
+        # orientations across SNR conditions.
+        source = (
+            Simulation(
+                n=n_images,
+                vols=volume,
+                offsets=0,
+                amplitudes=1,
+                seed=SEED,
+                noise_adder=WhiteNoiseAdder.from_snr(snr=snr),
+            )
+            .downsample(RESOLUTION)
+            .cache()
+        )
+
+        # Initialize the NUG orientation estimator, matching paramters
+        # to the referenced experiment.
+        estimator = CommonlineNUG(
+            source,
+            symmetry="C1",
+            max_shift=0,
+            verbose=False,
+        )
+
+        oriented_source = OrientedSource(source, estimator)
+        estimated_rotations = oriented_source.rotations
+
+        # Rotation.mse registers the estimates to ground truth before
+        # computing the mean squared error.
+        mse = Rotation(estimated_rotations).mse(Rotation(source.rotations))
+        results.append((n_images, snr, mse))
+
+        print(f"n={n_images}  SNR={snr}  rotation MSE={mse}", flush=True)
+
+        # The simulation has zero offsets. Use the estimated orientations
+        # with zero shifts to reconstruct from the cached noisy images.
+        oriented_source = oriented_source.update(offsets=0)
+
+        snr_tag = f"{snr:g}"
+        volume_filename = f"nug_2660_c1_n{n_images}_snr{snr_tag}_recon.mrc"
+
+        logger.info(f"Reconstructing and saving {volume_filename}")
+        estimated_volume = MeanEstimator(oriented_source).estimate()
+        estimated_volume.save(volume_filename, overwrite=True)
+
+# %%
+# Summary
+# -------
+print("\nC1 NUG rotation accuracy")
+print(f"{'Images':>6} {'SNR':>10} {'MSE':>14}")
+for n_images, snr, mse in results:
+    print(f"{n_images:6d} {snr:10g} {mse:14.6g}")

src/aspire/abinitio/__init__.py

@@ -2,10 +2,14 @@
 from .J_sync import JSync
 from .commonline_utils import (
     build_outer_products,
+    compare_rots_sym,
     g_sync,
+    g_sync_finite_group,
+    saff_kuijlaars,
 )
 from .commonline_base import Orient3D
 from .commonline_matrix import CLOrient3D
+from .commonline_nug import CommonlineNUG
 from .commonline_sdp import CommonlineSDP
 from .commonline_lud import CommonlineLUD
 from .commonline_irls import CommonlineIRLS