def convolve_volume(self, x): """ Convolve volume with kernel :param x: An N-by-N-by-N-by-... array of volumes to be convolved. :return: The original volumes convolved by the kernel with the same dimensions as before. """ N = x.shape[0] kernel_f = self.kernel N_ker = kernel_f.shape[0] x, sz_roll = unroll_dim(x, 4) ensure(x.shape[0] == x.shape[1] == x.shape[2] == N, "Volumes in x must be cubic") ensure(kernel_f.ndim == 3, "Convolution kernel must be cubic") ensure( len(set(kernel_f.shape)) == 1, "Convolution kernel must be cubic") is_singleton = x.ndim == 3 if is_singleton: x = fftn(x, (N_ker, N_ker, N_ker)) else: raise NotImplementedError('not yet') x = x * kernel_f if is_singleton: x = np.real(ifftn(x)) x = x[:N, :N, :N] else: raise NotImplementedError('not yet') x = roll_dim(x, sz_roll) return x
def evaluate(self, v): """ Evaluate coefficients in standard 2D coordinate basis from those in polar Fourier basis :param v: A coefficient vector (or an array of coefficient vectors) in polar Fourier basis to be evaluated. The first dimension must equal to `self.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`. """ v, sz_roll = unroll_dim(v, 2) nimgs = v.shape[1] half_size = self.nrad * self.ntheta // 2 v = m_reshape(v, (self.nrad, self.ntheta, nimgs)) v = (v[:, :self.ntheta // 2, :] + v[:, self.ntheta // 2:, :].conj()) v = m_reshape(v, (half_size, nimgs)) # finufftpy require it to be aligned in fortran order x = np.empty((self._sz_prod, nimgs), dtype='complex128', order='F') finufftpy.nufft2d1many(self.freqs[0, :half_size], self.freqs[1, :half_size], v, 1, 1e-15, self.sz[0], self.sz[1], x) x = m_reshape(x, (self.sz[0], self.sz[1], nimgs)) x = x.real # return coefficients whose first two dimensions equal to self.sz x = roll_dim(x, sz_roll) return x
def evaluate_t(self, x): """ Evaluate coefficient in polar Fourier grid 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 polar Fourier grid. This is an array of vectors whose first dimension is `self.count` and whose remaining dimensions correspond to higher dimensions of `x`. """ # ensure the first two dimensions with size of self.sz x, sz_roll = unroll_dim(x, self.ndim + 1) nimgs = x.shape[2] # finufftpy require it to be aligned in fortran order half_size = self.nrad * self.ntheta // 2 pf = np.empty((half_size, nimgs), dtype='complex128', order='F') finufftpy.nufft2d2many(self.freqs[0, :half_size], self.freqs[1, :half_size], pf, 1, 1e-15, x) pf = m_reshape(pf, (self.nrad, self.ntheta // 2, nimgs)) v = np.concatenate((pf, pf.conj()), axis=1) # return v coefficients with the first dimension size of self.count v = m_reshape(v, (self.nrad * self.ntheta, nimgs)) v = roll_dim(v, sz_roll) return v
def im_filter(im, filt, *args, **kwargs): # TODO: Move inside appropriate object L = im.shape[0] im, sz_roll = unroll_dim(im, 3) filter_vals = filt.evaluate_grid(L, *args, **kwargs) im_f = centered_fft2(im) if im_f.ndim > filter_vals.ndim: im_f = np.expand_dims(filter_vals, 2) * im_f else: im_f = filter_vals * im_f im = centered_ifft2(im_f) im = np.real(im) im = roll_dim(im, sz_roll) return im
def evaluate(self, v): """ Evaluate coefficients in standard coordinate basis from those in Dirac basis :param v: A coefficient vector (or an array of coefficient vectors) to be evaluated. The first dimension must equal `self.count`. :return: The evaluation of the coefficient vector(s) `v` for this basis. This is an array whose first dimensions equal `self.sz` and the remaining dimensions correspond to dimensions two and higher of `v`. """ v, sz_roll = unroll_dim(v, 2) x = np.zeros(shape=(self._sz_prod, ) + v.shape[1:]) x[self._mask, ...] = v x = m_reshape(x, self.sz + x.shape[1:]) x = roll_dim(x, sz_roll) return x
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 evaluate_t(self, x): """ Evaluate coefficient in Dirac basis from those in standard coordinate basis :param x: 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.count` and whose remaining dimensions correspond to higher dimensions of `v`. """ x, sz_roll = unroll_dim(x, self.ndim + 1) x = m_reshape(x, new_shape=(self._sz_prod, ) + x.shape[self.ndim:]) v = np.zeros(shape=(self.count, ) + x.shape[1:]) v = x[self._mask, ...] v = roll_dim(v, sz_roll) return v
def evaluate_t(self, v): """ Evaluate coefficient in FB basis from those in standard 2D coordinate 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.count` and whose remaining dimensions correspond to higher dimensions of `v`. """ x, sz_roll = unroll_dim(v, self.ndim + 1) x = m_reshape(x, new_shape=tuple([np.prod(self.sz)] + list(x.shape[self.ndim:]))) 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.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 evaluate(self, v): """ Evaluate coefficients in standard 2D coordinate basis from those in FB basis :param v: A coefficient vector (or an array of coefficient vectors) to be evaluated. The first dimension must equal `self.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 expand(self, x): """ Obtain coefficients in the basis from those in standard coordinate basis This is a similar function to evaluate_t but with more accuracy by using the cg optimizing of linear equation, Ax=b. :param x: An array whose first two or three dimensions are to be expanded the desired basis. These dimensions must equal `self.sz`. :return : The coefficients of `v` expanded in the desired basis. The first dimension of `v` is with size of `count` and the second and higher dimensions of the return value correspond to those higher dimensions of `x`. """ # ensure the first dimensions with size of self.sz x, sz_roll = unroll_dim(x, self.ndim + 1) ensure(x.shape[:self.ndim] == self.sz, f'First {self.ndim} dimensions of x must match {self.sz}.') operator = LinearOperator( shape=(self.count, self.count), matvec=lambda v: self.evaluate_t(self.evaluate(v))) # TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v tol = 10 * np.finfo(x.dtype).eps logger.info('Expanding array in basis') # number of image samples n_data = np.size(x, self.ndim) v = np.zeros((self.count, n_data), dtype=x.dtype) for isample in range(0, n_data): b = self.evaluate_t(x[..., isample]) # TODO: need check the initial condition x0 can improve the results or not. v[..., isample], info = cg(operator, b, tol=tol) if info != 0: raise RuntimeError('Unable to converge!') # return v coefficients with the first dimension of self.count v = roll_dim(v, sz_roll) return v
def expand(self, v): """ Expand array in basis 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 v = (B' * B)^(-1) * B' * x where the rows of `B` and columns of `x` are read as vectorized arrays. :param v: An array whose first few dimensions are to be expanded in this basis. These dimensions must equal `self.sz`. :return: The coefficients of `v` expanded in this basis. If more than one array of size `self.sz` is found in `v`, the second and higher dimensions of the return value correspond to those higher dimensions of `v`. .. seealso:: evaluate """ ensure(v.shape[:self.d] == self.sz, f'First {self.d} dimensions of v must match {self.sz}.') v, sz_roll = unroll_dim(v, self.d + 1) b = self.evaluate_t(v) operator = LinearOperator( shape=(self.basis_count, self.basis_count), matvec=lambda x: self.evaluate_t(self.evaluate(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 basis') v, info = cg(operator, b, tol=tol) if info != 0: raise RuntimeError('Unable to converge!') v = roll_dim(v, sz_roll) return v
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 = im_to_vec(self.evaluate(v)) operator = LinearOperator( shape=(self.N**2, self.N**2), matvec=lambda x: im_to_vec(self.evaluate(self.evaluate_t(vec_to_im(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_im(v) return x
def evaluate(self, v): """ Evaluate coefficients in standard 2D coordinate basis from those in FB basis :param v: A coefficient vector (or an array of coefficient vectors) in FB basis to be evaluated. The first dimension must equal `self.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.count v, sz_roll = unroll_dim(v, 2) v = m_reshape(v, (self.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 x = roll_dim(x, sz_roll) return x
def evaluate_t(self, x): """ Evaluate coefficient in FB 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 FB basis. This is an array of vectors whose first dimension equals `self.count` and whose remaining dimensions correspond to higher dimensions of `x`. """ # ensure the first two dimensions with size of self.sz x, sz_roll = unroll_dim(x, self.ndim + 1) 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.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.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.count v = roll_dim(v, sz_roll) return v
def evaluate_t(self, x): """ Evaluate coefficient in FB 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 FB basis. This is an array of vectors whose first dimension equals `self.count` and whose remaining dimensions correspond to higher dimensions of `x`. """ # ensure the first three dimensions with size of self.sz x, sz_roll = unroll_dim(x, self.ndim + 1) 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.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 v = roll_dim(v, sz_roll) return v
def evaluate(self, v): """ Evaluate coefficients in standard 3D coordinate basis from those in 3D FB basis :param v: A coefficient vector (or an array of coefficient vectors) in FB basis to be evaluated. The first dimension must equal `self.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.count v, sz_roll = unroll_dim(v, 2) v = m_reshape(v, (self.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 x = roll_dim(x, sz_roll) return x