Filter stack and radial covar optimizations

gallery/experiments/save_simulation_relion_reconstruct.py

@@ -51,7 +51,7 @@
 # that RELION will recover as optics groups.
 
 vol = emdb_2660()
-ctf_filters = [RadialCTFFilter(defocus=d) for d in defocus]
+ctf_filters = RadialCTFFilter(defocus=defocus)
 
 
 # %%
@@ -64,7 +64,7 @@
 sim = Simulation(
     n=n_particles,
     vols=vol,
-    unique_filters=ctf_filters,
+    filter_stack=ctf_filters,
     noise_adder=WhiteNoiseAdder.from_snr(snr),
 )
 sim.save(star_path, overwrite=True)

gallery/experiments/simulated_abinitio_pipeline.py

@@ -83,17 +83,19 @@ def noise_function(x, y):
 alpha = 0.1  # Amplitude contrast
 
 # Create filters
-ctf_filters = [
-    RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
-    for d in np.linspace(defocus_min, defocus_max, defocus_ct)
-]
+ctf_filters = RadialCTFFilter(
+    voltage,
+    defocus=np.linspace(defocus_min, defocus_max, defocus_ct),
+    Cs=2.0,
+    alpha=0.1,
+)
 
 # Finally create the Simulation
 src = Simulation(
     n=num_imgs,
     vols=og_v,
     noise_adder=custom_noise,
-    unique_filters=ctf_filters,
+    filter_stack=ctf_filters,
 )
 
 # Downsample

gallery/tutorials/aspire_introduction.py

@@ -570,10 +570,7 @@ def noise_function(x, y):
 defocus_ct = 7
 
 # Generate several CTFs.
-ctf_filters = [
-    RadialCTFFilter(defocus=d)
-    for d in np.linspace(defocus_min, defocus_max, defocus_ct)
-]
+ctf_filters = RadialCTFFilter(defocus=np.linspace(defocus_min, defocus_max, defocus_ct))
 
 # %%
 # Combining into a Simulation
@@ -586,7 +583,7 @@ def noise_function(x, y):
     amplitudes=1,
     offsets=0,
     noise_adder=white_noise_adder,
-    unique_filters=ctf_filters,
+    filter_stack=ctf_filters,
     seed=42,
 )
 

gallery/tutorials/pipeline_demo.py

@@ -66,10 +66,8 @@
 defocus_max = 25000
 defocus_ct = 7
 
-ctf_filters = [
-    RadialCTFFilter(defocus=d)
-    for d in np.linspace(defocus_min, defocus_max, defocus_ct)
-]
+ctf_filters = RadialCTFFilter(defocus=np.linspace(defocus_min, defocus_max, defocus_ct))
+
 
 # %%
 # Initialize Simulation Object
@@ -96,7 +94,7 @@
     n=2500,  # number of projections
     vols=original_vol,  # volume source
     offsets=0,  # Default: images are randomly shifted
-    unique_filters=ctf_filters,
+    filter_stack=ctf_filters,
     noise_adder=WhiteNoiseAdder(var=0.0002),  # desired noise variance
 ).cache()
 

gallery/tutorials/tutorials/cov2d_simulation.py

@@ -67,10 +67,13 @@
 
 print("Initialize simulation object and CTF filters.")
 # Create filters
-ctf_filters = [
-    RadialCTFFilter(voltage, defocus=d, Cs=2.0, alpha=0.1)
-    for d in np.linspace(defocus_min, defocus_max, defocus_ct)
-]
+ctf_filters = RadialCTFFilter(
+    voltage,
+    defocus=np.linspace(defocus_min, defocus_max, defocus_ct),
+    Cs=2.0,
+    alpha=0.1,
+)
+
 
 # Load the map file of a 70S Ribosome
 print(
@@ -89,7 +92,7 @@
     L=img_size,
     n=num_imgs,
     vols=vols,
-    unique_filters=ctf_filters,
+    filter_stack=ctf_filters,
     offsets=0.0,
     amplitudes=1.0,
     dtype=dtype,
@@ -109,9 +112,7 @@
 h_idx = sim.filter_indices
 
 # Evaluate CTF in the 8X8 FB basis
-h_ctf_fb = [
-    ffbbasis.filter_to_basis_mat(filt, pixel_size=pixel_size) for filt in ctf_filters
-]
+h_ctf_fb = ffbbasis.filter_stack_to_basis_mats(ctf_filters, pixel_size=pixel_size)
 
 # Get clean images from projections of 3D map.
 print("Apply CTF filters to clean images.")

gallery/tutorials/tutorials/cov3d_simulation.py

@@ -44,7 +44,7 @@
     L=img_size,
     n=num_imgs,
     vols=vols,
-    unique_filters=[RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)],
+    filter_stack=RadialCTFFilter(defocus=np.linspace(1.5e4, 2.5e4, 7)),
     dtype=dtype,
 )
 

