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
