class PolarFFT(FFTBase):
def __init__(self, nx, ny, nt, nr):
# Initialize the non-equispaced FFT
self.nfft = NFFT(N=(nx, ny), M=nx * ny, n=None, m=12, flags=None)
# Set up the polar sampling grid
theta = np.linspace(0, 2 * np.pi, nt)
r = np.linspace(0, 1., nr)
T, R = np.meshgrid(theta, r)
self.nfft.x = np.c_[T[..., None], R[..., None]].flatten()
self.nfft.precompute()
def analyze(self, f):
self.nfft.f_hat = f
f_hat = self.nfft.forward()
return f_hat
def synthesize(self, f_hat):
self.nfft.f = f_hat
f = self.nfft.adjoint()
return f_hat
def __init__(self, sz, fourier_pts, epsilon=1e-15, **kwargs):
"""
A plan for non-uniform FFT (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 3-by-K array. These need to be in the range [-pi, pi] in each dimension.
:param epsilon: The desired precision of the NUFFT
"""
self.sz = sz
self.dim = len(sz)
self.fourier_pts = fourier_pts
self.num_pts = fourier_pts.shape[1]
self.epsilon = epsilon
self.cutoff = PyNfftPlan.epsilon_to_nfft_cutoff(epsilon)
self.multi_bandwith = tuple(2 * 2**nextpow2(self.sz))
# TODO - no other flags used in the MATLAB code other than these 2 are supported by the PyNFFT wrapper
self._flags = ('PRE_PHI_HUT', 'PRE_PSI')
self._plan = NFFT(N=self.sz,
M=self.num_pts,
n=self.multi_bandwith,
m=self.cutoff,
flags=self._flags)
self._plan.x = ((1. / (2 * np.pi)) * self.fourier_pts).T
self._plan.precompute()
def __init__(self, domain, uv):
from pynfft.nfft import NFFT
npix = domain.shape[0]
assert npix == domain.shape[1]
assert len(domain.shape) == 2
assert type(npix) == int, "npix must be integer"
assert npix > 0 and (npix %
2) == 0, "npix must be an even, positive integer"
assert isinstance(uv, np.ndarray), "uv must be a Numpy array"
assert uv.dtype == np.float64, "uv must be an array of float64"
assert uv.ndim == 2, "uv must be a 2D array"
assert uv.shape[0] > 0, "at least one point needed"
assert uv.shape[1] == 2, "the second dimension of uv must be 2"
assert np.all(uv >= -0.5) and np.all(uv <= 0.5),\
"all coordinates must lie between -0.5 and 0.5"
self._domain = ift.DomainTuple.make(domain)
self._target = ift.DomainTuple.make(ift.UnstructuredDomain(
uv.shape[0]))
self._capability = self.TIMES | self.ADJOINT_TIMES
self.npt = uv.shape[0]
self.plan = NFFT(self.domain.shape, self.npt, m=6)
self.plan.x = uv
self.plan.precompute()
class PyNfftPlan(Plan):
@staticmethod
def epsilon_to_nfft_cutoff(epsilon):
# NOTE: These are obtained empirically. Should have a theoretical derivation.
rel_errs = [6e-2, 2e-3, 2e-5, 2e-7, 3e-9, 4e-11, 4e-13, 0]
return list(
filter(lambda i_err: i_err[1] < epsilon,
enumerate(rel_errs, start=1)))[0][0]
def __init__(self, sz, fourier_pts, epsilon=1e-15, **kwargs):
"""
A plan for non-uniform FFT (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 3-by-K array. These need to be in the range [-pi, pi] in each dimension.
:param epsilon: The desired precision of the NUFFT
"""
self.sz = sz
self.dim = len(sz)
self.fourier_pts = fourier_pts
self.num_pts = fourier_pts.shape[1]
self.epsilon = epsilon
self.cutoff = PyNfftPlan.epsilon_to_nfft_cutoff(epsilon)
self.multi_bandwith = tuple(2 * 2**nextpow2(self.sz))
# TODO - no other flags used in the MATLAB code other than these 2 are supported by the PyNFFT wrapper
self._flags = ('PRE_PHI_HUT', 'PRE_PSI')
self._plan = NFFT(N=self.sz,
M=self.num_pts,
n=self.multi_bandwith,
m=self.cutoff,
flags=self._flags)
self._plan.x = ((1. / (2 * np.pi)) * self.fourier_pts).T
self._plan.precompute()
def transform(self, signal):
ensure(signal.shape == self.sz,
f'Signal to be transformed must have shape {self.sz}')
self._plan.f_hat = signal.astype('complex64')
f = self._plan.trafo()
if signal.dtype == np.float32:
f = f.astype('complex64')
return f
def adjoint(self, signal):
self._plan.f = signal.astype('complex64')
f_hat = self._plan.adjoint()
if signal.dtype == np.float32:
f_hat = f_hat.astype('complex64')
return f_hat
def __init__(self, RandomField):
super().__init__(RandomField)
# Prepare NFFT objects
from pynfft.nfft import NFFT
x = RandomField.nodes
M = x.shape[1]
self.nfft_obj = NFFT(self.Nd, M)#, n=self.Nd, m=1)
self.nfft_obj.x = (x - 0.5) / np.tile(self.L, [M, 1]).T
self.nfft_obj.precompute()
def __init__(self, nx, ny, nt, nr):
# Initialize the non-equispaced FFT
self.nfft = NFFT(N=(nx, ny), M=nx * ny, n=None, m=12, flags=None)
# Set up the polar sampling grid
theta = np.linspace(0, 2 * np.pi, nt)
r = np.linspace(0, 1., nr)
T, R = np.meshgrid(theta, r)
self.nfft.x = np.c_[T[..., None], R[..., None]].flatten()
self.nfft.precompute()
class NFFT(ift.LinearOperator):
"""Performs a non-equidistant Fourier transform, i.e. a Fourier transform
followed by a degridding operation.
Parameters
----------
domain : RGSpace
Domain of the operator. It has to be two-dimensional and have shape
`(2N, 2N)`. The coordinates of the lower left pixel of the dirty image
are `(-N,-N)`, and of the upper right pixel `(N-1,N-1)`.
uv : numpy.ndarray
2D numpy array of type float64 and shape (M,2), where M is the number
of measurements. uv[i,0] and uv[i,1] contain the u and v coordinates
of measurement #i, respectively. All coordinates must lie in the range
`[-0.5; 0,5[`.
"""
def __init__(self, domain, uv):
from pynfft.nfft import NFFT
npix = domain.shape[0]
assert npix == domain.shape[1]
assert len(domain.shape) == 2
assert type(npix) == int, "npix must be integer"
assert npix > 0 and (npix %
2) == 0, "npix must be an even, positive integer"
assert isinstance(uv, np.ndarray), "uv must be a Numpy array"
assert uv.dtype == np.float64, "uv must be an array of float64"
assert uv.ndim == 2, "uv must be a 2D array"
assert uv.shape[0] > 0, "at least one point needed"
assert uv.shape[1] == 2, "the second dimension of uv must be 2"
assert np.all(uv >= -0.5) and np.all(uv <= 0.5),\
"all coordinates must lie between -0.5 and 0.5"
self._domain = ift.DomainTuple.make(domain)
self._target = ift.DomainTuple.make(ift.UnstructuredDomain(
uv.shape[0]))
self._capability = self.TIMES | self.ADJOINT_TIMES
self.npt = uv.shape[0]
self.plan = NFFT(self.domain.shape, self.npt, m=6)
self.plan.x = uv
self.plan.precompute()
def apply(self, x, mode):
self._check_input(x, mode)
if mode == self.TIMES:
self.plan.f_hat = x.to_global_data()
res = self.plan.trafo().copy()
else:
self.plan.f = x.to_global_data()
res = self.plan.adjoint().copy()
return ift.Field.from_global_data(self._tgt(mode), res)
class Sampling_NFFT(Sampling_method_freq):
def __init__(self, RandomField):
super().__init__(RandomField)
# Prepare NFFT objects
from pynfft.nfft import NFFT
x = RandomField.nodes
M = x.shape[1]
self.nfft_obj = NFFT(self.Nd, M)#, n=self.Nd, m=1)
self.nfft_obj.x = (x - 0.5) / np.tile(self.L, [M, 1]).T
self.nfft_obj.precompute()
def __call__(self, noise):
self.nfft_obj.f_hat = self.Spectrum * FourierOfGaussian(noise)
y = self.nfft_obj.trafo()
# assert (abs(y.imag) < 1.e-8).all(), np.amax(abs(y.imag))
y = np.array(y.real, dtype=np.float) / self.TransformNorm
return y
def test_pynfft():
pixels = 128
wavel = 1.5e-6
plate_scale = 1.2 # mas / pixel
x = np.arange(pixels) - pixels // 2
xx, yy = np.meshgrid(x, x)
image = np.zeros_like(xx)
rho = np.sqrt(xx ** 2 + yy ** 2)
image = image + 1.0 * (rho < 5)
mask = np.random.normal(0, 6, [100, 2])
def uv_samples(mask):
N = mask.shape[0]
uv = np.zeros([N * (N - 1) // 2, 2])
k = 0
for i in range(N):
for j in range(i + 1, N):
uv[k, 0] = mask[i, 0] - mask[j, 0]
uv[k, 1] = mask[i, 1] - mask[j, 1]
k += 1
return uv
uv = uv_samples(mask)
start = time.time()
ndft = NDFTM(uv, wavel, pixels, plate_scale)
vis1 = np.einsum("ij, j -> i", ndft, np.ravel(image))
end = time.time() - start
print(f"Took {end:.4f} second to compute NDFTM")
phase = np.exp(-1j * np.pi / wavel * mas2rad(plate_scale) * (uv[:, 0] + uv[:, 1]))
start = time.time()
plan = NFFT([pixels, pixels], uv.shape[0], n=[pixels, pixels])
plan.x = uv / wavel * mas2rad(plate_scale)
plan.f_hat = image
plan.precompute()
plan.trafo()
vis = plan.f.copy() * phase
end = time.time() - start
print(f"Took {end:.4f} seconds to compute NFFT")
assert np.allclose(np.abs(vis), np.abs(vis1), rtol=1e-5)
assert np.allclose(np.sin(np.angle(vis) - np.angle(vis1)), np.zeros_like(vis), atol=1e-5)
def test_adjoint_ndft():
for N, M, nfft_kwargs in tested_nfft_args:
plan = NFFT(N, M, **nfft_kwargs)
vrand_shifted_unit_double(plan.x.ravel())
plan.precompute()
yield check_adjoint_ndft, plan
DATA_PATH = pkg_resources.resource_filename('pynufft', './src/data/')
# om is normalized between [-pi, pi] but should be in [-0.5, 0.5[
om = np.load(DATA_PATH+'om2D.npz')['arr_0']
om = om/(2*np.pi)
assert(om.shape[1]==dim)
Kd = (Kd1,)*dim
Jd =(Jd1,)*dim
print('setting image dimension Nd...', Nd)
print('setting spectrum dimension Kd...', Kd)
print('setting interpolation size Jd...', Jd)
print('Fourier transform...')
## declaration and pre-computation
time_pre = time.clock()
plan = NFFT(image.shape,om.shape[0])
plan.x = om
plan.precompute()
plan.f_hat = image
## Compute 2D-Fourier transform
time_comp = time.clock()
y_nfft = plan.trafo()
time_end = time.clock()
time_preproc = time_comp - time_pre
time_proc = time_end - time_comp
time_total = time_preproc + time_proc
save_pynfft = {'y':y_nfft, 'Nd':Nd, 'Kd':Kd, 'Jd':Kd, 'om_path':om_path,\
'time_preproc':time_preproc, 'time_proc':time_proc, 'time_total':time_total,\
'adj':adj, 'title':title}
# From https://github.com/conda-forge/pynfft-feedstock/blob/7c29da63dbd4823b7911627488c1d36d25827feb/recipe/meta.yaml#L30-L34 import pynfft from pynfft.nfft import NFFT plan = NFFT(8, 8)
def execute(self):
try:
from pynfft.nfft import NFFT
except ImportError:
raise ImportError(errfmt("""\
Error: The 'pynfft' module seems to be missing. Please install it
through your OS's package manager or directly from PyPI:
https://pypi.python.org/pypi/pynfft/
"""))
if np.any(np.remainder(self.samples.shape, 2)):
raise ValueError(errfmt("""Number of samples must be even
in each direction (NFFT restriction)"""))
# Account for symmetry vs. asymmetry around center (even shape)
# (gfunc is symmetric, NFFT assumes asymmetric)
# TODO: check if this is really necessary
center = self.samples.center + self.samples.spacing / 2
# Compute shift factor here as `freqs` may be changed later
if ne:
freqs = self.freqs
dotp = ne.evaluate('sum(freqs * center, axis=1)')
if self.direction == 'forward':
shift_fac = ne.evaluate('exp(-1j * dotp)')
else:
shift_fac = ne.evaluate('exp(1j * dotp)')
else:
if self.direction == 'forward':
shift_fac = np.exp(-1j * np.dot(self.freqs, center))
else:
shift_fac = np.exp(1j * np.dot(self.freqs, center))
# The normalized frequencies must lie between -1/2 and 1/2
if ne:
spacing = self.samples.spacing[:]
norm_freqs = ne.evaluate('''freqs * spacing / (2 * pi)''')
else:
norm_freqs = self.freqs * self.samples.spacing / (2 * pi)
maxfreq = np.max(np.abs(norm_freqs))
if maxfreq > 0.5:
raise ValueError('''Frequencies must lie between -0.5 and 0.5
after normalization, got maximum absolute value: {}.
'''.format(maxfreq))
# Initialize the geometry and precompute
nfft_plan = NFFT(N=self.samples.shape, M=self.freqs.shape[0])
nfft_plan.x = norm_freqs
nfft_plan.precompute()
# Feed in the samples and compute the transform
if self.direction == 'forward':
# TODO: check ordering!
nfft_plan.f_hat = self.samples.fvals.flatten(order='F')
nfft_plan.trafo()
self.out = nfft_plan.f
else:
# TODO: check ordering!
nfft_plan.f = self.samples.fvals.flatten(order='F')
nfft_plan.adjoint()
self.out = nfft_plan.f_hat
# Account for shift and scaling
scaling = (np.prod(self.samples.spacing) /
(2 * pi)**(self.samples.dim / 2.))
if ne:
out = self.out
self.out = ne.evaluate('out * shift_fac * scaling')
else:
self.out *= shift_fac * scaling
return self.out
def compute_summations(x, y, err, H, ofac=5, hfac=1):
"""
Computes C, S, YC, YS, CC, CS, SS using
pyNFFT
"""
# convert errs to weights
w = weights(err)
# number of frequencies (+1 for 0 freq)
N = int(floor(0.5 * len(x) * ofac * hfac))
# shift times to [ -1/2, 1/2 ]
t = shift_t_for_nfft(x, ofac)
# compute angular frequencies
T = max(x) - min(x)
df = 1. / (ofac * T)
omegas = np.array([2 * np.pi * i * df for i in range(1, N)])
# compute weighted mean
ybar = np.dot(w, y)
# subtract off weighted mean
u = np.multiply(w, y - ybar)
# weighted variance
YY = np.dot(w, np.power(y - ybar, 2))
# plan NFFT's and precompute
plan = NFFT(4 * H * N, len(x))
plan.x = t
plan.precompute()
plan2 = NFFT(2 * H * N, len(x))
plan2.x = t
plan2.precompute()
# evaluate NFFT for w
plan.f = w
f_hat_w = plan.adjoint()[2 * H * N + 1:]
# evaluate NFFT for y - ybar
plan2.f = u
f_hat_u = plan2.adjoint()[H * N + 1:]
all_computed_sums = []
# Now compute the summation values at each frequency
for i in range(N - 1):
computed_sums = Summations(C=np.zeros(H),
S=np.zeros(H),
YC=np.zeros(H),
YS=np.zeros(H),
CCh=np.zeros((H, H)),
CSh=np.zeros((H, H)),
SSh=np.zeros((H, H)))
C_, S_ = np.zeros(2 * H), np.zeros(2 * H)
for j in range(2 * H):
# This sign factor is necessary
# but I don't know why.
s = (-1 if ((i % 2) == 0) and ((j % 2) == 0) else 1)
C_[j] = f_hat_w[(j + 1) * (i + 1) - 1].real * s
S_[j] = f_hat_w[(j + 1) * (i + 1) - 1].imag * s
if j < H:
computed_sums.YC[j] = f_hat_u[(j + 1) * (i + 1) - 1].real * s
computed_sums.YS[j] = f_hat_u[(j + 1) * (i + 1) - 1].imag * s
for j in range(H):
for k in range(H):
Sn, Cn = None, None
if j == k:
Sn = 0
Cn = 1
else:
Sn = np.sign(k - j) * S_[int(abs(k - j)) - 1]
Cn = C_[int(abs(k - j)) - 1]
Sp = S_[j + k + 1]
Cp = C_[j + k + 1]
computed_sums.CCh[j][k] = 0.5 * (Cn + Cp) - C_[j] * C_[k]
computed_sums.CSh[j][k] = 0.5 * (Sn + Sp) - C_[j] * S_[k]
computed_sums.SSh[j][k] = 0.5 * (Cn - Cp) - S_[j] * S_[k]
computed_sums.C[:] = C_[:H]
computed_sums.S[:] = S_[:H]
all_computed_sums.append(computed_sums)
return omegas, all_computed_sums, YY, w, ybar
def fast_summations(t, y, w, freqs, nh, eps=1E-5):
"""
Computes C, S, YC, YS, CC, CS, SS using
pyNFFT
"""
nf, df, dnf = inspect_freqs(freqs)
tmin = min(t)
# infer samples per peak
baseline = max(t) - tmin
samples_per_peak = 1. / (baseline * df)
eps = 1E-5
a = 0.5 - eps
r = 2 * a / df
tshift = a * (2 * (t - tmin) / r - 1)
# number of frequencies needed for NFFT
# need nf_nfft_u / 2 - 1 = H * (nf - 1 + dnf)
# nf_nfft_w / 2 - 1 = 2H * (nf - 1 + dnf)
nf_nfft_u = 2 * (nh * (nf + dnf - 1) + 1)
nf_nfft_w = 2 * (2 * nh * (nf + dnf - 1) + 1)
n_w0 = int(floor(nf_nfft_w / 2))
n_u0 = int(floor(nf_nfft_u / 2))
# transform y -> w_i * y_i - ybar
ybar = np.dot(w, y)
u = np.multiply(w, y - ybar)
# plan NFFT's and precompute
plan = NFFT(nf_nfft_w, len(tshift))
plan.x = tshift
plan.precompute()
plan2 = NFFT(nf_nfft_u, len(tshift))
plan2.x = tshift
plan2.precompute()
# NFFT(weights)
plan.f = w
f_hat_w = plan.adjoint()[n_w0:]
# NFFT(y - ybar)
plan2.f = u
f_hat_u = plan2.adjoint()[n_u0:]
# now correct for phase shift induced by transforming t -> (-1/2, 1/2)
beta = -a * (2 * tmin / r + 1)
I = 0. + 1j
twiddles = np.exp(-I * 2 * np.pi * np.arange(0, n_w0) * beta)
f_hat_u *= twiddles[:len(f_hat_u)]
f_hat_w *= twiddles[:len(f_hat_w)]
all_computed_sums = []
# Now compute the summation values at each frequency
for i in range(nf):
j = np.arange(2 * nh)
k = (j + 1) * (i + dnf)
C = f_hat_w[k].real
S = f_hat_w[k].imag
YC = f_hat_u[k[:nh]].real
YS = f_hat_u[k[:nh]].imag
#-------------------------------
# Note: redefining j and k here!
k = np.arange(nh)
j = k[:, np.newaxis]
Sn = np.sign(k - j) * S[abs(k - j) - 1]
Sn.flat[::nh + 1] = 0 # set diagonal to zero
Cn = C[abs(k - j) - 1]
Cn.flat[::nh + 1] = 1 # set diagonal to one
Sp = S[j + k + 1]
Cp = C[j + k + 1]
CC = 0.5 * (Cn + Cp) - C[j] * C[k]
CS = 0.5 * (Sn + Sp) - C[j] * S[k]
SS = 0.5 * (Cn - Cp) - S[j] * S[k]
all_computed_sums.append(
Summations(C=C[:nh], S=S[:nh], YC=YC, YS=YS, CC=CC, CS=CS, SS=SS))
return all_computed_sums
rm = 0.5 * pdim
return j1(rm * np.sqrt(x**2 + y**2)) / np.sqrt(x**2 + y**2) / rm**2
uv = uv_samples(mask)
start = time.time()
ndft = NDFTM(uv, wavel, pixels, plate_scale)
vis1 = np.einsum("ij, j -> i", ndft, np.ravel(image))
end = time.time() - start
print(f"Took {end:.4f} second to compute NDFTM")
m2pix = mas2rad(plate_scale) * pixels / wavel
phase = np.exp(-1j * np.pi / wavel * mas2rad(plate_scale) *
(uv[:, 0] + uv[:, 1]))
start = time.time()
plan = NFFT([pixels, pixels], uv.shape[0], n=[pixels, pixels])
plan.x = uv / wavel * mas2rad(plate_scale)
plan.f_hat = image
plan.precompute()
plan.trafo()
vis = plan.f.copy() * phase
end = time.time() - start
print(
f"Took {end:.4f} seconds to compute NFFT"
) # usually at least 5x faster, scales better with large number of pixels
# print(np.abs(vis))
# print(np.abs(vis1))
plt.figure()
plt.plot(sorted(np.abs(vis)), color="k")
plt.plot(sorted(np.abs(vis1)), color="r")
plt.show()