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(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 last dimension must equal to
`self.count`.
:return x: Image instance in standard 2D coordinate basis with
resolution of `self.sz`.
"""
if self.dtype != real_type(v.dtype):
msg = (f"Input data type, {v.dtype}, is not consistent with"
f" type defined in the class {self.dtype}.")
logger.error(msg)
raise TypeError(msg)
v = v.reshape(-1, self.ntheta, self.nrad)
nimgs = v.shape[0]
half_size = self.ntheta // 2
v = v[:, :half_size, :] + v[:, half_size:, :].conj()
v = v.reshape(nimgs, self.nrad * half_size)
x = anufft(v, self.freqs, self.sz, real=True)
return Image(x)
def testComplexCoversionErrorsToReal(self):
# Load a reasonable input
x = np.load(os.path.join(DATA_DIR, "fbbasis_coefficients_8_8.npy"))
# Express in an FB basis
cv1 = self.basis.to_complex(self.basis.expand(x.astype(self.dtype)))
# Test catching Errors
with raises(TypeError):
# Pass real into `to_real`
_ = self.basis.to_real(cv1.real.astype(np.float32))
# Test casting case, where basis and coef precision don't match
if self.basis.dtype == np.float32:
test_dtype = np.complex128
elif self.basis.dtype == np.float64:
test_dtype = np.complex64
# Result should be same precision as coef input, just real.
result_dtype = real_type(test_dtype)
v3 = self.basis.to_real(cv1.astype(test_dtype))
self.assertTrue(v3.dtype == result_dtype)
# Try a 0d vector, should not crash.
_ = self.basis.to_real(cv1.reshape(-1))
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 calculate_bispectrum(self,
coef,
flatten=False,
filter_nonzero_freqs=False,
freq_cutoff=None):
if coef.dtype == real_type(self.dtype):
coef = self.to_complex(coef)
return super().calculate_bispectrum(
coef,
flatten=flatten,
filter_nonzero_freqs=filter_nonzero_freqs,
freq_cutoff=freq_cutoff,
)
def to_real(self, complex_coef):
"""
Return real valued representation of complex coefficients.
This can be useful when comparing or implementing methods
from literature.
There is a corresponding method, to_complex.
:param complex_coef: Complex coefficients from this basis.
:return: Real coefficent representation from this basis.
"""
if complex_coef.ndim == 1:
complex_coef = complex_coef.reshape(1, -1)
if complex_coef.dtype not in (np.complex128, np.complex64):
raise TypeError("coef provided to to_real should be complex.")
# Pass through dtype precions, but check and warn if mismatched.
dtype = real_type(complex_coef.dtype)
if dtype != self.dtype:
logger.warning(
f"Complex coef dtype {complex_coef.dtype} does not match precision of basis.dtype {self.dtype}, returning {dtype}."
)
coef = np.zeros((complex_coef.shape[0], self.count), dtype=dtype)
ind = 0
idx = np.arange(self.k_max[0], dtype=int)
ind += np.size(idx)
ind_pos = ind
coef[:, idx] = complex_coef[:, idx].real
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]
c = complex_coef[:, idx]
coef[:, idx_pos] = 2.0 * np.real(c)
coef[:, idx_neg] = -2.0 * np.imag(c)
ind += np.size(idx)
ind_pos += 2 * self.k_max[ell]
return coef
def to_real(self, complex_coef):
"""
Return real valued representation of complex coefficients.
This can be useful when comparing or implementing methods
from literature.
There is a corresponding method, to_complex.
:param complex_coef: Complex coefficients from this basis.
:return: Real coefficent representation from this basis.
"""
if complex_coef.ndim == 1:
complex_coef = complex_coef.reshape(1, -1)
if complex_coef.dtype not in (np.complex128, np.complex64):
raise TypeError("coef provided to to_real should be complex.")
# Pass through dtype precions, but check and warn if mismatched.
dtype = real_type(complex_coef.dtype)
if dtype != self.dtype:
logger.warning(
f"Complex coef dtype {complex_coef.dtype} does not match precision of basis.dtype {self.dtype}, returning {dtype}."
)
coef = np.zeros((complex_coef.shape[0], self.count), dtype=dtype)
# map ordered index to (ell, q) key in dict
keymap = list(self.complex_indices_map.keys())
for i in range(self.complex_count):
# retreive index into reals
pos_i, neg_i = self.complex_indices_map[keymap[i]]
if self.complex_angular_indices[i] == 0:
coef[:, pos_i] = complex_coef[:, i].real
else:
coef[:, pos_i] = 2.0 * complex_coef[:, i].real
coef[:, neg_i] = -2.0 * complex_coef[:, i].imag
return coef