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