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 testPolarBasis2DAdjoint(self):
# The evaluate function should be the adjoint operator of evaluate_t.
# Namely, if A = evaluate, B = evaluate_t, and B=A^t, we will have
# (y, A*x) = (A^t*y, x) = (B*y, x)
x = randn(self.basis.count, seed=self.seed).astype(self.dtype)
x = m_reshape(x, (self.basis.nrad, self.basis.ntheta))
x = (1 / 2 * x[:, :self.basis.ntheta // 2] +
1 / 2 * x[:, :self.basis.ntheta // 2].conj())
x = np.concatenate((x, x.conj()), axis=1)
x = m_reshape(x, (self.basis.nrad * self.basis.ntheta, ))
x_t = self.basis.evaluate(x).asnumpy()
y = randn(np.prod(self.basis.sz), seed=self.seed).astype(self.dtype)
y_t = self.basis.evaluate_t(
Image(m_reshape(y, self.basis.sz)[np.newaxis, :])) # RCOPT
lhs = np.dot(y, m_reshape(x_t, (np.prod(self.basis.sz), )))
rhs = np.real(np.dot(y_t, x))
logging.debug(
f"lhs: {lhs} rhs: {rhs} absdiff: {np.abs(lhs-rhs)} atol: {utest_tolerance(self.dtype)}"
)
self.assertTrue(np.isclose(lhs, rhs, atol=utest_tolerance(self.dtype)))
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 precond_fun(S, x):
p = np.size(S, 0)
ensure(np.size(x) == p*p, 'The sizes of S and x are not consistent.')
x = m_reshape(x, (p, p))
y = S @ x @ S
y = m_reshape(y, (p**2,))
return y
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 _solve_covar(self, A_covar, b_covar, M, covar_est_opt):
ctf_fb = self.ctf_fb
def precond_fun(S, x):
p = np.size(S, 0)
ensure(np.size(x) == p*p, 'The sizes of S and x are not consistent.')
x = m_reshape(x, (p, p))
y = S @ x @ S
y = m_reshape(y, (p**2,))
return y
def apply(A, x):
p = np.size(A[0], 0)
x = m_reshape(x, (p, p))
y = np.zeros_like(x)
for k in range(0, len(A)):
y = y + A[k] @ x @ A[k].T
y = m_reshape(y, (p**2,))
return y
cg_opt = covar_est_opt
covar_coeff = BlkDiagMatrix.zeros_like(ctf_fb[0])
for ell in range(0, len(b_covar)):
A_ell = []
for k in range(0, len(A_covar)):
A_ell.append(A_covar[k][ell])
p = np.size(A_ell[0], 0)
b_ell = m_reshape(b_covar[ell], (p ** 2,))
S = inv(M[ell])
cg_opt['preconditioner'] = lambda x: precond_fun(S, x)
covar_coeff_ell, _, _ = conj_grad(lambda x: apply(A_ell, x), b_ell, cg_opt)
covar_coeff[ell] = m_reshape(covar_coeff_ell, (p, p))
return covar_coeff
def apply(A, x):
p = np.size(A[0], 0)
x = m_reshape(x, (p, p))
y = np.zeros_like(x)
for k in range(0, len(A)):
y = y + A[k] @ x @ A[k].T
y = m_reshape(y, (p**2, ))
return y
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 _getfbzeros(self):
"""
Generate zeros of Bessel functions
"""
# get upper_bound of zeros of Bessel functions
upper_bound = min(self.ell_max + 1, 2 * self.nres + 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.ndim - 2) / 2,
self.nres * 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 backproject(self, rot_matrices):
"""
Backproject images along rotation
:param im: An Image (stack) to backproject.
:param rot_matrices: An n-by-3-by-3 array of rotation matrices \
corresponding to viewing directions.
:return: Volume instance corresonding to the backprojected images.
"""
L = self.res
ensure(
self.n_images == rot_matrices.shape[0],
"Number of rotation matrices must match the number of images",
)
# TODO: rotated_grids might as well give us correctly shaped array in the first place
pts_rot = aspire.volume.rotated_grids(L, rot_matrices)
pts_rot = np.moveaxis(pts_rot, 1, 2)
pts_rot = m_reshape(pts_rot, (3, -1))
im_f = xp.asnumpy(fft.centered_fft2(xp.asarray(self.data))) / (L**2)
if L % 2 == 0:
im_f[:, 0, :] = 0
im_f[:, :, 0] = 0
im_f = im_f.flatten()
vol = anufft(im_f, pts_rot, (L, L, L), real=True) / L
return aspire.volume.Volume(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 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 compute_kernel(self):
_2L = 2 * self.L
kernel = np.zeros((_2L, _2L, _2L), dtype=self.dtype)
sq_filters_f = self.src.eval_filter_grid(self.L, power=2)
for i in range(0, self.n, self.batch_size):
_range = np.arange(i, min(self.n, i + self.batch_size), dtype=int)
pts_rot = rotated_grids(self.L, self.src.rots[_range, :, :])
weights = sq_filters_f[:, :, _range]
weights *= self.src.amplitudes[_range]**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) *
anufft(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 rand(size, seed=None):
"""
Note this is for MATLAB repro (see m_reshape).
Other uses prefer use of `random`.
"""
with Random(seed):
return m_reshape(np.random.random(np.prod(size)), size)
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 last 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 last dimension equals
`self.count` and whose first dimensions correspond to
first dimensions of `v`.
"""
if v.dtype != self.dtype:
logger.warning(
f"{self.__class__.__name__}::evaluate_t"
f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}"
)
if isinstance(v, Image):
v = v.asnumpy()
v = v.T # RCOPT
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:])), dtype=v.dtype)
for ell in range(0, self.ell_max + 1):
k_max = self.k_max[ell]
idx_radial = ind_radial + np.arange(0, k_max)
# include the normalization factor of angular part
ang_nrms = self.angular_norms[idx_radial]
radial = self._precomp["radial"][:, idx_radial]
radial = radial / ang_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.T # RCOPT
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 = self.n_r
n_theta = self.n_theta
r, w = lgwt(n_r, 0.0, self.kcut, dtype=self.dtype)
radial = np.zeros(shape=(np.sum(self.k_max), n_r), dtype=self.dtype)
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.kcut)
# NOTE: We need to remove the factor due to the discretization here
# since it is already included in our quadrature weights
# Only normalized by the radial part of basis function
nrm = 1 / (np.sqrt(np.prod(
self.sz))) * self.radial_norms[ind_radial]
radial[ind_radial] /= nrm
ind_radial += 1
# Only calculate "positive" frequencies in one half-plane.
freqs_x = m_reshape(r, (n_r, 1)) @ m_reshape(
np.cos(
np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)),
(1, n_theta),
)
freqs_y = m_reshape(r, (n_r, 1)) @ m_reshape(
np.sin(
np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)),
(1, n_theta),
)
freqs = np.vstack((freqs_y[np.newaxis, ...], freqs_x[np.newaxis, ...]))
return {
"gl_nodes": r,
"gl_weights": w,
"radial": radial,
"freqs": freqs
}
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 evaluate_grid(self, L, dtype=np.float32, *args, **kwargs):
grid2d = grid_2d(L, dtype=dtype)
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 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):
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 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_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 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 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 testPolarBasis2DAdjoint(self):
# The evaluate function should be the adjoint operator of evaluate_t.
# Namely, if A = evaluate, B = evaluate_t, and B=A^t, we will have
# (y, A*x) = (A^t*y, x) = (B*y, x)
x = np.random.randn(self.basis.count)
x = m_reshape(x, (self.basis.nrad, self.basis.ntheta))
x = (1 / 2 * x[:, :self.basis.ntheta // 2]
+ 1 / 2 * x[:, :self.basis.ntheta // 2].conj())
x = np.concatenate((x, x.conj()), axis=1)
x = m_reshape(x, (self.basis.nrad * self.basis.ntheta,))
x_t = self.basis.evaluate(x)
y = np.random.randn(np.prod(self.basis.sz))
y_t = self.basis.evaluate_t(m_reshape(y, self.basis.sz))
self.assertTrue(np.isclose(np.dot(y, m_reshape(x_t, (np.prod(self.basis.sz),))),
np.dot(y_t, x)))
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 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 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]
new_shape += (-1,)
Y = m_reshape(X, new_shape)
removed_dims = old_shape[dim:]
return Y, removed_dims
