cleanup some filter eval concerns

Author: garrettwrong

src/aspire/operators/filters.py

@@ -31,14 +31,18 @@ def evaluate_src_filters_on_grid(src, indices=None):
     grid2d = grid_2d(src.L, indexing="yx", dtype=src.dtype)
     omega = np.pi * np.vstack((grid2d["x"].flatten(), grid2d["y"].flatten()))
 
-    #  xxx filter opt (eval in bulk instead of loop here), remove branch
     # Initialize h as ones to mimic an IdentityFilter when src.filter_stack is None.
     h = np.ones((omega.shape[-1], len(indices)), dtype=src.dtype)
+    #### XXX I believe this might be what Tony reported ^
+
     if src.filter_stack is not None:
-        for i, filt in enumerate(src.filter_stack):
+        # Evaluate all filters in bulk
+        filter_stack_values = src.filter_stack.evaluate(
+            omega, pixel_size=src.pixel_size
+        )
+        for i, filter_values in enumerate(filter_stack_values):
             idx_k = np.where(src.filter_indices[indices] == i)[0]
             if len(idx_k) > 0:
-                filter_values = filt.evaluate(omega, pixel_size=src.pixel_size)
                 # convert filter_values row vector to column vector and tile broadcast
                 filter_values = filter_values.reshape(-1, 1)
                 h[:, idx_k] = np.tile(filter_values, len(idx_k))

src/aspire/source/image.py

@@ -792,7 +792,7 @@ def _apply_filters(
             return im
 
         # else evaluate filters
-        # XXXX broadcast filter eval
+        # TODO broadcast filter eval
         for i, filt in enumerate(filters):
             idx_k = np.where(indices == i)[0]
             if len(idx_k) > 0:
@@ -876,7 +876,6 @@ def downsample(self, L, zero_nyquist=True, centered_fft=True):
             )
         )
 
-        # XXXX sigh
         ds_factor = self.L / L
         if self.filter_stack is not None:
             self.filter_stack = self.filter_stack.scale(ds_factor)
@@ -1004,7 +1003,7 @@ def legacy_whiten(self, noise_response=None, delta=None, batch_size=512):
         if delta is None:
             delta = np.finfo(np.float32).eps
 
-        # # XXX This "should be better" but totally breaks things.
+        # # TODO This "should be better" but totally breaks things.
         # # First guess would be to check the strange normalization.
         # # Can't fix everything at once.
         # logger.info(f"Extending filter stack by legacy whitening Filter")
@@ -1027,7 +1026,6 @@ def phase_flip(self):
         logger.info("Perform phase flip on source object")
 
         if self.filter_stack is not None:
-            # XXXX
             unique_xforms = FilterXform(self.filter_stack.sign)
 
             logger.info("Adding Phase Flip Xform to end of generation pipeline")
@@ -1811,9 +1809,11 @@ def __init__(self, src, indices, memory=None):
             _unq, _inv = np.unique(_filter_indices, return_inverse=True)
             # Repack filter_stack
             self.filter_indices = _inv
-            self.filter_stack = copy.copy(
-                src.filter_stack[_unq]
-            )  # xxx, this might just work by slicing...
+            # This would work by slicing with current code,
+            #   but if future code mutated the filter objects, that would be a problem.
+            #   Copy for safety/intent.
+            #   Deep copy may be required if future code mutates underlying objects.
+            self.filter_stack = copy.copy(src.filter_stack[_unq])
         else:
             # Pass through the None case
             self.filter_stack = src.filter_stack

src/aspire/source/simulation.py

@@ -228,7 +228,8 @@ def _populate_ctf_metadata(self, filter_indices):
         ]
 
         # Unpack the `filter_stack` params across images using `filter_indices` mapping
-        # Note this does not include the B factor term (unique to ASPIRE?,xxx should we add to star if used?)
+        # Note this does not include the B factor term (hence the truncation)
+        #   B factor term looks unique to ASPIRE, should we add to star if used?
         filter_stack_params = self.filter_stack._ctf_params()[
             :, :6
         ]  # params per filter