Fixup the CWF experiment example (some dtypes issues)

Author: garrettwrong

gallery/experiments/simulated_abinitio_pipeline.py

@@ -19,8 +19,8 @@
 import numpy as np
 
 from aspire.abinitio import CLSyncVoting
-from aspire.basis import FFBBasis3D
-from aspire.classification import RIRClass2D
+from aspire.basis import FFBBasis2D, FFBBasis3D
+from aspire.classification import BFSReddyChatterjiAlign2D, RIRClass2D
 from aspire.denoising import DenoiserCov2D
 from aspire.noise import AnisotropicNoiseEstimator
 from aspire.operators import FunctionFilter, RadialCTFFilter
@@ -132,22 +132,32 @@ def noise_function(x, y):
 if interactive:
     src.images(0, 10).show()
 
+# # Optionally invert image contrast, depends on data.
 # logger.info("Invert the global density contrast")
 # src.invert_contrast()
 
-# # On Simulation data, better results so far were achieved without cov2d.
+# Cache to memory for some speedup
+src = ArrayImageSource(src.images(0, num_imgs).asnumpy(), angles=src.angles)
+
+# On Simulation data, better results so far were achieved without cov2d
+# However, we can demonstrate using CWF denoised images for classification.
+classification_src = src
+custom_aligner = None
 if do_cov2d:
     # Use CWF denoising
     cwf_denoiser = DenoiserCov2D(src)
-    src = cwf_denoiser.denoise()
-
-# Peek, what do the denoised images look like...
-if interactive:
-    src.images(0, 10).show()
+    # Use denoised src for classification
+    classification_src = cwf_denoiser.denoise()
+    # Peek, what do the denoised images look like...
+    if interactive:
+        classification_src.images(0, 10).show()
 
+    # Use regular `src` for the alignment and composition (averaging).
+    composite_basis = FFBBasis2D((src.L,) * 2, dtype=src.dtype)
+    custom_aligner = BFSReddyChatterjiAlign2D(
+        None, src, composite_basis, dtype=src.dtype
+    )
 
-# Cache to memory for some speedup
-src = ArrayImageSource(src.images(0, num_imgs).asnumpy(), angles=src.angles)
 
 # %%
 # Class Averaging
@@ -158,14 +168,15 @@ def noise_function(x, y):
 logger.info("Begin Class Averaging")
 
 rir = RIRClass2D(
-    src,
+    classification_src,  # Source used for classification
     fspca_components=400,
     bispectrum_components=300,  # Compressed Features after last PCA stage.
     n_nbor=n_nbor,
     n_classes=n_classes,
     large_pca_implementation="legacy",
     nn_implementation="sklearn",
     bispectrum_implementation="legacy",
+    aligner=custom_aligner,
 )
 
 classes, reflections, distances = rir.classify()

src/aspire/classification/align2d.py

@@ -421,9 +421,7 @@ def __init__(
         Instead see `alignment_source`.
         :param source: Source of original images.
         :param composite_basis:  Basis to be used during class average composition.
-        For current implementation of ReddyChatterjiAlign2D this should be `None`.
-        Instead this method uses `source` for composition of the averaged stack.
-        :param alignment_source:  Basis to be used during class average composition.
+        :param alignment_source: Optional, source to be used during class average alignment.
         Must be the same resolution as `source`.
         :param dtype: Numpy dtype to be used during alignment.
         """
@@ -888,9 +886,7 @@ def __init__(
         Instead see `alignment_source`.
         :param source: Source of original images.
         :param composite_basis:  Basis to be used during class average composition.
-        For current implementation of ReddyChatterjiAlign2D this should be `None`.
-        Instead this method uses `source` for composition of the averaged stack.
-        :param alignment_source:  Basis to be used during class average composition.
+        :param alignment_source: Optional, source to be used during class average alignment.
         Must be the same resolution as `source`.
         :param radius: Brute force translation search radius.
         Defaults to source.L//8.

src/aspire/covariance/covar2d.py

@@ -284,7 +284,7 @@ def identity(x):
 
         for k in np.unique(ctf_idx[:]):
 
-            coeff_k = coeffs[ctf_idx == k]
+            coeff_k = coeffs[ctf_idx == k].astype(self.dtype)
             weight = coeff_k.shape[0] / coeffs.shape[0]
 
             ctf_fb_k = ctf_fb[k]

src/aspire/denoising/denoiser_cov2d.py

@@ -123,7 +123,7 @@ def __init__(self, src, basis=None, var_noise=None):
         self.var_noise = var_noise
 
         if basis is None:
-            basis = FFBBasis2D((self.src.L, self.src.L))
+            basis = FFBBasis2D((self.src.L, self.src.L), dtype=src.dtype)
 
         if not isinstance(basis, FFBBasis2D):
             raise NotImplementedError("Currently only fast FB method is supported")