Update Pipeline Demo

gallery/tutorials/pipeline_demo.py

@@ -14,24 +14,28 @@
 # Ribosome, sourced from EMDB: https://www.ebi.ac.uk/emdb/EMD-2660.
 # This is one of several volume maps that can be downloaded with
 # ASPIRE's data downloading utility by using the following import.
+
 # sphinx_gallery_start_ignore
 # flake8: noqa
 # sphinx_gallery_end_ignore
+
 from aspire.downloader import emdb_2660
 
 # Load 80s Ribosome as a ``Volume`` object.
 original_vol = emdb_2660()
 
-# Downsample the volume
-res = 41
-vol = original_vol.downsample(res)
+# During the preprocessing stages of the pipeline we will downsample
+# the images to an image size of 64 pixels. Here, we also downsample the
+# volume so we can compare to our reconstruction later.
+res = 64
+vol_ds = original_vol.downsample(res)
 
 # %%
 # .. note::
 #     A ``Volume`` can be saved using the ``Volume.save()`` method as follows::
 #
 #         fn = f"downsampled_80s_ribosome_size{res}.mrc"
-#         vol.save(fn, overwrite=True)
+#         vol_ds.save(fn, overwrite=True)
 
 
 # %%
@@ -63,7 +67,7 @@
 defocus_ct = 7
 
 ctf_filters = [
-    RadialCTFFilter(pixel_size=vol.pixel_size, defocus=d)
+    RadialCTFFilter(pixel_size=original_vol.pixel_size, defocus=d)
     for d in np.linspace(defocus_min, defocus_max, defocus_ct)
 ]
 
@@ -72,8 +76,8 @@
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 # We feed our ``Volume`` and filters into ``Simulation`` to generate
 # the dataset of images.  When controlled white Gaussian noise is
-# desired, ``WhiteNoiseAdder.from_snr()`` can be used to generate a
-# simulation data set around a specific SNR.
+# desired, ``WhiteNoiseAdder(var=VAR)`` can be used to generate a
+# simulation data set around a specific noise variance.
 #
 # Alternatively, users can bring their own images using an
 # ``ArrayImageSource``, or define their own custom noise functions via
@@ -84,28 +88,25 @@
 from aspire.noise import WhiteNoiseAdder
 from aspire.source import Simulation
 
-# set parameters
-n_imgs = 2500
-
-# SNR target for white gaussian noise.
-snr = 0.5
-
-# %%
-# .. note::
-#   Note, the SNR value was chosen based on other parameters for this
-#   quick tutorial, and can be changed to adjust the power of the
-#   additive noise.
-
 # For this ``Simulation`` we set all 2D offset vectors to zero,
 # but by default offset vectors will be randomly distributed.
+# We cache the Simulation to prevent regenerating the projections
+# for each preprocessing stage.
 src = Simulation(
-    n=n_imgs,  # number of projections
-    vols=vol,  # volume source
+    n=2500,  # number of projections
+    vols=original_vol,  # volume source
     offsets=0,  # Default: images are randomly shifted
     unique_filters=ctf_filters,
-    noise_adder=WhiteNoiseAdder.from_snr(snr=snr),  # desired SNR
-)
+    noise_adder=WhiteNoiseAdder(var=0.0002),  # desired noise variance
+).cache()
 
+# %%
+# .. note::
+#   The noise variance value above was chosen based on other parameters for this
+#   quick tutorial, and can be changed to adjust the power of the additive noise.
+#   Alternatively, an SNR value can be prescribed as follows::
+#
+#       Simulation(..., noise_adder=WhiteNoiseAdder.from_snr(SNR))
 
 # %%
 # Several Views of the Projection Images
@@ -125,60 +126,77 @@
 # with noise and CTF corruption
 src.images[0:10].show()
 
+# %%
+# Image Preprocessing
+# -------------------
+# We apply some image preprocessing techniques to prepare the
+# the images for denoising via Class Averaging.
+
+# %%
+# Downsampling
+# ------------
+# We downsample the images. Reducing the image size will improve the
+# efficiency of subsequent pipeline stages. Metadata such as pixel size is
+# scaled accordingly to correspond correctly with the image resolution.
+
+src = src.downsample(res)
+src.images[:10].show()
 
 # %%
 # CTF Correction
 # --------------
 # We apply ``phase_flip()`` to correct for CTF effects.
 
 src = src.phase_flip()
+src.images[:10].show()
 
 # %%
-# Cache
-# -----
+# Normalize Background
+# --------------------
+# We apply ``normalize_background()`` to prepare the image class averaging.
+
+src = src.normalize_background()
+src.images[:10].show()
+
+# %%
+# Noise Whitening
+# ---------------
+# We apply ``whiten()`` to estimate and whiten the noise.
+
+src = src.whiten()
+src.images[:10].show()
+
+# %%
+# Contrast Inversion
+# ------------------
+# We apply ``invert_contrast()`` to ensure a positive valued signal.
+
+src = src.invert_contrast()
+
+# %%
+# Caching
+# -------
 # We apply ``cache`` to store the results of the ``ImageSource``
 # pipeline up to this point.  This is optional, but can provide
 # benefit when used intently on machines with adequate memory.
