def _calc_op(self):
src = self.src
ctf_fb = self.ctf_fb
ctf_idx = self.ctf_idx
A_mean = BlkDiagMatrix.zeros_like(ctf_fb[0])
A_covar = [None for _ in ctf_fb]
M_covar = BlkDiagMatrix.zeros_like(ctf_fb[0])
for k in np.unique(ctf_idx):
weight = np.count_nonzero(ctf_idx == k) / src.n
ctf_fb_k = ctf_fb[k]
ctf_fb_k_t = ctf_fb_k.T
ctf_fb_k_sq = ctf_fb_k_t @ ctf_fb_k
A_mean_k = weight * ctf_fb_k_sq
A_mean += A_mean_k
A_covar_k = np.sqrt(weight) * ctf_fb_k_sq
A_covar[k] = A_covar_k
M_covar += A_covar_k
self.A_mean = A_mean
self.A_covar = A_covar
self.M_covar = M_covar
def _calc_rhs(self):
src = self.src
basis = self.basis
ctf_fb = self.ctf_fb
ctf_idx = self.ctf_idx
zero_coeff = np.zeros((basis.count, ), dtype=self.dtype)
b_mean = [np.zeros(basis.count, dtype=self.dtype) for _ in ctf_fb]
b_covar = BlkDiagMatrix.zeros_like(ctf_fb[0])
for start in range(0, src.n, self.batch_size):
batch = np.arange(start, min(start + self.batch_size, src.n))
im = src.images(batch[0], len(batch))
coeff = basis.evaluate_t(im.data)
for k in np.unique(ctf_idx[batch]):
coeff_k = coeff[ctf_idx[batch] == k]
weight = np.size(coeff_k, 0) / src.n
mean_coeff_k = self._get_mean(coeff_k)
ctf_fb_k = ctf_fb[k]
ctf_fb_k_t = ctf_fb_k.T
b_mean_k = weight * ctf_fb_k_t.apply(mean_coeff_k)
b_mean[k] += b_mean_k
covar_coeff_k = self._get_covar(coeff_k, zero_coeff)
b_covar_k = ctf_fb_k_t @ covar_coeff_k
b_covar_k = b_covar_k @ ctf_fb_k
b_covar_k *= weight
b_covar += b_covar_k
self.b_mean = b_mean
self.b_covar = b_covar
def _solve_covar(self, A_covar, b_covar, M, covar_est_opt):
ctf_fb = self.ctf_fb
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
cg_opt = covar_est_opt
covar_coeff = BlkDiagMatrix.zeros_like(ctf_fb[0])
for ell in range(0, len(b_covar)):
A_ell = []
for k in range(0, len(A_covar)):
A_ell.append(A_covar[k][ell])
p = np.size(A_ell[0], 0)
b_ell = m_reshape(b_covar[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
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!")
# should assert we require none or both...
if (ctf_fb is None) or (ctf_idx is None):
ctf_idx = np.zeros(coeffs.shape[0], dtype=int)
ctf_fb = [
BlkDiagMatrix.eye_like(RadialCTFFilter().fb_mat(self.basis),
dtype=self.dtype)
]
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 = coeff_k.shape[0] / coeffs.shape[0]
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 get_covar(
self,
coeffs,
ctf_fb=None,
ctf_idx=None,
mean_coeff=None,
do_refl=True,
noise_var=1,
covar_est_opt=None,
make_psd=True,
):
"""
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.
:param make_psd: If True, make the covariance matrix positive semidefinite
: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[0], 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": self.dtype,
"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 _ in range(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 = coeff_k.shape[0] / coeffs.shape[0]
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 not b_coeff.check_psd():
logger.warning(
"Left side b in Cov2D is not positive semidefinite.")
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"],
)
if not b.check_psd():
logger.warning("Left side b after removing noise in Cov2D"
" is not positive semidefinite.")
# RCOPT okay, this looks like a big batch, come back later
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))
if not covar_coeff.check_psd():
logger.warning(
"Covariance matrix in Cov2D is not positive semidefinite.")
if make_psd:
logger.info("Convert matrices to positive semidefinite.")
covar_coeff = covar_coeff.make_psd()
return covar_coeff
def testBlkDiagMatrixZeros(self):
blk_zeros = BlkDiagMatrix.zeros(self.blk_partition)
self.allallfunc(blk_zeros, self.blk_zeros)
blk_zeros = BlkDiagMatrix.zeros_like(self.blk_a)
self.allallfunc(blk_zeros, self.blk_zeros)