gallery/tutorials/tutorials/ctf.py

@@ -154,9 +154,8 @@ def generate_example_image(L, noise_variance=0.1):
 
 # Construct a range of CTF filters.
 defoci = [2500, 5000, 10000, 20000]
-ctf_filters = [
-    RadialCTFFilter(voltage=200, defocus=d, Cs=2.26, alpha=0.07, B=0) for d in defoci
-]
+ctf_filters = RadialCTFFilter(voltage=200, defocus=defoci, Cs=2.26, alpha=0.07, B=0)
+
 
 # %%
 # Generate CTF corrupted Images
@@ -334,7 +333,7 @@ def generate_example_image(L, noise_variance=0.1):
 from aspire.source import Simulation
 
 # Create the Source.  ``ctf_filters`` are re-used from earlier section.
-src = Simulation(L=64, n=4, unique_filters=ctf_filters, pixel_size=1)
+src = Simulation(L=64, n=4, filter_stack=ctf_filters, pixel_size=1)
 src.images[:4].show()
 
 # %%

gallery/tutorials/tutorials/micrograph_source.py

@@ -181,9 +181,7 @@
 
 # Create our CTF Filter and add it to a list.
 # This configuration will apply the same CTF to all particles.
-ctfs = [
-    RadialCTFFilter(voltage=200, defocus=15000, Cs=2.26, alpha=0.07, B=0),
-]
+ctfs = RadialCTFFilter(voltage=200, defocus=15000, Cs=2.26, alpha=0.07, B=0)
 
 src = MicrographSimulation(
     vol,

gallery/tutorials/tutorials/orient3d_simulation.py

@@ -49,10 +49,13 @@
 
 print("Initialize simulation object and CTF filters.")
 # Create CTF filters
-filters = [
-    RadialCTFFilter(voltage, defocus=d, Cs=2.0, alpha=0.1)
-    for d in np.linspace(defocus_min, defocus_max, defocus_ct)
-]
+filters = RadialCTFFilter(
+    voltage,
+    defocus=np.linspace(defocus_min, defocus_max, defocus_ct),
+    Cs=2.0,
+    alpha=0.1,
+)
+
 
 # %%
 # Downsampling
@@ -74,7 +77,7 @@
 # Create a simulation object with specified filters and the downsampled 3D map
 print("Use downsampled map to creat simulation object.")
 sim = Simulation(
-    L=img_size, n=num_imgs, vols=vols, unique_filters=filters, pixel_size=5, dtype=dtype
+    L=img_size, n=num_imgs, vols=vols, filter_stack=filters, pixel_size=5, dtype=dtype
 )
 
 print("Get true rotation angles generated randomly by the simulation object.")

gallery/tutorials/tutorials/preprocess_imgs_sim.py

@@ -53,10 +53,13 @@
 
 print("Initialize simulation object and CTF filters.")
 # Create CTF filters
-ctf_filters = [
-    RadialCTFFilter(voltage, defocus=d, Cs=2.0, alpha=0.1)
-    for d in np.linspace(defocus_min, defocus_max, defocus_ct)
-]
+ctf_filters = RadialCTFFilter(
+    voltage,
+    defocus=np.linspace(defocus_min, defocus_max, defocus_ct),
+    Cs=2.0,
+    alpha=0.1,
+)
+
 
 # Load the map file of a 70S ribosome and downsample the 3D map to desired image size.
 print("Load 3D map from mrc file")
@@ -73,7 +76,7 @@
     L=img_size,
     n=num_imgs,
     vols=vols,
-    unique_filters=ctf_filters,
+    filter_stack=ctf_filters,
     noise_adder=noise_adder,
     pixel_size=pixel_size,
 )