-
 src = src.cache()
-src.images[0:10].show()
 
 # %%
 # Class Averaging
 # ---------------
-# We use ``RIRClass2D`` object to classify the images via the
-# rotationally invariant representation (RIR) algorithm. Class
-# selection is customizable. The classification module also includes a
-# set of protocols for selecting a set of images to be used after
-# classification.  Here we're using the simplest
-# ``DebugClassAvgSource`` which internally uses the
-# ``TopClassSelector`` to select the first ``n_classes`` images from
-# the source.  In practice, the selection is done by sorting class
-# averages based on some configurable notion of quality (contrast,
-# neighbor distance etc).
-
-from aspire.classification import RIRClass2D
-
-# set parameters
-n_classes = 200
-n_nbor = 6
-
-# We will customize our class averaging source. Note that the
-# ``fspca_components`` and ``bispectrum_components`` were selected for
-# this small tutorial.
-rir = RIRClass2D(
-    src,
-    fspca_components=40,
-    bispectrum_components=30,
-    n_nbor=n_nbor,
-)
+# For this tutorial we use the ``DebugClassAvgSource`` to generate an ``ImageSource``
+# of class averages. Internally, ``DebugClassAvgSource`` uses the ``RIRClass2D``
+# object to classify the source images via the rotationally invariant representation
+# (RIR) algorithm and the ``TopClassSelector`` object to select the first ``n_classes``
+# images in the original order from the source. In practice, class selection is commonly
+# done by sorting class averages based on some configurable notion of quality
+# (contrast, neighbor distance etc) which can be accomplished by providing a custom
+# class selector to ``ClassAverageSource``, which changes the ordering of the classes
+# returned by ``ClassAverageSource``.
 
 from aspire.denoising import DebugClassAvgSource
 
-avgs = DebugClassAvgSource(
-    src=src,
-    classifier=rir,
-)
+avgs = DebugClassAvgSource(src=src)
 
 # We'll continue our pipeline using only the first ``n_classes`` from
 # ``avgs``.  The ``cache()`` call is used here to precompute results
@@ -187,6 +205,7 @@
 # the following ``CLSyncVoting`` algorithm.  Outside of demonstration
 # purposes, where we are repeatedly peeking at various stage results,
 # such caching can be dropped allowing for more lazy evaluation.
+n_classes = 250
 avgs = avgs[:n_classes].cache()
 
 
@@ -216,17 +235,19 @@
 
 from aspire.abinitio import CLSyncVoting
 from aspire.source import OrientedSource
+from aspire.utils import Rotation
 
 # Stash true rotations for later comparison
-true_rotations = src.rotations[:n_classes]
+true_rotations = Rotation(src.rotations[:n_classes])
 
-# For this low resolution example we will customize the ``CLSyncVoting``
-# instance to use fewer theta points ``n_theta`` then the default value of 360.
-orient_est = CLSyncVoting(avgs, n_theta=72)
+# Instantiate a ``CLSyncVoting`` object for estimating orientations.
+orient_est = CLSyncVoting(avgs)
 
 # Instantiate an ``OrientedSource``.
 oriented_src = OrientedSource(avgs, orient_est)
 
+# Estimate Rotations.
+est_rotations = oriented_src.rotations
 
 # %%
 # Mean Error of Estimated Rotations
@@ -238,7 +259,7 @@
 from aspire.utils import mean_aligned_angular_distance
 
 # Compare with known true rotations
-mean_ang_dist = mean_aligned_angular_distance(oriented_src.rotations, true_rotations)
+mean_ang_dist = mean_aligned_angular_distance(est_rotations, true_rotations)
 print(f"Mean aligned angular distance: {mean_ang_dist} degrees")
 
 
@@ -269,7 +290,7 @@
 # Get the first 10 projections from the estimated volume using the
 # estimated orientations.  Recall that ``project`` returns an
 # ``Image`` instance, which we can peek at with ``show``.
-projections_est = estimated_volume.project(oriented_src.rotations[0:10])
+projections_est = estimated_volume.project(est_rotations[0:10])
 
 # We view the first 10 projections of the estimated volume.
 projections_est.show()
@@ -278,3 +299,24 @@
 
 # For comparison, we view the first 10 source projections.
 src.projections[0:10].show()
+
+
+# %%
+# Fourier Shell Correlation
+# -------------------------
+# Additionally, we can compare our reconstruction to the known source volume
+# by performing a Fourier shell correlation (FSC). We first find a rotation
+# matrix which best aligns the estimated rotations to the ground truth rotations
+# using the ``find_registration`` method. We then use that rotation to align
+# the reconstructed volume to the ground truth volume.
+
+# `find_registration` returns the best aligning rotation, `Q`, as well as
+# a `flag` which indicates if the volume needs to be reflected.
+Q, flag = Rotation(est_rotations).find_registration(true_rotations)
+aligned_vol = estimated_volume
+if flag == 1:
+    aligned_vol = aligned_vol.flip()
+aligned_vol = aligned_vol.rotate(Rotation(Q.T))
+
+# Compute the FSC.
+vol_ds.fsc(aligned_vol, cutoff=0.143, plot=True)