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 testUnrollDims(self):
m = np.arange(1, 1201).reshape((5, 2, 10, 3, 4), order="F")
m2, sz = unroll_dim(
m, 2
) # second argument is 1-indexed - all dims including and after this are unrolled
# m2 will now have shape (5, (2x10x3x4)) = (5, 240)
self.assertEqual(m2.shape, (5, 240))
# The values should still be filled in with the first axis values changing fastest
self.assertTrue(np.allclose(m2[:, 0], np.array([1, 2, 3, 4, 5])))
# sz are the dimensions that were unrolled
self.assertEqual(sz, (2, 10, 3, 4))
def evaluate_t(self, v):
"""
Evaluate coefficient in FB basis from those in standard 3D 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`.
"""
v = v.T
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=self.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)
nrms = self._norms[idx_radial]
radial = self._precomp["radial"][:, idx_radial]
radial = radial / nrms
for _ in range(-ell, ell + 1):
ang = self._precomp["ang"][:, ind_ang]
ang_radial = np.expand_dims(ang[ang_idx],
axis=1) * radial[r_idx]
idx = ind + np.arange(0, len(idx_radial))
v[idx] = np.real(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
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:], dtype=self.dtype)
x[self._mask, ...] = v
x = m_reshape(x, self.sz + x.shape[1:])
x = roll_dim(x, sz_roll)
return x
def evaluate(self, v):
"""
Evaluate coefficients in standard 3D 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 = v.T
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:])),
dtype=self.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)
nrms = self._norms[idx_radial]
radial = self._precomp["radial"][:, idx_radial]
radial = radial / nrms
for _ in range(-ell, ell + 1):
ang = self._precomp["ang"][:, ind_ang]
ang_radial = np.expand_dims(ang[ang_idx],
axis=1) * radial[r_idx]
idx = ind + np.arange(0, len(idx_radial))
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.T
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:], dtype=self.dtype)
v = x[self._mask, ...]
v = roll_dim(v, sz_roll)
return v
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[..., np.newaxis]
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.shape[3] == 1, "Convolution kernel must be cubic")
ensure(
len(set(kernel_f.shape[:3])) == 1,
"Convolution kernel must be cubic")
is_singleton = x.shape[3] == 1
if is_singleton:
x = fftn(x[..., 0], (N_ker, N_ker, N_ker))[..., np.newaxis]
else:
raise NotImplementedError("not yet")
x = x * kernel_f
if is_singleton:
x[..., 0] = np.real(ifftn(x[..., 0]))
x = x[:N, :N, :N, :]
else:
raise NotImplementedError("not yet")
x = roll_dim(x, sz_roll)
return np.real(x)
