def anufft(sig_f, fourier_pts, sz, real=False):
"""
Wrapper for 1, 2, and 3 dimensional Non Uniform FFT Adjoint.
Dimension is based on the dimension of fourier_pts and checked against sig_f.
Selects best available package from `nfft` `backends` configuration list.
:param sig_f: Array representing the signal(s) in Fourier space to be transformed. \
sig_f either matches length of fourier_pts or sig_f.shape is stack of (`ntransforms`, ...).
:param fourier_pts: The points in Fourier space where the Fourier transform is to be calculated,
arranged as a dimension-by-K array. These need to be in the range [-pi, pi] in each dimension.
:param sz: A tuple indicating the geometry of the signal.
:param real: Optional Bool indicating if you would like only the real components, Defaults False.
:return: The Non Uniform FFT adjoint transform.
"""
if fourier_pts.dtype != real_type(sig_f.dtype):
logger.warning("anufft passed inconsistent dtypes."
f" fourier_pts: {fourier_pts.dtype}"
f" forcing precision of signal data: {sig_f.dtype}.")
fourier_pts = fourier_pts.astype(real_type(sig_f.dtype))
if sig_f.dtype != complex_type(sig_f.dtype):
logger.debug("anufft passed real_type for signal, converting")
sig_f = sig_f.astype(complex_type(sig_f.dtype))
ntransforms = 1
if len(sig_f.shape) == 2:
ntransforms = sig_f.shape[0]
plan = Plan(sz=sz, fourier_pts=fourier_pts, ntransforms=ntransforms)
adjoint = plan.adjoint(sig_f)
return np.real(adjoint) if real else adjoint
def _evaluate_pswf2d_all(self, r, theta, max_ns):
"""
Evaluate the numerical values of PSWF functions for all N's, up to given n for each N
:param r: Radial part to evaluate
:param theta: Phase part to evaluate
:param max_ns: List of ints max_ns[i] is max n to to use for N=i, not included.
If max_ns[i]<1 N=i won't be used
:return: (len(r), sum(max_ns)) ndarray
Indices are corresponding to the list (N, n)
(0, 0),..., (0, max_ns[0]), (1, 0),..., (1, max_ns[1]),... , (len(max_ns)-1, 0),
(len(max_ns)-1, max_ns[-1])
"""
max_ns_ints = [int(max_n) for max_n in max_ns]
out_mat = []
for i, max_n in enumerate(max_ns_ints):
if max_n < 1:
continue
d_vec = self.d_vec_all[i]
phase_part = np.exp(1j * i * theta) / np.sqrt(2 * np.pi)
range_array = np.arange(len(d_vec))
r_radial_part_mat = t_radial_part_mat(
r, i, range_array, len(d_vec)).dot(d_vec[:, :max_n])
pswf_n_n_mat = phase_part * r_radial_part_mat.T
out_mat.extend(pswf_n_n_mat)
out_mat = np.array(out_mat, dtype=complex_type(self.dtype)).T
return out_mat
def evaluate(self, coefficients):
"""
Evaluate coefficients in standard 2D coordinate basis from those in PSWF basis
:param coeffcients: A coefficient vector (or an array of coefficient
vectors) in PSWF basis to be evaluated.
:return : The evaluation of the coefficient vector(s) in standard 2D
coordinate basis.
"""
coefficients = coefficients.T # RCOPT
# if we got only one vector
if len(coefficients.shape) == 1:
coefficients = coefficients[:, np.newaxis]
angular_is_zero = np.absolute(self.ang_freqs) == 0
flatten_images = self.samples[:, angular_is_zero].dot(
coefficients[angular_is_zero]) + 2.0 * np.real(
self.samples[:, ~angular_is_zero].dot(
coefficients[~angular_is_zero]))
n_images = int(flatten_images.shape[1])
images = np.zeros((self._image_height, self._image_height,
n_images)).astype(complex_type(self.dtype))
images[self._disk_mask, :] = flatten_images
images = np.transpose(images, axes=(1, 0, 2))
return np.real(images).T # RCOPT
def testComplexCoversionErrorsToComplex(self):
# Load a reasonable input
x = np.load(os.path.join(DATA_DIR, "fbbasis_coefficients_8_8.npy"))
# Express in an FB basis
v1 = self.basis.expand(x.astype(self.dtype))
# Test catching Errors
with raises(TypeError):
# Pass complex into `to_complex`
_ = self.basis.to_complex(v1.astype(np.complex64))
# Test casting case, where basis and coef don't match
if self.basis.dtype == np.float32:
test_dtype = np.float64
elif self.basis.dtype == np.float64:
test_dtype = np.float32
# Result should be same precision as coef input, just complex.
result_dtype = complex_type(test_dtype)
v3 = self.basis.to_complex(v1.astype(test_dtype))
self.assertTrue(v3.dtype == result_dtype)
# Try 0d vector, should not crash.
_ = self.basis.to_complex(v1.reshape(-1))
def evaluate(self, coefficients):
"""
Evaluate coefficients in standard 2D coordinate basis from those in PSWF basis
:param coefficients: A coefficient vector (or an array of coefficient vectors)
in PSWF basis to be evaluated.
:return : The evaluation of the coefficient vector(s) in standard 2D
coordinate basis.
"""
coefficients = coefficients.T # RCOPT
# if we got only one vector
if len(coefficients.shape) == 1:
coefficients = coefficients.reshape((len(coefficients), 1))
angular_is_zero = np.absolute(self.ang_freqs) == 0
flatten_images = self.samples[:, angular_is_zero].dot(
coefficients[angular_is_zero]
) + (
2.0
* np.real(
self.samples[:, ~angular_is_zero].dot(coefficients[~angular_is_zero])
)
)
n_images = int(flatten_images.shape[1])
images = np.zeros((self._image_height, self._image_height, n_images)).astype(
complex_type(self.dtype)
)
images[self._disk_mask, :] = flatten_images
# TODO: no need to switch x and y any more, need to make consistent with direct method
return np.real(images).T # RCOPT
def nufft(sig_f, fourier_pts, real=False):
"""
Wrapper for 1, 2, and 3 dimensional Non Uniform FFT
Dimension is based on the dimension of fourier_pts and checked against sig_f.
Selects best available package from `nfft` `backends` configuration list.
:param sig_f: Array representing the signal(s) in real space to be transformed. \
sig_f either matches `sz` or sig_f.shape is stack of (..., `ntransforms`).
:param fourier_pts: The points in Fourier space where the Fourier transform is to be calculated,
arranged as a dimension-by-K array. These need to be in the range [-pi, pi] in each dimension.
:param real: Optional Bool indicating if you would like only the real components, Defaults False.
:return: The Non Uniform FFT transform.
"""
if fourier_pts.dtype != real_type(sig_f.dtype):
logger.warning("nufft passed inconsistent dtypes."
f" fourier_pts: {fourier_pts.dtype}"
f" forcing precision of signal data: {sig_f.dtype}.")
fourier_pts = fourier_pts.astype(real_type(sig_f.dtype))
if sig_f.dtype != complex_type(sig_f.dtype):
logger.debug("nufft passed real_type for signal, converting")
sig_f = sig_f.astype(complex_type(sig_f.dtype))
# Unpack the dimension of the signal
# Note, infer dimension from fourier_pts because signal might be a stack.
dimension = fourier_pts.shape[0]
# Unpack the resolution of the signal
resolution = sig_f.shape[1]
# Construct tuple describing geometry of signal
sz = (resolution, ) * dimension
ntransforms = 1
if len(sig_f.shape) == dimension + 1:
ntransforms = sig_f.shape[0]
plan = Plan(sz=sz, fourier_pts=fourier_pts, ntransforms=ntransforms)
transform = plan.transform(sig_f)
return np.real(transform) if real else transform
def estimate_psd(self, blocks, tapers_1d):
"""
Estimate the power spectrum of the micrograph using the multi-taper method
:param blocks: 3-D NumPy array containing windows extracted from the micrograph in the preprocess function.
:param tapers_1d: NumPy array of data tapers.
:return: NumPy array of estimated power spectrum.
"""
num_1d_tapers = tapers_1d.shape[-1]
tapers_1d = tapers_1d.astype(complex_type(self.dtype), copy=False)
blocks_mt = np.zeros(blocks[0, :, :].shape, dtype=self.dtype)
blocks_tapered = np.zeros(blocks[0, :, :].shape,
dtype=complex_type(self.dtype))
taper_2d = np.zeros((blocks.shape[1], blocks.shape[2]),
dtype=complex_type(self.dtype))
for ax1 in range(num_1d_tapers):
for ax2 in range(num_1d_tapers):
np.matmul(
tapers_1d[:, ax1, np.newaxis],
tapers_1d[:, ax2, np.newaxis].T,
out=taper_2d,
)
for m in range(blocks.shape[0]):
np.multiply(blocks[m, :, :], taper_2d, out=blocks_tapered)
blocks_mt_post_fft = fft.fftn(blocks_tapered,
axes=(-2, -1))
blocks_mt += abs2(blocks_mt_post_fft)
blocks_mt /= blocks.shape[0]**2
blocks_mt /= tapers_1d.shape[0]**2
amplitude_spectrum = fft.fftshift(
blocks_mt) # max difference 10^-13, max relative difference 10^-14
return Image(amplitude_spectrum)
def testSmallFitTransform(self):
# Reals
pca = PCA(n_components=self.components_small)
Y1 = pca.fit_transform(self.X_small)
# Complex
cpca = ComplexPCA(n_components=self.components_small)
Y2 = cpca.fit_transform(self.X_small.astype(complex_type(self.dtype)))
# Real part should be the same.
self.assertTrue(np.allclose(np.real(Y2), Y1))
# Imag part should be zero.
self.assertTrue(np.allclose(np.imag(Y2), 0))
def _pswf_integration(self, images_nufft):
"""
Perform integration part for rotational invariant property.
"""
num_images = images_nufft.shape[1]
n_max_float = float(self.n_max) / 2
r_n_eval_mat = np.zeros(
(len(self.radial_quad_pts), self.n_max, num_images),
dtype=complex_type(self.dtype),
)
for i in range(len(self.radial_quad_pts)):
curr_r_mat = images_nufft[
self.r_quad_indices[i] : self.r_quad_indices[i]
+ self.num_angular_pts[i],
:,
]
curr_r_mat = np.concatenate((curr_r_mat, np.conj(curr_r_mat)))
fft_plan = xp.asnumpy(fft.fft(xp.asarray(curr_r_mat), axis=0))
angular_eval = fft_plan * self.quad_rule_radial_wts[i]
r_n_eval_mat[i, :, :] = np.tile(
angular_eval,
(int(max(1, np.ceil(n_max_float / self.num_angular_pts[i]))), 1),
)[: self.n_max, :]
r_n_eval_mat = r_n_eval_mat.reshape(
(len(self.radial_quad_pts) * self.n_max, num_images), order="F"
)
coeff_vec_quad = np.zeros(
(len(self.ang_freqs), num_images), dtype=complex_type(self.dtype)
)
m = len(self.pswf_radial_quad)
for i in range(self.n_max):
coeff_vec_quad[
self.indices_for_n[i] + np.arange(self.numel_for_n[i]), :
] = np.dot(self.blk_r[i], r_n_eval_mat[i * m : (i + 1) * m, :])
return coeff_vec_quad
def to_complex(self, coef):
"""
Return complex valued representation of coefficients.
This can be useful when comparing or implementing methods
from literature.
There is a corresponding method, to_real.
:param coef: Coefficients from this basis.
:return: Complex coefficent representation from this basis.
"""
if coef.ndim == 1:
coef = coef.reshape(1, -1)
if coef.dtype not in (np.float64, np.float32):
raise TypeError("coef provided to to_complex should be real.")
# Pass through dtype precions, but check and warn if mismatched.
dtype = complex_type(coef.dtype)
if coef.dtype != self.dtype:
logger.warning(
f"coef dtype {coef.dtype} does not match precision of basis.dtype {self.dtype}, returning {dtype}."
)
# Return the same precision as coef
imaginary = dtype(1j)
ccoef = np.zeros((coef.shape[0], self.complex_count), dtype=dtype)
ccoef_d = OrderedDict()
for i in range(self.count):
ell = self.angular_indices[i]
q = self.radial_indices[i]
sgn = self.signs_indices[i]
ccoef_d.setdefault((ell, q), 0 + 0j)
if ell == 0:
ccoef_d[(ell, q)] = coef[:, i]
elif sgn == 1:
ccoef_d[(ell, q)] += coef[:, i] / 2.0
elif sgn == -1:
ccoef_d[(ell, q)] -= imaginary * coef[:, i] / 2.0
else:
raise ValueError("sgns should be +-1")
for i, k in enumerate(ccoef_d.keys()):
ccoef[:, i] = ccoef_d[k]
return ccoef
def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
"""
A plan for non-uniform FFT in 2D or 3D.
:param sz: A tuple indicating the geometry of the signal.
:param fourier_pts: The points in Fourier space where the Fourier
transform is to be calculated, arranged as a dimension-by-K array.
These need to be in the range [-pi, pi] in each dimension.
:param epsilon: The desired precision of the NUFFT.
:param ntransforms: Optional integer indicating if you would like
to compute a batch of `ntransforms`.
transforms. Implies vol_f.shape is (`ntransforms`, ...).
"""
self.ntransforms = ntransforms
self.sz = sz
self.dim = len(sz)
self.dtype = fourier_pts.dtype
self.complex_dtype = complex_type(self.dtype)
self.fourier_pts = np.ascontiguousarray(
np.mod(fourier_pts + np.pi, 2 * np.pi) - np.pi)
self.num_pts = fourier_pts.shape[1]
self.epsilon = max(epsilon, np.finfo(self.dtype).eps)
if self.epsilon != epsilon:
logger.debug(f"FinufftPlan adjusted eps={self.epsilon}"
f" from requested {epsilon}.")
self._transform_plan = finufft.Plan(
nufft_type=2,
n_modes_or_dim=self.sz,
eps=self.epsilon,
n_trans=self.ntransforms,
dtype=self.dtype,
)
self._adjoint_plan = finufft.Plan(
nufft_type=1,
n_modes_or_dim=self.sz,
eps=self.epsilon,
n_trans=self.ntransforms,
dtype=self.dtype,
)
self._transform_plan.setpts(*self.fourier_pts)
self._adjoint_plan.setpts(*self.fourier_pts)
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 to_complex(self, coef):
"""
Return complex valued representation of coefficients.
This can be useful when comparing or implementing methods
from literature.
There is a corresponding method, to_real.
:param coef: Coefficients from this basis.
:return: Complex coefficent representation from this basis.
"""
if coef.ndim == 1:
coef = coef.reshape(1, -1)
if coef.dtype not in (np.float64, np.float32):
raise TypeError("coef provided to to_complex should be real.")
# Pass through dtype precions, but check and warn if mismatched.
dtype = complex_type(coef.dtype)
if coef.dtype != self.dtype:
logger.warning(
f"coef dtype {coef.dtype} does not match precision of basis.dtype {self.dtype}, returning {dtype}."
)
# Return the same precision as coef
imaginary = dtype(1j)
ccoef = np.zeros((coef.shape[0], self.complex_count), dtype=dtype)
ind = 0
idx = np.arange(self.k_max[0], dtype=int)
ind += np.size(idx)
ccoef[:, idx] = coef[:, idx]
for ell in range(1, self.ell_max + 1):
idx = ind + np.arange(self.k_max[ell], dtype=int)
ccoef[:, idx] = (
coef[:, self._pos[idx]] - imaginary * coef[:, self._neg[idx]]
) / 2.0
ind += np.size(idx)
return ccoef
def evaluate(self, v):
"""
Evaluate coefficients in standard 2D coordinate basis from those in FB basis
:param v: A coefficient vector (or an array of coefficient vectors)
in FB basis to be evaluated. The last dimension must equal `self.count`.
:return x: The evaluation of the coefficient vector(s) `x` in standard 2D
coordinate basis. This is Image instance with resolution of `self.sz`
and the first dimension correspond to remaining dimension of `v`.
"""
if v.dtype != self.dtype:
logger.debug(
f"{self.__class__.__name__}::evaluate"
f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}")
sz_roll = v.shape[:-1]
v = v.reshape(-1, self.count)
# number of 2D image samples
n_data = v.shape[0]
# get information on polar grids from precomputed data
n_theta = np.size(self._precomp["freqs"], 2)
n_r = np.size(self._precomp["freqs"], 1)
# go through each basis function and find corresponding coefficient
pf = np.zeros((n_data, 2 * n_theta, n_r),
dtype=complex_type(self.dtype))
mask = self._indices["ells"] == 0
ind = 0
idx = ind + np.arange(self.k_max[0], dtype=int)
# include the normalization factor of angular part into radial part
radial_norm = self._precomp["radial"] / np.expand_dims(
self.angular_norms, 1)
pf[:, 0, :] = v[:, mask] @ radial_norm[idx]
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], dtype=int)
idx_pos = ind_pos + np.arange(self.k_max[ell], dtype=int)
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 = v_ell @ radial_norm[idx]
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 * xp.asnumpy(fft.ifft(xp.asarray(pf), axis=1))
# 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 = np.reshape(pf, (n_data, n_r * n_theta))
# perform inverse non-uniformly FFT transform back to 2D coordinate basis
freqs = m_reshape(self._precomp["freqs"], (2, n_r * n_theta))
x = 2 * anufft(pf, 2 * pi * freqs, self.sz, real=True)
# Return X as Image instance with the last two dimensions as *self.sz
x = x.reshape((*sz_roll, *self.sz))
return Image(x)
def testPolarBasis2DEvaluate(self):
v = np.array(
[
0.38243133 - 6.66608316e-18j,
0.3249317 + 1.47839074e-01j,
0.14819172 - 3.78171168e-03j,
-0.22808599 - 5.29338933e-02j,
0.38243133 - 6.66608316e-18j,
0.34595014 + 1.06355385e-01j,
0.15519289 + 4.75602164e-02j,
-0.22401193 - 4.33128746e-03j,
0.38243133 - 6.66608316e-18j,
0.36957165 + 5.69575709e-02j,
0.17389327 + 5.53498385e-02j,
-0.11601473 + 1.35405676e-02j,
0.38243133 - 6.66608316e-18j,
0.39045046 + 2.17911945e-03j,
0.18146449 + 1.37089189e-02j,
-0.02110144 - 6.65071497e-03j,
0.38243133 - 6.66608316e-18j,
0.4063995 - 5.21354967e-02j,
0.15674204 - 3.85815662e-02j,
-0.02886296 - 3.91489615e-02j,
0.38243133 - 6.66608316e-18j,
0.41872477 - 9.98946906e-02j,
0.11862477 - 5.15231952e-02j,
-0.05298751 - 1.95319478e-02j,
0.38243133 - 6.66608316e-18j,
0.43013599 - 1.38307796e-01j,
0.10075763 - 1.25689289e-02j,
-0.04052728 + 5.66863498e-02j,
0.38243133 - 6.66608316e-18j,
0.44144497 - 1.68826980e-01j,
0.11446016 + 4.53003874e-02j,
-0.03546515 + 1.13544145e-01j,
0.38243133 - 6.66608316e-18j,
0.44960099 - 1.94794929e-01j,
0.15053714 + 8.11915305e-02j,
-0.04800556 + 1.15828804e-01j,
0.38243133 - 6.66608316e-18j,
0.44872328 - 2.17957567e-01j,
0.19116871 + 7.99536373e-02j,
-0.05683092 + 9.72225058e-02j,
0.38243133 - 6.66608316e-18j,
0.43379428 - 2.36681249e-01j,
0.21025378 + 5.48466438e-02j,
-0.05318826 + 8.54948014e-02j,
0.38243133 - 6.66608316e-18j,
0.40485577 - 2.47073481e-01j,
0.18680217 + 3.31766116e-02j,
-0.06674163 + 7.94216591e-02j,
0.38243133 - 6.66608316e-18j,
0.36865853 - 2.45913767e-01j,
0.13660805 + 3.68947359e-02j,
-0.11467046 + 8.49198927e-02j,
0.38243133 - 6.66608316e-18j,
0.33597018 - 2.32971425e-01j,
0.1072859 + 6.24686168e-02j,
-0.12932565 + 1.06139634e-01j,
0.38243133 - 6.66608316e-18j,
0.31616666 - 2.10791785e-01j,
0.11876919 + 7.93812474e-02j,
-0.1094488 + 1.20159845e-01j,
0.38243133 - 6.66608316e-18j,
0.31313975 - 1.82190396e-01j,
0.14075481 + 5.85637416e-02j,
-0.15198775 + 1.02156797e-01j,
0.38243133 + 6.66608316e-18j,
0.3249317 - 1.47839074e-01j,
0.14819172 + 3.78171168e-03j,
-0.22808599 + 5.29338933e-02j,
0.38243133 + 6.66608316e-18j,
0.34595014 - 1.06355385e-01j,
0.15519289 - 4.75602164e-02j,
-0.22401193 + 4.33128746e-03j,
0.38243133 + 6.66608316e-18j,
0.36957165 - 5.69575709e-02j,
0.17389327 - 5.53498385e-02j,
-0.11601473 - 1.35405676e-02j,
0.38243133 + 6.66608316e-18j,
0.39045046 - 2.17911945e-03j,
0.18146449 - 1.37089189e-02j,
-0.02110144 + 6.65071497e-03j,
0.38243133 + 6.66608316e-18j,
0.4063995 + 5.21354967e-02j,
0.15674204 + 3.85815662e-02j,
-0.02886296 + 3.91489615e-02j,
0.38243133 + 6.66608316e-18j,
0.41872477 + 9.98946906e-02j,
0.11862477 + 5.15231952e-02j,
-0.05298751 + 1.95319478e-02j,
0.38243133 + 6.66608316e-18j,
0.43013599 + 1.38307796e-01j,
0.10075763 + 1.25689289e-02j,
-0.04052728 - 5.66863498e-02j,
0.38243133 + 6.66608316e-18j,
0.44144497 + 1.68826980e-01j,
0.11446016 - 4.53003874e-02j,
-0.03546515 - 1.13544145e-01j,
0.38243133 + 6.66608316e-18j,
0.44960099 + 1.94794929e-01j,
0.15053714 - 8.11915305e-02j,
-0.04800556 - 1.15828804e-01j,
0.38243133 + 6.66608316e-18j,
0.44872328 + 2.17957567e-01j,
0.19116871 - 7.99536373e-02j,
-0.05683092 - 9.72225058e-02j,
0.38243133 + 6.66608316e-18j,
0.43379428 + 2.36681249e-01j,
0.21025378 - 5.48466438e-02j,
-0.05318826 - 8.54948014e-02j,
0.38243133 + 6.66608316e-18j,
0.40485577 + 2.47073481e-01j,
0.18680217 - 3.31766116e-02j,
-0.06674163 - 7.94216591e-02j,
0.38243133 + 6.66608316e-18j,
0.36865853 + 2.45913767e-01j,
0.13660805 - 3.68947359e-02j,
-0.11467046 - 8.49198927e-02j,
0.38243133 + 6.66608316e-18j,
0.33597018 + 2.32971425e-01j,
0.1072859 - 6.24686168e-02j,
-0.12932565 - 1.06139634e-01j,
0.38243133 + 6.66608316e-18j,
0.31616666 + 2.10791785e-01j,
0.11876919 - 7.93812474e-02j,
-0.1094488 - 1.20159845e-01j,
0.38243133 + 6.66608316e-18j,
0.31313975 + 1.82190396e-01j,
0.14075481 - 5.85637416e-02j,
-0.15198775 - 1.02156797e-01j,
],
dtype=complex_type(self.dtype),
)
x = self.basis.evaluate(v)
result = np.array(
[
[
9.8593804,
7.94242903,
7.23336975,
7.33314303,
7.41260132,
7.59483694,
7.94830958,
9.47324547,
],
[
7.27801941,
8.29797686,
7.17234599,
7.31082685,
7.04347376,
6.91956664,
8.12234596,
8.36258646,
],
[
8.76188511,
10.69546884,
8.37029969,
9.87512737,
9.73946157,
6.56646752,
5.69555713,
8.77758976,
],
[
10.42069436,
12.3649092,
14.23951952,
20.41736454,
22.32664939,
18.11535113,
7.95059873,
8.79515046,
],
[
11.23152882,
12.61468396,
17.92585027,
25.82097043,
26.4633412,
25.11167661,
11.90634511,
9.05131389,
],
[
10.7048523,
11.73534566,
16.53838035,
25.13242621,
23.58037996,
21.37129485,
12.1024389,
10.26313743,
],
[
8.24162377,
11.90490143,
14.82292441,
19.50174891,
17.69291969,
15.06781768,
10.4669263,
10.2082326,
],
[
5.26532858,
9.60999648,
12.68642275,
12.42354237,
10.87648517,
10.60647963,
9.11026567,
8.53250276,
],
],
dtype=complex_type(self.dtype),
).T # RCOPT
self.assertTrue(np.allclose(x.asnumpy(), result))
def testPolarBasis2DEvaluate_t(self):
x = Image(
np.array(
[
[
0.00000000e00,
0.00000000e00,
0.00000000e00,
0.00000000e00,
-1.08106869e-17,
0.00000000e00,
0.00000000e00,
0.00000000e00,
],
[
0.00000000e00,
0.00000000e00,
-6.40456062e-03,
-3.32961020e-03,
-1.36887927e-02,
-5.42770488e-03,
7.63680861e-03,
0.00000000e00,
],
[
0.00000000e00,
3.16377602e-03,
-9.31273350e-03,
9.46128404e-03,
1.93239220e-02,
3.79891953e-02,
1.06841173e-02,
-2.36467925e-03,
],
[
0.00000000e00,
1.72736955e-03,
-1.00710814e-02,
4.93520304e-02,
3.77702656e-02,
6.57365438e-02,
3.94739462e-03,
-4.41228496e-03,
],
[
4.01551066e-18,
-3.08071647e-03,
-1.61670565e-02,
8.66886286e-02,
5.09898409e-02,
7.19313349e-02,
1.68313715e-02,
5.19180892e-03,
],
[
0.00000000e00,
2.87262215e-03,
-3.37732956e-02,
4.51706505e-02,
5.72215879e-02,
4.63553081e-02,
1.86552175e-03,
1.12608805e-02,
],
[
0.00000000e00,
2.77905016e-03,
-2.77499404e-02,
-4.02645374e-02,
-1.54969139e-02,
-1.66229153e-02,
-2.07389259e-02,
6.64060546e-03,
],
[
0.00000000e00,
0.00000000e00,
5.20080934e-03,
-1.06788196e-02,
-1.14761672e-02,
-1.27443126e-02,
-1.15563484e-02,
0.00000000e00,
],
],
dtype=self.dtype,
).T) # RCOPT
pf = self.basis.evaluate_t(x)
result = np.array(
[
0.38243133 + 6.66608316e-18j,
0.3249317 - 1.47839074e-01j,
0.14819172 + 3.78171168e-03j,
-0.22808599 + 5.29338933e-02j,
0.38243133 + 6.66608316e-18j,
0.34595014 - 1.06355385e-01j,
0.15519289 - 4.75602164e-02j,
-0.22401193 + 4.33128746e-03j,
0.38243133 + 6.66608316e-18j,
0.36957165 - 5.69575709e-02j,
0.17389327 - 5.53498385e-02j,
-0.11601473 - 1.35405676e-02j,
0.38243133 + 6.66608316e-18j,
0.39045046 - 2.17911945e-03j,
0.18146449 - 1.37089189e-02j,
-0.02110144 + 6.65071497e-03j,
0.38243133 + 6.66608316e-18j,
0.4063995 + 5.21354967e-02j,
0.15674204 + 3.85815662e-02j,
-0.02886296 + 3.91489615e-02j,
0.38243133 + 6.66608316e-18j,
0.41872477 + 9.98946906e-02j,
0.11862477 + 5.15231952e-02j,
-0.05298751 + 1.95319478e-02j,
0.38243133 + 6.66608316e-18j,
0.43013599 + 1.38307796e-01j,
0.10075763 + 1.25689289e-02j,
-0.04052728 - 5.66863498e-02j,
0.38243133 + 6.66608316e-18j,
0.44144497 + 1.68826980e-01j,
0.11446016 - 4.53003874e-02j,
-0.03546515 - 1.13544145e-01j,
0.38243133 + 6.66608316e-18j,
0.44960099 + 1.94794929e-01j,
0.15053714 - 8.11915305e-02j,
-0.04800556 - 1.15828804e-01j,
0.38243133 + 6.66608316e-18j,
0.44872328 + 2.17957567e-01j,
0.19116871 - 7.99536373e-02j,
-0.05683092 - 9.72225058e-02j,
0.38243133 + 6.66608316e-18j,
0.43379428 + 2.36681249e-01j,
0.21025378 - 5.48466438e-02j,
-0.05318826 - 8.54948014e-02j,
0.38243133 + 6.66608316e-18j,
0.40485577 + 2.47073481e-01j,
0.18680217 - 3.31766116e-02j,
-0.06674163 - 7.94216591e-02j,
0.38243133 + 6.66608316e-18j,
0.36865853 + 2.45913767e-01j,
0.13660805 - 3.68947359e-02j,
-0.11467046 - 8.49198927e-02j,
0.38243133 + 6.66608316e-18j,
0.33597018 + 2.32971425e-01j,
0.1072859 - 6.24686168e-02j,
-0.12932565 - 1.06139634e-01j,
0.38243133 + 6.66608316e-18j,
0.31616666 + 2.10791785e-01j,
0.11876919 - 7.93812474e-02j,
-0.1094488 - 1.20159845e-01j,
0.38243133 + 6.66608316e-18j,
0.31313975 + 1.82190396e-01j,
0.14075481 - 5.85637416e-02j,
-0.15198775 - 1.02156797e-01j,
0.38243133 - 6.66608316e-18j,
0.3249317 + 1.47839074e-01j,
0.14819172 - 3.78171168e-03j,
-0.22808599 - 5.29338933e-02j,
0.38243133 - 6.66608316e-18j,
0.34595014 + 1.06355385e-01j,
0.15519289 + 4.75602164e-02j,
-0.22401193 - 4.33128746e-03j,
0.38243133 - 6.66608316e-18j,
0.36957165 + 5.69575709e-02j,
0.17389327 + 5.53498385e-02j,
-0.11601473 + 1.35405676e-02j,
0.38243133 - 6.66608316e-18j,
0.39045046 + 2.17911945e-03j,
0.18146449 + 1.37089189e-02j,
-0.02110144 - 6.65071497e-03j,
0.38243133 - 6.66608316e-18j,
0.4063995 - 5.21354967e-02j,
0.15674204 - 3.85815662e-02j,
-0.02886296 - 3.91489615e-02j,
0.38243133 - 6.66608316e-18j,
0.41872477 - 9.98946906e-02j,
0.11862477 - 5.15231952e-02j,
-0.05298751 - 1.95319478e-02j,
0.38243133 - 6.66608316e-18j,
0.43013599 - 1.38307796e-01j,
0.10075763 - 1.25689289e-02j,
-0.04052728 + 5.66863498e-02j,
0.38243133 - 6.66608316e-18j,
0.44144497 - 1.68826980e-01j,
0.11446016 + 4.53003874e-02j,
-0.03546515 + 1.13544145e-01j,
0.38243133 - 6.66608316e-18j,
0.44960099 - 1.94794929e-01j,
0.15053714 + 8.11915305e-02j,
-0.04800556 + 1.15828804e-01j,
0.38243133 - 6.66608316e-18j,
0.44872328 - 2.17957567e-01j,
0.19116871 + 7.99536373e-02j,
-0.05683092 + 9.72225058e-02j,
0.38243133 - 6.66608316e-18j,
0.43379428 - 2.36681249e-01j,
0.21025378 + 5.48466438e-02j,
-0.05318826 + 8.54948014e-02j,
0.38243133 - 6.66608316e-18j,
0.40485577 - 2.47073481e-01j,
0.18680217 + 3.31766116e-02j,
-0.06674163 + 7.94216591e-02j,
0.38243133 - 6.66608316e-18j,
0.36865853 - 2.45913767e-01j,
0.13660805 + 3.68947359e-02j,
-0.11467046 + 8.49198927e-02j,
0.38243133 - 6.66608316e-18j,
0.33597018 - 2.32971425e-01j,
0.1072859 + 6.24686168e-02j,
-0.12932565 + 1.06139634e-01j,
0.38243133 - 6.66608316e-18j,
0.31616666 - 2.10791785e-01j,
0.11876919 + 7.93812474e-02j,
-0.1094488 + 1.20159845e-01j,
0.38243133 - 6.66608316e-18j,
0.31313975 - 1.82190396e-01j,
0.14075481 + 5.85637416e-02j,
-0.15198775 + 1.02156797e-01j,
],
dtype=complex_type(self.dtype),
)
self.assertTrue(np.allclose(pf, result))
def calculate_bispectrum(self,
complex_coef,
flatten=False,
filter_nonzero_freqs=False,
freq_cutoff=None):
"""
Calculate bispectrum for a set of coefs in this basis.
The Bispectum matrix is of shape:
(count, count, unique_radial_indices)
where count is the number of complex coefficients.
:param coef: Coefficients representing a (single) image expanded in this basis.
:param flatten: Optionally extract symmetric values (tril) and then flatten.
:param filter_nonzero_freqs: Remove indices corresponding to zero frequency (defaults False).
:param freq_cutoff: Truncate (zero) high k frequecies above (int) value, defaults off (None).
:return: Bispectum matrix (complex valued).
"""
# Check shape
if complex_coef.shape[0] != 1:
raise ValueError(
"Due to potentially large sizes, bispectrum is limited to a single set of coefs."
f" Passed shape {complex_coef.shape}")
if complex_coef.shape[1] != self.complex_count:
raise ValueError(
"Basis.calculate_bispectrum coefs expected"
f" to have (complex) count {self.complex_count}, received {complex_coef.shape}."
)
# From here just treat complex_coef as 1d vector instead of 1 by count 2d array.
complex_coef = complex_coef[0]
if freq_cutoff and freq_cutoff > np.max(self.complex_angular_indices):
logger.warning(
f"Bispectrum frequency cutoff {freq_cutoff} outside max {np.max(self.complex_angular_indices)}"
)
# Notes, regarding the naming:
# radial freq indices q in paper/slides, _indices["ks"] in code
radial_indices = self.complex_radial_indices # q
# angular freq indices k in paper/slides, _indices["ells"] in code
angular_indices = self.complex_angular_indices # k
# Compute the set of all unique q in the compressed basis
# Note that np.unique is also sorted.
unique_radial_indices = np.unique(radial_indices)
# When compressed, we need maps between the basis and uncompressed set of q
# to a reduced set of q that remain after compression.
# One map is provided by self.complex_radial_indices
# which maps an index in the basis remaining after compression to a q value.
# The following code computes a similar but inverted map,
# given a q value, find an index into the set of unique q after compression.
# Also, it is known that the set of q gets sparser with increasing k
# but that's ignored here, instead construct a dense
# array and filter it later.
# The plan is to revisit this code after appropriate coef classes are derived.
# Default array to fill_value, we can use a value
# k should never achieve..
fill_value = self.complex_count**2
compressed_radial_map = (
np.ones(np.max(unique_radial_indices) + 1, dtype=int) * fill_value)
for uniq_q_index, q_value in enumerate(unique_radial_indices):
compressed_radial_map[q_value] = uniq_q_index
B = np.zeros(
(self.complex_count, self.complex_count,
unique_radial_indices.shape[0]),
dtype=complex_type(self.dtype),
)
logger.info(f"Calculating bispectrum matrix with shape {B.shape}.")
for ind1 in range(self.complex_count):
k1 = angular_indices[ind1]
if freq_cutoff and k1 > freq_cutoff:
continue
coef1 = complex_coef[ind1]
for ind2 in range(self.complex_count):
k2 = angular_indices[ind2]
if freq_cutoff and k2 > freq_cutoff:
continue
coef2 = complex_coef[ind2]
k3 = k1 + k2
intermodulated_coef_inds = angular_indices == k3
if np.any(intermodulated_coef_inds):
# Get the specific q indices related to feasible k3 angular_indices
Q3_ind = radial_indices[intermodulated_coef_inds]
if hasattr(self, "compressed") and self.compressed:
# Map those Q3_ind values to indices into compressed unique_radial_indices
# by using the compressed_radial_map prepared above.
Q3_ind = compressed_radial_map[Q3_ind]
Coef3 = complex_coef[intermodulated_coef_inds]
B[ind1, ind2, Q3_ind] = coef1 * coef2 * np.conj(Coef3)
if filter_nonzero_freqs:
non_zero_freqs = angular_indices != 0
B = B[non_zero_freqs][:, non_zero_freqs]
if flatten:
# B is sym, start by taking lower triangle.
tril = np.tri(B.shape[0], dtype=bool)
B = B[tril, :]
# Then flatten
B = B.flatten()
return B
