Simplify Whiten syntax

gallery/experiments/cov2d_experiment.py.dontrun

@@ -47,7 +47,7 @@ logger.info(f"var_noise before whitening {var_noise}")
 
 # Whiten the noise of images
 logger.info(f"Whiten the noise of images from the noise estimator")
-source.whiten(noise_estimator.filter)
+source.whiten(noise_estimator)
 #   Note this changes the noise variance,
 #   flattening spectrum and converging towards 1.
 #   Noise variance will be recomputed in DenoiserCov2D by default.

gallery/experiments/cov3d_experiment.dontrun

@@ -40,7 +40,7 @@ noise_estimator = AnisotropicNoiseEstimator(source, batchSize=512)
 
 # Whiten the noise of images
 print("Whiten the noise of images from the noise estimator")
-source.whiten(noise_estimator.filter)
+source.whiten(noise_estimator)
 # Estimate the noise variance. This is needed for the covariance estimation step below.
 noise_variance = noise_estimator.estimate()
 print(f"Noise Variance = {noise_variance}")

gallery/experiments/experimental_abinitio_pipeline.py

@@ -78,7 +78,7 @@
 
 # Estimate the noise and `Whiten` based on the estimated noise
 aiso_noise_estimator = AnisotropicNoiseEstimator(src)
-src.whiten(aiso_noise_estimator.filter)
+src.whiten(aiso_noise_estimator)
 
 # Plot the noise profile for inspection
 if interactive:

gallery/experiments/preprocess_imgs_exp.py.dontrun

@@ -45,7 +45,7 @@ imgs_nb = source.images[:nimgs_ext]
 
 print('Whiten noise of images')
 noise_estimator = WhiteNoiseEstimator(source)
-source.whiten(noise_estimator.filter)
+source.whiten(noise_estimator)
 imgs_wt = source.images[:nimgs_ext]
 
 print('Invert global density contrast')

gallery/experiments/simulated_abinitio_pipeline.py

@@ -118,7 +118,7 @@ def noise_function(x, y):
 
 # Estimate the noise and `Whiten` based on the estimated noise
 aiso_noise_estimator = AnisotropicNoiseEstimator(src)
-src.whiten(aiso_noise_estimator.filter)
+src.whiten(aiso_noise_estimator)
 
 # Plot the noise profile for inspection
 if interactive:

gallery/tutorials/basic_image_array.py

@@ -139,7 +139,7 @@ def noise_function(x, y):
 
 # Once we have the estimator instance,
 #   we can use it in a transform applied to our Source.
-imgs_src.whiten(noise_estimator.filter)
+imgs_src.whiten(noise_estimator)
 
 
 # Peek at two whitened images and their corresponding spectrum.

gallery/tutorials/lecture_feature_demo.py

@@ -272,14 +272,14 @@ def noise_function(x, y):
 # Noise Whitening
 # ---------------
 #
-# Applying the ``Simulation.whiten()`` method just requires passing the filter corresponding to the estimated noise instance.
+# Applying the ``Simulation.whiten()`` method just requires passing a `NoiseEstimator` instance.
 # Then we can inspect some of the whitened images.  While noise is still present, we can see a dramatic change.
 
 # Estimate noise.
 aiso_noise_estimator = AnisotropicNoiseEstimator(sim4)
 
 # Whiten based on the estimated noise
-sim4.whiten(aiso_noise_estimator.filter)
+sim4.whiten(aiso_noise_estimator)
 
 # What do the whitened images look like...
 sim4.images[:10].show()

gallery/tutorials/preprocess_imgs_sim.py

@@ -97,7 +97,7 @@
 
 logger.info("Whiten noise of images")
 noise_estimator = WhiteNoiseEstimator(source)
-source.whiten(noise_estimator.filter)
+source.whiten(noise_estimator)
 imgs_wt = source.images[0].asnumpy()
 
 logger.info("Invert the global density contrast if need")

gallery/tutorials/starfile.py

@@ -37,7 +37,7 @@
 # Estimate noise in the ImageSource instance
 noise_estimator = AnisotropicNoiseEstimator(source)
 # Apply whitening to ImageSource
-source.whiten(noise_estimator.filter)
+source.whiten(noise_estimator)
 
 # Display subset of the images
 images = source.images[:10]

