def testMSE(self):
# need larger numbers of images and higher resolution for good MSE
dtype = np.float32
img_size = 64
num_imgs = 1024
noise_var = 0.1848
noise_filter = ScalarFilter(dim=2, value=noise_var)
filters = [
RadialCTFFilter(5, 200, defocus=d, Cs=2.0, alpha=0.1)
for d in np.linspace(1.5e4, 2.5e4, 7)
]
# set simulation object
sim = Simulation(
L=img_size,
n=num_imgs,
unique_filters=filters,
offsets=0.0,
amplitudes=1.0,
dtype=dtype,
noise_filter=noise_filter,
)
imgs_clean = sim.projections()
# Specify the fast FB basis method for expending the 2D images
ffbbasis = FFBBasis2D((img_size, img_size), dtype=dtype)
denoiser = DenoiserCov2D(sim, ffbbasis, noise_var)
denoised_src = denoiser.denoise(batch_size=64)
imgs_denoised = denoised_src.images(0, num_imgs)
# Calculate the normalized RMSE of the estimated images.
nrmse_ims = (imgs_denoised - imgs_clean).norm() / imgs_clean.norm()
self.assertTrue(nrmse_ims < 0.25)
def setUp(self):
L = 8
n = 32
C = 1
SNR = 1
pixel_size = 5
voltage = 200
defocus_min = 1.5e4
defocus_max = 2.5e4
defocus_ct = 7
Cs = 2.0
alpha = 0.1
filters = [
RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
for d in np.linspace(defocus_min, defocus_max, defocus_ct)
]
# Since FFBBasis2D doesn't yet implement dtype, we'll set this to double to match its built in types.
sim = Simulation(n=n, C=C, filters=filters, dtype='double')
vols = np.load(os.path.join(DATA_DIR, 'clean70SRibosome_vol.npy'))
vols = vols[..., np.newaxis]
vols = downsample(vols, (L * np.ones(3, dtype=int)))
sim.vols = vols
self.basis = FFBBasis2D((L, L))
# use new methods to generate random rotations and clean images
sim.rots = qrand_rots(n, seed=0)
self.imgs_clean = vol2img(vols[..., 0], sim.rots)
self.h_idx = np.array([filters.index(f) for f in sim.filters])
self.filters = filters
self.h_ctf_fb = [filt.fb_mat(self.basis) for filt in self.filters]
self.imgs_ctf_clean = sim.eval_filters(self.imgs_clean)
sim.cache(self.imgs_ctf_clean)
power_clean = anorm(self.imgs_ctf_clean)**2 / np.size(
self.imgs_ctf_clean)
self.noise_var = power_clean / SNR
self.imgs_ctf_noise = self.imgs_ctf_clean + np.sqrt(
self.noise_var) * randn(L, L, n, seed=0)
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 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(f'Estimate the noise of images using isotropic method')
noise_estimator = WhiteNoiseEstimator(source)
else:
logger.info(f'Estimate the noise of images using anisotropic method')
noise_estimator = AnisotropicNoiseEstimator(source)
# Whiten the noise of images
logger.info(f'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(f'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 _build(self):
src = self.src
if self.basis is None:
from aspire.basis.ffb_2d import FFBBasis2D
self.basis = FFBBasis2D((src.L, src.L))
if src.filters is None:
logger.info(f'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(f'Represent CTF filters in FB basis')
unique_filters = list(set(src.filters))
self.ctf_idx = np.array([unique_filters.index(f) for f in src.filters])
self.ctf_fb = [f.fb_mat(self.basis) for f in unique_filters]
def setUp(self):
n = 32
L = 8
noise_var = 0.1848
pixel_size = 5
voltage = 200
defocus_min = 1.5e4
defocus_max = 2.5e4
defocus_ct = 7
Cs = 2.0
alpha = 0.1
filters = [
RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
for d in np.linspace(defocus_min, defocus_max, defocus_ct)
]
# Since FFBBasis2D doesn't yet implement dtype, we'll set this to double to match its built in types.
src = Simulation(L, n, filters=filters, dtype='double')
basis = FFBBasis2D((L, L))
unique_filters = list(set(src.filters))
ctf_idx = np.array([unique_filters.index(f) for f in src.filters])
ctf_fb = [f.fb_mat(basis) for f in unique_filters]
im = src.images(0, src.n)
coeff = basis.evaluate_t(im.data).astype(src.dtype)
cov2d = RotCov2D(basis)
bcov2d = BatchedRotCov2D(src, basis, batch_size=7)
self.src = src
self.basis = basis
self.ctf_fb = ctf_fb
self.ctf_idx = ctf_idx
self.cov2d = cov2d
self.bcov2d = bcov2d
self.coeff = coeff
nrmse_ims = anorm(fb_images - org_images) / anorm(org_images)
logger.info(f'FB estimated images normalized RMSE: {nrmse_ims}')
# plot the first images using the normal FB method
plt.subplot(3, 4, 1)
plt.imshow(np.real(org_images[..., 0]), cmap='gray')
plt.title('Original')
plt.subplot(3, 4, 5)
plt.imshow(np.real(fb_images[..., 0]), cmap='gray')
plt.title('FB Image')
plt.subplot(3, 4, 9)
plt.imshow(np.real(org_images[..., 0] - fb_images[..., 0]), cmap='gray')
plt.title('Differences')
# Specify the fast FB basis method for expanding the 2D images
ffb_basis = FFBBasis2D((img_size, img_size))
# 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.')
def estimate_ctf(
data_folder,
pixel_size,
cs,
amplitude_contrast,
voltage,
num_tapers,
psd_size,
g_min,
g_max,
output_dir,
dtype=np.float32,
):
"""
Given paramaters estimates CTF from experimental data
and returns CTF as a mrc file.
"""
dtype = np.dtype(dtype)
assert dtype in (np.float32, np.float64)
dir_content = os.scandir(data_folder)
mrc_files = [
f.name for f in dir_content if os.path.splitext(f)[1] == ".mrc"
]
mrcs_files = [
f.name for f in dir_content if os.path.splitext(f)[1] == ".mrcs"
]
file_names = mrc_files + mrcs_files
amp = amplitude_contrast
amplitude_contrast = np.arctan(amplitude_contrast /
np.sqrt(1 - amplitude_contrast**2))
lmbd = voltage_to_wavelength(voltage) / 10 # (Angstrom)
ctf_object = CtfEstimator(pixel_size,
cs,
amplitude_contrast,
voltage,
psd_size,
num_tapers,
dtype=dtype)
# Note for repro debugging, suggest use of doubles,
# closer to original code.
ffbbasis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype)
results = []
for name in file_names:
with mrcfile.open(os.path.join(data_folder, name),
mode="r",
permissive=True) as mrc:
micrograph = mrc.data
# Try to match dtype used in Basis instance
micrograph = micrograph.astype(dtype, copy=False)
micrograph_blocks = ctf_object.preprocess_micrograph(
micrograph, psd_size)
tapers_1d = ctf_object.tapers(psd_size, num_tapers / 2, num_tapers)
signal_observed = ctf_object.estimate_psd(micrograph_blocks, tapers_1d)
amplitude_spectrum, _ = ctf_object.elliptical_average(
ffbbasis, signal_observed,
True) # absolute differenceL 10^-14. Relative error: 10^-7
# Optionally changing to: linprog_method='simplex',
# will more deterministically repro results in exchange for speed.
signal_1d, background_1d = ctf_object.background_subtract_1d(
amplitude_spectrum, linprog_method="interior-point")
avg_defocus, low_freq_skip = ctf_object.opt1d(
signal_1d,
pixel_size,
cs,
lmbd, # (Angstrom)
amplitude_contrast,
signal_observed.shape[-1],
)
low_freq_skip = 12
signal, background_2d = ctf_object.background_subtract_2d(
signal_observed, background_1d, low_freq_skip)
ratio = ctf_object.pca(signal_observed, pixel_size, g_min, g_max)
signal, additional_background = ctf_object.elliptical_average(
ffbbasis, signal.sqrt(), False)
background_2d = background_2d + additional_background
initial_df1 = (avg_defocus * 2) / (1 + ratio)
initial_df2 = (avg_defocus * 2) - initial_df1
grid = grid_2d(psd_size, normalized=True, dtype=dtype)
rb = np.sqrt((grid["x"] / 2)**2 + (grid["y"] / 2)**2).T
r_ctf = rb * (10 / pixel_size)
theta = grid["phi"].T
angle = -5 / 12 * np.pi # Radians (-75 degrees)
cc_array = np.zeros((6, 4))
for a in range(0, 6):
df1, df2, angle_ast, p = ctf_object.gd(
signal,
initial_df1,
initial_df2,
angle + a * np.pi / 6.0, # Radians, + a*30degrees
r_ctf,
theta,
pixel_size,
g_min,
g_max,
amplitude_contrast,
lmbd, # (Angstrom)
cs,
)
cc_array[a, 0] = df1
cc_array[a, 1] = df2
cc_array[a, 2] = angle_ast # Radians
cc_array[a, 3] = p
ml = np.argmax(cc_array[:, 3], -1)
result = (
cc_array[ml, 0],
cc_array[ml, 1],
cc_array[ml, 2], # Radians
cs,
voltage,
pixel_size,
amp,
name,
)
ctf_object.write_star(*result, output_dir)
results.append(result)
ctf_object.set_df1(cc_array[ml, 0])
ctf_object.set_df2(cc_array[ml, 1])
ctf_object.set_angle(cc_array[ml, 2]) # Radians
ctf_object.generate_ctf()
with mrcfile.new(output_dir + "/" + os.path.splitext(name)[0] +
"_noise.mrc",
overwrite=True) as mrc:
mrc.set_data(background_2d[0].astype(np.float32))
mrc.voxel_size = pixel_size
mrc.close()
df = (cc_array[ml, 0] + cc_array[ml, 1]) * np.ones(
theta.shape, theta.dtype) + (cc_array[ml, 0] - cc_array[
ml, 1]) * np.cos(2 * theta - 2 * cc_array[ml, 2] *
np.ones(theta.shape, theta.dtype))
ctf_im = -np.sin(np.pi * lmbd * r_ctf**2 / 2 *
(df - lmbd**2 * r_ctf**2 * cs * 1e6) +
amplitude_contrast)
ctf_signal = np.zeros(ctf_im.shape, ctf_im.dtype)
ctf_signal[:ctf_im.shape[0] // 2, :] = ctf_im[:ctf_im.shape[0] // 2, :]
ctf_signal[ctf_im.shape[0] // 2 +
1:, :] = signal[:, :, ctf_im.shape[0] // 2 + 1]
with mrcfile.new(output_dir + "/" + os.path.splitext(name)[0] + ".ctf",
overwrite=True) as mrc:
mrc.set_data(np.float32(ctf_signal))
mrc.voxel_size = pixel_size
mrc.close()
return results
def setUp(self):
self.basis = FFBBasis2D((8, 8))