def evaluate_t(self, x): """ Evaluate coefficient in polar Fourier grid from those in standard 2D coordinate basis :param x: The Image instance representing coefficient array in the standard 2D coordinate basis to be evaluated. :return v: The evaluation of the coefficient array `v` in the polar Fourier grid. This is an array of vectors whose first dimension corresponds to x.n_images, and last dimension equals `self.count`. """ assert isinstance(x, Image) if self.dtype != x.dtype: msg = (f"Input data type, {x.dtype}, is not consistent with" f" type defined in the class {self.dtype}.") logger.error(msg) raise TypeError(msg) nimgs = x.n_images half_size = self.ntheta // 2 pf = nufft(x.asnumpy(), self.freqs) pf = pf.reshape((nimgs, self.nrad, half_size)) v = np.concatenate((pf, pf.conj()), axis=1) # return v coefficients with the last dimension size of self.count v = v.reshape(nimgs, -1) return v
def _compute_nfft_potts(self, images, start, finish): """ Perform NuFFT transform for images in rectangular coordinates """ x = self.us_fft_pts num_images = finish - start m = x.shape[0] images_nufft = np.zeros((m, num_images), dtype=complex_type(self.dtype)) for i in range(start, finish): images_nufft[:, i - start] = nufft(images[..., i], 2 * pi * x.T) return images_nufft
def project(self, vol_idx, rot_matrices): """ Using the stack of rot_matrices, project images of Volume[vol_idx]. :param vol_idx: Volume index :param rot_matrices: Stack of rotations. Rotation or ndarray instance. :return: `Image` instance. """ # If we are an ASPIRE Rotation, get the numpy representation. if isinstance(rot_matrices, Rotation): rot_matrices = rot_matrices.matrices if rot_matrices.dtype != self.dtype: logger.warning( f"{self.__class__.__name__}" f" rot_matrices.dtype {rot_matrices.dtype}" f" != self.dtype {self.dtype}." " In the future this will raise an error." ) data = self[vol_idx].T # RCOPT n = rot_matrices.shape[0] pts_rot = np.moveaxis(rotated_grids(self.resolution, rot_matrices), 1, 2) # TODO: rotated_grids might as well give us correctly shaped array in the first place pts_rot = m_reshape(pts_rot, (3, self.resolution ** 2 * n)) im_f = nufft(data, pts_rot) / self.resolution im_f = im_f.reshape(-1, self.resolution, self.resolution) if self.resolution % 2 == 0: im_f[:, 0, :] = 0 im_f[:, :, 0] = 0 im_f = xp.asnumpy(fft.centered_ifft2(xp.asarray(im_f))) return aspire.image.Image(np.real(im_f))
def evaluate_t(self, x): """ Evaluate coefficient in FB basis from those in standard 2D coordinate basis :param x: The Image instance representing coefficient array in the standard 2D coordinate basis to be evaluated. :return v: The evaluation of the coefficient array `v` in the FB basis. This is an array of vectors whose last dimension equals `self.count` and whose first dimension correspond to `x.n_images`. """ if x.dtype != self.dtype: logger.warning( f"{self.__class__.__name__}::evaluate_t" f" Inconsistent dtypes v: {x.dtype} self: {self.dtype}") if not isinstance(x, Image): logger.warning(f"{self.__class__.__name__}::evaluate_t" " passed numpy array instead of Image.") x = Image(x) # 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 = np.reshape(self._precomp["freqs"], (2, n_r * n_theta)) # number of 2D image samples n_images = x.n_images x_data = x.data # resamping x in a polar Fourier gird using nonuniform discrete Fourier transform pf = nufft(x_data, 2 * pi * freqs) pf = np.reshape(pf, (n_images, n_r, n_theta)) # Recover "negative" frequencies from "positive" half plane. pf = np.concatenate((pf, pf.conjugate()), axis=2) # 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) * xp.asnumpy(fft.fft(xp.asarray(pf))) # This only makes it easier to slice the array later. v = np.zeros((n_images, self.count), 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 # include the normalization factor of angular part into radial part radial_norm = self._precomp["radial"] / np.expand_dims( self.angular_norms, 1) v[:, mask] = pf[:, :, 0].real @ radial_norm[idx].T 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 = pf[:, :, ell] @ radial_norm[idx].T 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
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 last 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 last dimension equals `self.count` and whose remaining dimensions correspond to higher dimensions of `x`. """ # roll dimensions sz_roll = x.shape[:-3] x = x.reshape((-1, *self.sz)) n_data = x.shape[0] 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 = nufft(x, self._precomp["fourier_pts"]) pf = m_reshape(pf.T, (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((n_data, self.count), 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.T # Roll dimensions, last dimension should be self.count, # Higher dimensions like x. v = v.reshape((*sz_roll, self.count)) return v