def qr_vols_forward(sim, s, n, vols, k):
"""
TODO: Write docstring
TODO: Find a better place for this!
:param sim:
:param s:
:param n:
:param vols:
:param k:
:return:
"""
ims = np.zeros((sim.L, sim.L, n, k), dtype=vols.dtype)
for ell in range(k):
ims[:, :, :, ell] = sim.vol_forward(vols[:, :, :, ell], s, n)
ims = np.swapaxes(ims, 2, 3)
Q_vecs = np.zeros((sim.L**2, k, n), dtype=vols.dtype)
Rs = np.zeros((k, k, n), dtype=vols.dtype)
im_vecs = mat_to_vec(ims)
for i in range(n):
Q_vecs[:, :, i], Rs[:, :, i] = qr(im_vecs[:, :, i])
Qs = vec_to_mat(Q_vecs)
return Qs, Rs
def src_wiener_coords(sim, mean_vol, eig_vols, lambdas=None, noise_var=0, batch_size=512):
"""
Calculate coordinates using Wiener filter
:param sim: A simulation object containing the images whose coordinates we want.
:param mean_vol: The mean volume of the source in an L-by-L-by-L array.
:param eig_vols: The eigenvolumes of the source in an L-by-L-by-L-by-K array.
:param lambdas: The eigenvalues in a K-by-K diagonal matrix (default `eye(K)`).
:param noise_var: The variance of the noise in the images (default 0).
:param batch_size: The size of the batches in which to compute the coordinates (default 512).
:return: A K-by-`src.n` array of coordinates corresponding to the Wiener filter coordinates of each image in sim.
The coordinates are obtained by the formula
alpha_s = eig_vols^T H_s ( y_s - P_s mean_vol ) ,
where P_s is the forward image mapping and y_s is the sth image,
H_s = Sigma * P_s^T ( P_s Sigma P_s^T + noise_var I )^(-1) ,
and Sigma is the covariance matrix eig_vols * lambdas * eig_vols^T.
Note that when noise_var is zero, this reduces to the projecting y_s onto the span of P_s eig_vols.
# TODO: Find a better place for this functionality other than in utils
"""
k = eig_vols.shape[-1]
if lambdas is None:
lambdas = np.eye(k)
coords = np.zeros((k, sim.n))
covar_noise = noise_var * np.eye(k)
for i in range(0, sim.n, batch_size):
ims = sim.images(i, batch_size)
batch_n = ims.shape[-1]
ims -= sim.vol_forward(mean_vol, i, batch_n)
Qs, Rs = qr_vols_forward(sim, i, batch_n, eig_vols, k)
Q_vecs = mat_to_vec(Qs)
im_vecs = mat_to_vec(ims)
for j in range(batch_n):
im_coords = Q_vecs[:, :, j].T @ im_vecs[:, j]
covar_im = (Rs[:, :, j] @ lambdas @ Rs[:, :, j].T) + covar_noise
xx = solve(covar_im, im_coords)
im_coords = lambdas @ Rs[:, :, j].T @ xx
coords[:, i+j] = im_coords
return coords
def testMatToVecSymm1(self):
# We create an unsymmetric matrix and pass it to the functions as a symmetric matrix,
# just so we can closely inspect the returned values without confusion
m = np.array([[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14],
[3, 7, 11, 15]])
v = mat_to_vec(m, is_symmat=True)
# Notice the order of the elements in symmetric matrix - axis 0 first, then axis 1
self.assertTrue(
np.allclose(v, np.array([0, 1, 2, 3, 5, 6, 7, 10, 11, 15])))
def testMatToVecSymm2(self):
# We create an unsymmetric matrix and pass it to the functions as a symmetric matrix,
# just so we can closely inspect the returned values without confusion
m = np.array([[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14],
[3, 7, 11, 15]])
# Make 2 copies depthwise
m = np.dstack((m, m))
v = mat_to_vec(m, is_symmat=True)
# Notice the order of the elements in symmetric matrix - axis 0 first, then axis 1
self.assertTrue(
np.allclose(
v,
np.array([[0, 0], [1, 1], [2, 2], [3, 3], [5, 5], [6, 6],
[7, 7], [10, 10], [11, 11], [15, 15]])))
def testMatToVec2(self):
m = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# Make 2 copies depthwise
m = np.dstack((m, m))
v = mat_to_vec(m)
self.assertTrue(
np.allclose(
v,
np.array([
[1, 1],
[4, 4],
[7, 7],
[2, 2],
[5, 5],
[8, 8],
[3, 3],
[6, 6],
[9, 9],
])))
def testMatToVec1(self):
m = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
v = mat_to_vec(m)
self.assertTrue(np.allclose(v, np.array([1, 4, 7, 2, 5, 8, 3, 6, 9])))