def get_cwf_coeffs(
self,
coeffs,
ctf_fb=None,
ctf_idx=None,
mean_coeff=None,
covar_coeff=None,
noise_var=1,
):
"""
Estimate the expansion coefficients using the Covariance Wiener Filtering (CWF) method.
: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 covar_coeff: The block diagonal covariance matrix of the clean coefficients represented by a cell array.
:param noise_var: The estimated variance of noise. The value should be zero for `coeffs`
from clean images of simulation data.
:return: The estimated coefficients of the unfiltered images in certain math basis.
These are obtained using a Wiener filter with the specified covariance for the clean images
and white noise of variance `noise_var` for the noise.
"""
if mean_coeff is None:
mean_coeff = self.get_mean(coeffs, ctf_fb, ctf_idx)
if covar_coeff is None:
covar_coeff = self.get_covar(coeffs,
ctf_fb,
ctf_idx,
mean_coeff,
noise_var=noise_var)
# should be 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(covar_coeff)]
noise_covar_coeff = noise_var * BlkDiagMatrix.eye_like(covar_coeff)
coeffs_est = np.zeros_like(coeffs)
for k in np.unique(ctf_idx[:]):
coeff_k = coeffs[ctf_idx == k]
ctf_fb_k = ctf_fb[k]
ctf_fb_k_t = ctf_fb_k.T
sig_covar_coeff = ctf_fb_k @ covar_coeff @ ctf_fb_k_t
sig_noise_covar_coeff = sig_covar_coeff + noise_covar_coeff
mean_coeff_k = ctf_fb_k.apply(mean_coeff)
coeff_est_k = coeff_k - mean_coeff_k
coeff_est_k = sig_noise_covar_coeff.solve(coeff_est_k.T).T
coeff_est_k = (covar_coeff @ ctf_fb_k_t).apply(coeff_est_k.T).T
coeff_est_k = coeff_est_k + mean_coeff
coeffs_est[ctf_idx == k] = coeff_est_k
return coeffs_est
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 testBlkDiagMatrixCompat(self):
""" Check incompatible matrix raises exception. """
# Create a differently shaped matrix
x = BlkDiagMatrix.from_list(self.blk_a[1:-1])
# code should raise
with pytest.raises(RuntimeError):
_ = x + self.blk_a
def testBlkDiagMatrixInPlace(self):
""" Tests sequence of in place optimized arithmetic (add, sub, mul) """
_ = [x + x + 10.0 for x in self.blk_a]
_ = [np.ones(x.shape) * 5.0 for x in self.blk_a]
# make a block diagonal object to mutate
blk_c = self.blk_a.copy()
# store the python object ids of each array in blk_c
# we want to ensure the actual object refs are _not_ changing
# for in place operations.
id0 = [id(x) for x in blk_c]
blk_c += self.blk_a
self.allallid(blk_c, id0)
blk_c += 10.0
self.allallid(blk_c, id0)
blk_a5 = BlkDiagMatrix.ones(self.blk_partition)
id1 = [id(x) for x in blk_a5]
blk_a5 *= 5.0
self.allallid(blk_a5, id1)
blk_c -= blk_a5
blk_c -= blk_a5
self.allallid(blk_c, id0)
blk_c -= self.blk_a
self.allallid(blk_c, id0)
self.allallfunc(blk_c, self.blk_a)
def _get_covar(self, coeffs, mean_coeff=None, do_refl=True):
"""
Calculate the covariance matrix from the expansion coefficients without CTF information.
:param coeffs: A coefficient vector (or an array of coefficient vectors) calculated from 2D images.
:param mean_coeff: The mean vector calculated from the `coeffs`.
:param do_refl: If true, enforce invariance to reflection (default false).
:return: The covariance matrix of coefficients for all images.
"""
if coeffs.size == 0:
raise RuntimeError("The coefficients need to be calculated first!")
if mean_coeff is None:
mean_coeff = self._get_mean(coeffs)
# Initialize a totally empty BlkDiagMatrix, build incrementally.
covar_coeff = BlkDiagMatrix.empty(0, dtype=coeffs.dtype)
ell = 0
mask = self.basis._indices["ells"] == ell
coeff_ell = coeffs[..., mask] - mean_coeff[mask]
covar_ell = np.array(coeff_ell.T @ coeff_ell / coeffs.shape[0])
covar_coeff.append(covar_ell)
for ell in range(1, self.basis.ell_max + 1):
mask = self.basis._indices["ells"] == ell
mask_pos = [
mask[i] and (self.basis._indices["sgns"][i] == +1)
for i in range(len(mask))
]
mask_neg = [
mask[i] and (self.basis._indices["sgns"][i] == -1)
for i in range(len(mask))
]
covar_ell_diag = np.array(
coeffs[:, mask_pos].T @ coeffs[:, mask_pos] +
coeffs[:, mask_neg].T @ coeffs[:, mask_neg]) / (
2 * coeffs.shape[0])
if do_refl:
covar_coeff.append(covar_ell_diag)
covar_coeff.append(covar_ell_diag)
else:
covar_ell_off = np.array(
(coeffs[:, mask_pos] @ coeffs[:, mask_neg].T /
coeffs.shape[0] -
coeffs[:, mask_pos].T @ coeffs[:, mask_neg]) /
(2 * coeffs.shape[0]))
hsize = covar_ell_diag.shape[0]
covar_coeff_blk = np.zeros((2, hsize, 2, hsize))
covar_coeff_blk[0:2, :,
0:2, :] = covar_ell_diag[:hsize, :hsize]
covar_coeff_blk[0, :, 1, :] = covar_ell_off[:hsize, :hsize]
covar_coeff_blk[1, :, 0, :] = covar_ell_off.T[:hsize, :hsize]
covar_coeff.append(
covar_coeff_blk.reshape(2 * hsize, 2 * hsize))
return covar_coeff
def setUp(self):
self.num_blks = 10
self.blk_partition = [(i, i) for i in range(self.num_blks, 0, -1)]
self.dense_shape = np.sum(self.blk_partition, axis=0)
n = np.sum(np.prod(np.array(self.blk_partition), axis=1))
self.flat = np.arange(n)
self.revflat = self.flat[::-1].copy()
diag_ind = np.array([0, 0])
ind = 0
zeros = []
ones = []
eyes = []
A = []
B = []
self.dense = np.zeros(self.dense_shape)
for blk_shp in self.blk_partition:
zeros.append(np.zeros(blk_shp))
ones.append(np.ones(blk_shp))
eyes.append(np.eye(blk_shp[0]))
offt = np.prod(blk_shp)
blk = self.flat[ind : ind + offt].reshape(blk_shp)
A.append(blk)
B.append(self.revflat[ind : ind + offt].reshape(blk_shp))
ind += offt
# Also build a dense array.
self.dense[
diag_ind[0] : diag_ind[0] + blk_shp[0],
diag_ind[1] : diag_ind[1] + blk_shp[1],
] = blk
diag_ind += blk_shp
self.blk_a = BlkDiagMatrix.from_list(A)
self.blk_b = BlkDiagMatrix.from_list(B)
self.blk_zeros = BlkDiagMatrix.from_list(zeros)
self.blk_ones = BlkDiagMatrix.from_list(ones)
self.blk_eyes = BlkDiagMatrix.from_list(eyes)
def filter_to_fb_mat(h_fun, fbasis):
"""
Convert a nonradial function in k space into a basis representation.
:param h_fun: The function form in k space.
:param fbasis: The basis object for expanding.
:return: a BlkDiagMatrix instance representation using the
`fbasis` expansion.
"""
# These form a circular dependence, import locally until time to clean up.
from aspire.basis import FFBBasis2D
from aspire.basis.basis_utils import lgwt
if not isinstance(fbasis, FFBBasis2D):
raise NotImplementedError("Currently only fast FB method is supported")
# Set same dimensions as basis object
n_k = fbasis.n_r
n_theta = fbasis.n_theta
radial = fbasis.get_radial()
# get 2D grid in polar coordinate
k_vals, wts = lgwt(n_k, 0, 0.5, dtype=fbasis.dtype)
k, theta = np.meshgrid(k_vals,
np.arange(n_theta) * 2 * np.pi / (2 * n_theta),
indexing="ij")
# Get function values in polar 2D grid and average out angle contribution
omegax = k * np.cos(theta)
omegay = k * np.sin(theta)
omega = 2 * np.pi * np.vstack((omegax.flatten("C"), omegay.flatten("C")))
h_vals2d = h_fun(omega).reshape(n_k, n_theta).astype(fbasis.dtype)
h_vals = np.sum(h_vals2d, axis=1) / n_theta
# Represent 1D function values in fbasis
h_fb = BlkDiagMatrix.empty(2 * fbasis.ell_max + 1, dtype=fbasis.dtype)
ind_ell = 0
for ell in range(0, fbasis.ell_max + 1):
k_max = fbasis.k_max[ell]
rmat = 2 * k_vals.reshape(n_k, 1) * fbasis.r0[0:k_max, ell].T
fb_vals = np.zeros_like(rmat)
ind_radial = np.sum(fbasis.k_max[0:ell])
fb_vals[:, 0:k_max] = radial[ind_radial:ind_radial + k_max].T
h_fb_vals = fb_vals * h_vals.reshape(n_k, 1)
h_fb_ell = fb_vals.T @ (h_fb_vals * k_vals.reshape(n_k, 1) *
wts.reshape(n_k, 1))
h_fb[ind_ell] = h_fb_ell
ind_ell += 1
if ell > 0:
h_fb[ind_ell] = h_fb[ind_ell - 1]
ind_ell += 1
return h_fb
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 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)
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 _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 _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 testBlkDiagMatrixOnes(self):
blk_ones = BlkDiagMatrix.ones(self.blk_partition)
self.allallfunc(blk_ones, self.blk_ones)
def get_cwf_coeffs(self,
coeffs,
ctf_fb,
ctf_idx,
mean_coeff,
covar_coeff,
noise_var=0):
"""
Estimate the expansion coefficients using the Covariance Wiener Filtering (CWF) method.
: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 covar_coeff: The block diagonal covariance matrix of the clean coefficients represented by a cell array.
:param noise_var: The estimated variance of noise. The value should be zero for `coeffs`
from clean images of simulation data.
:return: The estimated coefficients of the unfiltered images in certain math basis.
These are obtained using a Wiener filter with the specified covariance for the clean images
and white noise of variance `noise_var` for the noise.
"""
if mean_coeff is None:
mean_coeff = self.get_mean()
if covar_coeff is None:
covar_coeff = self.get_covar(noise_var=noise_var,
mean_coeff=mean_coeff)
# Handle CTF arguments.
if (ctf_fb is None) ^ (ctf_idx is None):
raise RuntimeError(
"Both `ctf_fb` and `ctf_idx` should be provided,"
" or both should be `None`."
f' Given {"ctf_fb" if ctf_idx is None else "ctf_idx"}')
elif ctf_fb is None:
# Setup defaults for CTF
ctf_idx = np.zeros(coeffs.shape[0], dtype=int)
ctf_fb = [BlkDiagMatrix.eye_like(covar_coeff)]
noise_covar_coeff = noise_var * BlkDiagMatrix.eye_like(covar_coeff)
coeffs_est = np.zeros_like(coeffs)
for k in np.unique(ctf_idx[:]):
coeff_k = coeffs[ctf_idx == k]
ctf_fb_k = ctf_fb[k]
ctf_fb_k_t = ctf_fb_k.T
mean_coeff_k = ctf_fb_k.apply(mean_coeff)
coeff_est_k = coeff_k - mean_coeff_k
if noise_var == 0:
coeff_est_k = ctf_fb_k.solve(coeff_est_k.T).T
else:
sig_covar_coeff = ctf_fb_k @ covar_coeff @ ctf_fb_k_t
sig_noise_covar_coeff = sig_covar_coeff + noise_covar_coeff
coeff_est_k = sig_noise_covar_coeff.solve(coeff_est_k.T).T
coeff_est_k = (covar_coeff @ ctf_fb_k_t).apply(coeff_est_k.T).T
coeff_est_k = coeff_est_k + mean_coeff
coeffs_est[ctf_idx == k] = coeff_est_k
return coeffs_est
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 testBlkDiagMatrixEye(self):
blk_eye = BlkDiagMatrix.eye(self.blk_partition)
self.allallfunc(blk_eye, self.blk_eyes)
blk_eye = BlkDiagMatrix.eye_like(self.blk_a)
self.allallfunc(blk_eye, self.blk_eyes)
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)
def _compute_spca(self, coef):
"""
Algorithm 2 from paper.
It has been adopted to use ASPIRE-Python's
cov2d (real) covariance estimation.
"""
# Compute coefficient vector of mean image at zeroth component
self.mean_coef_zero = self.mean_coef_est[self.angular_indices == 0]
# Make the Data matrix (A_k)
# # Construct A_k, matrix of expansion coefficients a^i_k_q
# # for image i, angular index k, radial index q,
# # (around eq 31-33)
# # Rows radial indices, columns image i.
# #
# # We can extract this directly (up to transpose) from
# # fb coef matrix where ells == angular_index
# # then use the transpose so image stack becomes columns.
# Initialize a totally empty BlkDiagMatrix, then build incrementally.
A = BlkDiagMatrix.empty(0, dtype=coef.dtype)
# Zero angular index is special case of indexing.
mask = self.basis._indices["ells"] == 0
A_0 = coef[:, mask] - self.mean_coef_zero
A.append(A_0)
# Remaining angular indices have postive and negative entries in real representation.
for ell in range(1, self.basis.ell_max +
1): # `ell` in this code is `k` from paper
mask = self.basis._indices["ells"] == ell
mask_pos = [
mask[i] and (self.basis._indices["sgns"][i] == +1)
for i in range(len(mask))
]
mask_neg = [
mask[i] and (self.basis._indices["sgns"][i] == -1)
for i in range(len(mask))
]
A.append(coef[:, mask_pos])
A.append(coef[:, mask_neg])
if len(A) != len(self.covar_coef_est):
raise RuntimeError(
"Data matrix A should have same number of blocks as Covar matrix.",
f" {len(A)} != {len(self.covar_coef_est)}",
)
# For each angular frequency (`ells` in FB code, `k` from paper)
# we use the properties of Block Diagonal Matrices to work
# on the correspong block.
eigval_index = 0
for angular_index, C_k in enumerate(self.covar_coef_est):
# # Eigen/SVD, covariance block C_k should be symmetric.
eigvals_k, eigvecs_k = np.linalg.eigh(C_k)
# Determistically enforce eigen vector sign convention
eigvecs_k = fix_signs(eigvecs_k)
# Sort eigvals_k
sorted_indices = np.argsort(-eigvals_k)
eigvals_k, eigvecs_k = (
eigvals_k[sorted_indices],
eigvecs_k[:, sorted_indices],
)
# These are the dense basis indices for this block.
basis_inds = np.arange(eigval_index, eigval_index + len(eigvals_k))
# Store the eigvals for this block, note this is a flat array.
self.eigvals[basis_inds] = eigvals_k
# Store the eigvecs, note this is a BlkDiagMatrix and is assigned incrementally.
self.eigvecs[angular_index] = eigvecs_k
# To compute new expansion coefficients using spca basis
# we combine the basis coefs using the eigen decomposition.
# Note image stack slow moving axis, otherwise this is just a
# block by block matrix multiply.
self.spca_coef[:, basis_inds] = A[angular_index] @ eigvecs_k
eigval_index += len(eigvals_k)
# Sanity check we have same dimension of eigvals and basis coefs.
if eigval_index != self.basis.count:
raise RuntimeError(
f"eigvals dimension {eigval_index} != basis coef count {self.basis.count}."
)
# Store a map of indices sorted by eigenvalue.
# We don't resort then now because this would destroy the block diagonal structure.
#
# sorted_indices[i] is the ith most powerful eigendecomposition index
#
# We can pass a full or truncated slice of sorted_indices to any array indexed by
# the coefs. This is used later for compression and index re-generation.
self.sorted_indices = np.argsort(-np.abs(self.eigvals))