def eigs(self):
"""
Eigendecomposition of volume covariance matrix of simulation
:return: A 2-tuple:
eigs_true: The eigenvectors of the volume covariance matrix in the form of an L-by-L-by-L-by-(C-1) array,
where C is the number of distinct states in the simulation
lambdas_true: The eigenvalues of the covariance matrix in the form of a (C-1)-by-(C-1) diagonal matrix.
"""
C = self.C
vols_c = self.vols - np.expand_dims(self.mean_true(), 3)
p = np.ones(C) / C
vols_c = vol_to_vec(vols_c)
Q, R = qr(vols_c, mode='economic')
# Rank is at most C-1, so remove last vector
Q = Q[:, :-1]
R = R[:-1, :]
w, v = eigh(make_symmat(R @ np.diag(p) @ R.T))
eigs_true = vec_to_vol(Q @ v)
# Arrange in descending order (flip column order in eigenvector matrix)
w = w[::-1]
eigs_true = np.flip(eigs_true, axis=-1)
return eigs_true, np.diag(w)
def src_backward(self, mean_vol, noise_variance, shrink_method=None):
"""
Apply adjoint mapping to source
:return: The sum of the outer products of the mean-subtracted images in `src`, corrected by the expected noise
contribution and expressed as coefficients of `basis`.
"""
covar_b = np.zeros((self.L, self.L, self.L, self.L, self.L, self.L),
dtype=self.as_type)
for i in range(0, self.n, self.batch_size):
im = self.src.images(i, self.batch_size)
batch_n = im.shape[-1]
im_centered = im - self.src.vol_forward(mean_vol, i,
self.batch_size)
im_centered_b = np.zeros((self.L, self.L, self.L, batch_n),
dtype=self.as_type)
for j in range(batch_n):
im_centered_b[:, :, :, j] = self.src.im_backward(
Image(im_centered[:, :, j]), i + j)
im_centered_b = vol_to_vec(im_centered_b)
covar_b += vecmat_to_volmat(
im_centered_b @ im_centered_b.T) / self.n
covar_b_coeff = self.basis.mat_evaluate_t(covar_b)
return self._shrink(covar_b_coeff, noise_variance, shrink_method)
def covar_true(self):
eigs_true, lamdbas_true = self.eigs()
eigs_true = vol_to_vec(eigs_true)
covar_true = eigs_true @ lamdbas_true @ eigs_true.T
covar_true = vecmat_to_volmat(covar_true)
return covar_true
def expand_t(self, v):
"""
Expand array in dual basis
This is a similar function to `evaluate` but with more accuracy by
using the cg optimizing of linear equation, Ax=b.
If `v` is a matrix of size `basis.ct`-by-..., `B` is the change-of-basis
matrix of this basis, and `x` is a matrix of size `self.sz`-by-...,
the function calculates x = (B * B')^(-1) * B * v, where the rows of `B`
and columns of `x` are read as vectorized arrays.
:param v: An array whose first dimension is to be expanded in this
basis's dual. This dimension must be equal to `self.count`.
:return: The coefficients of `v` expanded in the dual of `basis`. If more
than one vector is supplied in `v`, the higher dimensions of the return
value correspond to second and higher dimensions of `v`.
.. seealso:: expand
"""
ensure(v.shape[0] == self.count,
f'First dimension of v must be {self.count}')
v, sz_roll = unroll_dim(v, 2)
b = vol_to_vec(self.evaluate(v))
operator = LinearOperator(
shape=(self.nres ** 3, self.nres ** 3),
matvec=lambda x: vol_to_vec(self.evaluate(self.evaluate_t(vec_to_vol(x))))
)
# TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v
tol = 10 * np.finfo(v.dtype).eps
logger.info('Expanding array in dual basis')
v, info = cg(operator, b, tol=tol)
v = v[..., np.newaxis]
if info != 0:
raise RuntimeError('Unable to converge!')
v = roll_dim(v, sz_roll)
x = vec_to_vol(v)
return x
def vol_coords(self, mean_vol=None, eig_vols=None):
"""
Coordinates of simulation volumes in a given basis
:param mean_vol: A mean volume in the form of an L-by-L-by-L array (default `mean_true`).
:param eig_vols: A set of eigenvolumes in an L-by-L-by-L-by-K array (default `eigs`).
:return:
"""
if mean_vol is None:
mean_vol = self.mean_true()
if eig_vols is None:
eig_vols = self.eigs()[0]
vols = self.vols - np.expand_dims(mean_vol, 3)
coords = vol_to_vec(eig_vols).T @ vol_to_vec(vols)
res = vols - vec_to_vol(vol_to_vec(eig_vols) @ coords)
res_norms = np.diag(anorm(res, (0, 1, 2)))
res_inners = vol_to_vec(mean_vol).T @ vol_to_vec(res)
return coords.squeeze(), res_norms, res_inners
def compute_kernel(self):
# TODO: Most of this stuff is duplicated in MeanEstimator - move up the hierarchy?
n = self.n
L = self.L
_2L = 2 * self.L
kernel = np.zeros((_2L, _2L, _2L, _2L, _2L, _2L), dtype=self.as_type)
filters_f = self.src.filters.evaluate_grid(L)
sq_filters_f = np.array(filters_f**2, dtype=self.as_type)
for i in tqdm(range(0, n, self.batch_size)):
pts_rot = rotated_grids(L, self.src.rots[:, :,
i:i + self.batch_size])
weights = sq_filters_f[:, :,
self.src.filters.indices[i:i +
self.batch_size]]
weights *= self.src.amplitudes[i:i + self.batch_size]**2
if L % 2 == 0:
weights[0, :, :] = 0
weights[:, 0, :] = 0
# TODO: This is where this differs from MeanEstimator
pts_rot = m_reshape(pts_rot, (3, L**2, -1))
weights = m_reshape(weights, (L**2, -1))
batch_n = weights.shape[-1]
factors = np.zeros((_2L, _2L, _2L, batch_n), dtype=self.as_type)
# TODO: Numpy has got to have a functional shortcut to avoid looping like this!
for j in range(batch_n):
factors[:, :, :, j] = anufft3(weights[:, j],
pts_rot[:, :, j],
(_2L, _2L, _2L),
real=True)
factors = vol_to_vec(factors)
kernel += vecmat_to_volmat(factors @ factors.T) / (n * L**8)
# Ensure symmetric kernel
kernel[0, :, :, :, :, :] = 0
kernel[:, 0, :, :, :, :] = 0
kernel[:, :, 0, :, :, :] = 0
kernel[:, :, :, 0, :, :] = 0
kernel[:, :, :, :, 0, :] = 0
kernel[:, :, :, :, :, 0] = 0
logger.info('Computing non-centered Fourier Transform')
kernel = mdim_ifftshift(kernel, range(0, 6))
kernel_f = fftn(kernel)
# Kernel is always symmetric in spatial domain and therefore real in Fourier
kernel_f = np.real(kernel_f)
return FourierKernel(kernel_f, centered=False)
def eval_eigs(self, eigs_est, lambdas_est):
"""
Evaluate covariance eigendecomposition accuracy
:param eigs_est: The estimated volume eigenvectors in an L-by-L-by-L-by-K array.
:param lambdas_est: The estimated eigenvalues in a K-by-K diagonal matrix (default `diag(ones(K, 1))`).
:return:
"""
eigs_true, lambdas_true = self.eigs()
B = vol_to_vec(eigs_est).T @ vol_to_vec(eigs_true)
norm_true = anorm(lambdas_true)
norm_est = anorm(lambdas_est)
inner = ainner(B @ lambdas_true, lambdas_est @ B)
err = np.sqrt(norm_true**2 + norm_est**2 - 2 * inner)
rel_err = err / norm_true
corr = inner / (norm_true * norm_est)
# TODO: Determine Principal Angles and return as a dict value
return {'err': err, 'rel_err': rel_err, 'corr': corr}
def eval_coords(self, mean_vol, eig_vols, coords_est):
"""
Evaluate coordinate estimation
:param mean_vol: A mean volume in the form of an L-by-L-by-L array.
:param eig_vols: A set of eigenvolumes in an L-by-L-by-L-by-K array.
:param coords_est: The estimated coordinates in the affine space defined centered at `mean_vol` and spanned
by `eig_vols`.
:return:
"""
coords_true, res_norms, res_inners = self.vol_coords(
mean_vol, eig_vols)
# 0-indexed states vector
states = self.states - 1
coords_true = coords_true[states]
res_norms = res_norms[states]
res_inners = res_inners[states]
mean_eigs_inners = np.asscalar(
vol_to_vec(mean_vol).T @ vol_to_vec(eig_vols))
coords_err = coords_true - coords_est
err = np.hypot(res_norms, coords_err)
mean_vol_norm2 = anorm(mean_vol)**2
norm_true = np.sqrt(coords_true**2 + mean_vol_norm2 + 2 * res_inners +
2 * mean_eigs_inners * coords_true)
norm_true = np.hypot(res_norms, norm_true)
rel_err = err / norm_true
inner = mean_vol_norm2 + mean_eigs_inners * (
coords_true + coords_est) + coords_true * coords_est + res_inners
norm_est = np.sqrt(coords_est**2 + mean_vol_norm2 +
2 * mean_eigs_inners * coords_est)
corr = inner / (norm_true * norm_est)
return {'err': err, 'rel_err': rel_err, 'corr': corr}
def expand_t(self, v):
ensure(v.shape[0] == self.basis_count,
f'First dimension of v must be {self.basis_count}')
v, sz_roll = unroll_dim(v, 2)
b = vol_to_vec(self.evaluate(v))
operator = LinearOperator(
shape=(self.N**3, self.N**3),
matvec=lambda x: vol_to_vec(
self.evaluate(self.evaluate_t(vec_to_vol(x)))))
# TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v
tol = 10 * np.finfo(v.dtype).eps
logger.info('Expanding array in dual basis')
v, info = cg(operator, b, tol=tol)
if info != 0:
raise RuntimeError('Unable to converge!')
v = roll_dim(v, sz_roll)
x = vec_to_vol(v)
return x
def toeplitz(self, L=None):
"""
Compute the 3D Toeplitz matrix corresponding to this Fourier Kernel
:param L: The size of the volumes to be convolved (default M/2, where the dimensions of this Fourier Kernel
are MxMxM
:return: An six-dimensional Toeplitz matrix of size L describing the convolution of a volume with this kernel
"""
if L is None:
L = int(self.M / 2)
A = np.eye(L**3, dtype=self.as_type)
for i in range(L**3):
A[:, i] = vol_to_vec(self.convolve_volume(vec_to_vol(A[:, i])))
A = vecmat_to_volmat(A)
return A
