Candidate fix for #533: Whitening applied twice to Simulation objects (#537)

* override ImageSource.whiten() so as to not whiten CTF filters in sim

* make _apply_filters accept filters and filter_indices as parameters

* formatting imports

* removed _apply_unique_filters from sim._images()

* create default value for filter_indices

* filter_indices must have a default of [0,0,...]

* apply_unique_filters -> apply_source_filters

* np.zeros() instead of np.array()

* requested change and 549 bugfix

* formatting

* cleaner solution: np.at_least1D

* spelling

* comments

* format

Author: chris-langfield

src/aspire/source/image.py

@@ -145,7 +145,7 @@ def states(self, values):
 
     @property
     def filter_indices(self):
-        return self.get_metadata("__filter_indices")
+        return np.atleast_1d(self.get_metadata("__filter_indices"))
 
     @filter_indices.setter
     def filter_indices(self, indices):
@@ -345,17 +345,18 @@ def _images(self, start=0, num=np.inf, indices=None):
             "Subclasses should implement this and return an Image object"
         )
 
-    def _apply_filters(self, im_orig, start=0, num=np.inf, indices=None):
+    def _apply_filters(
+        self,
+        im_orig,
+        filters,
+        indices,
+    ):
         """
-        For each image in `im_orig` specified by start, num, or indices,
-        the unique_filters associated with the corresponding index in the
-        `ImageSource` are applied. The images are then returned as an `Image`
-        stack.
+        For images in `im_orig`, `filters` associated with the corresponding
+        index in the supplied `indices` are applied. The images are then returned as an `Image` stack.
         :param im_orig: An `Image` object
-        :param start: Starting index of images in `im_orig`.
-        :param num: Number of images to work on, starting at `start`.
-        :param indices: A numpy array of image indices. If specified,`start` and `num` are ignored.
-        :return: An `Image` instance with the unique filters of the source applied at the given indices.
+        :param filters: A list of `Filter` objects
+        :param indices: A list of indices indicating the corresponding filter in `filters`
         """
         if not isinstance(im_orig, Image):
             logger.warning(
@@ -366,16 +367,20 @@ def _apply_filters(self, im_orig, start=0, num=np.inf, indices=None):
 
         im = im_orig.copy()
 
-        if indices is None:
-            indices = np.arange(start, min(start + num, self.n))
-
-        for i, filt in enumerate(self.unique_filters):
-            idx_k = np.where(self.filter_indices[indices] == i)[0]
+        for i, filt in enumerate(filters):
+            idx_k = np.where(indices == i)[0]
             if len(idx_k) > 0:
                 im[idx_k] = Image(im[idx_k]).filter(filt).asnumpy()
 
         return im
 
+    def _apply_source_filters(self, im_orig, indices):
+        return self._apply_filters(
+            im_orig,
+            self.unique_filters,
+            self.filter_indices[indices],
+        )
+
     def cache(self):
         logger.info("Caching source images")
         self._cached_im = self.images(start=0, num=np.inf)
@@ -531,7 +536,7 @@ def im_backward(self, im, start):
         all_idx = np.arange(start, min(start + num, self.n))
         im *= self.amplitudes[all_idx, np.newaxis, np.newaxis]
         im = im.shift(-self.offsets[all_idx, :])
-        im = self._apply_filters(im, start=start, num=num)
+        im = self._apply_source_filters(im, all_idx)
 
         vol = im.backproject(self.rots[start : start + num, :, :])[0]
 
@@ -553,7 +558,7 @@ def vol_forward(self, vol, start, num):
             logger.warning(f"Volume.dtype {vol.dtype} inconsistent with {self.dtype}")
 
         im = vol.project(0, self.rots[all_idx, :, :])
-        im = self._apply_filters(im, start, num)
+        im = self._apply_source_filters(im, all_idx)
         im = im.shift(self.offsets[all_idx, :])
         im *= self.amplitudes[all_idx, np.newaxis, np.newaxis]
         return im

src/aspire/source/simulation.py

@@ -1,3 +1,4 @@
+import copy
 import logging
 
 import numpy as np
@@ -84,12 +85,17 @@ def __init__(
         if unique_filters is None:
             unique_filters = []
         self.unique_filters = unique_filters
+        # sim_filters must be a deep copy so that it is not changed
+        # when unique_filters is changed
+        self.sim_filters = copy.deepcopy(unique_filters)
 
         # Create filter indices and fill the metadata based on unique filters
         if unique_filters:
             if filter_indices is None:
                 filter_indices = randi(len(unique_filters), n, seed=seed) - 1
             self.filter_indices = filter_indices
+        else:
+            self.filter_indices = np.zeros(n)
 
         self.offsets = offsets
         self.amplitudes = amplitudes
@@ -134,7 +140,9 @@ def _images(self, start=0, num=np.inf, indices=None, enable_noise=True):
 
         im = self.projections(start=start, num=num, indices=indices)
 
-        im = self._apply_filters(im, start=start, num=num, indices=indices)
+        # apply original CTF distortion to image
+        im = self._apply_sim_filters(im, indices)
+
         im = im.shift(self.offsets[indices, :])
 
         im *= self.amplitudes[indices].reshape(len(indices), 1, 1).astype(self.dtype)
@@ -144,6 +152,13 @@ def _images(self, start=0, num=np.inf, indices=None, enable_noise=True):
 
         return im
 
+    def _apply_sim_filters(self, im, indices):
+        return self._apply_filters(
+            im,
+            self.sim_filters,
+            self.filter_indices[indices],
+        )
+
     def vol_coords(self, mean_vol=None, eig_vols=None):
         """
         Coordinates of simulation volumes in a given basis