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
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.ndim == 3, "Convolution kernel must be cubic")
ensure(
len(set(kernel_f.shape)) == 1, "Convolution kernel must be cubic")
is_singleton = x.ndim == 3
if is_singleton:
x = fftn(x, (N_ker, N_ker, N_ker))
else:
raise NotImplementedError('not yet')
x = x * kernel_f
if is_singleton:
x = np.real(ifftn(x))
x = x[:N, :N, :N]
else:
raise NotImplementedError('not yet')
x = roll_dim(x, sz_roll)
return x
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 im_filter(im, filter, *args, **kwargs):
# TODO: Move inside appropriate object
L = im.shape[0]
im, sz_roll = unroll_dim(im, 3)
filter_vals = filter.evaluate_grid(L, *args, **kwargs)
im_f = centered_fft2(im)
if im_f.ndim > filter_vals.ndim:
im_f = np.expand_dims(filter_vals, 2) * im_f
else:
im_f = filter_vals * im_f
im = centered_ifft2(im_f)
im = np.real(im)
im = roll_dim(im, sz_roll)
return im
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 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 expand(self, v):
"""
Expand array in basis
If `v` is a matrix of size `basis.ct`-by-..., `B` is the change-of-basis matrix of this basis, and `x` is a
matrix of size `self.sz`-by-..., the function calculates
v = (B' * B)^(-1) * B' * x
where the rows of `B` and columns of `x` are read as vectorized arrays.
:param v: An array whose first few dimensions are to be expanded in this basis.
These dimensions must equal `self.sz`.
:return: The coefficients of `v` expanded in this basis. If more than one array of size `self.sz` is found in
`v`, the second and higher dimensions of the return value correspond to those higher dimensions of `v`.
.. seealso:: evaluate
"""
ensure(v.shape[:self.d] == self.sz,
f'First {self.d} dimensions of v must match {self.sz}.')
v, sz_roll = unroll_dim(v, self.d + 1)
b = self.evaluate_t(v)
operator = LinearOperator(
shape=(self.basis_count, self.basis_count),
matvec=lambda x: self.evaluate_t(self.evaluate(x)))
# TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v
tol = 10 * np.finfo(v.dtype).eps
logger.info('Expanding array in basis')
v, info = cg(operator, b, tol=tol)
if info != 0:
raise RuntimeError('Unable to converge!')
v = roll_dim(v, sz_roll)
return v
def expand_t(self, v):
ensure(v.shape[0] == self.basis_count,
f'First dimension of v must be {self.basis_count}')
v, sz_roll = unroll_dim(v, 2)
b = im_to_vec(self.evaluate(v))
operator = LinearOperator(
shape=(self.N**2, self.N**2),
matvec=lambda x: im_to_vec(
self.evaluate(self.evaluate_t(vec_to_im(x)))))
# TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v
tol = 10 * np.finfo(v.dtype).eps
logger.info('Expanding array in dual basis')
v, info = cg(operator, b, tol=tol)
if info != 0:
raise RuntimeError('Unable to converge!')
v = roll_dim(v, sz_roll)
x = vec_to_im(v)
return x