src/aspire/commands/denoise.py

@@ -96,7 +96,7 @@ def denoise(
 
     # Whiten the noise of images
     logger.info("Whiten the noise of images from the noise estimator")
-    source.whiten(noise_estimator.filter)
+    source.whiten(noise_estimator)
 
     if denoise_method == "CWF":
         logger.info("Denoise the images using CWF cov2D method.")

src/aspire/commands/extract_particles.py

@@ -148,7 +148,7 @@ def extract_particles(
         src.normalize_background()
     if whiten:
         estimator = WhiteNoiseEstimator(src)
-        src.whiten(estimator.filter)
+        src.whiten(estimator)
     if invert_contrast:
         src.invert_contrast()
 

src/aspire/commands/preprocess.py

@@ -110,7 +110,7 @@ def preprocess(
     if whiten:
         logger.info("Whiten noise of images")
         noise_estimator = WhiteNoiseEstimator(source)
-        source.whiten(noise_estimator.filter)
+        source.whiten(noise_estimator)
 
     if invert_contrast:
         logger.info("Invert global density contrast")

src/aspire/source/image.py

@@ -17,9 +17,10 @@
     Multiply,
     Pipeline,
 )
-from aspire.noise import WhiteNoiseEstimator
+from aspire.noise import NoiseEstimator, WhiteNoiseEstimator
 from aspire.operators import (
     CTFFilter,
+    Filter,
     IdentityFilter,
     MultiplicativeFilter,
     PowerFilter,
@@ -454,14 +455,29 @@ def downsample(self, L):
 
         self.L = L
 
-    def whiten(self, noise_filter):
+    def whiten(self, noise_estimate):
         """
         Modify the `ImageSource` in-place by appending a whitening filter to the generation pipeline.
 
-        :param noise_filter: The noise psd of the images as a `Filter` object. Typically determined by a
-            NoiseEstimator class, and available as its `filter` attribute.
+        :param noise_estimate: `NoiseEstimator` or `Filter`. When
+            passed a `NoiseEstimator` the `filter` attribute will be
+            queried.  Alternatively, the noise PSD may be passed
+            directly as a `Filter` object.
         :return: On return, the `ImageSource` object has been modified in place.
         """
+
+        if isinstance(noise_estimate, NoiseEstimator):
+            # Any NoiseEstimator instance should have a `filter`.
+            noise_filter = noise_estimate.filter
+        elif isinstance(noise_estimate, Filter):
+            # We were given a `Filter` object.
+            noise_filter = noise_estimate
+        else:
+            raise TypeError(
+                f"Whiten passed {noise_estimate}"
+                " instead of `NoiseEstimator` or `Filter`."
+            )
+
         logger.info("Whitening source object")
         whiten_filter = PowerFilter(noise_filter, power=-0.5)
 

tests/test_coordinate_source.py

@@ -515,7 +515,7 @@ def testPreprocessing(self):
         src.downsample(60)
         src.normalize_background()
         noise_estimator = WhiteNoiseEstimator(src)
-        src.whiten(noise_estimator.filter)
+        src.whiten(noise_estimator)
         src.invert_contrast()
         # call .images() to ensure the filters are applied
         # and not just added to pipeline

tests/test_preprocess_pipeline.py

@@ -94,7 +94,7 @@ def testWhiten(dtype):
     L = 64
     sim = get_sim_object(L, dtype)
     noise_estimator = AnisotropicNoiseEstimator(sim)
-    sim.whiten(noise_estimator.filter)
+    sim.whiten(noise_estimator)
     imgs_wt = sim.images[:num_images].asnumpy()
 
     # calculate correlation between two neighboring pixels from background
@@ -110,8 +110,10 @@ def testWhiten(dtype):
 def testWhiten2(dtype):
     # Excercises missing cases using odd image resolutions with filter.
     #  Relates to GitHub issue #401.
-    # Otherwise this is the same as testWhiten, though the accuracy
+    # Otherwise this is the similar to testWhiten, though the accuracy
     #  (atol) for odd resolutions seems slightly worse.
+    # Note, we also use this test to excercise calling `whiten`
+    #  directly with a `Filter`.
     L = 63
     sim = get_sim_object(L, dtype)
     noise_estimator = AnisotropicNoiseEstimator(sim)