src/aspire/basis/fb_2d.py

@@ -291,6 +291,6 @@ def calculate_bispectrum(
 
     def filter_to_basis_mat(self, *args, **kwargs):
         """
-        See `SteerableBasis2D.filter_to_basis_mat`.
+        See `SteerableBasis2D.filter_stack_to_basis_mat`.
         """
         return super().filter_to_basis_mat(*args, **kwargs)

src/aspire/basis/ffb_2d.py

@@ -67,6 +67,13 @@ def _build(self):
             self._precomp["gl_nodes"]
         )
 
+        # Generate radial filter point set for radial optimized eval
+        # Weights appear a little sensitive to dtype, otherwise could use self._precomp["gl_nodes"]
+        k_vals, _ = lgwt(self.n_r, 0, self.kcut, dtype=np.float64)
+        self._filter_pts = np.pad(
+            2 * np.pi * k_vals.reshape(1, -1), ((0, 1), (0, 0))
+        ).astype(self.dtype)
+
     def _precomp(self):
         """
         Precomute the basis functions on a polar Fourier grid
@@ -236,63 +243,127 @@ def _evaluate_t(self, x):
 
         return xp.asnumpy(v)
 
-    def filter_to_basis_mat(self, f, **kwargs):
+    def _filter_stack_to_basis_mats(self, f, **kwargs):
         """
         See `SteerableBasis2D.filter_to_basis_mat`.
         """
-        # Note 'method' and 'truncate' not relevant for this optimized FFB code.
-        if kwargs.get("method", None) is not None:
+        # Note 'truncate' not relevant for this specific FFB code.
+        # `expand_method=radial` should have already been diverted
+        # by the wrapping code in `super().filter_stack_to_basis_mats`.
+        # Permits forcing 2d calc on a radial filter via `expand_method=evaluate_t`
+        expand_method = kwargs.get("expand_method", None)
+        if expand_method not in [None, "evaluate_t"]:
             raise NotImplementedError(
-                "`FFBBasis2D.filter_to_basis_mat` method {method} not supported."
-                "  Use `method=None`."
+                f"`FFBBasis2D.filter_to_basis_mat` expand_method '{expand_method}' not supported."
+                "  Use `expand_method=` `None` or `evaluate_t`."
             )
 
         pixel_size = kwargs.get("pixel_size", None)
 
-        # These form a circular dependence, import locally until time to clean up.
-        from aspire.basis.basis_utils import lgwt
-
         # Get the filter's evaluate function.
         h_fun = f.evaluate
 
         # Set same dimensions as basis object
         n_k = self.n_r
         n_theta = self.n_theta
-        radial = self._precomp["radial"]
+        k_vals = self._precomp["gl_nodes"]
 
         # get 2D grid in polar coordinate
-        k_vals, wts = lgwt(n_k, 0, 0.5, dtype=self.dtype)
-        k, theta = np.meshgrid(
-            k_vals, np.arange(n_theta) * 2 * np.pi / (2 * n_theta), indexing="ij"
+        k, theta = xp.meshgrid(
+            xp.asarray(k_vals),
+            xp.arange(n_theta) * 2 * np.pi / (2 * n_theta),
+            indexing="ij",
         )
 
         # Get function values in polar 2D grid and average out angle contribution
-        omegax = k * np.cos(theta)
-        omegay = k * np.sin(theta)
-        omega = 2 * np.pi * np.vstack((omegax.flatten("C"), omegay.flatten("C")))
+        omegax = k * xp.cos(theta)
+        omegay = k * xp.sin(theta)
+        omega = 2 * np.pi * xp.vstack((omegax.flatten("C"), omegay.flatten("C")))
 
+        # This should return either a single 2d array, or stack of 2d arrays
+        # Reshape singleton to stack of 1.
         h_vals2d = (
-            h_fun(omega, pixel_size=pixel_size).reshape(n_k, n_theta).astype(self.dtype)
+            h_fun(omega, pixel_size=pixel_size)
+            .reshape(len(f), n_k, n_theta)
+            .astype(self.dtype)
         )
-        h_vals = np.sum(h_vals2d, axis=1) / n_theta
+        h_vals = h_vals2d.sum(axis=-1) / n_theta
+
+        h_basis = self.expand_radial_vec(h_vals, **kwargs)
+        return h_basis
+
+    def filter_to_basis_mat(self, f, **kwargs):
+        """
+        See `SteerableBasis2D.filter_stack_to_basis_mats`.
+        """
+        if len(f) != 1:
+            raise RuntimeError("Unexpected filter length.")
+        return self._filter_stack_to_basis_mats(f, **kwargs)[0]
+
+    def expand_radial_vec(self, radial_vec, **kwargs):
+        """
+        Expands radial vector or stack of vetors `radial_vec` to basis matrix.
+
+        See `_filter_pts` for point set.
+
+        :param radial_vec: Array holding radial vector,
+            shaped (n_radial_pts) or (n_vectors, n_radial_pts)
+        :force_diag: Optionally flush off-diagonal elements to zero and return `DiagMatrix`
+        :return: List of `BlkDiagMatrix`, or list of `DiagMatrix`
+        """
+        force_diag = kwargs.get("force_diag", False)
+
+        # Convert vector to (1,...)
+        if radial_vec.ndim == 1:
+            radial_vec = radial_vec.reshape(1, *radial_vec.shape)
+        # Optionally transfer to GPU
+        radial_vec = xp.asarray(radial_vec)
+
+        # Set same dimensions as basis object
+        n_k = self.n_r
+        radial = self._precomp["radial"]
+        k_vals = xp.asarray(self._precomp["gl_nodes"])
+        wts = xp.asarray(self._precomp["gl_weights"])
 
         # Represent 1D function values in basis
-        h_basis = BlkDiagMatrix.empty(2 * self.ell_max + 1, dtype=self.dtype)
+        h_basis = [
+            BlkDiagMatrix.empty(2 * self.ell_max + 1, dtype=self.dtype)
+            for _ in radial_vec
+        ]
+
         ind_ell = 0
+        # Reshapes for broadcasting
+        radial_vec = radial_vec.reshape(len(h_basis), n_k, 1)
+        k_vals = k_vals.reshape(1, n_k, 1)
+        wts = wts.reshape(1, n_k, 1)
         for ell in range(0, self.ell_max + 1):
             k_max = self.k_max[ell]
-            rmat = 2 * k_vals.reshape(n_k, 1) * self.r0[ell][0:k_max].T
-            basis_vals = np.zeros_like(rmat)
+            basis_vals = xp.zeros((n_k, k_max), dtype=self.dtype)
             ind_radial = np.sum(self.k_max[0:ell])
-            basis_vals[:, 0:k_max] = radial[ind_radial : ind_radial + k_max].T
-            h_basis_vals = basis_vals * h_vals.reshape(n_k, 1)
-            h_basis_ell = basis_vals.T @ (
-                h_basis_vals * k_vals.reshape(n_k, 1) * wts.reshape(n_k, 1)
+            basis_vals[:, 0:k_max] = xp.asarray(
+                radial[ind_radial : ind_radial + k_max]
+            ).T
+            h_basis_vals = basis_vals * radial_vec
+            h_basis_ell = basis_vals.T @ (h_basis_vals * k_vals * wts)
+            h_basis_ell = xp.asnumpy(h_basis_ell)
+
+            # loop over filters in `radial_vec`,
+            #   assigning current `ind_ell`.
+            for _filter in range(len(radial_vec)):
+                _tmp = h_basis[_filter][ind_ell] = h_basis_ell[_filter]
+                if ell > 0:
+                    # for non-zero `elle` also assign `ind_ell+1`
+                    h_basis[_filter][ind_ell + 1] = _tmp
+                # On the last filter, iterate the outer `ind_ell`
+                if _filter == len(radial_vec) - 1:
+                    ind_ell += 1
+                    #   For non-zero `ell`, iterate twice.
+                    if ell > 0:
+                        ind_ell += 1
+        if force_diag:
+            logger.warning(
+                "Forcing block diagonal to diagonal. Zeroing all off diagonal values."
             )
-            h_basis[ind_ell] = h_basis_ell
-            ind_ell += 1
-            if ell > 0:
-                h_basis[ind_ell] = h_basis[ind_ell - 1]
-                ind_ell += 1
+            h_basis = [h.diag() for h in h_basis]
 
         return h_basis