def eigs(A, k): """ Multidimensional partial eigendecomposition :param A: An array of size `sig_sz`-by-`sig_sz`, where `sig_sz` is a size containing d dimensions. The array represents a matrix with d indices for its rows and columns. :param k: The number of eigenvalues and eigenvectors to calculate (default 6). :return: A 2-tuple of values V: An array of eigenvectors of size `sig_sz`-by-k. D: A matrix of size k-by-k containing the corresponding eigenvalues in the diagonals. """ sig_sz = A.shape[:int(A.ndim / 2)] sig_len = np.prod(sig_sz) A = m_reshape(A, (sig_len, sig_len)) dtype = A.dtype w, v = eigh(A.astype('float64'), eigvals=(sig_len - 1 - k + 1, sig_len - 1)) # Arrange in descending order (flip column order in eigenvector matrix) and typecast to proper type w = w[::-1].astype(dtype) v = np.fliplr(v) v = m_reshape(v, sig_sz + (k, )).astype(dtype) return v, np.diag(w)
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 im_backproject(im, rot_matrices): """ Backproject images along rotation :param im: An L-by-L-by-n array of images to backproject. :param rot_matrices: An 3-by-3-by-n array of rotation matrices corresponding to viewing directions. :return: An L-by-L-by-L volumes corresponding to the sum of the backprojected images. """ L, _, n = im.shape ensure(L == im.shape[1], "im must be LxLxK") ensure(n == rot_matrices.shape[2], "No. of rotation matrices must match the number of images") pts_rot = rotated_grids(L, rot_matrices) pts_rot = m_reshape(pts_rot, (3, -1)) im_f = centered_fft2(im) / (L**2) if L % 2 == 0: im_f[0, :, :] = 0 im_f[:, 0, :] = 0 im_f = m_flatten(im_f) plan = Plan(sz=(L, L, L), fourier_pts=pts_rot) vol = np.real(plan.adjoint(im_f)) / L return vol
def compute_kernel(self): _2L = 2 * self.L kernel = np.zeros((_2L, _2L, _2L), dtype=self.as_type) filters_f = self.src.filters.evaluate_grid(self.L) sq_filters_f = np.array(filters_f ** 2, dtype=self.as_type) for i in range(0, self.n, self.batch_size): pts_rot = rotated_grids(self.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 self.L % 2 == 0: weights[0, :, :] = 0 weights[:, 0, :] = 0 pts_rot = m_reshape(pts_rot, (3, -1)) weights = m_flatten(weights) kernel += 1 / (self.n * self.L ** 4) * anufft3(weights, pts_rot, (_2L, _2L, _2L), real=True) # Ensure symmetric kernel kernel[0, :, :] = 0 kernel[:, 0, :] = 0 kernel[:, :, 0] = 0 logger.info('Computing non-centered Fourier Transform') kernel = mdim_ifftshift(kernel, range(0, 3)) kernel_f = fft2(kernel, axes=(0, 1, 2)) kernel_f = np.real(kernel_f) return FourierKernel(kernel_f, centered=False)
def _getfbzeros(self): # get upper_bound of zeros of Bessel functions upper_bound = min(self.ell_max + 1, 2 * self.N + 1) # List of number of zeros n = [] # List of zero values (each entry is an ndarray; all of possibly different lengths) zeros = [] # generate zeros of Bessel functions for each ell for ell in range(upper_bound): _n, _zeros = num_besselj_zeros(ell + (self.d - 2) / 2, self.N * np.pi / 2) if _n == 0: break else: n.append(_n) zeros.append(_zeros) # get maximum number of ell self.ell_max = len(n) - 1 # set the maximum of k for each ell self.k_max = np.array(n, dtype=int) max_num_zeros = max(len(z) for z in zeros) for i, z in enumerate(zeros): zeros[i] = np.hstack((z, np.zeros(max_num_zeros - len(z)))) self.r0 = m_reshape(np.hstack(zeros), (-1, self.ell_max + 1))
def vec_to_symmat(vec): """ Convert packed lower triangular vector to symmetric matrix :param vec: A vector of size N*(N+1)/2-by-... describing a symmetric (or Hermitian) matrix. :return: An array of size N-by-N-by-... which indexes symmetric/Hermitian matrices that occupy the first two dimensions. The lower triangular parts of these matrices consists of the corresponding vectors in vec. """ # TODO: Handle complex values in vec if np.iscomplex(vec).any(): raise NotImplementedError('Coming soon') # M represents N(N+1)/2 M = vec.shape[0] N = int(round(np.sqrt(2 * M + 0.25) - 0.5)) ensure((M == 0.5 * N * (N + 1)) and N != 0, "Vector must be of size N*(N+1)/2 for some N>0.") vec, sz_roll = unroll_dim(vec, 2) index_matrix = np.empty((N, N)) i_upper = np.triu_indices_from(index_matrix) index_matrix[i_upper] = np.arange( M) # Incrementally populate upper triangle in row major order index_matrix.T[i_upper] = index_matrix[i_upper] # Copy to lower triangle mat = vec[index_matrix.flatten('F').astype('int')] mat = m_reshape(mat, (N, N) + mat.shape[1:]) mat = roll_dim(mat, sz_roll) return mat
def evaluate_grid(self, L, *args, **kwargs): grid2d = grid_2d(L) omega = np.pi * np.vstack((grid2d['x'].flatten('F'), grid2d['y'].flatten('F'))) h = self.evaluate(omega, *args, **kwargs) h = m_reshape(h, grid2d['x'].shape) return h
def roll_dim(X, dim): # TODO: dim is still 1-indexed like in MATLAB to reduce headaches for now if len(dim) > 0: old_shape = X.shape new_shape = old_shape[:-1] + dim Y = m_reshape(X, new_shape) return Y else: return X
def vol_project(vol, rot_matrices): L = vol.shape[0] n = rot_matrices.shape[-1] pts_rot = rotated_grids(L, rot_matrices) # TODO: rotated_grids might as well give us correctly shaped array in the first place pts_rot = m_reshape(pts_rot, (3, L**2 * n)) im_f = 1. / L * Plan(vol.shape, pts_rot).transform(vol) im_f = m_reshape(im_f, (L, L, -1)) if L % 2 == 0: im_f[0, :, :] = 0 im_f[:, 0, :] = 0 im = centered_ifft2(im_f) return np.real(im)
def evaluate_grid(self, L, *args, **kwargs): # Todo: remove redundancy wrt a single Filter's evaluate_grid grid2d = grid_2d(L) omega = np.pi * np.vstack( (grid2d['x'].flatten('F'), grid2d['y'].flatten('F'))) h = self.evaluate(omega, *args, **kwargs) h = m_reshape(h, grid2d['x'].shape + (len(self.filters), )) return h
def circularize_1d(self, kernel, dim): ndim = kernel.ndim sz = kernel.shape N = int(sz[dim] / 2) top, bottom = np.split(kernel, 2, axis=dim) # Multiplier for weighted average mult_shape = [1] * ndim mult_shape[dim] = N mult_shape = tuple(mult_shape) mult = m_reshape((np.arange(N) / N), mult_shape) kernel_circ = mult * top mult = m_reshape((np.arange(N, 0, -1) / N), mult_shape) kernel_circ += mult * bottom return fftshift(kernel_circ, dim)
def vec_to_im(X): """ Unroll vectors to images :param X: N^2-by-... array. :return: An N-by-N-by-... array. """ shape = X.shape N = round(shape[0]**(1 / 2)) ensure(N**2 == shape[0], "First dimension of X must be square") return m_reshape(X, (N, N) + (shape[1:]))
def vol_to_vec(X): """ Roll up volumes into vectors :param X: N-by-N-by-N-by-... array. :return: An N^3-by-... array. """ shape = X.shape ensure(X.ndim >= 3, "Array should have at least 3 dimensions") ensure(shape[0] == shape[1] == shape[2], "Array should have first 3 dimensions identical") return m_reshape(X, (shape[0]**3,) + (shape[3:]))
def im_to_vec(im): """ Roll up images into vectors :param im: An N-by-N-by-... array. :return: An N^2-by-... array. """ shape = im.shape ensure(im.ndim >= 2, "Array should have at least 2 dimensions") ensure(shape[0] == shape[1], "Array should have first 2 dimensions identical") return m_reshape(im, (shape[0]**2,) + (shape[2:]))
def vec_to_vol(X): """ Unroll vectors to volumes :param X: N^3-by-... array. :return: An N-by-N-by-N-by-... array. """ shape = X.shape N = round(shape[0]**(1 / 3)) ensure(N**3 == shape[0], "First dimension of X must be cubic") return m_reshape(X, (N, N, N) + (shape[1:]))
def precomp(self): """ Precomute the basis functions on a polar Fourier grid. Gaussian quadrature points and weights are also generated. The sampling criterion requires n_r=4*c*R and n_theta= 16*c*R. """ n_r = int(np.ceil(4 * self.R * self.c)) r, w = lgwt(n_r, 0.0, self.c) radial = np.zeros(shape=(n_r, np.sum(self.k_max))) ind_radial = 0 for ell in range(0, self.ell_max + 1): for k in range(1, self.k_max[ell] + 1): radial[:, ind_radial] = jv(ell, self.r0[k - 1, ell] * r / self.c) # NOTE: We need to remove the factor due to the discretization here # since it is already included in our quadrature weights nrm = 1 / (np.sqrt(np.prod(self.sz))) * self.basis_norm_2d( ell, k) radial[:, ind_radial] /= nrm ind_radial += 1 n_theta = np.ceil(16 * self.c * self.R) n_theta = int((n_theta + np.mod(n_theta, 2)) / 2) # Only calculate "positive" frequencies in one half-plane. freqs_x = m_reshape(r, (n_r, 1)) @ m_reshape( np.cos(np.arange(n_theta) * 2 * pi / (2 * n_theta)), (1, n_theta)) freqs_y = m_reshape(r, (n_r, 1)) @ m_reshape( np.sin(np.arange(n_theta) * 2 * pi / (2 * n_theta)), (1, n_theta)) freqs = np.vstack((freqs_x[np.newaxis, ...], freqs_y[np.newaxis, ...])) return { 'gl_nodes': r, 'gl_weights': w, 'radial': radial, 'freqs': freqs }
def vec_to_mat(vec, is_symmat=False): """ Converts a vectorized matrix into a matrix :param vec: The vectorized representations. If the matrix is non-symmetric, this array has the dimensions N^2-by-..., but if the matrix is symmetric, the dimensions are N*(N+1)/2-by-... . :param is_symmat: True if the vectors represent symmetric matrices (default False) :return: The array of size N-by-N-by-... representing the matrices. """ if not is_symmat: sz = vec.shape N = int(round(np.sqrt(sz[0]))) ensure(sz[0] == N**2, "Vector must represent square matrix.") return m_reshape(vec, (N, N) + sz[1:]) else: return vec_to_symmat(vec)
def unroll_dim(X, dim): # TODO: dim is still 1-indexed like in MATLAB to reduce headaches for now # TODO: unroll/roll are great candidates for a context manager since they're always used in conjunction. dim = dim - 1 old_shape = X.shape new_shape = old_shape[:dim] if dim < len(old_shape): new_shape += (-1, ) if old_shape != new_shape: Y = m_reshape(X, new_shape) else: Y = X removed_dims = old_shape[dim:] return Y, removed_dims
def mat_to_vec(mat, is_symmat=False): """ Converts a matrix into vectorized form :param mat: An array of size N-by-N-by-... containing the matrices to be vectorized. :param is_symmat: Specifies whether the matrices are symmetric/Hermitian, in which case they are stored in packed form using symmat_to_vec (default False). :return: The vectorized form of the matrices, with dimension N^2-by-... or N*(N+1)/2-by-... depending on the value of is_symmat. """ if not is_symmat: sz = mat.shape N = sz[0] ensure(sz[1] == N, "Matrix must be square") return m_reshape(mat, (N**2, ) + sz[2:]) else: return symmat_to_vec(mat)
def volmat_to_vecmat(X): """ Unroll volume matrices to vector matrices :param X: A volume "matrix" of size L1-by-L1-by-L1-by-L2-by-L2-by-L2-by-... :return: A vector matrix of size L1^3-by-L2^3-by-... """ # TODO: Use context manager? shape = X.shape ensure(X.ndim >= 6, "Array should have at least 6 dimensions") ensure(shape[0] == shape[1] == shape[2], "Dimensions 1-3 should be identical") ensure(shape[3] == shape[4] == shape[5], "Dimensions 4-6 should be identical") l1 = shape[0] l2 = shape[3] return m_reshape(X, (l1**3, l2**3) + (shape[6:]))
def vecmat_to_volmat(X): """ Roll up vector matrices into volume matrices :param X: A vector matrix of size L1^3-by-L2^3-by-... :return: A volume "matrix" of size L1-by-L1-by-L1-by-L2-by-L2-by-L2-by-... """ # TODO: Use context manager? shape = X.shape ensure(X.ndim >= 2, "Array should have at least 2 dimensions") L1 = round(shape[0]**(1 / 3)) L2 = round(shape[1]**(1 / 3)) ensure(L1**3 == shape[0], "First dimension of X must be cubic") ensure(L2**3 == shape[1], "Second dimension of X must be cubic") return m_reshape(X, (L1, L1, L1, L2, L2, L2) + (shape[2:]))
def evaluate_t(self, v): """ Evaluate coefficient in dual basis :param v: The coefficient array to be evaluated. The first dimensions must equal `self.sz`. :return: The evaluation of the coefficient array `v` in the dual basis of `basis`. This is an array of vectors whose first dimension equals `self.basis_count` and whose remaining dimensions correspond to higher dimensions of `v`. """ x, sz_roll = unroll_dim(v, self.d + 1) x = m_reshape(x, new_shape=tuple([np.prod(self.sz)] + list(x.shape[self.d:]))) r_idx = self.basis_coords['r_idx'] ang_idx = self.basis_coords['ang_idx'] mask = m_flatten(self.basis_coords['mask']) ind = 0 ind_radial = 0 ind_ang = 0 v = np.zeros(shape=tuple([self.basis_count] + list(x.shape[1:]))) for ell in range(0, self.ell_max + 1): k_max = self.k_max[ell] idx_radial = ind_radial + np.arange(0, k_max) nrms = self._norms[idx_radial] radial = self._precomp['radial'][:, idx_radial] radial = radial / nrms sgns = (1, ) if ell == 0 else (1, -1) for _ in sgns: ang = self._precomp['ang'][:, ind_ang] ang_radial = np.expand_dims(ang[ang_idx], axis=1) * radial[r_idx] idx = ind + np.arange(0, k_max) v[idx] = ang_radial.T @ x[mask] ind += len(idx) ind_ang += 1 ind_radial += len(idx_radial) v = roll_dim(v, sz_roll) return v
def eval_gaussian_blobs(self, L, Q, D, mu): g = grid_3d(L) # Migration Note - Matlab (:) flattens in column-major order, so specify 'F' with flatten() coords = np.array( [g['x'].flatten('F'), g['y'].flatten('F'), g['z'].flatten('F')]) K = Q.shape[-1] vol = np.zeros(shape=(1, coords.shape[-1])).astype(self.dtype) for k in range(K): coords_k = coords - mu[:, k, np.newaxis] coords_k = Q[:, :, k] / np.sqrt(np.diag( D[:, :, k])) @ Q[:, :, k].T @ coords_k vol += np.exp(-0.5 * np.sum(np.abs(coords_k)**2, axis=0)) vol = m_reshape(vol, g['x'].shape) return vol
def evaluate(self, v): """ Evaluate coefficient vector in basis :param v: A coefficient vector (or an array of coefficient vectors) to be evaluated. The first dimension must equal `self.basis_count`. :return: The evaluation of the coefficient vector(s) `v` for this basis. This is an array whose first dimensions equal `self.z` and the remaining dimensions correspond to dimensions two and higher of `v`. """ v, sz_roll = unroll_dim(v, 2) r_idx = self.basis_coords['r_idx'] ang_idx = self.basis_coords['ang_idx'] mask = m_flatten(self.basis_coords['mask']) ind = 0 ind_radial = 0 ind_ang = 0 x = np.zeros(shape=tuple([np.prod(self.sz)] + list(v.shape[1:]))) for ell in range(0, self.ell_max + 1): k_max = self.k_max[ell] idx_radial = ind_radial + np.arange(0, k_max) nrms = self._norms[idx_radial] radial = self._precomp['radial'][:, idx_radial] radial = radial / nrms sgns = (1, ) if ell == 0 else (1, -1) for _ in sgns: ang = self._precomp['ang'][:, ind_ang] ang_radial = np.expand_dims(ang[ang_idx], axis=1) * radial[r_idx] idx = ind + np.arange(0, k_max) x[mask] += ang_radial @ v[idx] ind += len(idx) ind_ang += 1 ind_radial += len(idx_radial) x = m_reshape(x, self.sz + x.shape[1:]) x = roll_dim(x, sz_roll) return x
def rotated_grids(L, rot_matrices): """ Generate rotated Fourier grids in 3D from rotation matrices :param L: The resolution of the desired grids. :param rot_matrices: An array of size 3-by-3-by-K containing K rotation matrices :return: A set of rotated Fourier grids in three dimensions as specified by the rotation matrices. Frequencies are in the range [-pi, pi]. """ # TODO: Flattening and reshaping at end may not be necessary! grid2d = grid_2d(L) num_pts = L**2 num_rots = rot_matrices.shape[-1] pts = np.pi * np.vstack([ grid2d['x'].flatten('F'), grid2d['y'].flatten('F'), np.zeros(num_pts) ]) pts_rot = np.zeros((3, num_pts, num_rots)) for i in range(num_rots): pts_rot[:, :, i] = rot_matrices[:, :, i] @ pts pts_rot = m_reshape(pts_rot, (3, L, L, num_rots)) return pts_rot
def evaluate(self, v): """ Evaluate coefficients in standard 2D coordinate basis from those in Fourier-Bessel basis :param v: A coefficient vector (or an array of coefficient vectors) in FB basis to be evaluated. The first dimension must equal `self.basis_count`. :return x: The evaluation of the coefficient vector(s) `x` in standard 2D coordinate basis. This is an array whose first two dimensions equal `self.sz` and the remaining dimensions correspond to dimensions two and higher of `v`. """ # make should the first dimension of v is self.basis_count v = m_reshape(v, (self.basis_count, -1)) # get information on polar grids from precomputed data n_theta = np.size(self._precomp["freqs"], 2) n_r = np.size(self._precomp["freqs"], 1) # number of 2D image samples n_data = np.size(v, 1) # go through each basis function and find corresponding coefficient pf = np.zeros((n_r, 2 * n_theta, n_data), dtype=np.complex) mask = self._indices["ells"] == 0 ind = 0 idx = ind + np.arange(self.k_max[0]) pf[:, 0, :] = self._precomp["radial"][:, idx] @ v[mask, ...] ind = ind + np.size(idx) ind_pos = ind for ell in range(1, self.ell_max + 1): idx = ind + np.arange(self.k_max[ell]) idx_pos = ind_pos + np.arange(self.k_max[ell]) idx_neg = idx_pos + self.k_max[ell] v_ell = (v[idx_pos, :] - 1j * v[idx_neg, :]) / 2.0 if np.mod(ell, 2) == 1: v_ell = 1j * v_ell pf_ell = self._precomp["radial"][:, idx] @ v_ell pf[:, ell, :] = pf_ell if np.mod(ell, 2) == 0: pf[:, 2 * n_theta - ell, :] = pf_ell.conjugate() else: pf[:, 2 * n_theta - ell, :] = -pf_ell.conjugate() ind = ind + np.size(idx) ind_pos = ind_pos + 2 * self.k_max[ell] # 1D inverse FFT in the degree of polar angle pf = 2 * pi * ifft(pf, axis=1, overwrite_x=True) # Only need "positive" frequencies. hsize = int(np.size(pf, 1) / 2) pf = pf[:, 0:hsize, :] for i_r in range(0, n_r): pf[i_r, ...] = pf[i_r, ...] * (self._precomp["gl_weights"][i_r] * self._precomp["gl_nodes"][i_r]) pf = m_reshape(pf, (n_r * n_theta, n_data)) # perform inverse non-uniformly FFT transform back to 2D coordinate basis freqs = m_reshape(self._precomp["freqs"], (2, n_r * n_theta)) x = np.zeros((self.sz[0], self.sz[1], n_data), dtype=v.dtype) for isample in range(0, n_data): x[..., isample] = 2 * np.real( anufft3(pf[:, isample], 2 * pi * freqs, self.sz)) # return the x with the first two dimensions of self.sz return x
def evaluate_t(self, x): """ Evaluate coefficient in Fourier Bessel basis from those in standard 3D coordinate basis :param x: The coefficient array in the standard 3D coordinate basis to be evaluated. The first three dimensions must equal `self.sz`. :return v: The evaluation of the coefficient array `v` in the Fourier Bessel basis. This is an array of vectors whose first dimension equals `self.basis_count` and whose remaining dimensions correspond to higher dimensions of `x`. """ # ensure the first three dimensions with size of self.sz x = m_reshape(x, (self.sz[0], self.sz[1], self.sz[2], -1)) n_data = np.size(x, 3) n_r = np.size(self._precomp['radial_wtd'], 0) n_phi = np.size(self._precomp['ang_phi_wtd_even'][0], 0) n_theta = np.size(self._precomp['ang_theta_wtd'], 0) # resamping x in a polar Fourier gird using nonuniform discrete Fourier transform pf = np.zeros((n_theta * n_phi * n_r, n_data), dtype=complex) for isample in range(0, n_data): pf[..., isample] = nufft3(x[..., isample], self._precomp['fourier_pts'], self.sz) pf = m_reshape(pf, (n_theta, n_phi * n_r * n_data)) # evaluate the theta parts u_even = self._precomp['ang_theta_wtd'].T @ np.real(pf) u_odd = self._precomp['ang_theta_wtd'].T @ np.imag(pf) u_even = m_reshape(u_even, (2 * self.ell_max + 1, n_phi, n_r, n_data)) u_odd = m_reshape(u_odd, (2 * self.ell_max + 1, n_phi, n_r, n_data)) u_even = np.transpose(u_even, (1, 2, 3, 0)) u_odd = np.transpose(u_odd, (1, 2, 3, 0)) w_even = np.zeros((int(np.floor(self.ell_max / 2) + 1), n_r, 2 * self.ell_max + 1, n_data), dtype=x.dtype) w_odd = np.zeros((int(np.ceil( self.ell_max / 2)), n_r, 2 * self.ell_max + 1, n_data), dtype=x.dtype) # evaluate the phi parts for m in range(0, self.ell_max + 1): ang_phi_wtd_m_even = self._precomp['ang_phi_wtd_even'][m] ang_phi_wtd_m_odd = self._precomp['ang_phi_wtd_odd'][m] n_even_ell = np.size(ang_phi_wtd_m_even, 1) n_odd_ell = np.size(ang_phi_wtd_m_odd, 1) if m == 0: sgns = (1, ) else: sgns = (1, -1) for sgn in sgns: u_m_even = u_even[:, :, :, self.ell_max + sgn * m] u_m_odd = u_odd[:, :, :, self.ell_max + sgn * m] u_m_even = m_reshape(u_m_even, (n_phi, n_r * n_data)) u_m_odd = m_reshape(u_m_odd, (n_phi, n_r * n_data)) w_m_even = ang_phi_wtd_m_even.T @ u_m_even w_m_odd = ang_phi_wtd_m_odd.T @ u_m_odd w_m_even = m_reshape(w_m_even, (n_even_ell, n_r, n_data)) w_m_odd = m_reshape(w_m_odd, (n_odd_ell, n_r, n_data)) end = np.size(w_even, 0) w_even[end - n_even_ell:end, :, self.ell_max + sgn * m, :] = w_m_even end = np.size(w_odd, 0) w_odd[end - n_odd_ell:end, :, self.ell_max + sgn * m, :] = w_m_odd w_even = np.transpose(w_even, (1, 2, 3, 0)) w_odd = np.transpose(w_odd, (1, 2, 3, 0)) # evaluate the radial parts v = np.zeros((self.basis_count, n_data), dtype=x.dtype) for ell in range(0, self.ell_max + 1): k_max_ell = self.k_max[ell] radial_wtd = self._precomp['radial_wtd'][:, 0:k_max_ell, ell] if np.mod(ell, 2) == 0: v_ell = w_even[:, int(self.ell_max - ell):int(self.ell_max + 1 + ell), :, int(ell / 2)] else: v_ell = w_odd[:, int(self.ell_max - ell):int(self.ell_max + 1 + ell), :, int((ell - 1) / 2)] v_ell = m_reshape(v_ell, (n_r, (2 * ell + 1) * n_data)) v_ell = radial_wtd.T @ v_ell v_ell = m_reshape(v_ell, (k_max_ell * (2 * ell + 1), n_data)) # TODO: Fix this to avoid lookup each time. ind = self._indices['ells'] == ell v[ind, :] = v_ell return v
def precomp(self): """ Precomute the basis functions on a polar Fourier 3D grid. Gaussian quadrature points and weights are also generated in radical and phi dimensions. """ n_r = int(self.ell_max + 1) n_theta = int(2 * self.sz[0]) n_phi = int(self.ell_max + 1) r, wt_r = lgwt(n_r, 0.0, self.c) z, wt_z = lgwt(n_phi, -1, 1) r = m_reshape(r, (n_r, 1)) wt_r = m_reshape(wt_r, (n_r, 1)) z = m_reshape(z, (n_phi, 1)) wt_z = m_reshape(wt_z, (n_phi, 1)) phi = np.arccos(z) wt_phi = wt_z theta = 2 * pi * np.arange(n_theta).T / (2 * n_theta) theta = m_reshape(theta, (n_theta, 1)) # evaluate basis function in the radial dimension radial_wtd = np.zeros(shape=(n_r, np.max(self.k_max), self.ell_max + 1)) for ell in range(0, self.ell_max + 1): k_max_ell = self.k_max[ell] rmat = r * self.r0[0:k_max_ell, ell].T / self.c radial_ell = np.zeros_like(rmat) for ik in range(0, k_max_ell): radial_ell[:, ik] = sph_bessel(ell, rmat[:, ik]) nrm = np.abs(sph_bessel(ell + 1, self.r0[0:k_max_ell, ell].T) / 4) radial_ell = radial_ell / nrm radial_ell_wtd = r**2 * wt_r * radial_ell radial_wtd[:, 0:k_max_ell, ell] = radial_ell_wtd # evaluate basis function in the phi dimension ang_phi_wtd_even = [] ang_phi_wtd_odd = [] for m in range(0, self.ell_max + 1): n_even_ell = int( np.floor((self.ell_max - m) / 2) + 1 - np.mod(self.ell_max, 2) * np.mod(m, 2)) n_odd_ell = int(self.ell_max - m + 1 - n_even_ell) phi_wtd_m_even = np.zeros((n_phi, n_even_ell), dtype=phi.dtype) phi_wtd_m_odd = np.zeros((n_phi, n_odd_ell), dtype=phi.dtype) ind_even = 0 ind_odd = 0 for ell in range(m, self.ell_max + 1): phi_m_ell = norm_assoc_legendre(ell, m, z) nrm_inv = np.sqrt(0.5 / pi) phi_m_ell = nrm_inv * phi_m_ell phi_wtd_m_ell = wt_phi * phi_m_ell if np.mod(ell, 2) == 0: phi_wtd_m_even[:, ind_even] = phi_wtd_m_ell[:, 0] ind_even = ind_even + 1 else: phi_wtd_m_odd[:, ind_odd] = phi_wtd_m_ell[:, 0] ind_odd = ind_odd + 1 ang_phi_wtd_even.append(phi_wtd_m_even) ang_phi_wtd_odd.append(phi_wtd_m_odd) # evaluate basis function in the theta dimension ang_theta = np.zeros((n_theta, 2 * self.ell_max + 1), dtype=theta.dtype) ang_theta[:, 0:self.ell_max] = np.sqrt(2) * np.sin( theta @ m_reshape(np.arange(self.ell_max, 0, -1), (1, self.ell_max))) ang_theta[:, self.ell_max] = np.ones(n_theta, dtype=theta.dtype) ang_theta[:, self.ell_max + 1:2 * self.ell_max + 1] = np.sqrt(2) * np.cos( theta @ m_reshape(np.arange(1, self.ell_max + 1), (1, self.ell_max))) ang_theta_wtd = (2 * pi / n_theta) * ang_theta theta_grid, phi_grid, r_grid = np.meshgrid(theta, phi, r, sparse=False, indexing='ij') fourier_x = m_flatten(r_grid * np.cos(theta_grid) * np.sin(phi_grid)) fourier_y = m_flatten(r_grid * np.sin(theta_grid) * np.sin(phi_grid)) fourier_z = m_flatten(r_grid * np.cos(phi_grid)) fourier_pts = 2 * pi * np.vstack( (fourier_x[np.newaxis, ...], fourier_y[np.newaxis, ...], fourier_z[np.newaxis, ...])) return { 'radial_wtd': radial_wtd, 'ang_phi_wtd_even': ang_phi_wtd_even, 'ang_phi_wtd_odd': ang_phi_wtd_odd, 'ang_theta_wtd': ang_theta_wtd, 'fourier_pts': fourier_pts }
def evaluate(self, v): """ Evaluate coefficients in standard 3D coordinate basis from those in 3D Fourier-Bessel basis :param v: A coefficient vector (or an array of coefficient vectors) in FB basis to be evaluated. The first dimension must equal `self.basis_count`. :return x: The evaluation of the coefficient vector(s) `x` in standard 3D coordinate basis. This is an array whose first three dimensions equal `self.sz` and the remaining dimensions correspond to dimensions two and higher of `v`. """ # make should the first dimension of v is self.basis_count v = m_reshape(v, (self.basis_count, -1)) # get information on polar grids from precomputed data n_theta = np.size(self._precomp['ang_theta_wtd'], 0) n_phi = np.size(self._precomp['ang_phi_wtd_even'][0], 0) n_r = np.size(self._precomp['radial_wtd'], 0) # number of 3D image samples n_data = np.size(v, 1) u_even = np.zeros((n_r, int(2 * self.ell_max + 1), n_data, int(np.floor(self.ell_max / 2) + 1)), dtype=v.dtype) u_odd = np.zeros((n_r, int(2 * self.ell_max + 1), n_data, int(np.ceil(self.ell_max / 2))), dtype=v.dtype) # go through each basis function and find corresponding coefficient # evaluate the radial parts for ell in range(0, self.ell_max + 1): k_max_ell = self.k_max[ell] radial_wtd = self._precomp['radial_wtd'][:, 0:k_max_ell, ell] ind = self._indices['ells'] == ell v_ell = m_reshape(v[ind, :], (k_max_ell, (2 * ell + 1) * n_data)) v_ell = radial_wtd @ v_ell v_ell = m_reshape(v_ell, (n_r, 2 * ell + 1, n_data)) if np.mod(ell, 2) == 0: u_even[:, int(self.ell_max - ell):int(self.ell_max + ell + 1), :, int(ell / 2)] = v_ell else: u_odd[:, int(self.ell_max - ell):int(self.ell_max + ell + 1), :, int((ell - 1) / 2)] = v_ell u_even = np.transpose(u_even, (3, 0, 1, 2)) u_odd = np.transpose(u_odd, (3, 0, 1, 2)) w_even = np.zeros((n_phi, n_r, n_data, 2 * self.ell_max + 1), dtype=v.dtype) w_odd = np.zeros((n_phi, n_r, n_data, 2 * self.ell_max + 1), dtype=v.dtype) # evaluate the phi parts for m in range(0, self.ell_max + 1): ang_phi_wtd_m_even = self._precomp['ang_phi_wtd_even'][m] ang_phi_wtd_m_odd = self._precomp['ang_phi_wtd_odd'][m] n_even_ell = np.size(ang_phi_wtd_m_even, 1) n_odd_ell = np.size(ang_phi_wtd_m_odd, 1) if m == 0: sgns = (1, ) else: sgns = (1, -1) for sgn in sgns: end = np.size(u_even, 0) u_m_even = u_even[end - n_even_ell:end, :, self.ell_max + sgn * m, :] end = np.size(u_odd, 0) u_m_odd = u_odd[end - n_odd_ell:end, :, self.ell_max + sgn * m, :] u_m_even = m_reshape(u_m_even, (n_even_ell, n_r * n_data)) u_m_odd = m_reshape(u_m_odd, (n_odd_ell, n_r * n_data)) w_m_even = ang_phi_wtd_m_even @ u_m_even w_m_odd = ang_phi_wtd_m_odd @ u_m_odd w_m_even = m_reshape(w_m_even, (n_phi, n_r, n_data)) w_m_odd = m_reshape(w_m_odd, (n_phi, n_r, n_data)) w_even[:, :, :, self.ell_max + sgn * m] = w_m_even w_odd[:, :, :, self.ell_max + sgn * m] = w_m_odd w_even = np.transpose(w_even, (3, 0, 1, 2)) w_odd = np.transpose(w_odd, (3, 0, 1, 2)) u_even = w_even u_odd = w_odd u_even = m_reshape(u_even, (2 * self.ell_max + 1, n_phi * n_r * n_data)) u_odd = m_reshape(u_odd, (2 * self.ell_max + 1, n_phi * n_r * n_data)) # evaluate the theta parts w_even = self._precomp['ang_theta_wtd'] @ u_even w_odd = self._precomp['ang_theta_wtd'] @ u_odd pf = w_even + 1j * w_odd pf = m_reshape(pf, (n_theta * n_phi * n_r, n_data)) # perform inverse non-uniformly FFT transformation back to 3D rectangular coordinates freqs = m_reshape(self._precomp['fourier_pts'], (3, n_r * n_theta * n_phi, -1)) x = np.zeros((self.sz[0], self.sz[1], self.sz[2], n_data), dtype=v.dtype) for isample in range(0, n_data): x[..., isample] = np.real(anufft3(pf[:, isample], freqs, self.sz)) # return the x with the first three dimensions of self.sz return x
def evaluate_t(self, x): """ Evaluate coefficient in Fourier Bessel basis from those in standard 2D coordinate basis :param x: The coefficient array in the standard 2D coordinate basis to be evaluated. The first two dimensions must equal `self.sz`. :return v: The evaluation of the coefficient array `v` in the Fourier Bessel basis. This is an array of vectors whose first dimension equals `self.basis_count` and whose remaining dimensions correspond to higher dimensions of `x`. """ # ensure the first two dimensions with size of self.sz x = m_reshape(x, (self.sz[0], self.sz[1], -1)) # get information on polar grids from precomputed data n_theta = np.size(self._precomp["freqs"], 2) n_r = np.size(self._precomp["freqs"], 1) freqs = m_reshape(self._precomp["freqs"], new_shape=(2, n_r * n_theta)) # number of 2D image samples n_data = np.size(x, 2) pf = np.zeros((n_r * n_theta, n_data), dtype=complex) # resamping x in a polar Fourier gird using nonuniform discrete Fourier transform for isample in range(0, n_data): pf[..., isample] = nufft3(x[..., isample], 2 * pi * freqs, self.sz) pf = m_reshape(pf, new_shape=(n_r, n_theta, n_data)) # Recover "negative" frequencies from "positive" half plane. pf = np.concatenate((pf, pf.conjugate()), axis=1) # evaluate radial integral using the Gauss-Legendre quadrature rule for i_r in range(0, n_r): pf[i_r, ...] = pf[i_r, ...] * (self._precomp["gl_weights"][i_r] * self._precomp["gl_nodes"][i_r]) # 1D FFT on the angular dimension for each concentric circle pf = 2 * pi / (2 * n_theta) * fft(pf, 2 * n_theta, 1) # This only makes it easier to slice the array later. v = np.zeros((self.basis_count, n_data), dtype=x.dtype) # go through each basis function and find the corresponding coefficient ind = 0 idx = ind + np.arange(self.k_max[0]) mask = self._indices["ells"] == 0 v[mask, :] = self._precomp["radial"][:, idx].T @ pf[:, 0, :].real v = m_reshape(v, (self.basis_count, -1)) ind = ind + np.size(idx) ind_pos = ind for ell in range(1, self.ell_max + 1): idx = ind + np.arange(self.k_max[ell]) idx_pos = ind_pos + np.arange(self.k_max[ell]) idx_neg = idx_pos + self.k_max[ell] v_ell = self._precomp["radial"][:, idx].T @ pf[:, ell, :] if np.mod(ell, 2) == 0: v_pos = np.real(v_ell) v_neg = -np.imag(v_ell) else: v_pos = np.imag(v_ell) v_neg = np.real(v_ell) v[idx_pos, :] = v_pos v[idx_neg, :] = v_neg ind = ind + np.size(idx) ind_pos = ind_pos + 2 * self.k_max[ell] # return v coefficients with the first dimension of self.basis_count return v