def setUp(self):
self.dtype = np.float32
# Test Volume
v = Volume(
np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype(
self.dtype)).downsample(32)
# Create Sim object.
# Creates 10 projects so there is something to feed FSPCABasis.
self.src = Simulation(L=v.resolution, n=10, vols=v, dtype=v.dtype)
# Original projection image to transform.
self.orig_img = self.src.images(0, 1)
# Rotate 90 degrees in cartesian coordinates using third party tool.
self.rt90_img = Image(np.rot90(self.orig_img.asnumpy(), axes=(1, 2)))
# Prepare a Fourier Bessel Basis
self.basis = FFBBasis2D((self.orig_img.res, ) * 2, dtype=self.dtype)
self.v1 = self.basis.evaluate_t(self.orig_img)
self.v2 = self.basis.evaluate_t(self.rt90_img)
# These should _not_ be equal or the test is pointless.
self.assertFalse(np.allclose(self.v1, self.v2))
# Prepare a FSPCA Basis too.
self.fspca_basis = FSPCABasis(self.src, self.basis)
def setUp(self):
n = 32
L = 8
filters = [
RadialCTFFilter(5, 200, defocus=d, Cs=2.0, alpha=0.1)
for d in np.linspace(1.5e4, 2.5e4, 7)
]
self.dtype = np.float32
self.noise_var = 0.1848
# Initial noise filter to generate noise images.
# Noise variance is set to a value far away that is used to calculate
# covariance matrix and CWF coefficients in order to check the function
# for rebuilding positive definite covariance matrix.
noise_filter = ScalarFilter(dim=2, value=self.noise_var * 0.001)
self.src = Simulation(L,
n,
unique_filters=filters,
dtype=self.dtype,
noise_filter=noise_filter)
self.basis = FFBBasis2D((L, L), dtype=self.dtype)
self.coeff = self.basis.evaluate_t(self.src.images(0, self.src.n))
self.ctf_idx = self.src.filter_indices
self.ctf_fb = [f.fb_mat(self.basis) for f in self.src.unique_filters]
self.cov2d = RotCov2D(self.basis)
self.bcov2d = BatchedRotCov2D(self.src, self.basis, batch_size=7)
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)
else:
logger.warn(
f"Unable to downsample to {max_resolution}, using {source.L}")
source.cache()
# Specify the fast FB basis method for expending the 2D images
basis = FFBBasis2D((source.L, source.L))
# Estimate the noise of images
noise_estimator = None
if noise_type == "White":
logger.info("Estimate the noise of images using white noise method")
noise_estimator = WhiteNoiseEstimator(source)
elif noise_type == "Anisotropic":
logger.info("Estimate the noise of images using anisotropic method")
noise_estimator = AnisotropicNoiseEstimator(source)
else:
raise RuntimeError(f"Unsupported noise_type={noise_type}")
# Whiten the noise of images
logger.info("Whiten the noise of images from the noise estimator")
source.whiten(noise_estimator.filter)
if denoise_method == "CWF":
logger.info("Denoise the images using CWF cov2D method.")
denoiser = DenoiserCov2D(source, basis)
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 setUp(self):
self.dtype = np.float32
L = 8
n = 32
pixel_size = 5.0 * 65 / L
voltage = 200
defocus_min = 1.5e4
defocus_max = 2.5e4
defocus_ct = 7
self.noise_var = 1.3957e-4
noise_filter = ScalarFilter(dim=2, value=self.noise_var)
unique_filters = [
RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
for d in np.linspace(defocus_min, defocus_max, defocus_ct)
]
vols = Volume(
np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype(
self.dtype
)
) # RCOPT
vols = vols.downsample((L * np.ones(3, dtype=int))) * 1.0e3
# Since FFBBasis2D doesn't yet implement dtype, we'll set this to double to match its built in types.
sim = Simulation(
n=n,
L=L,
vols=vols,
unique_filters=unique_filters,
offsets=0.0,
amplitudes=1.0,
dtype=self.dtype,
noise_filter=noise_filter,
)
self.basis = FFBBasis2D((L, L), dtype=self.dtype)
self.h_idx = sim.filter_indices
self.h_ctf_fb = [filt.fb_mat(self.basis) for filt in unique_filters]
self.imgs_clean = sim.projections()
self.imgs_ctf_clean = sim.clean_images()
self.imgs_ctf_noise = sim.images(start=0, num=n)
self.cov2d = RotCov2D(self.basis)
self.coeff_clean = self.basis.evaluate_t(self.imgs_clean)
self.coeff = self.basis.evaluate_t(self.imgs_ctf_noise)
def setUp(self):
self.vols = Volume(
np.load(os.path.join(DATA_DIR,
"clean70SRibosome_vol.npy"))).downsample(17)
self.resolution = self.vols.resolution
self.n_img = 3
self.dtype = np.float64
# Create a Basis to use in alignment.
self.basis = FFBBasis2D((self.resolution, self.resolution),
dtype=self.dtype)
# This sets up a trivial class, where there is one group having all images.
self.classes = np.arange(self.n_img, dtype=int).reshape(1, self.n_img)
self.reflections = np.zeros(self.classes.shape, dtype=bool)
def testRotate(self):
# Now low res (8x8) had problems;
# better with odd (7x7), but still not good.
# We'll use a higher res test image.
# fh = np.load(os.path.join(DATA_DIR, 'ffbbasis2d_xcoeff_in_8_8.npy'))[:7,:7]
# Use a real data volume to generate a clean test image.
v = Volume(
np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype(
np.float64))
src = Simulation(L=v.resolution, n=1, vols=v, dtype=v.dtype)
# Extract, this is the original image to transform.
x1 = src.images(0, 1)
# Rotate 90 degrees in cartesian coordinates.
x2 = Image(np.rot90(x1.asnumpy(), axes=(1, 2)))
# Express in an FB basis
basis = FFBBasis2D((x1.res, ) * 2, dtype=x1.dtype)
v1 = basis.evaluate_t(x1)
v2 = basis.evaluate_t(x2)
v3 = basis.evaluate_t(x1)
v4 = basis.evaluate_t(x1)
# Reflect in the FB basis space
v4 = basis.rotate(v1, 0, refl=[True])
# Rotate in the FB basis space
v3 = basis.rotate(v1, 2 * np.pi)
v1 = basis.rotate(v1, -np.pi / 2)
# Evaluate back into cartesian
y1 = basis.evaluate(v1)
y2 = basis.evaluate(v2)
y3 = basis.evaluate(v3)
y4 = basis.evaluate(v4)
# Rotate 90
self.assertTrue(np.allclose(y1[0], y2[0], atol=1e-4))
# 2*pi Identity
self.assertTrue(
np.allclose(x1[0], y3[0], atol=utest_tolerance(self.dtype)))
# Refl (flipped using flipud)
self.assertTrue(np.allclose(np.flipud(x1[0]), y4[0], atol=1e-4))
def _build(self):
src = self.src
if self.basis is None:
from aspire.basis import FFBBasis2D
self.basis = FFBBasis2D((src.L, src.L), dtype=self.dtype)
if src.unique_filters is None:
logger.info("CTF filters are not included in Cov2D denoising")
# set all CTF filters to an identity filter
self.ctf_idx = np.zeros(src.n, dtype=int)
self.ctf_fb = [
BlkDiagMatrix.eye_like(RadialCTFFilter().fb_mat(self.basis))
]
else:
logger.info("Represent CTF filters in FB basis")
unique_filters = src.unique_filters
self.ctf_idx = src.filter_indices
self.ctf_fb = [f.fb_mat(self.basis) for f in unique_filters]
def __init__(self, src, basis=None, noise_var=None, components=400):
"""
:param src: Source instance
:param basis: Optional Fourier Bessel Basis (usually FFBBasis2D)
:param noise_var: None estimates noise (default).
0 forces "clean" treatment (no weighting).
Other values assigned to noise_var.
"""
self.src = src
# Automatically generate basis if needed.
if basis is None:
basis = FFBBasis2D((self.src.L, ) * 2, dtype=self.src.dtype)
self.basis = basis
# Components are used for `compress` during `build`.
self.components = components
# check/warn dtypes
self.dtype = self.src.dtype
if self.basis.dtype != self.dtype:
logger.warning(f"basis.dtype {self.basis.dtype} does not match"
f" source {self.src.dtype}, using {self.dtype}.")
self.count = self.basis.count
self.complex_count = self.basis.complex_count
self.angular_indices = self.basis.angular_indices
self.radial_indices = self.basis.radial_indices
self.signs_indices = self.basis._indices["sgns"]
self.complex_angular_indices = self.basis.complex_angular_indices
self.complex_radial_indices = self.basis.complex_radial_indices
self.complex_indices_map = self._get_complex_indices_map()
assert (len(self.complex_indices_map) == self.complex_count
), f"{len(self.complex_indices_map)} != {self.complex_count}"
self.noise_var = noise_var # noise_var is handled during `build` call.
self.build()
def setUp(self):
self.resolution = 16
self.dtype = np.float64
# Create some projections
v = Volume(
np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype(
self.dtype))
v = v.downsample(self.resolution)
# Clean
self.clean_src = Simulation(
L=self.resolution,
n=321,
vols=v,
dtype=self.dtype,
)
# With Noise
noise_var = 0.01 * np.var(np.sum(v[0], axis=0))
noise_filter = ScalarFilter(dim=2, value=noise_var)
self.noisy_src = Simulation(
L=self.resolution,
n=123,
vols=v,
dtype=self.dtype,
noise_filter=noise_filter,
)
# Set up FFB
# Setup a Basis
self.basis = FFBBasis2D((self.resolution, self.resolution),
dtype=self.dtype)
# Create Basis, use precomputed Basis
self.clean_fspca_basis = FSPCABasis(
self.clean_src, self.basis, noise_var=0
) # Note noise_var assigned zero, skips eigval filtering.
# Ceate another fspca_basis, use autogeneration FFB2D Basis
self.noisy_fspca_basis = FSPCABasis(self.noisy_src)
# Create a simulation object with specified filters and the downsampled 3D map
logger.info("Use downsampled map to creat simulation object.")
sim = Simulation(
L=img_size,
n=num_imgs,
vols=vols,
unique_filters=ctf_filters,
offsets=0.0,
amplitudes=1.0,
dtype=dtype,
noise_filter=noise_filter,
)
# Specify the fast FB basis method for expending the 2D images
ffbbasis = FFBBasis2D((img_size, img_size), dtype=dtype)
# Assign the CTF information and index for each image
h_idx = sim.filter_indices
# Evaluate CTF in the 8X8 FB basis
h_ctf_fb = [filt.fb_mat(ffbbasis) for filt in ctf_filters]
# Get clean images from projections of 3D map.
logger.info("Apply CTF filters to clean images.")
imgs_clean = sim.projections()
imgs_ctf_clean = sim.clean_images()
power_clean = imgs_ctf_clean.norm() ** 2 / imgs_ctf_clean.size
sn_ratio = power_clean / noise_var
logger.info(f"Signal to noise ratio is {sn_ratio}.")
def setUp(self):
self.dtype = np.float32 # Required for convergence of this test
self.L = 8
self.basis = FFBBasis2D((self.L, self.L), dtype=self.dtype)
plt.subplot(1, 3, 2)
plt.imshow(np.real(fb_images[0]), cmap="gray")
plt.title("FB Image")
plt.subplot(1, 3, 3)
plt.imshow(np.real(org_images[0] - fb_images[0]), cmap="gray")
plt.title("Differences")
plt.tight_layout()
# %%
# Expand Images with Fast FB Basis Method
# ---------------------------------------
# Specify the fast FB basis method for expanding the 2D images
# Note, we'll set the Basis dtype to be the same as the `org_image` data,
# as good practice.
ffb_basis = FFBBasis2D((img_size, img_size), dtype=org_images.dtype)
# Get the expansion coefficients based on fast FB basis
logger.info("start fast FB expansion of original images.")
tstart = timeit.default_timer()
ffb_coeffs = ffb_basis.evaluate_t(org_images)
tstop = timeit.default_timer()
dtime = tstop - tstart
logger.info(
f"Finish fast FB expansion of original images in {dtime:.4f} seconds.")
# Reconstruct images from the expansion coefficients based on fast FB basis
ffb_images = ffb_basis.evaluate(ffb_coeffs)
logger.info(
"Finish reconstruction of images from fast FB expansion coefficients.")