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 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((k, n, sim.L, sim.L), dtype=vols.dtype) for ell in range(k): ims[ell] = sim.vol_forward(Volume(vols[ell]), s, n).asnumpy() ims = np.swapaxes(ims, 1, 3) ims = np.swapaxes(ims, 0, 2) 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 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 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])))
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 as Volume instance. :param eig_vols: The eigenvolumes of the source as Volume instance. :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 """ if not isinstance(mean_vol, Volume): logger.debug( "src_wiener_coords mean_vol should be a Volume instance. Attempt correction." ) if len(mean_vol.shape) == 4 and mean_vol.shape[3] != 1: msg = ( f"Cannot naively convert {mean_vol.shape} to Volume instance." "Please change calling code." ) logger.error(msg) raise RuntimeError(msg) mean_vol = Volume(mean_vol) if not isinstance(eig_vols, Volume): logger.debug( "src_wiener_coords eig_vols should be a Volume instance. Correcting for now." ) eig_vols = Volume(eig_vols) if not sim.dtype == mean_vol.dtype == eig_vols.dtype: logger.warning( "Inconsistent types in src_wiener_coords" f" sim {sim.dtype}," f" mean_vol {mean_vol.dtype}," f" eig_vols {eig_vols.dtype}" ) k = eig_vols.n_vols 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[0] 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) # RCOPT ims = np.moveaxis(ims.data, 0, 2) 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