def cov3d(starfile, data_folder, pixel_size, max_rows, max_resolution, cg_tol):
"""Estimate mean volume and covariance from a starfile."""
source = RelionSource(starfile,
data_folder=data_folder,
pixel_size=pixel_size,
max_rows=max_rows)
source.downsample(max_resolution)
source.cache()
source.whiten()
basis = FBBasis3D((max_resolution, max_resolution, max_resolution))
mean_estimator = MeanEstimator(source, basis, batch_size=8192)
mean_est = mean_estimator.estimate()
noise_estimator = WhiteNoiseEstimator(source, batchSize=500)
# Estimate the noise variance. This is needed for the covariance estimation step below.
noise_variance = noise_estimator.estimate()
logger.info(f"Noise Variance = {noise_variance}")
# Passing in a mean_kernel argument to the following constructor speeds up some calculations
covar_estimator = CovarianceEstimator(source,
basis,
mean_kernel=mean_estimator.kernel)
covar_estimator.estimate(mean_est, noise_variance, tol=cg_tol)
def __init__(self, src, basis, var_noise=None):
"""
Initialize an object for denoising 2D images using Cov2D method
:param src: The source object of 2D images with metadata
:param basis: The basis method to expand 2D images
:param var_noise: The estimated variance of noise
"""
super().__init__(src)
# When var_noise is not specfically over-ridden,
# recompute it now. See #496.
if var_noise is None:
logger.info("Estimating noise of images using WhiteNoiseEstimator")
noise_estimator = WhiteNoiseEstimator(src)
var_noise = noise_estimator.estimate()
logger.info(f"Estimated Noise Variance: {var_noise}")
self.var_noise = var_noise
if not isinstance(basis, FFBBasis2D):
raise NotImplementedError("Currently only fast FB method is supported")
self.basis = basis
self.cov2d = None
self.mean_est = None
self.covar_est = None
def denoise(
data_folder,
starfile_in,
starfile_out,
pixel_size,
max_rows,
max_resolution,
noise_type,
denoise_method,
):
"""
Denoise the images and output the clean images using the default CWF method.
"""
# Create a source object for 2D images
logger.info(
f"Read in images from {starfile_in} and preprocess the images.")
source = RelionSource(starfile_in,
data_folder,
pixel_size=pixel_size,
max_rows=max_rows)
logger.info(f"Set the resolution to {max_resolution} X {max_resolution}")
if max_resolution < source.L:
# Downsample the images
source.downsample(max_resolution)
source.cache()
# Specify the fast FB basis method for expending the 2D images
basis = FFBBasis2D((max_resolution, max_resolution))
# Estimate the noise of images
noise_estimator = None
if noise_type == "Isotropic":
logger.info("Estimate the noise of images using isotropic method")
noise_estimator = WhiteNoiseEstimator(source)
else:
logger.info("Estimate the noise of images using anisotropic method")
noise_estimator = AnisotropicNoiseEstimator(source)
# Whiten the noise of images
logger.info("Whiten the noise of images from the noise estimator")
source.whiten(noise_estimator.filter)
var_noise = noise_estimator.estimate()
if denoise_method == "CWF":
logger.info("Denoise the images using CWF cov2D method.")
denoiser = DenoiserCov2D(source, basis, var_noise)
denoised_src = denoiser.denoise(batch_size=512)
denoised_src.save(starfile_out,
batch_size=512,
save_mode="single",
overwrite=False)
else:
raise NotImplementedError(
"Currently only covariance Wiener filtering method is supported")
def preprocess(
data_folder,
starfile_in,
starfile_out,
pixel_size,
max_rows,
flip_phase,
max_resolution,
normalize_background,
whiten_noise,
invert_contrast,
batch_size,
save_mode,
overwrite,
):
"""
Preprocess the raw images and output desired images for future analysis
"""
# Create a source object for 2D images
logger.info(
f"Read in images from {starfile_in} and preprocess the images.")
source = RelionSource(starfile_in,
data_folder,
pixel_size=pixel_size,
max_rows=max_rows)
if flip_phase:
logger.info("Perform phase flip to input images")
source.phase_flip()
if max_resolution < source.L:
logger.info(
f"Downsample resolution to {max_resolution} X {max_resolution}")
source.downsample(max_resolution)
if normalize_background:
logger.info("Normalize images to noise background")
source.normalize_background()
if whiten_noise:
logger.info("Whiten noise of images")
noise_estimator = WhiteNoiseEstimator(source)
source.whiten(noise_estimator.filter)
if invert_contrast:
logger.info("Invert global density contrast")
source.invert_contrast()
source.save(starfile_out,
batch_size=batch_size,
save_mode=save_mode,
overwrite=overwrite)
def testParseval(self):
"""
Here we construct a source of white noise.
Then code tests that the average noise power in the real domain,
is equivalent to the sum of the magnitudes squared
of all frequency coefficients in the Fourier domain.
These are calculated by WhiteNoiseEstimator and
AnisotropicNoiseEstimator respectively.
See Parseval/Plancherel's Theorem.
"""
wht_noise = np.random.randn(1024, 128, 128).astype(self.dtype)
src = ArrayImageSource(wht_noise)
wht_noise_estimator = WhiteNoiseEstimator(src, batchSize=512)
wht_noise_variance = wht_noise_estimator.estimate()
noise_estimator = AnisotropicNoiseEstimator(src, batchSize=512)
noise_variance = noise_estimator.estimate()
self.assertTrue(np.allclose(noise_variance, wht_noise_variance))
def build(self):
"""
Computes the FSPCA basis.
This may take some time for large image stacks.
"""
coef = self.basis.evaluate_t(self.src.images(0, self.src.n))
if self.noise_var is None:
from aspire.noise import WhiteNoiseEstimator
logger.info("Estimating the noise of images.")
self.noise_var = WhiteNoiseEstimator(self.src).estimate()
logger.info(f"Setting noise_var={self.noise_var}")
cov2d = RotCov2D(self.basis)
covar_opt = {
"shrinker": "frobenius_norm",
"verbose": 0,
"max_iter": 250,
"iter_callback": [],
"store_iterates": False,
"rel_tolerance": 1e-12,
"precision": "float64",
"preconditioner": "identity",
}
self.mean_coef_est = cov2d.get_mean(coef)
self.covar_coef_est = cov2d.get_covar(
coef,
mean_coeff=self.mean_coef_est,
noise_var=self.noise_var,
covar_est_opt=covar_opt,
)
# Create the arrays to be packed by _compute_spca
self.eigvals = np.zeros(self.basis.count, dtype=self.dtype)
self.eigvecs = BlkDiagMatrix.empty(2 * self.basis.ell_max + 1,
dtype=self.dtype)
self.spca_coef = np.zeros((self.src.n, self.basis.count),
dtype=self.dtype)
self._compute_spca(coef)
self._compress(self.components)
logger.info("Perform phase flip to input images.")
source.phase_flip()
imgs_pf = source.images(start=0, num=1).asnumpy()
max_resolution = 15
logger.info(f"Downsample resolution to {max_resolution} X {max_resolution}")
if max_resolution < source.L:
source.downsample(max_resolution)
imgs_ds = source.images(start=0, num=1).asnumpy()
logger.info("Normalize images to background noise.")
source.normalize_background()
imgs_nb = source.images(start=0, num=1).asnumpy()
logger.info("Whiten noise of images")
noise_estimator = WhiteNoiseEstimator(source)
source.whiten(noise_estimator.filter)
imgs_wt = source.images(start=0, num=1).asnumpy()
logger.info("Invert the global density contrast if need")
source.invert_contrast()
imgs_rc = source.images(start=0, num=1).asnumpy()
# plot the first images
logger.info("Plot first image from each preprocess steps")
idm = 0
plt.subplot(2, 3, 1)
plt.imshow(imgs_od[idm], cmap="gray")
plt.colorbar(orientation="horizontal")
plt.title("original image")
def testWhiteNoise(self):
noise_estimator = WhiteNoiseEstimator(self.sim, batchSize=512)
noise_variance = noise_estimator.estimate()
self.assertAlmostEqual(noise_variance, 0.00307627)
num_imgs = 1024 # number of images
num_eigs = 16 # number of eigen-vectors to keep
# Create a simulation object with specified filters
sim = Simulation(
L=img_size,
n=num_imgs,
C=num_vols,
unique_filters=[RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)],
)
# Specify the normal FB basis method for expending the 2D images
basis = FBBasis3D((img_size, img_size, img_size))
# Estimate the noise variance. This is needed for the covariance estimation step below.
noise_estimator = WhiteNoiseEstimator(sim, batchSize=500)
noise_variance = noise_estimator.estimate()
logger.info(f"Noise Variance = {noise_variance}")
# %%
# Estimate Mean Volume and Covariance
# -----------------------------------
#
# Estimate the mean. This uses conjugate gradient on the normal equations for
# the least-squares estimator of the mean volume. The mean volume is represented internally
# using the basis object, but the output is in the form of an
# L-by-L-by-L array.
mean_estimator = MeanEstimator(sim, basis)
mean_est = mean_estimator.estimate()
# The `aspire.noise <https://computationalcryoem.github.io/ASPIRE-Python/aspire.noise.html>`_
# package contains several useful classes for generating and estimating different types of noise.
#
# In this case, we know the noise to be white, so we can proceed directly to
# `WhiteNoiseEstimator <https://computationalcryoem.github.io/ASPIRE-Python/aspire.noise.html#aspire.noise.noise.WhiteNoiseEstimator>`_. The noise estimators consume from a ``Source``.
#
# The white noise estimator should log a diagnostic variance value. How does this compare with the known noise variance above?
# Create another Simulation source to tinker with.
sim_wht = Simulation(L=v2.resolution,
n=num_imgs,
vols=v2,
noise_filter=white_noise_filter)
# Estimate the white noise.
noise_estimator = WhiteNoiseEstimator(sim_wht)
# %%
# A Custom ``FunctionFilter``
# ---------------------------
#
# We will now apply some more interesting noise, using a custom function, and then apply a ``whitening`` process to our data.
#
# Using ``FunctionFilter`` we can create our own custom functions to apply in a pipeline.
# Here we want to apply a custom filter as a noise adder. We can use a function of two variables for example.
def noise_function(x, y):
return 1e-7 * np.exp(-(x * x + y * y) / (2 * 0.3**2))
def adaptive_support(img_src, energy_threshold=0.99):
"""
Determine size of the compact support in both real and Fourier Space.
Returns c_limit (support radius in Fourier space),
and R_limit (support radius in real space).
Fourier c_limit is scaled in range [0, 0.5].
R_limit is in pixels [0, Image.res/2].
:param img_src: Input `Source` of images.
:param energy_threshold: [0, 1] threshold limit
:return: (c_limit, R_limit)
"""
if not isinstance(img_src, ImageSource):
raise RuntimeError(
"adaptive_support expects `Source` instance or subclass.")
# Sanity Check Threshold is in range
if energy_threshold <= 0 or energy_threshold > 1:
raise ValueError(
f"Given energy_threshold {energy_threshold} outside sane range [0,1]"
)
L = img_src.L
N = L // 2
r = grid_2d(L, shifted=False, normalized=False, dtype=img_src.dtype)["r"]
# Estimate noise
noise_est = WhiteNoiseEstimator(img_src)
noise_var = noise_est.estimate()
# Transform to Fourier space
img = img_src.images(0, img_src.n).asnumpy()
imgf = fft.centered_fft2(img)
# Compute the Variance and Power Spectrum
# Mean along image stack.
variance_map = np.mean(np.abs(img)**2, axis=0)
pspec = np.mean(np.abs(imgf)**2, axis=0)
# Compute the Radial Variance and Radial Power Spectrum
radial_var = np.zeros(N)
radial_pspec = np.zeros(N)
for i in range(N):
mask = (r >= i) & (r < i + 1)
# Mean along radial track defined by mask
radial_var[i] = np.mean(variance_map[mask])
radial_pspec[i] = np.mean(pspec[mask])
# Subtract the noise variance
radial_pspec -= noise_var
radial_var -= noise_var
# Lower bound variance and power by 0
np.clip(radial_pspec, 0, a_max=None, out=radial_pspec)
np.clip(radial_var, 0, a_max=None, out=radial_var)
# Construct range of Fourier limits. We need a half-sample correction
# since each ring is centered between two integer radii. Same for spatial
# domain (R).
c = (np.arange(N) + 0.5) / (2 * N)
R = np.arange(N) + 0.5
# Calculate cumulative energy
cum_pspec = np.cumsum(radial_pspec * c)
cum_var = np.cumsum(radial_var * R)
# Normalize energies [0,1]
# Multiply threshold to avoid unstable division
c_energy_threshold = energy_threshold * cum_pspec[-1]
R_energy_threshold = energy_threshold * cum_var[-1]
# First note legacy code *=L for Fourier limit,
# but then only uses divided by L... so just removed here.
# This makes it consistent with Nyquist, ie [0, .5]
# Second note, we attempt to find the cutoff,
# but when a search fails returns the last (-1) element,
# essentially the maximal radius.
# Third note, to increase accuracy, we take a weighted average of the two
# points around the cutoff. This mostly affects c since R is rounded.
ind = np.argmax(cum_pspec > c_energy_threshold)
if ind > 0:
c_limit = (cum_pspec[ind - 1] * c[ind - 1] + cum_pspec[ind] *
c[ind]) / (cum_pspec[ind - 1] + cum_pspec[ind])
else:
c_limit = c[-1]
ind = np.argmax(cum_var > R_energy_threshold)
if ind > 0:
R_limit = round(
(cum_var[ind - 1] * R[ind - 1] + cum_var[ind] * R[ind]) /
(cum_var[ind - 1] + cum_var[ind]))
else:
R_limit = R[-1]
return c_limit, R_limit