def SweepSpectraCython(h, N, M, xdtype=np.complex_):
# implements Doppler filter:
# y[n, p] = SUM_k exp(2*pi*j*n*(k - (L-1))/N) * (h[k] * x[p - k])
# = SUM_k exp(-2*pi*j*n*k/N) * (s*[k] * x[p - (L-1) + k])
L = len(h)
outlen = M + L - 1
# when N < L, still need to take FFT with nfft >= L so we don't lose data
# then subsample to get our N points that we desire
step = L // N + 1
nfft = N*step
# ensure that h is C-contiguous as required by the Cython function
h = np.asarray(h, order='C')
# make sure xdtype is a dtype object
xdtype = np.dtype(xdtype)
demodpad = np.zeros((outlen, nfft), np.result_type(xdtype, h.dtype, np.complex64))
demodpad = pyfftw.n_byte_align(demodpad, 16)
y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)
sweepspectra_cython = libdopplerbanks.SweepSpectraCython(h, demodpad, y, fft,
step, N, M, xdtype)
sweepspectra_cython = dopplerbank_dec(h, N, M)(sweepspectra_cython)
return sweepspectra_cython
def SweepSpectraNumba(h, N, M, xdtype=np.complex_):
# implements Doppler filter:
# y[n, p] = SUM_k exp(2*pi*j*n*(k - (L-1))/N) * (h[k] * x[p - k])
# = SUM_k exp(-2*pi*j*n*k/N) * (s*[k] * x[p - (L-1) + k])
L = len(h)
outlen = M + L - 1
# when N < L, still need to take FFT with nfft >= L so we don't lose data
# then subsample to get our N points that we desire
step = L // N + 1
nfft = N*step
hrev = h[::-1]
xpad = np.zeros(M + 2*(L - 1), xdtype) # x[0] at xpad[L - 1]
demodpad = np.zeros((outlen, nfft), np.result_type(xdtype, h.dtype, np.complex64))
demodpad = pyfftw.n_byte_align(demodpad, 16)
y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)
xtype = numba.__getattribute__(str(np.dtype(xdtype)))
#@jit(argtypes=[xtype[::1]])
@jit
def sweepspectra_numba(x):
xpad[(L - 1):outlen] = x
for p in range(outlen):
demodpad[p, :L] = hrev*xpad[p:(p + L)]
fft.execute() # input is demodpad, output is y
yc = np.array(y[:, ::step].T) # we need a copy, which np.array provides
return yc
sweepspectra_numba = dopplerbank_dec(h, N, M)(sweepspectra_numba)
return sweepspectra_numba
def ShiftConvNumbaFFT(h, N, M, xdtype=np.complex_, powerof2=True):
# implements Doppler filter:
# y[n, p] = SUM_k (exp(2*pi*j*n*(k - (L-1))/N) * h[k]) * x[p - k]
# = SUM_k (exp(-2*pi*j*n*k/N) * s*[k]) * x[p - (L-1) + k]
L = len(h)
outlen = M + L - 1
nfft = outlen
if powerof2:
nfft = pow2(nfft)
dopplermat = np.exp(2*np.pi*1j*np.arange(N)[:, np.newaxis]*(np.arange(L) - (L - 1))/N)
dopplermat.astype(np.result_type(h.dtype, np.complex64)) # cast to complex type with precision of h
hbank = h*dopplermat
# speed not critical here, just use numpy fft
hbankpad = zero_pad(hbank, nfft)
H = np.fft.fft(hbankpad) / nfft # divide by nfft b/c FFTW's ifft does not do this
xcdtype = np.result_type(xdtype, np.complex64) # cast to complex type with precision of x
xpad = pyfftw.n_byte_align(np.zeros(nfft, xcdtype), 16)
X = pyfftw.n_byte_align(np.zeros(nfft, xcdtype), 16)
xfft = pyfftw.FFTW(xpad, X, threads=_THREADS)
ydtype = np.result_type(H.dtype, xcdtype)
Y = pyfftw.n_byte_align_empty(H.shape, 16, ydtype)
y = pyfftw.n_byte_align_empty(H.shape, 16, ydtype)
ifft = pyfftw.FFTW(Y, y, direction='FFTW_BACKWARD', threads=_THREADS)
xtype = numba.__getattribute__(str(np.dtype(xdtype)))
#htype = numba.__getattribute__(str(H.dtype))
#xctype = numba.__getattribute__(str(X.dtype))
#ytype = numba.__getattribute__(str(Y.dtype))
#@jit(argtypes=[htype[:, ::1], xctype[::1], ytype[:, ::1], xtype[::1]])
#def fun(H, X, Y, x):
#xpad[:M] = x
#xfft.execute() # input is xpad, output is X
#Y[:, :] = H*X # need expression optimized by numba but that writes into Y
#ifft.execute() # input is Y, output is y
#yc = np.array(y)[:, :outlen] # need a copy, which np.array provides
#return yc
#@dopplerbank_dec(h, N, M, nfft=nfft, H=H)
#def shiftconv_numba_fft(x):
#return fun(H, X, Y, x)
#@jit(argtypes=[xtype[::1]])
@jit
def shiftconv_numba_fft(x):
xpad[:M] = x
xfft.execute() # input is xpad, output is X
Y[:, :] = X*H # need expression optimized by numba but that writes into Y
ifft.execute() # input is Y, output is y
yc = np.array(y[:, :outlen]) # need a copy, which np.array provides
return yc
shiftconv_numba_fft = dopplerbank_dec(h, N, M, nfft=nfft, H=H)(shiftconv_numba_fft)
return shiftconv_numba_fft
def SweepSpectraStridedTapsMod(h, N, M, xdtype=np.complex_):
"""Doppler bank where the signal is downshift modulated before filtering."""
# implements Doppler filter:
# y[n, p] = SUM_k (h[p - k] * x[k]) * exp(-2*pi*j*n*k/N)
# = SUM_k (s*[k + (L-1) - p] * x[k]) * exp(-2*pi*j*n*k/N)
L = len(h)
# when N < M, still need to take FFT with nfft >= M so we don't lose data
# then subsample to get our N points that we desire
step = M // N + 1
nfft = N*step
hpad = np.hstack((np.zeros(M - 1, h.dtype), h, np.zeros(M - 1, h.dtype)))
hmat = np.lib.stride_tricks.as_strided(hpad[(M - 1):], (M + L - 1, M),
(hpad.itemsize, -hpad.itemsize))
demodpad = np.zeros((M + L - 1, nfft), np.result_type(xdtype, h.dtype, np.complex64))
demodpad = pyfftw.n_byte_align(demodpad, 16)
y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)
@dopplerbank_dec(h, N, M, hmat=hmat)
def sweepspectra_strided_taps_mod(x):
np.multiply(hmat, x, demodpad[:, :M])
fft.execute() # input is demodpad, output is y
yc = np.array(y[:, ::step].T) # need a copy, which np.array provides
return yc
return sweepspectra_strided_taps_mod
def acquire(self, n):
# Leave all array allocation beyond storage to subclasses
# Set up counters
self.n = n
self.i = -1 # First thing it does is get incremented
self.start_countup = 0
# Set up storage
self.storage = np.zeros((self.n, self.x), dtype=np.int32)
# Set up fft arrays
self.ft_in = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
self.ft_out = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
self.ift_in = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
self.ift_out = fftw.n_byte_align(np.zeros((n, self.x), dtype=np.complex128), 16)
self.high_pass = np.zeros((n, self.x), dtype=np.complex128)
self.high_pass[0,:] = 1
self.fft = fftw.FFTW(self.ft_in, self.ft_out, axes=(0,), direction='FFTW_FORWARD',
flags=('FFTW_MEASURE',), threads=1, planning_timelimit=None)
self.ifft = fftw.FFTW(self.ift_in, self.ift_out, axes=(0,), direction='FFTW_BACKWARD',
flags=('FFTW_MEASURE',), threads=1, planning_timelimit=None)
# self.fft = fftw.FFTW(self.ft_in, self.ft_out, axes=(0), direction='FFTW_FORWARD')
# self.ifft = fftw.FFTW(self.ift_in, self.ift_out, axes=(0), direction='FFTW_BACKWARD')
# Get ready to calculate fps
self.plot_times = []
# Go
self.new_image.connect(self.send_plot)
self.running = True
def createFFTArrays(ftbuf, axis_len, buf_len):
# Set up fft arrays
ftbuf.ft_in = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.ft_out = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.ift_in = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.ift_out = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.fft = fftw.FFTW(ftbuf.ft_in,
ftbuf.ft_out,
axes=(1, ),
direction='FFTW_FORWARD',
flags=('FFTW_MEASURE', ),
threads=1,
planning_timelimit=None)
ftbuf.ifft = fftw.FFTW(ftbuf.ift_in,
ftbuf.ift_out,
axes=(1, ),
direction='FFTW_BACKWARD',
flags=('FFTW_MEASURE', ),
threads=1,
planning_timelimit=None)
return ftbuf
def SweepSpectraStridedInput(h, N, M, xdtype=np.complex_):
# implements Doppler filter:
# y[n, p] = SUM_k exp(2*pi*j*n*(k - (L-1))/N) * (h[k] * x[p - k])
# = SUM_k exp(-2*pi*j*n*k/N) * (s*[k] * x[p - (L-1) + k])
L = len(h)
outlen = M + L - 1
# when N < L, still need to take FFT with nfft >= L so we don't lose data
# then subsample to get our N points that we desire
step = L // N + 1
nfft = N*step
hrev = h[::-1]
xpad = np.zeros(M + 2*(L - 1), xdtype) # x[0] at xpad[L - 1]
xshifted = np.lib.stride_tricks.as_strided(xpad,
(outlen, L),
(xpad.itemsize, xpad.itemsize))
demodpad = np.zeros((outlen, nfft), np.result_type(xdtype, h.dtype, np.complex64))
demodpad = pyfftw.n_byte_align(demodpad, 16)
y = pyfftw.n_byte_align(np.zeros_like(demodpad), 16)
fft = pyfftw.FFTW(demodpad, y, threads=_THREADS)
@dopplerbank_dec(h, N, M)
def sweepspectra_strided_input(x):
xpad[(L - 1):outlen] = x
np.multiply(xshifted, hrev, demodpad[:, :L])
fft.execute() # input is demodpad, output is y
yc = np.array(y[:, ::step].T) # we need a copy, which np.array provides
return yc
return sweepspectra_strided_input
def phrt(im, plan, method=4, nr_threads=2):
"""Process a tomographic projection image with the selected phase retrieval algorithm.
Parameters
----------
im : array_like
Flat corrected image data as numpy array.
plan : structure
Structure with pre-computed data (see prepare_plan function)
method : int
Phase retrieval filter {1 = TIE (default), 2 = CTF, 3 = CTF first-half sine, 4 = Quasiparticle,
5 = Quasiparticle first half sine}.
nr_threads : int
Number of threads to be used in the computation of FFT by PyFFTW
Credits
-------
Julian Moosmann, KIT (Germany) is acknowledged for this code
"""
# Extract plan values:
dim0_o = plan['dim0']
dim1_o = plan['dim1']
n_pad0 = plan['npad0']
n_pad1 = plan['npad1']
marg0 = (n_pad0 - dim0_o) / 2
marg1 = (n_pad1 - dim1_o) / 2
filter = plan['filter']
# Pad image (if required):
im = padImage(im, n_pad0, n_pad1)
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = fft2(im - 1, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = fft2(im - 1)
# Apply phase retrieval filter:
im = filter * im
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = real(ifft2(im, threads=nr_threads))
# (Un)comment the following line to use NumPy:
#im = real(ifft2(im))
# Return the negative:
im = - im
# Return cropped output:
return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def phrt(im, plan, method=4, nr_threads=2):
"""Process a tomographic projection image with the selected phase retrieval algorithm.
Parameters
----------
im : array_like
Flat corrected image data as numpy array.
plan : structure
Structure with pre-computed data (see prepare_plan function)
method : int
Phase retrieval filter {1 = TIE (default), 2 = CTF, 3 = CTF first-half sine, 4 = Quasiparticle,
5 = Quasiparticle first half sine}.
nr_threads : int
Number of threads to be used in the computation of FFT by PyFFTW
Credits
-------
Julian Moosmann, KIT (Germany) is acknowledged for this code
"""
# Extract plan values:
dim0_o = plan['dim0']
dim1_o = plan['dim1']
n_pad0 = plan['npad0']
n_pad1 = plan['npad1']
marg0 = (n_pad0 - dim0_o) / 2
marg1 = (n_pad1 - dim1_o) / 2
filter = plan['filter']
# Pad image (if required):
im = padImage(im, n_pad0, n_pad1)
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = fft2(im - 1, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = fft2(im - 1)
# Apply phase retrieval filter:
im = filter * im
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = real(ifft2(im, threads=nr_threads))
# (Un)comment the following line to use NumPy:
#im = real(ifft2(im))
# Return the negative:
im = -im
# Return cropped output:
return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def NumbaFFTW(h, M, xdtype=np.complex_, powerof2=True):
L = len(h)
outlen = M + L - 1
nfft = outlen
if powerof2:
nfft = pow2(nfft)
outdtype = np.result_type(h.dtype, xdtype)
fftdtype = np.result_type(outdtype, np.complex64) # output is always complex, promote using smallest
# speed not critical here, just use numpy fft
# cast to outdtype so we use same type of fft as when transforming x
hpad = zero_pad(h, nfft).astype(outdtype)
if np.iscomplexobj(hpad):
H = np.fft.fft(hpad)
else:
H = np.fft.rfft(hpad)
H = (H / nfft).astype(fftdtype) # divide by nfft b/c FFTW's ifft does not do this
xpad = pyfftw.n_byte_align(np.zeros(nfft, outdtype), 16) # outdtype so same type fft as h->H
X = pyfftw.n_byte_align(np.zeros(len(H), fftdtype), 16) # len(H) b/c rfft may be used
xfft = pyfftw.FFTW(xpad, X, threads=_THREADS)
y = pyfftw.n_byte_align_empty(nfft, 16, outdtype)
ifft = pyfftw.FFTW(X, y, direction='FFTW_BACKWARD', threads=_THREADS)
xtype = numba.__getattribute__(str(np.dtype(xdtype)))
outtype = numba.__getattribute__(str(outdtype))
ffttype = numba.__getattribute__(str(fftdtype))
#@jit(restype=outtype[::1],
#argtypes=[outtype[::1], ffttype[::1], ffttype[::1], outtype[::1], xtype[::1]])
#def filt(xpad, X, H, y, x):
#xpad[:M] = x
#xfft.execute() # input in xpad, result in X
#X[:] = H*X
#ifft.execute() # input in X, result in y
#yc = y[:outlen].copy()
#return yc
#@filter_dec(h, M, nfft=nfft, H=H)
#def numba_fftw(x):
#return filt(xpad, X, H, y, x)
#@jit(argtypes=[xtype[::1]])
@jit
def numba_fftw(x):
xpad[:M] = x
xfft.execute() # input in xpad, result in X
X[:] = H*X # want expression that is optimized by numba but writes into X
ifft.execute() # input in X, result in y
yc = y[:outlen].copy()
return yc
numba_fftw = filter_dec(h, M, nfft=nfft, H=H)(numba_fftw)
return numba_fftw
def test_call_with_unaligned(self):
'''Make sure the right thing happens with unaligned data.
'''
input_array = (numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape))
output_array = self.fft(
input_array=n_byte_align(input_array.copy(), 16)).copy()
input_array = n_byte_align(input_array, 16)
output_array = n_byte_align(output_array, 16)
# Offset by one from 16 byte aligned to guarantee it's not
# 16 byte aligned
a = n_byte_align(input_array.copy(), 16)
a__ = n_byte_align_empty(
numpy.prod(a.shape)*a.itemsize+1, 16, dtype='int8')
a_ = a__[1:].view(dtype=a.dtype).reshape(*a.shape)
a_[:] = a
# Create a different second array the same way
b = n_byte_align(output_array.copy(), 16)
b__ = n_byte_align_empty(
numpy.prod(b.shape)*a.itemsize+1, 16, dtype='int8')
b_ = b__[1:].view(dtype=b.dtype).reshape(*b.shape)
b_[:] = a
# Set up for the first array
fft = FFTW(input_array, output_array)
a_[:] = a
output_array = fft().copy()
# Check a_ is not aligned...
self.assertRaisesRegex(ValueError, 'Invalid input alignment',
self.fft.update_arrays, *(a_, output_array))
# and b_ too
self.assertRaisesRegex(ValueError, 'Invalid output alignment',
self.fft.update_arrays, *(input_array, b_))
# But it should still work with the a_
fft(a_)
# However, trying to update the output will raise an error
self.assertRaisesRegex(ValueError, 'Invalid output alignment',
self.fft.update_arrays, *(input_array, b_))
# Same with SIMD off
fft = FFTW(input_array, output_array, flags=('FFTW_UNALIGNED',))
fft(a_)
self.assertRaisesRegex(ValueError, 'Invalid output alignment',
self.fft.update_arrays, *(input_array, b_))
def _create_ifft_plan(self):
x = pyfftw.n_byte_align(
np.zeros((self.P, self.N), self.xydtype),
pyfftw.simd_alignment
)
ifft_out = pyfftw.n_byte_align(
np.zeros_like(x),
pyfftw.simd_alignment
)
ifft = pyfftw.FFTW(x, ifft_out, direction='FFTW_BACKWARD')
return ifft
def _create_fft_plan(self):
delaymult_out = pyfftw.n_byte_align(
np.zeros((self.P, self.N), self.xydtype),
pyfftw.simd_alignment
)
fft_out = pyfftw.n_byte_align(
np.zeros_like(delaymult_out),
pyfftw.simd_alignment
)
fft = pyfftw.FFTW(delaymult_out, fft_out)
return fft
def test_update_data_with_alignment_error(self):
in_shape = self.input_shapes['2d']
out_shape = self.output_shapes['2d']
byte_error = 1
axes=(-1,)
a, b = self.create_test_arrays(in_shape, out_shape)
a = n_byte_align(a, 16)
b = n_byte_align(b, 16)
fft, ifft = self.run_validate_fft(a, b, axes)
a, b = self.create_test_arrays(in_shape, out_shape)
# Offset from 16 byte aligned to guarantee it's not
# 16 byte aligned
a__ = n_byte_align_empty(
numpy.prod(in_shape)*a.itemsize+byte_error,
16, dtype='int8')
a_ = (a__[byte_error:]
.view(dtype=self.input_dtype).reshape(*in_shape))
a_[:] = a
b__ = n_byte_align_empty(
numpy.prod(out_shape)*b.itemsize+byte_error,
16, dtype='int8')
b_ = (b__[byte_error:]
.view(dtype=self.output_dtype).reshape(*out_shape))
b_[:] = b
with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft,
create_array_copies=False)
with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft,
create_array_copies=False)
# Should also be true for the unaligned case
fft, ifft = self.run_validate_fft(a, b, axes,
force_unaligned_data=True)
with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft,
create_array_copies=False)
with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft,
create_array_copies=False)
def get_fftw3_speed(arr, iters=10, direction=None, dtype=np.float32, **kwargs):
""" measure the fftw3 speed for various data sizes by using
plan with estimate, and running one instance. If arr is int,
then the elements are different array sizes, otherwise use the
array.
direction default - just fwd - use 'both' for both
To "train": - the print allows you to ^c out.
from pyfusion.utils.primefactors import get_fftw3_speed
from pyfusion.utils.fftw3_bdb_utils import save_wisdom, load_wisdom
from pyfusion.data.filters import next_nice_number
from pyfusion.utils.primefactors import get_fftw3_speed
for n in next_nice_number(None):
print(n);
get_fftw3_speed(n, flags=['FFTW_MEASURE'],direction='both',dtype=np.float32)
save_wisdom()
Accepts all pyfftw.FFTW args e.g. planning_timelimit
"""
import pyfftw
from time import time as seconds
if np.isscalar(arr): arr = np.array([arr])
if np.issubdtype(arr.dtype, int):
atimes = []
for n in arr:
atimes.append([
n,
get_fftw3_speed(np.ones(n, dtype=dtype),
direction=direction,
iters=iters,
**kwargs)
])
return (np.array(atimes))
else: # do one example
build_kwargs = dict(flags=['FFTW_ESTIMATE'])
build_kwargs.update(kwargs)
print(build_kwargs)
simd_align = pyfftw.simd_alignment # 16 at the moment.
arr = pyfftw.n_byte_align(arr, simd_align)
out = pyfftw.n_byte_align(
np.ones(len(arr) / 2 + 1, dtype=np.complex64), simd_align)
fwd = pyfftw.FFTW(arr, out, **build_kwargs)
if direction == 'both':
rev = pyfftw.FFTW(out,
arr,
direction='FFTW_BACKWARD',
**build_kwargs)
st = seconds()
for i in range(iters):
fwd.execute()
if direction == 'both':
rev.execute()
return ((seconds() - st) / iters)
def tiehom2020(im, plan, nr_threads=2):
"""Process a tomographic projection image with the Generalized TIE-HOM (Paganin et al 2020)
phase retrieval algorithm.
Parameters
----------
im : array_like
Flat corrected image data as numpy array.
plan : structure
Structure with pre-computed data (see tiehom_plan function).
nr_threads : int
Number of threads to be used in the computation of FFT by PyFFTW (default = 2).
"""
# Extract plan values:
dim0_o = plan['dim0']
dim1_o = plan['dim1']
n_pad0 = plan['npad0']
n_pad1 = plan['npad1']
marg0 = (n_pad0 - dim0_o) / 2
marg1 = (n_pad1 - dim1_o) / 2
den = plan['den']
mu = plan['mu']
# Pad image (if required):
im = padImage(im, n_pad0, n_pad1)
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = rfft2(im, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = rfft2(im)
# Apply Paganin's (pre-computed) formula:
im = im / den
# (Un)comment the following two lines to use PyFFTW:
n_byte_align(im, simd_alignment)
im = irfft2(im, threads=nr_threads)
# (Un)comment the following line to use NumPy:
#im = irfft2(im)
im = im.astype(float32)
im = -1 / mu * nplog(im)
# Return cropped output:
return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def tiehom(im, plan, nr_threads=2): """Process a tomographic projection image with the TIE-HOM (Paganin's) phase retrieval algorithm. Parameters ---------- im : array_like Flat corrected image data as numpy array. plan : structure Structure with pre-computed data (see tiehom_plan function). nr_threads : int Number of threads to be used in the computation of FFT by PyFFTW (default = 2). """ # Extract plan values: dim0_o = plan['dim0'] dim1_o = plan['dim1'] n_pad0 = plan['npad0'] n_pad1 = plan['npad1'] marg0 = (n_pad0 - dim0_o) / 2 marg1 = (n_pad1 - dim1_o) / 2 den = plan['den'] mu = plan['mu'] # Pad image (if required): im = padImage(im, n_pad0, n_pad1) # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = rfft2(im, threads=nr_threads) # (Un)comment the following line to use NumPy: #im = rfft2(im) # Apply Paganin's (pre-computed) formula: im = im / den # (Un)comment the following two lines to use PyFFTW: n_byte_align(im, simd_alignment) im = irfft2(im, threads=nr_threads) # (Un)comment the following line to use NumPy: #im = irfft2(im) im = im.astype(float32) im = -1 / mu * nplog(im) # Return cropped output: return im[marg0:dim0_o + marg0, marg1:dim1_o + marg1]
def test_call_with_keyword_input_update(self):
'''Test the class call with a keyword input update.
'''
input_array = n_byte_align(
numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape), 16)
output_array = self.fft(
input_array=n_byte_align(input_array.copy(), 16)).copy()
self.update_arrays(input_array, self.output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(output_array == self.output_array))
def test_call_with_keyword_output_update(self):
"""Test the class call with a keyword output update.
"""
output_array = n_byte_align(
(numpy.random.randn(*self.output_array.shape) + 1j * numpy.random.randn(*self.output_array.shape)),
simd_alignment,
)
returned_output_array = self.fft(output_array=n_byte_align(output_array.copy(), simd_alignment)).copy()
self.update_arrays(self.input_array, output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(returned_output_array == output_array))
def createFFTArrays(ftbuf, axis_len, buf_len):
# Set up fft arrays
ftbuf.ft_in = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.ft_out = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.ift_in = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.ift_out = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
ftbuf.fft = fftw.FFTW(ftbuf.ft_in, ftbuf.ft_out,
axes=(1,), direction='FFTW_FORWARD',
flags=('FFTW_MEASURE',), threads=1, planning_timelimit=None)
ftbuf.ifft = fftw.FFTW(ftbuf.ift_in, ftbuf.ift_out,
axes=(1,), direction='FFTW_BACKWARD',
flags=('FFTW_MEASURE',), threads=1, planning_timelimit=None)
return ftbuf
def test_call_with_positional_input_update(self):
"""Test the class call with a positional input update.
"""
input_array = n_byte_align(
(numpy.random.randn(*self.input_array.shape) + 1j * numpy.random.randn(*self.input_array.shape)),
simd_alignment,
)
output_array = self.fft(n_byte_align(input_array.copy(), simd_alignment)).copy()
self.update_arrays(input_array, self.output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(output_array == self.output_array))
def test_n_byte_align_integer_shape(self):
shape = 100
a = numpy.random.randn(shape)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
b = n_byte_align(a, n)
self.assertTrue(b.ctypes.data % n == 0)
def fftImage(image, use_pyfftw=True):
if use_pyfftw:
pyfftw.interfaces.cache.enable()
im = pyfftw.n_byte_align(image, 16)
return pyfftw.interfaces.numpy_fft.fftshift(
pyfftw.interfaces.numpy_fft.fft2(image))
return np.fft.fftshift(np.fft.fft2(image))
def test_avoid_copy(self):
'''Test the avoid_copy flag
'''
dtype_tuple = io_dtypes[functions[self.func]]
for dtype in dtype_tuple[0]:
for test_shape, s, kwargs in self.test_data:
_kwargs = kwargs.copy()
_kwargs['avoid_copy'] = True
s2 = copy.copy(s)
try:
for each_axis, length in enumerate(s):
s2[each_axis] += 2
except TypeError:
s2 += 2
input_array = dtype_tuple[1](test_shape, dtype)
self.assertRaisesRegex(ValueError,
'Cannot avoid copy.*transform shape.*',
getattr(builders, self.func),
input_array, s2, **_kwargs)
non_contiguous_shape = [
each_dim * 2 for each_dim in test_shape]
non_contiguous_slices = (
[slice(None, None, 2)] * len(test_shape))
misaligned_input_array = dtype_tuple[1](
non_contiguous_shape, dtype)[non_contiguous_slices]
self.assertRaisesRegex(ValueError,
'Cannot avoid copy.*not contiguous.*',
getattr(builders, self.func),
misaligned_input_array, s, **_kwargs)
# Offset by one from 16 byte aligned to guarantee it's not
# 16 byte aligned
_input_array = n_byte_align_empty(
numpy.prod(test_shape)*input_array.itemsize+1,
16, dtype='int8')
misaligned_input_array = _input_array[1:].view(
dtype=input_array.dtype).reshape(*test_shape)
self.assertRaisesRegex(ValueError,
'Cannot avoid copy.*not aligned.*',
getattr(builders, self.func),
misaligned_input_array, s, **_kwargs)
_input_array = n_byte_align(input_array.copy(), 16)
FFTW_object = getattr(builders, self.func)(
_input_array, s, **_kwargs)
# A catch all to make sure the internal array
# is not a copy
self.assertTrue(FFTW_object.get_input_array() is
_input_array)
def test_call_with_different_input_dtype(self):
'''Test the class call with an array with a different input dtype
'''
input_array = n_byte_align(numpy.complex64(
numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape)), 16)
output_array = self.fft(n_byte_align(input_array.copy(), 16)).copy()
_input_array = numpy.asarray(input_array,
dtype=self.input_array.dtype)
self.update_arrays(_input_array, self.output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(output_array == self.output_array))
def test_n_byte_align_consistent_data(self):
shape = (10, 10)
a = numpy.int16(numpy.random.randn(*shape) * 16000)
b = numpy.float64(numpy.random.randn(*shape))
c = numpy.int8(numpy.random.randn(*shape) * 255)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
d = n_byte_align(a, n)
self.assertTrue(numpy.array_equal(a, d))
d = n_byte_align(b, n)
self.assertTrue(numpy.array_equal(b, d))
d = n_byte_align(c, n)
self.assertTrue(numpy.array_equal(c, d))
def test_call_with_different_striding(self):
"""Test the input update with different strides to internal array.
"""
input_array_shape = self.input_array.shape + (2,)
internal_array_shape = self.internal_array.shape
internal_array = n_byte_align(
numpy.random.randn(*internal_array_shape) + 1j * numpy.random.randn(*internal_array_shape), simd_alignment
)
fft = utils._FFTWWrapper(
internal_array,
self.output_array,
input_array_slicer=self.input_array_slicer,
FFTW_array_slicer=self.FFTW_array_slicer,
)
test_output_array = fft().copy()
new_input_array = n_byte_align_empty(input_array_shape, simd_alignment, dtype=internal_array.dtype)
new_input_array[:] = 0
new_input_array[:, :, 0][self.input_array_slicer] = internal_array[self.FFTW_array_slicer]
new_output = fft(new_input_array[:, :, 0]).copy()
# Test the test!
self.assertTrue(new_input_array[:, :, 0].strides != internal_array.strides)
self.assertTrue(numpy.alltrue(test_output_array == new_output))
def test_n_byte_align_consistent_data(self):
shape = (10, 10)
a = numpy.int16(numpy.random.randn(*shape) * 16000)
b = numpy.float64(numpy.random.randn(*shape))
c = numpy.int8(numpy.random.randn(*shape) * 255)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
d = n_byte_align(a, n)
self.assertTrue(numpy.array_equal(a, d))
d = n_byte_align(b, n)
self.assertTrue(numpy.array_equal(b, d))
d = n_byte_align(c, n)
self.assertTrue(numpy.array_equal(c, d))
def test_n_byte_align_integer_shape(self):
shape = 100
a = numpy.random.randn(shape)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
b = n_byte_align(a, n)
self.assertTrue(b.ctypes.data % n == 0)
def main_loop(self):
if stream.is_active() and data_queue != None:
try:
in_data = data_queue.pop()
except IndexError:
return
# calculate fourier transform
input = struct.unpack(str(2 * chunk_size) + "B", in_data)
input = (numpy.array(input, dtype='b')[::2] + 128)
input = pyfftw.n_byte_align(input, n=16, dtype='float64')
# # reduce volume/amplitude
# volmax = input.max()
# for i in range(len(input)):
# # if i > 5 and i < len(input) - 5:
# input[i] /= 10
# input[0] = volmax
print("max=" + str(input.max()))
print("min=" + str(input.min()))
# print("med=" + str(numpy.median(input)))
print("")
fourier = numpy.absolute(fftw_obj(input), dtype='float64')
fourier = fourier[0 : len(fourier) // 2]
# zero out fourier edges
for i in range(5):
fourier[i] = 0
fourier = fourier[:len(fourier) - i - 1]
# scale values to log scale
scaled_fourier = fourier # todo
# scaled_fourier = input # test wave visualizer
# calculate percentages
max = scaled_fourier.max()
if max == 0:
max = 1
percentages = numpy.empty(len(scaled_fourier), dtype='float64')
for i in range(len(percentages)):
percentages[i] = scaled_fourier[i] / max
# modify depth of percentage transform to custom rectangle depth by averaging segments
depth = len(percentages)
step = math.ceil(float(depth) / float(numrects))
moddepth = numrects * step
if (moddepth > depth):
percentages = numpy.append(percentages, numpy.zeros(moddepth - depth, dtype='float64'))
i = 0
r = 0
while i < moddepth:
avg = 0.0
for j in range(step):
avg += percentages[i + j]
avg /= step
rheights[r] = avg
i += step
r += 1
# for r in range(numrects):
# rheights[r] = percentages[r]
# when taking averages, average one more value each time (increment step each time)
gui.update()
def acquire(self, n):
# Leave all array allocation beyond storage to subclasses
# Set up counters
self.n = n
self.i = -1 # First thing it does is get incremented
self.start_countup = 0
# Set up storage
self.storage = np.zeros((self.n, self.x), dtype=np.int32)
# Set up fft arrays
self.ft_in = fftw.n_byte_align(
np.zeros((n, self.x), dtype=np.complex128), 16)
self.ft_out = fftw.n_byte_align(
np.zeros((n, self.x), dtype=np.complex128), 16)
self.ift_in = fftw.n_byte_align(
np.zeros((n, self.x), dtype=np.complex128), 16)
self.ift_out = fftw.n_byte_align(
np.zeros((n, self.x), dtype=np.complex128), 16)
self.high_pass = np.zeros((n, self.x), dtype=np.complex128)
self.high_pass[0, :] = 1
self.fft = fftw.FFTW(self.ft_in,
self.ft_out,
axes=(0, ),
direction='FFTW_FORWARD',
flags=('FFTW_MEASURE', ),
threads=1,
planning_timelimit=None)
self.ifft = fftw.FFTW(self.ift_in,
self.ift_out,
axes=(0, ),
direction='FFTW_BACKWARD',
flags=('FFTW_MEASURE', ),
threads=1,
planning_timelimit=None)
# self.fft = fftw.FFTW(self.ft_in, self.ft_out, axes=(0), direction='FFTW_FORWARD')
# self.ifft = fftw.FFTW(self.ift_in, self.ift_out, axes=(0), direction='FFTW_BACKWARD')
# Get ready to calculate fps
self.plot_times = []
# Go
self.new_image.connect(self.send_plot)
self.running = True
def _maybe_align(a):
global planned
if not planned:
print_update("Note: FFTW is configuring itself. This will take " +
"several seconds, but subsequent calls will run " +
"*much* faster.")
planned = True
result = pyfftw.n_byte_align(a, a.dtype.alignment)
return result
def _maybe_align(a):
global planned
if not planned:
warnings.warn("FFTW is configuring itself. This will take " +
"several sections, but subsequent calls will run " +
"*much* faster.", UserWarning)
sys.stderr.flush() # Show that warning immediately.
planned = True
return pyfftw.n_byte_align(a, a.dtype.alignment)
def _maybe_align(a):
global planned
if not planned:
print_update("Note: FFTW is configuring itself. This will take " +
"several seconds, but subsequent calls will run " +
"*much* faster.")
planned = True
result = pyfftw.n_byte_align(a, a.dtype.alignment)
return result
def test_n_byte_align_set_dtype(self):
shape = (10, 10)
a = numpy.int16(numpy.random.randn(*shape) * 16000)
b = numpy.float64(numpy.random.randn(*shape))
c = numpy.int8(numpy.random.randn(*shape) * 255)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
d = n_byte_align(a, n, dtype='float32')
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.dtype == 'float32')
d = n_byte_align(b, n, dtype='float32')
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.dtype == 'float32')
d = n_byte_align(c, n, dtype='float64')
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.dtype == 'float64')
def __init__(self, axis_len, buf_len):
self.ft_in = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.ft_out = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.ift_in = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.ift_out = fftw.n_byte_align(np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.fft = fftw.FFTW(self.ft_in, self.ft_out,
axes=(1,),
direction='FFTW_FORWARD',
flags=('FFTW_MEASURE',),
threads=1,
planning_timelimit=None)
self.ifft = fftw.FFTW(self.ift_in, self.ift_out,
axes=(1,),
direction='FFTW_BACKWARD',
flags=('FFTW_MEASURE',),
threads=1,
planning_timelimit=None)
def test_n_byte_align_different_dtypes(self):
shape = (10, 10)
a = numpy.int16(numpy.random.randn(*shape) * 16000)
b = numpy.float64(numpy.random.randn(*shape))
c = numpy.int8(numpy.random.randn(*shape) * 255)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
d = n_byte_align(a, n)
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.__class__ == a.__class__)
d = n_byte_align(b, n)
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.__class__ == b.__class__)
d = n_byte_align(c, n)
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.__class__ == c.__class__)
def test_call_with_copy_with_missized_array_error(self):
"""Force an input copy with a missized array.
"""
shape = list(self.input_array.shape + (2,))
shape[0] += 1
input_array = n_byte_align(numpy.random.randn(*shape) + 1j * numpy.random.randn(*shape), simd_alignment)
self.assertRaisesRegex(ValueError, "Invalid input shape", self.fft, **{"input_array": input_array[:, :, 0]})
def test_n_byte_align_set_dtype(self):
shape = (10, 10)
a = numpy.int16(numpy.random.randn(*shape) * 16000)
b = numpy.float64(numpy.random.randn(*shape))
c = numpy.int8(numpy.random.randn(*shape) * 255)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
d = n_byte_align(a, n, dtype="float32")
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.dtype == "float32")
d = n_byte_align(b, n, dtype="float32")
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.dtype == "float32")
d = n_byte_align(c, n, dtype="float64")
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.dtype == "float64")
def fftn(a):
global planned
if not planned:
logger.info("Note: FFTW is configuring itself. This will take " +
"several seconds, but subsequent calls will run " +
"*much* faster.")
planned = True
a = pyfftw.n_byte_align(a, a.dtype.alignment)
return pyfftw.interfaces.numpy_fft.fftn(a).astype(np.complex128)
def test_call_with_positional_updates(self):
'''Test the class call with a positional array updates.
'''
input_array = n_byte_align((numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape)), 16)
output_array = n_byte_align((numpy.random.randn(*self.output_array.shape)
+ 1j*numpy.random.randn(*self.output_array.shape)), 16)
returned_output_array = self.fft(
n_byte_align(input_array.copy(), 16),
n_byte_align(output_array.copy(), 16)).copy()
self.update_arrays(input_array, output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(returned_output_array == output_array))
def test_n_byte_align_different_dtypes(self):
shape = (10, 10)
a = numpy.int16(numpy.random.randn(*shape) * 16000)
b = numpy.float64(numpy.random.randn(*shape))
c = numpy.int8(numpy.random.randn(*shape) * 255)
# Test a few alignments
for n in [3, 7, 9, 16, 24, 23, 63, 64]:
d = n_byte_align(a, n)
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.__class__ == a.__class__)
d = n_byte_align(b, n)
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.__class__ == b.__class__)
d = n_byte_align(c, n)
self.assertTrue(d.ctypes.data % n == 0)
self.assertTrue(d.__class__ == c.__class__)
def fftn(a):
global planned
if not planned:
print_update("Note: FFTW is configuring itself. This will take " +
"several seconds, but subsequent calls will run " +
"*much* faster.")
planned = True
a = pyfftw.n_byte_align(a, a.dtype.alignment)
return pyfftw.interfaces.numpy_fft.fftn(a).astype(np.complex128)
def pf2psf(self, PF, zs, intensity=True, verbose=False, use_pyfftw=False):
"""
Computes the point spread function for a given pupil function.
Parameters
----------
PF: array
The complex pupil function.
zs: number or iteratable
The axial position or a list of axial positions which should be computed. Focus is at z=0.
intensity: bool
Specifies if the intensity or the complex field should be returned.
Returns
-------
PSF: array or memmap
The complex PSF. If the memory is to small, a memmap will be
returned instead of an array.
"""
ny, nx = self.ny, self.nx
if np.isscalar(zs):
zs = [zs]
print("The z positions:", zs)
nz = len(zs)
kz = self.kz
# The normalization for ifft2:
N = np.sqrt(self.nx * self.ny)
# Preallocating memory for PSF:
if intensity:
PSF = np.zeros((nz, nx, nx))
else:
PSF = np.zeros((nz, nx, nx)) + 1j * np.zeros((nz, nx, nx))
for i in range(nz):
if verbose:
print('Calculating PSF slice for z={0}um.'.format(zs[i]))
U = np.zeros((nx, nx)) + 1j * np.zeros((nx, nx))
for j in range(0, self.kzs.shape[0]):
if use_pyfftw:
aligned = pyfftw.n_byte_align(
_fftpack.ifftshift(
np.exp(2 * np.pi * 1j * self.kzs[j] * zs[i]) * PF,
16))
U = U + N * pyfftw.interfaces.numpy_fft.ifft2(aligned)
else:
U = U + N * _fftpack.ifft2(
np.exp(2 * np.pi * 1j * self.kzs[j] * zs[i]) * PF)
U = U / self.numWavelengths
_slice_ = _fftpack.ifftshift(U) # move the x0 to the center
if intensity:
_slice_ = np.abs(_slice_)**2
PSF[i] = _slice_
if nz == 1:
PSF = PSF[0]
return PSF
def _maybe_align(a):
global planned
if not planned:
warnings.warn(
"FFTW is configuring itself. This will take " +
"several sections, but subsequent calls will run " +
"*much* faster.", UserWarning)
sys.stderr.flush() # Show that warning immediately.
planned = True
return pyfftw.n_byte_align(a, a.dtype.alignment)
def main_loop():
if stream.is_active() and data_queue != None:
# try:
# in_data = data_queue.pop()
# except IndexError:
# return
if (current_in_data == None):
return
# calculate fourier transform
input = numpy.fromstring(current_in_data, dtype='int16')
amplitude = pyfftw.n_byte_align(numpy.asarray(input,
dtype='float64'),
n=16,
dtype='float64')
fourier = numpy.absolute(fftw_obj(amplitude), dtype='float64')
# zero out fourier edges
# for i in range(10):
# fourier[i] = 0
# fourier[len(fourier) - (i + 1)] = 0
fourier = fourier[5:len(fourier) - 5]
if split_fourier == "left":
fourier = fourier[:len(fourier) // 2]
elif split_fourier == "right":
fourier = fourier[len(fourier) // 2:]
# calculate percentages
max = 1.0
for x in fourier:
if x > max:
max = x
percentages = numpy.empty(len(fourier), dtype='float64')
for i in range(len(percentages)):
percentages[i] = fourier[i] / max
# # zero out percentage edges
# for i in range(10):
# percentages[i] = 0
# percentages[len(percentages) - (i + 1)] = 0
# modify depth of percentage transform to custom rectangle depth by averaging segments
depth = len(percentages)
step = math.ceil(float(depth) / float(numrects))
moddepth = numrects * step
if (moddepth > depth):
percentages = numpy.append(
percentages, numpy.zeros(moddepth - depth,
dtype='float64'))
i = 0
r = 0
while i < moddepth:
avg = 0.0
for j in range(step):
avg += percentages[i + j]
avg /= step
rheights[r] = avg * h_mult
i += step
r += 1
serial_update(rheights)
def test_1d():
for use_pyfftw in (False, True):
buf = allocate(10, dtype=np.float32, use_pyfftw=use_pyfftw)
assert buf.shape == (10,)
assert buf.dtype == np.float32
if use_pyfftw:
try:
import pyfftw
assert pyfftw.n_byte_align(buf, 16) is buf
except ImportError:
pass
def test_3d():
for use_pyfftw in (False, True):
buf = allocate((4, 6, 8), dtype=np.complex64, use_pyfftw=use_pyfftw)
assert buf.shape == (4, 6, 8)
assert buf.dtype == np.complex64
if use_pyfftw:
try:
import pyfftw
assert pyfftw.n_byte_align(buf, 16) is buf
except ImportError:
pass
def __init__(self, A):
"""Set up arrays and FFTs for convolution.
Parameters
----------
A : ndarray (3-d)
PSF, assumed to be centered in the array at the
"reference wavelength."
"""
self.nw, self.ny, self.nx = A.shape
# The attribute `fftconv` stores the Fourier-space array
# necessary to convolve another array by the PSF. This is done
# by mulitiplying the input array by `fftconv` in fourier
# space.
#
# We shift the PSF so that instead of being exactly centered
# in the array, it is exactly centered on the lower left
# pixel. (For convolution in Fourier space, the (0, 0)
# element of the kernel is effectively the "center.")
# Note that this shifting is different than simply
# creating the PSF centered at the lower left pixel to begin
# with, due to wrap-around.
#
#`ifft2(fftconv).real` would be the PSF in
# real space, shifted to be centered on the lower-left pixel.
shift = -(self.ny - 1) / 2., -(self.nx - 1) / 2.
fshift = fft_shift_phasor_2d((self.ny, self.nx), shift)
fftconv = fft2(A) * fshift
# align on SIMD boundary.
self.fftconv = pyfftw.n_byte_align(fftconv,
pyfftw.simd_alignment,
dtype=np.complex128)
# set up input and output arrays for FFTs.
self.fftin = pyfftw.n_byte_align_empty(A.shape,
pyfftw.simd_alignment,
dtype=np.complex128)
self.fftout = pyfftw.n_byte_align_empty(A.shape,
pyfftw.simd_alignment,
dtype=np.complex128)
# Set up forward and backward FFTs.
self.fft = pyfftw.FFTW(self.fftin, self.fftout, axes=(1, 2), threads=1)
self.ifft = pyfftw.FFTW(self.fftout,
self.fftin,
axes=(1, 2),
threads=1,
direction='FFTW_BACKWARD')
self.fftnorm = 1. / (self.ny * self.nx)
def align_array(arr):
"""Return memory aligned copy of arr. This may be speed up pyfftw calls.
Parameters
----------
arr : array
Returns
-------
arr_aligned : array
"""
return pyfftw.n_byte_align(arr, pyfftw_simd_alignment)
def filtering(proj,geo,angles,parker):
if parker:
proj = parkerweight(proj.transpose(0,2,1),geo,angles,parker).transpose(0,2,1)
# proj=parkerweight(proj,geo,angles,parker)
filt_len = max(64,2 ** nextpow2(2 * geo.nDetector[0]))
ramp_kernel = ramp_flat(filt_len)
d = 1
filt = filter(geo.filter,ramp_kernel[0],filt_len,d)
filt = np.kron(np.ones((geo.nDetector[1],1)),filt).transpose()
for i in range(len(angles)):
# Zero-padding:
fproj = np.zeros((filt_len,geo.nDetector[1]),dtype=np.float32)
fproj[int(filt_len / 2 - geo.nDetector[0] / 2):int(filt_len / 2 + geo.nDetector[0] / 2),:] = proj[:,:,i]
# Fourier-transform:
n_byte_align(fproj, simd_alignment)
fproj = rfft(fproj, axis=0, threads=2)
# Filter:
fproj = fproj * filt[0:fproj.shape[0],:]
# Inverse-Fourier:
n_byte_align(fproj, simd_alignment)
fproj = np.real(irfft(fproj, axis=0, threads=2))
end = len(fproj)
# Crop:
fproj = fproj[int(end / 2 - geo.nDetector[0] / 2):int(end / 2 + geo.nDetector[0] / 2),:] / 2 / geo.dDetector[0] * \
(2 * np.pi / len(angles)) / 2 * (geo.DSD / geo.DSO)
proj[:,:,i] = fproj.astype(np.float32)
return proj
def __init__(self, axis_len, buf_len):
self.ft_in = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.ft_out = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.ift_in = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.ift_out = fftw.n_byte_align(
np.zeros((axis_len, buf_len), dtype=np.complex128), 16)
self.fft = fftw.FFTW(self.ft_in,
self.ft_out,
axes=(1, ),
direction='FFTW_FORWARD',
flags=('FFTW_MEASURE', ),
threads=1,
planning_timelimit=None)
self.ifft = fftw.FFTW(self.ift_in,
self.ift_out,
axes=(1, ),
direction='FFTW_BACKWARD',
flags=('FFTW_MEASURE', ),
threads=1,
planning_timelimit=None)
def main_loop(self):
if stream.is_active() and data_queue != None:
try:
in_data = data_queue.pop()
except IndexError:
return
# calculate fourier transform
input = struct.unpack(str(2 * chunk_size) + "B", in_data)
input = numpy.array(input, dtype='b')[::2] + 128
input = pyfftw.n_byte_align(input, n=16, dtype='float64')
# fourier = numpy.absolute(fftw_obj(input), dtype='float64')
fourier = input # no transform actually occurs
# zero out fourier edges
# for i in range(10):
# fourier[i] = 0
# fourier[len(fourier) - (i + 1)] = 0
# fourier = fourier[5:len(fourier) - 5]
# calculate percentages
max = 1.0
for x in fourier:
if x > max:
max = x
percentages = numpy.empty(len(fourier), dtype='float64')
for i in range(len(percentages)):
percentages[i] = fourier[i] / max
# modify depth of percentage transform to custom rectangle depth by averaging segments
depth = len(percentages)
step = math.ceil(float(depth) / float(numrects))
moddepth = numrects * step
if (moddepth > depth):
percentages = numpy.append(percentages, numpy.zeros(moddepth - depth, dtype='float64'))
i = 0
r = 0
while i < moddepth:
avg = 0.0
for j in range(step):
avg += percentages[i + j]
avg /= step
rheights[r] = avg
i += step
r += 1
gui.update()
def MASS(y, n, cum_sumx, cum_sumx2, sumx2, sumx, meanx, sigmax2, sigmax, X,
value):
'''
Implementation of the MASS algorithm
input: a time_series x, and query y
output:
'''
#normlaize line below
y = (y - np.mean(y)) / np.std(y)
m = len(y)
y = y[::-1]
sumy = np.sum(y)
sumy2 = np.sum(np.power(y, 2))
y = np.append(y, ((n) - (m)) * [0])
Y = fft(y)
Z = X * Y
Z = pyfftw.n_byte_align(Z, None)
z = pyfftw.interfaces.scipy_fftpack.ifft(Z)
dist = value - 2 * (z[m:n] - sumy * meanx) / sigmax + sumy2
return np.abs(np.lib.scimath.sqrt(dist))
def psf2pf(self,
PSF,
dz,
mu,
A,
nIterations=10,
z_offset=0,
use_pyfftw=True,
resetAmp=True,
symmeterize=False):
'''
Retrieves the complex pupil function from an intensity-only
PSF stack by relative entropy minimization. The algorithm is
based on Kner et al., 2010, doi:10.1117/12.840943, which in turn
is based on Deming, 2007, J Opt Soc Am A, Vol 24, No 11, p.3666.
Parameters
---------
PSF: 3D numpy.array
An intensity PSF stack. PSF.shape has to be
(nz, Simulation.nx, Simulation.nx), where nz is the arbitrary
number of z slices.
dz: float
The distance between two PSF slices.
mu: float
The noise level of the PSF.
A: 2D numpy.array
The initial guess for the complex pupil function with shape
(Simulation.nx, Simulation.nx).
'''
# z spacing:
dz = float(dz)
# Number of z slices:
nz = PSF.shape[0]
# Noise level:
mu = float(mu)
kz = self.kz
k = self.k
k_max = self.k_max
# Z position of slices:
upper = 0.5 * (nz - 1) * dz
zs = _np.linspace(-upper, upper, nz) - z_offset
# Normalization for fft2:
N = _np.sqrt(self.nx * self.ny)
if use_pyfftw:
pyfftw.interfaces.cache.enable()
Ue = _np.ones_like(PSF).astype(_np.complex128)
U = _np.ones_like(PSF).astype(_np.complex128)
Uconj = _np.ones_like(PSF).astype(_np.complex128)
Ic = _np.ones_like(PSF).astype(_np.complex128)
expr1 = "Ue = (PSF/Ic)*U"
expr2 = "Ic = mu + (U * Uconj)"
for ii in range(nIterations):
print('Iteration', ii + 1)
# Calculate PSF field from given PF:
U = self.pf2psf(A, zs, intensity=False)
# Calculated PSF intensity with noise:
Uconj = _np.conj(U)
#weave.blitz(expr2)
Ic = mu + (U * Uconj)
minFunc = _np.mean(PSF * _np.log(PSF / Ic))
print('Relative entropy per pixel:', minFunc)
redChiSq = mean((PSF - Ic)**2)
print('Reduced Chi square:', redChiSq)
# Comparing measured with calculated PSF by entropy minimization:
Ue = (PSF / Ic) * U
#weave.blitz(expr1)
# New PF guess:
A = _np.zeros_like(Ue) + 1j * _np.zeros_like(Ue)
for i in range(len(zs)):
#Ue[i] = _fftpack.fftshift(Ue[i])
if use_pyfftw:
Ue_aligned = pyfftw.n_byte_align(_fftpack.fftshift(Ue[i]),
16)
fted_ue = pyfftw.interfaces.numpy_fft.fft2(Ue_aligned)
A[i] = fted_ue / _np.exp(2 * _np.pi * 1j * kz * zs[i]) / N
else:
fted_ue = _fftpack.fft2(_fftpack.fftshift(Ue[i]))
A[i] = fted_ue / _np.exp(2 * _np.pi * 1j * kz * zs[i]) / N
for j in range(0, self.kzs.shape[0]):
A[i] = A[i] + fted_ue / _np.exp(
2 * _np.pi * 1j * self.kzs[j] * zs[i]) / N
A[i] = A[i] / (1 + self.kzs.shape[0])
A = _np.mean(A, axis=0)
#mean(abs(A))*_np.exp(1j*_np.angle(A))
# NA restriction:
A[k > k_max] = 0
if resetAmp:
amp = ndimage.gaussian_filter(np.abs(A), 15)
A = amp * np.nan_to_num(A / np.abs(A))
if symmeterize:
if ii > (nIterations / 2):
A = 0.5 * (A + _np.flipud(A))
#counts = sum(abs(A))/self.pupil_npxl
#A = counts*_np.exp(1j*angle(A))
#A[k>k_max] = 0
return A
def ifftn(a):
a = pyfftw.n_byte_align(a, a.dtype.alignment)
return pyfftw.interfaces.numpy_fft.ifftn(a)
i += step
r += 1
serial_update(rheights)
# set up and open stream
stream = p.open(format=p.get_format_from_width(2),
channels=int(
p.get_device_info_by_host_api_device_index(
0, input_device_id).get('maxInputChannels')),
rate=int(
p.get_device_info_by_host_api_device_index(
0, input_device_id).get('defaultSampleRate')),
input=True,
output=True,
input_device_index=input_device_id,
output_device_index=output_device_id,
stream_callback=callback)
stream.start_stream()
# fftw setup
amplitude = pyfftw.n_byte_align(numpy.zeros(2048, dtype='float64'),
n=16,
dtype='float64')
fftw_obj = pyfftw.builders.fft(amplitude)
# run main loop
refresh_rate = 0.1
while True:
main_loop()
time.sleep(refresh_rate)
r += 1
# for r in range(numrects):
# rheights[r] = percentages[r]
# when taking averages, average one more value each time (increment step each time)
gui.update()
# set up and open stream
stream_channels = int(pa.get_device_info_by_host_api_device_index(0, input_device_id).get('maxInputChannels'))
stream_rate = int(pa.get_device_info_by_host_api_device_index(0, input_device_id).get('defaultSampleRate'))
stream = pa.open(format=pa.get_format_from_width(2),
channels=stream_channels,
rate=stream_rate,
input=True,
output=True,
input_device_index=input_device_id,
output_device_index=output_device_id,
stream_callback=callback,
frames_per_buffer=int(chunk_size / stream_channels))
stream.start_stream()
# fftw setup
amplitude = pyfftw.n_byte_align(numpy.zeros(chunk_size, dtype='float64'), n=16, dtype='float64')
fftw_obj = pyfftw.builders.fft(amplitude)
# display window
Window.size = (w, h)
gui = AV_GUI()
gui.run()
def filter_fourier_bandpass(input_data,
passband,
stopband,
taper=None,
debug=None):
"""
Note: Is MUCH (2.2x faster) more efficient to use real ffts, (implemented April)
Use a Fourier space taper/tophat or pseudo gaussian filter to perform
narrowband filtering (much narrower than butterworth).
Problem is that bursts may generate ringing.
This should be better with taper=2, but it is not clear
debug = None (follow pyfusion debug), 2 plot responses
See the __main__ code below for nice test facilities
twid is the width of the transition from stop to pass (not impl.?)
>>> tb = timebase(np.linspace(0,20,512))
>>> w = 2*np.pi* 1 # 1 Hertz
>>> dat = dummysig(tb,np.sin(w*tb)*(tb<np.max(tb)/3))
>>> fop = filter_fourier_bandpass(dat,[0.9,1.1],[0.8,1.2],debug=1).signal[0]
Testing can be done on the dummy data set generated after running
filters.py
e.g. (with pyfusion,DEBUG=2
make_mask(512, [0.8,.93], [0.9,.98],dat,2)
# medium sharp shoulder
fopmed = filter_fourier_bandpass(dat,[9.5,10.5],[9,11],debug=1).signal[0]
# very sharp shoulders
fopsharp = filter_fourier_bandpass(dat,[9.1,10.9],[9,11],debug=1)
"""
if debug is None: debug = pyfusion.DBG()
# normalising makes it easier to think about - also for But'w'h
if (passband[0] < stopband[0]) or (passband[1] > stopband[1]):
raise ValueError('passband {pb} outside stopband {sb}'.format(
pb=passband, sb=stopband))
norm_passband = input_data.timebase.normalise_freq(np.array(passband))
norm_stopband = input_data.timebase.normalise_freq(np.array(stopband))
NS = len(input_data.timebase)
NA = next_nice_number(NS)
input_data.history += str(" {fftt} : nice number: {NA} cf {NS}\n".format(
fftt=pyfusion.fft_type, NA=NA, NS=NS))
# take a little more to speed up FFT
mask = make_mask(NA, norm_passband, norm_stopband, input_data, taper)
output_data = deepcopy(input_data) # was output_data = input_data
if (pyfusion.fft_type == 'fftw3'):
# should migrate elsewhere, but the import is only 6us
# the setup time seems about 150-250us even if size is in wisdom
# for 384, numpy is 20us, fftw3 is 4us, so fftw3 slower for less than
# 10 channels (unless we cache the plan)
#time# st=seconds()
import pyfftw
#time# im=seconds()
tdtype = np.float32
fdtype = np.complex64
# this could be useful to cache.
simd_align = pyfftw.simd_alignment # 16 at the moment.
tdom = pyfftw.n_byte_align(np.zeros(NA, dtype=tdtype), simd_align)
FT = pyfftw.n_byte_align_empty(NA / 2 + 1, simd_align, fdtype)
ids = [[id(tdom), id(FT)]] # check to see if it moves out of alignment
#time# alloc_t = seconds()
fwd = pyfftw.FFTW(tdom,
FT,
direction='FFTW_FORWARD',
**pyfusion.fftw3_args)
rev = pyfftw.FFTW(FT,
tdom,
direction='FFTW_BACKWARD',
**pyfusion.fftw3_args)
#time# pl=seconds()
#time# print("import {im:.2g}, alloc {al:.2g}, planboth {pl:.2g}"
#time# .format(im=im-st, al=alloc_t-im, pl=pl-alloc_t))
else:
tdtype = np.float32
tdom = np.zeros(NA, dtype=tdtype)
# example of tuning
#pyfusion.fftw3_args= {'planning_timelimit': 50.0, 'threads':1, 'flags':['FFTW_MEASURE']}
singl = not isinstance(output_data.channels, (list, tuple, np.ndarray))
if singl:
## this is a fudge - need to set the value part to a list.
output_data.signal = [
output_data.signal
] # not right - need to use same fudge as acq/bas
# output_data.signal.n_channels = fudgey_n_channels #bdb fudge for single channel diag - doesn't work, because we fudged output_data.signal to be a list (or because n_channels is a func)
for i, s in enumerate(output_data.signal):
#if len(output_data.signal) == 1: print('bug for a single signal')
#time run -i pyfusion/examples/plot_svd.py "dev_name='LHD'" start_time=.497 "normalise='r'" shot_number=90091 numpts=512 diag_name=MP2010HMPno612 "filter=dict(centre=8e3,bw=5e3,taper=2)" plot_mag=1 plot_phase=1 separate=1 closed=0 time_range=[0.0000,4.]
# 4.5 cf 15.8diag_name=MP2010HMPno612, time_range=[0.0000,2.80000]
# 0, 4.194304 2**21 samples, 21.8 cf 6.8 1thr
# (0,2)secs 90091 =2000000 samples 17 np, 5.73 2thread, nosimd, 6.1 1thread (mem bw?) 3.2 sec no filt
# note - the above are on an intermeittently loaded E4300 2 processor, below on 4 core 5/760
# 0, 4.194304 2**21 samples, 10.9 cf 3.16 thr2 3.47 1thr and 2.0 secs no filter
# for 17 fft/ifft takes about 1.16 sec 2 threads - should be (27.5ms+28.6)*17 = 952ms (14.2 2thr) OK
# duplicate the fft execute lines 4.3(3.47) 2thr 3.7(3.16) extra 810ms (expect 14.2ms * 2 * 17) =482
# the difference between 2 and 1thr should be 14*2*17 ms 500ms.
# orignall - 90ms/channel extra in reverse trasnform - maybe the 50 sec limit stopped optimization
# next _nice: 5.74 for 10 sec lenny
# E4300: 1thr 9.3 (39np) for 10 sec 90091; 5.5 for 4 sec (19.6 np)
if (pyfusion.fft_type == 'fftw3'
): # fftw3 nosim, no thread 2.8s cf 10s
#time# sst = seconds()
tdom[
0:
NS] = s # indexed to make sure tdom is in the right part of memory
if NS != NA: tdom[NS:] = 0.
fwd.execute()
FT[:] = FT * mask[0:NA / 2 + 1] # 12ms
rev.execute()
output_data.signal[i] = tdom[0:NS] / NA # doco says NA
ids.append([id(tdom), id(FT)])
#time# print("{dt:.1f}us".format(dt=(seconds()-sst)/1e-6)),
else: # default to numpy
tdom[0:NS] = s
FT = np.fft.fft(tdom)
IFT = np.fft.ifft(mask * FT)
if np.max(np.abs(IFT.imag)) > 1e-6 * np.max(np.abs(IFT.real)):
pyfusion.logger.warning("inverse fft imag part > 1e-6")
output_data.signal[i] = IFT.real[0:NS]
if debug > 2: print('ids of fftw3 input and output: {t}'.format(t=ids))
if debug > 1:
fig = plt.figure()
#fplot = host_subplot(111)
fplot = fig.add_subplot(111)
tplot = fplot.twinx()
tplot.plot(input_data.signal[0], 'c', label='input')
# for a while I needed a factor of 3 here too for fftw - why ????
#plt.plot(output_data.signal[0]/(3*NA),'m',label='output/{N}'.format(N=3*NA))
tplot.plot(output_data.signal[0], 'm', label='output')
tplot.set_ylim(-2.4, 1.1)
fplot.plot(mask, 'r.-', label='mask')
fplot.plot(np.abs(FT) / len(mask), label='FT')
#fplot.set_ylim(-.2,3.8)
#fplot.set_yscale('log', subsy=[2,5])
#fplot.set_ylim(1e-7,1e5)
fplot.set_yscale('symlog', linthreshy=1e-6)
fplot.set_ylim(0, 1e8)
fig.suptitle('Passband {pbl}...{pbh}'.format(pbl=passband[0],
pbh=passband[1]))
fplot.legend(loc=4) # bottom right
tplot.legend(loc=0)
plt.show()
debug_(debug, 3, key='filter_fourier')
if np.max(mask) == 0: raise ValueError('Filter blocks all signals')
from pyfusion.data.timeseries import Signal
if singl:
output_data.signal = Signal(output_data.signal[0])
# output_data.signal.n_channels = fudgey_n_channels would work if signal was propoer not just an array of data
return output_data
def homomorphic(im, args):
"""Process the input image with an homomorphic filtering.
Parameters
----------
im : array_like
Image data as numpy array.
d0 : float
Cutoff in the range [0.01, 0.99] of the high pass Gaussian filter.
Higher values means more high pass effect. [Suggested for default: 0.80].
alpha : float
Offset to preserve the zero frequency where. Higher values means less
high pass effect. [Suggested for default: 0.2]
(Parameters d0 and alpha have to passed as a string separated by ;)
Example (using tiffile.py)
--------------------------
>>> im = imread('im_orig.tif')
>>> im = homomorphic(im, '0.5;0.2')
>>> imsave('im_flt.tif', im)
References
----------
"""
# Get args:
param1, param2 = args.split(";")
d0 = (1.0 - float(param1)) # Simpler for user
alpha = float(param2)
# Internal parameters for Gaussian low-pass filtering:
d0 = d0 * (im.shape[1] / 2.0)
# Take the log:
im = nplog(1 + im)
# Compute FFT:
n_byte_align(im, simd_alignment)
im = rfft2(im, threads=2)
# Prepare the frequency coordinates:
u = arange(0, im.shape[0], dtype=float32)
v = arange(0, im.shape[1], dtype=float32)
# Compute the indices for meshgrid:
u[(u > im.shape[0] / 2.0)] = u[(u > im.shape[0] / 2.0)] - im.shape[0]
v[(v > im.shape[1] / 2.0)] = v[(v > im.shape[1] / 2.0)] - im.shape[1]
# Compute the meshgrid arrays:
V, U = meshgrid(v, u)
# Compute the distances D(U, V):
D = sqrt(U**2 + V**2)
# Prepare Guassian filter:
H = npexp(-(D**2) / (2 * (d0**2)))
H = (1 - H) + alpha
# Do the filtering:
im = H * im
# Compute inverse FFT of the filtered data:
n_byte_align(im, simd_alignment)
#im = real(irfft2(im, threads=2))
im = irfft2(im, threads=2)
# Take the exp:
im = npexp(im) - 1
# Return image according to input type:
return im.astype(float32)
