def testRadialCTFFilterMultiplierGrid(self):
filter = RadialCTFFilter(defocus=2.5e4) * RadialCTFFilter(
defocus=2.5e4)
result = filter.evaluate_grid(8)
self.assertEqual(result.shape, (8, 8))
self.assertTrue(
np.allclose(
result,
np.array([[
0.461755701877834, -0.995184514498978, 0.063120922443392,
0.833250206225063, 0.961464660252150, 0.833250206225063,
0.063120922443392, -0.995184514498978
],
[
-0.995184514498978, 0.626977423649552,
0.799934516166400, 0.004814348317439,
-0.298096205735759, 0.004814348317439,
0.799934516166400, 0.626977423649552
],
[
0.063120922443392, 0.799934516166400,
-0.573061561512667, -0.999286510416273,
-0.963805291282899, -0.999286510416273,
-0.573061561512667, 0.799934516166400
],
[
0.833250206225063, 0.004814348317439,
-0.999286510416273, -0.633095739808868,
-0.368890743119366, -0.633095739808868,
-0.999286510416273, 0.004814348317439
],
[
0.961464660252150, -0.298096205735759,
-0.963805291282899, -0.368890743119366,
-0.070000000000000, -0.368890743119366,
-0.963805291282899, -0.298096205735759
],
[
0.833250206225063, 0.004814348317439,
-0.999286510416273, -0.633095739808868,
-0.368890743119366, -0.633095739808868,
-0.999286510416273, 0.004814348317439
],
[
0.063120922443392, 0.799934516166400,
-0.573061561512667, -0.999286510416273,
-0.963805291282899, -0.999286510416273,
-0.573061561512667, 0.799934516166400
],
[
-0.995184514498978, 0.626977423649552,
0.799934516166400, 0.004814348317439,
-0.298096205735759, 0.004814348317439,
0.799934516166400, 0.626977423649552
]])**2))
def get_mean(self, coeffs, ctf_fb=None, ctf_idx=None):
"""
Calculate the mean vector from the expansion coefficients with CTF information.
:param coeffs: A coefficient vector (or an array of coefficient vectors) to be averaged.
:param ctf_fb: The CFT functions in the FB expansion.
:param ctf_idx: An array of the CFT function indices for all 2D images.
If ctf_fb or ctf_idx is None, the identity filter will be applied.
:return: The mean value vector for all images.
"""
if coeffs.size == 0:
raise RuntimeError('The coefficients need to be calculated!')
if (ctf_fb is None) or (ctf_idx is None):
ctf_idx = np.zeros(coeffs.shape[1], dtype=int)
ctf_fb = [BlkDiagMatrix.eye_like(RadialCTFFilter().fb_mat(self.basis))]
b = np.zeros(self.basis.count, dtype=coeffs.dtype)
A = BlkDiagMatrix.zeros_like(ctf_fb[0])
for k in np.unique(ctf_idx[:]).T:
coeff_k = coeffs[:, ctf_idx == k]
weight = np.size(coeff_k, 1)/np.size(coeffs, 1)
mean_coeff_k = self._get_mean(coeff_k)
ctf_fb_k = ctf_fb[k]
ctf_fb_k_t = ctf_fb_k.T
b += weight * ctf_fb_k_t.apply(mean_coeff_k)
A += weight * (ctf_fb_k_t @ ctf_fb_k)
mean_coeff = A.solve(b)
return mean_coeff
def setUp(self):
sim = Simulation(n=1024,
filters=[
RadialCTFFilter(defocus=d)
for d in np.linspace(1.5e4, 2.5e4, 7)
])
basis = FBBasis3D((8, 8, 8))
self.estimator = MeanEstimator(sim, basis, preconditioner='none')
self.estimator_with_preconditioner = MeanEstimator(
sim, basis, preconditioner='circulant')
def setUp(self):
self.sim = Simulation(n=1024,
L=8,
filters=[
RadialCTFFilter(defocus=d)
for d in np.linspace(1.5e4, 2.5e4, 7)
],
seed=0,
noise_filter=IdentityFilter(),
dtype='single')
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 _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
def setUpClass(cls):
cls.sim = Simulation(n=1024,
filters=[
RadialCTFFilter(defocus=d)
for d in np.linspace(1.5e4, 2.5e4, 7)
])
basis = FBBasis3D((8, 8, 8))
cls.noise_variance = 0.0030762743633643615
cls.mean_estimator = MeanEstimator(cls.sim, basis)
cls.mean_est = np.load(os.path.join(DATA_DIR, 'mean_8_8_8.npy'))
# Passing in a mean_kernel argument to the following constructor speeds up some calculations
cls.covar_estimator = CovarianceEstimator(
cls.sim,
basis,
mean_kernel=cls.mean_estimator.kernel,
preconditioner='none')
cls.covar_estimator_with_preconditioner = CovarianceEstimator(
cls.sim,
basis,
mean_kernel=cls.mean_estimator.kernel,
preconditioner='circulant')
# Set the signal-noise ratio
sn_ratio = 1
# Specify the CTF parameters
pixel_size = 5 # Pixel size of the images (in angstroms).
voltage = 200 # Voltage (in KV)
defocus_min = 1.5e4 # Minimum defocus value (in angstroms).
defocus_max = 2.5e4 # Maximum defocus value (in angstroms).
defocus_ct = 7 # Number of defocus groups.
Cs = 2.0 # Spherical aberration
alpha = 0.1 # Amplitude contrast
logger.info('Initialize simulation object and CTF filters.')
# Create filters
filters = [
RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
for d in np.linspace(defocus_min, defocus_max, defocus_ct)
]
# Load the map file of a 70S Ribosome and downsample the 3D map to desired resolution.
# The downsampling should be done by the internal function of sim object in future.
# Below we use alternative implementation to obtain the exact result with Matlab version.
logger.info(f'Load 3D map and downsample 3D map to desired grids '
f'of {img_size} x {img_size} x {img_size}.')
infile = mrcfile.open(os.path.join(DATA_DIR, 'clean70SRibosome_vol_65p.mrc'))
vols = infile.data
vols = vols[..., np.newaxis]
vols = downsample(vols, (img_size * np.ones(3, dtype=int)))
# Create a simulation object with specified filters and the downsampled 3D map
logger.info('Use downsampled map to creat simulation object.')
parser.add_argument('--num_volumes', default=2, type=int)
parser.add_argument('--image_size', default=8, type=int)
parser.add_argument('--num_images', default=1024, type=int)
parser.add_argument('--num_eigs', default=16, type=int)
with parser.parse_args() as args:
C = args.num_volumes
L = args.image_size
n = args.num_images
sim = Simulation(
n=n,
C=C,
filters=SourceFilter(
[RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)],
n=n
)
)
basis = FBBasis3D((L, L, L))
noise_estimator = WhiteNoiseEstimator(sim, batchSize=500)
# Estimate the noise variance. This is needed for the covariance estimation step below.
noise_variance = noise_estimator.estimate()
print(f'Noise Variance = {noise_variance}')
"""
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.
"""
def setUp(self):
self.sim = Simulation(n=1024,
filters=[
RadialCTFFilter(defocus=d)
for d in np.linspace(1.5e4, 2.5e4, 7)
])
def testRadialCTFFilter(self):
filter = RadialCTFFilter(defocus=2.5e4)
result = filter.evaluate(self.omega)
self.assertEqual(result.shape, (256, ))
def get_covar(self, coeffs, ctf_fb=None, ctf_idx=None, mean_coeff=None,
do_refl=True, noise_var=1, covar_est_opt=None):
"""
Calculate the covariance matrix from the expansion coefficients and CTF information.
:param coeffs: A coefficient vector (or an array of coefficient vectors) to be calculated.
:param ctf_fb: The CFT functions in the FB expansion.
:param ctf_idx: An array of the CFT function indices for all 2D images.
If ctf_fb or ctf_idx is None, the identity filter will be applied.
:param mean_coeff: The mean value vector from all images.
:param noise_var: The estimated variance of noise. The value should be zero for `coeffs`
from clean images of simulation data.
:param covar_est_opt: The optimization parameter list for obtaining the Cov2D matrix.
:return: The basis coefficients of the covariance matrix in
the form of cell array representing a block diagonal matrix. These
block diagonal matrices are implemented as BlkDiagMatrix instances.
The covariance is calculated from the images represented by the coeffs array,
along with all possible rotations and reflections. As a result, the computed covariance
matrix is invariant to both reflection and rotation. The effect of the filters in ctf_fb
are accounted for and inverted to yield a covariance estimate of the unfiltered images.
"""
if coeffs.size == 0:
raise RuntimeError('The coefficients need to be calculated!')
if (ctf_fb is None) or (ctf_idx is None):
ctf_idx = np.zeros(coeffs.shape[1], dtype=int)
ctf_fb = [BlkDiagMatrix.eye_like(RadialCTFFilter().fb_mat(self.basis))]
def identity(x):
return x
default_est_opt = {'shrinker': 'None', 'verbose': 0, 'max_iter': 250, 'iter_callback': [],
'store_iterates': False, 'rel_tolerance': 1e-12, 'precision': 'float64',
'preconditioner': identity}
covar_est_opt = fill_struct(covar_est_opt, default_est_opt)
if mean_coeff is None:
mean_coeff = self.get_mean(coeffs, ctf_fb, ctf_idx)
b_coeff = BlkDiagMatrix.zeros_like(ctf_fb[0])
b_noise = BlkDiagMatrix.zeros_like(ctf_fb[0])
A = []
for k in range(0, len(ctf_fb)):
A.append(BlkDiagMatrix.zeros_like(ctf_fb[0]))
M = BlkDiagMatrix.zeros_like(ctf_fb[0])
for k in np.unique(ctf_idx[:]):
coeff_k = coeffs[:, ctf_idx == k]
weight = np.size(coeff_k, 1)/np.size(coeffs, 1)
ctf_fb_k = ctf_fb[k]
ctf_fb_k_t = ctf_fb_k.T
mean_coeff_k = ctf_fb_k.apply(mean_coeff)
covar_coeff_k = self._get_covar(coeff_k, mean_coeff_k)
b_coeff += weight * (ctf_fb_k_t @ covar_coeff_k @ ctf_fb_k)
ctf_fb_k_sq = ctf_fb_k_t @ ctf_fb_k
b_noise += weight * ctf_fb_k_sq
A[k] = np.sqrt(weight) * ctf_fb_k_sq
M += A[k]
if covar_est_opt['shrinker'] == 'None':
b = b_coeff - noise_var * b_noise
else:
b = self.shrink_covar_backward(b_coeff, b_noise, np.size(coeffs, 1),
noise_var, covar_est_opt['shrinker'])
cg_opt = covar_est_opt
covar_coeff = BlkDiagMatrix.zeros_like(ctf_fb[0])
def precond_fun(S, x):
p = np.size(S, 0)
ensure(np.size(x) == p*p, 'The sizes of S and x are not consistent.')
x = m_reshape(x, (p, p))
y = S @ x @ S
y = m_reshape(y, (p ** 2,))
return y
def apply(A, x):
p = np.size(A[0], 0)
x = m_reshape(x, (p, p))
y = np.zeros_like(x)
for k in range(0, len(A)):
y = y + A[k] @ x @ A[k].T
y = m_reshape(y, (p ** 2,))
return y
for ell in range(0, len(b)):
A_ell = []
for k in range(0, len(A)):
A_ell.append(A[k][ell])
p = np.size(A_ell[0], 0)
b_ell = m_reshape(b[ell], (p ** 2,))
S = inv(M[ell])
cg_opt["preconditioner"] = lambda x: precond_fun(S, x)
covar_coeff_ell, _, _ = conj_grad(lambda x: apply(A_ell, x), b_ell, cg_opt)
covar_coeff[ell] = m_reshape(covar_coeff_ell, (p, p))
return covar_coeff
