diff --git a/gallery/experiments/legacy_preprocessing.py b/gallery/experiments/legacy_preprocessing.py new file mode 100644 index 0000000000..e3cd441365 --- /dev/null +++ b/gallery/experiments/legacy_preprocessing.py @@ -0,0 +1,47 @@ +""" +Legacy Preprocessing Pipeline +============================= + +This notebook demonstrates reproducing preprocessing results from the +legacy ASPIRE MATLAB workflow system for the EMPIAR 10028 ribosome +dataset. + +https://www.ebi.ac.uk/empiar/EMPIAR-10028 +""" + +# %% +# Source data and Preprocessing +# ----------------------------- +# +# Load the data and apply preprocessing steps corresponding to the MATLAB prompts: +# +# >>> Phaseflip projections? (Y/N)? [Y] Y +# >>> Enter pixel size in Angstrom (-1 to read from STAR file): [-1.000000] +# >>> Number of projections to read? [105247] +# >>> Crop? (Y/N)? [Y] Y +# >>> Crop to size? [360]359 +# >>> Downsample? (Y/N)? [Y] Y +# >>> Downsample to size? [360]179 +# >>> Normalize background of images to variance 1? (Y/N)? [Y] +# >>> Prewhiten? (Y/N)? [Y] +# >>> Split data into groups? (Y/N)? [Y] N + +from aspire.source import RelionSource + +# Inputs +# Note the published ``shiny_2sets.star`` requires removal of a stray '9' character on line 5476. +starfile_in = "10028/data/shiny_2sets_fixed9.star" +# Caching, while not required, will increase speed in exchange for potentially increased memory usage. +src = RelionSource(starfile_in).cache() + +# Define preprocessing steps. +src = src.phase_flip().cache() +src = src.crop_pad(359).cache() +src = src.legacy_downsample(179).cache() +src = src.legacy_normalize_background().cache() +src = src.legacy_whiten().cache() +src = src.invert_contrast().cache() + +# Save the preprocessed images. +# `save_mode=single` will save a STAR file and single mrcs holding the image stack. +src.save("10028_legacy_preprocessed_179px.star", save_mode="single", overwrite=True) diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py index eb3aed3856..d9e136bf71 100644 --- a/src/aspire/noise/noise.py +++ b/src/aspire/noise/noise.py @@ -361,7 +361,9 @@ class LegacyNoiseEstimator(NoiseEstimator): Isotropic noise estimator ported from MATLAB `cryo_noise_estimation`. """ - def __init__(self, src, bg_radius=None, max_d=None, batch_size=512): + def __init__( + self, src, bg_radius=None, max_d=None, batch_size=512, normalize_psd=False + ): """ Given an `ImageSource`, prepares for estimation of noise spectrum. @@ -375,6 +377,8 @@ def __init__(self, src, bg_radius=None, max_d=None, batch_size=512): :param max_d: Max computed correlation distance as a proportion of `src.L`. Default of `None` uses `np.floor(src.L/3) / L`. :param batch_size: The size of the batches in which to compute the variance estimate. + :param normalize_psd: Optionally normalize PSD in way that reproduces MATLAB intermediates. + Only useful if utilizing the `PSD` for applications outside of built-in legacy whitening. """ if bg_radius is None: @@ -386,6 +390,8 @@ def __init__(self, src, bg_radius=None, max_d=None, batch_size=512): if self.max_d is None: self.max_d = np.floor(src.L / 3) / src.L + self.normalize_psd = bool(normalize_psd) + def estimate(self): """ :return: The estimated noise variance of the images. @@ -414,6 +420,7 @@ def _estimate_noise_psd(self): samples_idx, max_d_pixels, batch_size=self.batch_size, + normalize_psd=self.normalize_psd, )[0] return psd @@ -552,7 +559,7 @@ def _estimate_power_spectrum_distribution_2d( Default of `None` yields `np.floor(L / 3)`. :param batch_size: The size of the batches in which to compute the variance estimate. :normalize_psd: Optionally normalize returned PSD. - Disabled by default because it will typiccally be + Disabled by default because it will typically be renormalized later in preperation for the convolution application in `Image.legacy_whiten`. Enable to reproduce legacy PSD. @@ -567,9 +574,6 @@ def _estimate_power_spectrum_distribution_2d( L = samples_idx.shape[-1] batch_size = min(batch_size, n_img) - # Migrate mask to GPU as needed - _samples_idx = xp.asarray(samples_idx) - # Correlations more than `max_d` pixels apart are not computed. if max_d is None: max_d = np.floor(L / 3) @@ -612,25 +616,21 @@ def _estimate_power_spectrum_distribution_2d( # to estimate it. E = 0.0 # Total energy of the noise samples used to estimate the power spectrum. - samples = xp.zeros((batch_size, L, L), dtype=np.float64) + n_samples_per_img = int(np.count_nonzero(samples_idx)) + samples = xp.zeros((batch_size, n_samples_per_img), dtype=np.float64) for start in trange( 0, n_img, batch_size, desc="Estimating image noise energy" ): end = min(n_img, start + batch_size) cnt = end - start - samples[:cnt, _samples_idx] = images[start:end].asnumpy()[0][ - samples_idx - ] + samples[:cnt] = xp.asarray(images[start:end].asnumpy()[:, samples_idx]) E += xp.sum( - ( - samples[:cnt] - - xp.mean(samples[:cnt], axis=(1, 2)).reshape(cnt, 1, 1) - ) + (samples[:cnt] - xp.mean(samples[:cnt], axis=-1).reshape(cnt, 1)) ** 2 ) + # Mean energy of the noise samples - n_samples_per_img = xp.count_nonzero(_samples_idx) meanE = E / (n_samples_per_img * n_img) # Normalize P2 such that its mean energy is preserved and is equal to