def test_call_with_normalisation_precision(self):
'''The normalisation should use a double precision scaling.
'''
# Should be the case for double inputs...
_input_array = empty_aligned((256, 512), dtype='complex128', n=16)
self.fft()
ifft = FFTW(self.output_array, _input_array,
direction='FFTW_BACKWARD')
ref_output = ifft(normalise_idft=False).copy()/numpy.float64(ifft.N)
test_output = ifft(normalise_idft=True).copy()
self.assertTrue(numpy.alltrue(ref_output == test_output))
# ... and single inputs.
_input_array = empty_aligned((256, 512), dtype='complex64', n=16)
ifft = FFTW(numpy.array(self.output_array, _input_array.dtype),
_input_array,
direction='FFTW_BACKWARD')
ref_output = ifft(normalise_idft=False).copy()/numpy.float64(ifft.N)
test_output = ifft(normalise_idft=True).copy()
self.assertTrue(numpy.alltrue(ref_output == test_output))
def test_incorrect_byte_alignment_fails(self):
in_shape = self.input_shapes["2d"]
out_shape = self.output_shapes["2d"]
input_dtype_alignment = self.get_input_dtype_alignment()
axes = (-1,)
a, b = self.create_test_arrays(in_shape, out_shape)
a = byte_align(a, n=16)
b = byte_align(b, n=16)
fft, ifft = self.run_validate_fft(a, b, axes, force_unaligned_data=True)
a, b = self.create_test_arrays(in_shape, out_shape)
# Offset from 16 byte aligned to guarantee it's not
# 16 byte aligned
a__ = empty_aligned(numpy.prod(in_shape) * a.itemsize + 1, dtype="int8", n=16)
a_ = a__[1:].view(dtype=self.input_dtype).reshape(*in_shape)
a_[:] = a
b__ = empty_aligned(numpy.prod(out_shape) * b.itemsize + 1, dtype="int8", n=16)
b_ = b__[1:].view(dtype=self.output_dtype).reshape(*out_shape)
b_[:] = b
self.assertRaisesRegex(ValueError, "Invalid output alignment", FFTW, *(a, b_))
self.assertRaisesRegex(ValueError, "Invalid input alignment", FFTW, *(a_, b))
self.assertRaisesRegex(ValueError, "Invalid input alignment", FFTW, *(a_, b_))
def test_flags(self):
'''Test to see if the flags are correct
'''
fft = FFTW(self.input_array, self.output_array)
self.assertEqual(fft.flags, ('FFTW_MEASURE',))
fft = FFTW(self.input_array, self.output_array,
flags=('FFTW_DESTROY_INPUT', 'FFTW_UNALIGNED'))
self.assertEqual(fft.flags, ('FFTW_DESTROY_INPUT', 'FFTW_UNALIGNED'))
# Test an implicit flag
_input_array = empty_aligned(256, dtype='complex64', n=16)
_output_array = empty_aligned(256, dtype='complex64', n=16)
# These are guaranteed to be misaligned (due to dtype size == 8)
input_array = _input_array[:-1]
output_array = _output_array[:-1]
u_input_array = _input_array[1:]
u_output_array = _output_array[1:]
fft = FFTW(input_array, u_output_array)
self.assertEqual(fft.flags, ('FFTW_MEASURE', 'FFTW_UNALIGNED'))
fft = FFTW(u_input_array, output_array)
self.assertEqual(fft.flags, ('FFTW_MEASURE', 'FFTW_UNALIGNED'))
fft = FFTW(u_input_array, u_output_array)
self.assertEqual(fft.flags, ('FFTW_MEASURE', 'FFTW_UNALIGNED'))
def generate_wisdom(self):
for each_dtype in (numpy.complex128, numpy.complex64,
numpy.clongdouble):
a = empty_aligned((1,1024), each_dtype, n=16)
b = empty_aligned(a.shape, dtype=a.dtype, n=16)
fft = FFTW(a,b)
def test_update_data_with_unaligned_original(self):
in_shape = self.input_shapes['2d']
out_shape = self.output_shapes['2d']
input_dtype_alignment = self.get_input_dtype_alignment()
axes=(-1,)
a, b = self.create_test_arrays(in_shape, out_shape)
# Offset from 16 byte aligned to guarantee it's not
# 16 byte aligned
a__ = empty_aligned(
numpy.prod(in_shape)*a.itemsize + input_dtype_alignment,
dtype='int8', n=16)
a_ = a__[input_dtype_alignment:].view(dtype=self.input_dtype).reshape(*in_shape)
a_[:] = a
b__ = empty_aligned(
numpy.prod(out_shape)*b.itemsize + input_dtype_alignment,
dtype='int8', n=16)
b_ = b__[input_dtype_alignment:].view(dtype=self.output_dtype).reshape(*out_shape)
b_[:] = b
fft, ifft = self.run_validate_fft(a_, b_, axes,
force_unaligned_data=True)
self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft)
self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft)
self.run_validate_fft(a_, b_, axes, fft=fft, ifft=ifft)
def __init__(self,mask, tccList):
self.mask = mask
self.tcc = tccList
self.order = tccList.order
self.kernelList = tccList.kernelList
self.coefList = tccList.coefList
self.focusList = tccList.focusList
self.focusCoef = tccList.focusCoef
self.doseList = [1.0]
self.doseCoef = [1.0]
self.AIList = []
self.RIList = []
self.resist_a = 80
self.resist_tRef = 0.5
self.norm = self.mask.y_gridnum*self.mask.x_gridnum
self.x1 = np.floor(self.mask.x_gridnum/2) - self.tcc.s.fnum
self.x2 = np.floor(self.mask.x_gridnum/2) + self.tcc.s.fnum + 1
self.y1 = np.floor(self.mask.y_gridnum/2) - self.tcc.s.gnum
self.y2 = np.floor(self.mask.y_gridnum/2) + self.tcc.s.gnum + 1
self.spat_part = pyfftw.empty_aligned((self.mask.y_gridnum,self.mask.x_gridnum),\
dtype='complex128')
self.freq_part = pyfftw.empty_aligned((self.mask.y_gridnum,self.mask.x_gridnum),\
dtype='complex128')
self.ifft_image = pyfftw.FFTW(self.freq_part,self.spat_part,axes=(0,1),\
direction='FFTW_BACKWARD')
def test_misaligned_data_doesnt_clobber_cache(self):
'''A bug was highlighted in #197 in which misaligned data causes
an overwrite of an FFTW internal array which is also the same as
an output array. The correct behaviour is for the cache to have
alignment as a key to stop this happening.
'''
interfaces.cache.enable()
N = 64
pyfftw.interfaces.cache.enable()
np.random.seed(12345)
Um = pyfftw.empty_aligned((N, N+1), dtype=np.float32, order='C')
Vm = pyfftw.empty_aligned((N, N+1), dtype=np.float32, order='C')
U = np.ndarray((N, N), dtype=Um.dtype, buffer=Um.data, offset=0)
V = np.ndarray(
(N, N), dtype=Vm.dtype, buffer=Vm.data, offset=Vm.itemsize)
U[:] = np.random.randn(N, N).astype(np.float32)
V[:] = np.random.randn(N, N).astype(np.float32)
uh = hashlib.md5(U).hexdigest()
vh = hashlib.md5(V).hexdigest()
x = interfaces.numpy_fft.rfftn(
U, None, axes=(0, 1), overwrite_input=False)
y = interfaces.numpy_fft.rfftn(
V, None, axes=(0, 1), overwrite_input=False)
self.assertTrue(uh == hashlib.md5(U).hexdigest())
self.assertTrue(vh == hashlib.md5(V).hexdigest())
interfaces.cache.disable()
def __init__(self, size):
self.size = size
self._time = pyfftw.empty_aligned(size, 'float64')
self._freq = pyfftw.empty_aligned(size//2 + 1, 'complex128')
self.fft = pyfftw.FFTW(self._time, self._freq, threads=os.cpu_count(),
direction='FFTW_FORWARD')
self.ifft = pyfftw.FFTW(self._freq, self._time, threads=os.cpu_count(),
direction='FFTW_BACKWARD')
def setUp(self):
self.input_array = empty_aligned((256, 512), dtype='complex128', n=16)
self.output_array = empty_aligned((256, 512), dtype='complex128', n=16)
self.fft = FFTW(self.input_array, self.output_array)
self.input_array[:] = (numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape))
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=byte_align(input_array.copy(), n=16)).copy()
input_array = byte_align(input_array, n=16)
output_array = byte_align(output_array, n=16)
# Offset by one from 16 byte aligned to guarantee it's not
# 16 byte aligned
a = byte_align(input_array.copy(), n=16)
a__ = empty_aligned(numpy.prod(a.shape)*a.itemsize+1, dtype='int8',
n=16)
a_ = a__[1:].view(dtype=a.dtype).reshape(*a.shape)
a_[:] = a
# Create a different second array the same way
b = byte_align(output_array.copy(), n=16)
b__ = empty_aligned(numpy.prod(b.shape)*a.itemsize+1, dtype='int8',
n=16)
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 pyfftw_container(ny, nx, bwd = False):
'''
construct a fftw container to perform fftw.
'''
a = pyfftw.empty_aligned((ny,nx),dtype = 'complex128')
b = pyfftw.empty_aligned((ny,nx),dtype = 'complex128')
if bwd:
container = pyfftw.FFTW(a,b,axes = (0,1),direction = 'FFTW_BACKWARD')
else:
container = pyfftw.FFTW(a,b,axes = (0,1),direction = 'FFTW_FORWARD')
return container
def test_failure(self):
for dtype, npdtype in zip(['32', '64', 'ld'], [np.complex64, np.complex128, np.clongdouble]):
if dtype == 'ld' and np.dtype(np.clongdouble) == np.dtype(np.complex128):
# skip this test on systems where clongdouble is complex128
continue
if dtype not in _supported_types:
a = empty_aligned((1,1024), npdtype, n=16)
b = empty_aligned(a.shape, dtype=a.dtype, n=16)
msg = "Rebuild pyFFTW with support for %s precision!" % _all_types_human_readable[dtype]
with self.assertRaisesRegex(NotImplementedError, msg):
FFTW(a,b)
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 = byte_align(a, n=16)
b = byte_align(b, n=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__ = empty_aligned(
numpy.prod(in_shape)*a.itemsize+byte_error,
dtype='int8', n=16)
a_ = (a__[byte_error:]
.view(dtype=self.input_dtype).reshape(*in_shape))
a_[:] = a
b__ = empty_aligned(
numpy.prod(out_shape)*b.itemsize+byte_error,
dtype='int8', n=16)
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_fft(self):
try:
import pyfftw
if not hasattr(self,'__fftw_ffts'):
a = pyfftw.empty_aligned((self._nx, self._ny), dtype='complex128')
b = pyfftw.empty_aligned((self._nx, self._ny), dtype='complex128')
fft2 = pyfftw.FFTW(a, b, axes=(0,1), threads=self._thread_count, direction = 'FFTW_FORWARD')
ifft2 = pyfftw.FFTW(b, a, axes=(0,1), threads=self._thread_count, direction = 'FFTW_BACKWARD')
self.__fftw_ffts = fft2,ifft2
return self.__fftw_ffts
except ImportError:
from numpy.fft import fft2,ifft2
return fft2,ifft2
def test_is_byte_aligned(self):
a = empty_aligned(100)
self.assertTrue(is_byte_aligned(a, get_expected_alignment(None)))
a = empty_aligned(100, n=16)
self.assertTrue(is_byte_aligned(a, n=16))
a = empty_aligned(100, n=5)
self.assertTrue(is_byte_aligned(a, n=5))
a = empty_aligned(100, dtype="float32", n=16)[1:]
self.assertFalse(is_byte_aligned(a, n=16))
self.assertTrue(is_byte_aligned(a, n=4))
def test_call_with_ortho_on(self):
_input_array = empty_aligned((256, 512), dtype='complex128', n=16)
ifft = FFTW(self.output_array, _input_array,
direction='FFTW_BACKWARD')
self.fft(ortho=True, normalise_idft=False)
# ortho case preserves the norm in forward direction
self.assertTrue(
numpy.allclose(numpy.linalg.norm(self.input_array),
numpy.linalg.norm(self.output_array)))
ifft(ortho=True, normalise_idft=False)
# ortho case preserves the norm in backward direction
self.assertTrue(
numpy.allclose(numpy.linalg.norm(_input_array),
numpy.linalg.norm(self.output_array)))
self.assertTrue(numpy.allclose(self.input_array, _input_array))
# cant select both ortho and normalise_idft
self.assertRaisesRegex(ValueError, 'Invalid options',
self.fft, normalise_idft=True, ortho=True)
# cant specify orth=True with default normalise_idft=True
self.assertRaisesRegex(ValueError, 'Invalid options',
self.fft, ortho=True)
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 = byte_align(
numpy.random.randn(*internal_array_shape)
+ 1j*numpy.random.randn(*internal_array_shape))
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 = empty_aligned(input_array_shape,
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_avoid_copy(self):
'''Test the avoid_copy flag
'''
dtype_tuple = input_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 = empty_aligned(
numpy.prod(test_shape)*input_array.itemsize+1,
dtype='int8', n=16)
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 = byte_align(input_array.copy())
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.input_array is
_input_array)
def conversion(self, missing, alt1, alt2):
'''If the ``missing`` precision is not available, the builder should convert to
``alt1`` precision. If that isn't available either, it should fall back to
``alt2``. If input precision is lost, a warning should be emitted.
'''
missing, alt1, alt2 = [np.dtype(x) for x in (missing, alt1, alt2)]
if _all_types_np[missing] in _supported_types:
return
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
itemsize = alt1.itemsize
a = empty_aligned((1, 512), dtype=missing)
b = interfaces.numpy_fft.fft(a)
res = _rc_dtype_pairs.get(alt1.char, None)
if res is not None:
self.assertEqual(b.dtype, res)
else:
itemsize = alt2.itemsize
self.assertEqual(b.dtype, _rc_dtype_pairs[alt2.char])
if itemsize < missing.itemsize:
print(itemsize, missing.itemsize)
assert len(w) == 1
assert "Narrowing conversion" in str(w[-1].message)
print("Found narrowing conversion from %d to %d bytes" % (missing.itemsize, itemsize))
else:
assert len(w) == 0
def readCoherent(self, size=2**25):
"""Return coherently dedispersed timestream
Read size number of samples, coherently dedisperse, take first
step samples to chop off wraparound
"""
if size != self.size or not 'dd' in self.__dict__ :
# Only compute dedispersion phases and fft plans once for given size
print("Calculating de-dispersion phase factors for size {0}".format(size))
self.f = self.fedge + self.forder*np.fft.rfftfreq(size, self.dt1)[:, np.newaxis]
self.dd = self.dm.phase_factor(self.f, self.fref)
for j in range(len(self.forder)):
if self.forder[j] == 1:
self.dd[...,j] = np.conj(self.dd[...,j])
self.size = size
self.step = int(size - 2**(np.ceil(np.log2(self.samploss))))
a = pyfftw.empty_aligned((self.size, self.npol), dtype='float32', n=16)
b = pyfftw.empty_aligned((self.size//2+1, self.npol), dtype='complex64', n=16)
print("planning FFTs for coherent dedispersion...")
self.fft_ts = pyfftw.FFTW(a,b, axes=(0,), direction='FFTW_FORWARD',
planning_timelimit=1.0, threads=8 )
print("...")
self.ifft_ts = pyfftw.FFTW(b,a, axes=(0,), direction='FFTW_BACKWARD',
planning_timelimit=1.0, threads=8 )
d = pyfftw.empty_aligned((size,self.npol), dtype='float32')
#d = self.fh.read(size)
# need better solution...
if self.dtype == 'vdif':
d[:] = self.fh.read(size)[self.thread_ids]
else:
d[:] = self.fh.read(size)
#ft = np.fft.rfft(d, axis=0)
ft = self.fft_ts(d)
ft *= self.dd
dift = pyfftw.empty_aligned((size//2+1,self.npol), dtype='complex64')
dift[:] = ft
d = self.ifft_ts(dift)
#d = np.fft.irfft(ft, axis=0)[:self.step]
return d
def _fftw(self, a, func, nthreads=ncpu):
if 0 in a.shape:
raise ValueError('This array cannot be transformed, shape: %s' % str(a.shape))
axes = [i - a.ndim for i in range(a.ndim)]
af = pyfftw.empty_aligned(a.shape, dtype=a.dtype.type.__name__)
plan = func(af, axes=axes, threads=nthreads)
af[:] = a[:]
return plan()
def compute_motion_shifts(scan, template, in_place=True, num_threads=8):
""" Compute shifts in y and x for rigid subpixel motion correction.
Returns the number of pixels that each image in the scan was to the right (x_shift)
or below (y_shift) the template. Negative shifts mean the image was to the left or
above the template.
:param np.array scan: 2 or 3-dimensional scan (image_height, image_width[, num_frames]).
:param np.array template: 2-d template image. Each frame in scan is aligned to this.
:param bool in_place: Whether the scan can be overwritten.
:param int num_threads: Number of threads used for the ffts.
:returns: (y_shifts, x_shifts) Two arrays (num_frames) with the y, x motion shifts.
..note:: Based in imreg_dft.translation().
"""
import pyfftw
from imreg_dft import utils
# Add third dimension if scan is a single image
if scan.ndim == 2:
scan = np.expand_dims(scan, -1)
# Get some params
image_height, image_width, num_frames = scan.shape
taper = np.outer(signal.tukey(image_height, 0.2), signal.tukey(image_width, 0.2))
# Prepare fftw
frame = pyfftw.empty_aligned((image_height, image_width), dtype='complex64')
fft = pyfftw.builders.fft2(frame, threads=num_threads, overwrite_input=in_place,
avoid_copy=True)
ifft = pyfftw.builders.ifft2(frame, threads=num_threads, overwrite_input=in_place,
avoid_copy=True)
# Get fourier transform of template
template_freq = fft(template * taper).conj() # we only need the conjugate
abs_template_freq = abs(template_freq)
eps = abs_template_freq.max() * 1e-15
# Compute subpixel shifts per image
y_shifts = np.empty(num_frames)
x_shifts = np.empty(num_frames)
for i in range(num_frames):
# Compute correlation via cross power spectrum
image_freq = fft(scan[:, :, i] * taper)
cross_power = (image_freq * template_freq) / (abs(image_freq) * abs_template_freq + eps)
shifted_cross_power = np.fft.fftshift(abs(ifft(cross_power)))
# Get best shift
shifts = np.unravel_index(np.argmax(shifted_cross_power), shifted_cross_power.shape)
shifts = utils._interpolate(shifted_cross_power, shifts, rad=3)
# Map back to deviations from center
y_shifts[i] = shifts[0] - image_height // 2
x_shifts[i] = shifts[1] - image_width // 2
return y_shifts, x_shifts
def test_call_with_normalisation_on(self):
_input_array = empty_aligned((256, 512), dtype="complex128", n=16)
ifft = FFTW(self.output_array, _input_array, direction="FFTW_BACKWARD")
self.fft(normalise_idft=True) # Shouldn't make any difference
ifft(normalise_idft=True)
self.assertTrue(numpy.allclose(self.input_array, _input_array))
def energy_spectrum(self,nx,ny,dx,dy,w):
epsilon = 1.0e-6
kx = np.empty(nx)
ky = np.empty(ny)
kx[0:int(nx/2)] = 2*np.pi/(np.float64(nx)*dx)*np.float64(np.arange(0,int(nx/2)))
kx[int(nx/2):nx] = 2*np.pi/(np.float64(nx)*dx)*np.float64(np.arange(-int(nx/2),0))
ky[0:ny] = kx[0:ny]
kx[0] = epsilon
ky[0] = epsilon
kx, ky = np.meshgrid(kx, ky, indexing='ij')
a = pyfftw.empty_aligned((nx,ny),dtype= 'complex128')
b = pyfftw.empty_aligned((nx,ny),dtype= 'complex128')
fft_object = pyfftw.FFTW(a, b, axes = (0,1), direction = 'FFTW_FORWARD')
wf = fft_object(w[1:nx+1,1:ny+1])
es = np.empty((nx,ny))
kk = np.sqrt(kx[:,:]**2 + ky[:,:]**2)
es[:,:] = np.pi*((np.abs(wf[:,:])/(nx*ny))**2)/kk
n = int(np.sqrt(nx*nx + ny*ny)/2.0)-1
en = np.zeros(n+1)
for k in range(1,n+1):
en[k] = 0.0
ic = 0
ii,jj = np.where((kk[1:,1:]>(k-0.5)) & (kk[1:,1:]<(k+0.5)))
ic = ii.size
ii = ii+1
jj = jj+1
en[k] = np.sum(es[ii,jj])
en[k] = en[k]/ic
return en, n
def Get_Bands_Matrix(cls,
Ground: bool = False,
Cluster: bool = False) -> np.ndarray:
Mminous, Mplus = cls.Sample_State(Ground=Ground)
x = np.arange(-(cls.N_size - 1) / 2, (cls.N_size - 1) / 2 + 1)
if Cluster:
M_plus = np.fft.ifftshift(Mplus)
M_minous = np.fft.ifftshift(Mminous)
Fourier_minous = np.fft.fft(M_minous)
Fourier_plus = np.fft.fft(M_plus)
else:
M_plus = pyfftw.empty_aligned(cls.N_size, dtype='complex128')
M_plus[:] = np.fft.ifftshift(Mplus)
M_minous = pyfftw.empty_aligned(cls.N_size, dtype='complex128')
M_minous[:] = np.fft.ifftshift(Mminous)
Fourier_minous = pyfftw.interfaces.numpy_fft.fft(M_minous)
Fourier_plus = pyfftw.interfaces.numpy_fft.fft(M_plus)
return Fourier_minous / cls.N_size, Fourier_plus / cls.N_size
def test_call_with_normalisation_default(self):
_input_array = empty_aligned((256, 512), dtype="complex128", n=16)
ifft = FFTW(self.output_array, _input_array, direction="FFTW_BACKWARD")
self.fft()
ifft()
# Scaling is performed by default
self.assertTrue(numpy.allclose(self.input_array, _input_array))
def fftw_test(input_data):
pyfftw.forget_wisdom() # This is just here to keep the tests honest
outLength = len(
input_data
) // 2 + 1 # For a real FFT, the output is symetrical. fftw returns only half the data in this case
a = pyfftw.empty_aligned(
len(input_data), dtype='float32'
) # This is the input array. It will be cleared when the fft object is created
outData = pyfftw.empty_aligned(
outLength, dtype='complex64'
) # This is the output array. Not that the size and type must be appropriate
fft_obj = pyfftw.FFTW(
a, outData, flags=('FFTW_ESTIMATE', ),
planning_timelimit=1.0) # NB: The flags tuple has a comma
a[:] = np.array(
input_data, dtype='float32'
) # We have to fill the array after fft_obj is created. NB: 'a[:] =' puts data into the existing array, 'a =' creates a new array
return fft_obj(
) # Calling the object returns the FFT of the data now in a. The result is also in outData
def ifft(self, a, axis=-1, workers=-1):
mutex.acquire()
comp_type = a.dtype
try:
a_ = pyfftw.empty_aligned(a.shape, dtype=comp_type)
b = pyfftw.empty_aligned(a.shape, dtype=comp_type)
cls = pyfftw.FFTW(
a_,
b,
axes=(axis, ),
direction="FFTW_BACKWARD",
threads=_workers(workers),
)
cls(a, b)
finally:
mutex.release()
return b
def get_length(image):
distortion_angle = get_direction(image)
image_size = image.shape
h_center, w_center = image_size[0] // 2, image_size[1] // 2
float_placeholder = pyfftw.empty_aligned(image_size, dtype='float32')
complex_placeholder = pyfftw.empty_aligned(image_size, dtype='complex64')
float_placeholder[:, :] = image
image_cepstrum = pyfftw.builders.fftn(float_placeholder)()
image_cepstrum = np.log(np.absolute(image_cepstrum))
complex_placeholder[:, :] = image_cepstrum
image_cepstrum = np.absolute(pyfftw.builders.ifftn(complex_placeholder)())
image_cepstrum = np.roll(image_cepstrum, h_center, axis=0)
image_cepstrum = np.roll(image_cepstrum, w_center, axis=1)
smooth_mask = gauss_2d(image_size, sigma=10)
smooth_mask = np.max(smooth_mask) - smooth_mask
sector_mask = np.zeros(image_size)
if 0 <= distortion_angle < 90 or 180 <= distortion_angle < 270:
sector_mask[:h_center, w_center:] = 1
sector_mask[h_center:, :w_center] = 1
else:
sector_mask[:h_center, :w_center:] = 1
sector_mask[h_center:, w_center:] = 1
masked_image_cepstrum = image_cepstrum * smooth_mask * sector_mask
max_index = np.argmax(masked_image_cepstrum)
h_max_indexes, w_max_indexes = np.abs(max_index // image_size[1] - h_center), \
np.abs(max_index % image_size[1] - w_center)
lenght = np.sqrt(h_max_indexes**2 + w_max_indexes**2)
lenght += (1 - lenght % 2)
return lenght, distortion_angle
def __init__(self, shape, float_precision, complex_precision, threads=2):
"""
Parameters
----------
shape : tuple
Shape of the arrays which you will take the Fourier transforms of.
float_precision : `~numpy.dtype`
complex_precision : `~numpy.dtype`
threads : int, optional
This FFT implementation uses multithreading, with
two threads by default.
"""
# Allocate byte-aligned
self.buffer_float = pyfftw.empty_aligned(shape,
dtype=float_precision.__name__)
self.buffer_complex = pyfftw.empty_aligned(shape,
dtype=complex_precision.__name__)
self._fft2 = pyfftw.builders.fft2(self.buffer_float, threads=threads)
self._ifft2 = pyfftw.builders.ifft2(self.buffer_complex,
threads=threads)
def test_call_with_ortho_off(self):
_input_array = empty_aligned((256, 512), dtype='complex128', n=16)
ifft = FFTW(self.output_array, _input_array,
direction='FFTW_BACKWARD')
self.fft(ortho=False)
ifft(ortho=False)
# Scaling by normalise_idft is performed by default
self.assertTrue(numpy.allclose(self.input_array, _input_array))
def test_scipy_overwrite(self):
new_style_scipy_fftn = False
try:
scipy_fftn = scipy.signal.signaltools.fftn
scipy_ifftn = scipy.signal.signaltools.ifftn
except AttributeError:
scipy_fftn = scipy.fftpack.fftn
scipy_ifftn = scipy.fftpack.ifftn
new_style_scipy_fftn = True
a = pyfftw.empty_aligned((128, 64), dtype='complex128', n=16)
b = pyfftw.empty_aligned((128, 64), dtype='complex128', n=16)
a[:] = (numpy.random.randn(*a.shape) +
1j*numpy.random.randn(*a.shape))
b[:] = (numpy.random.randn(*b.shape) +
1j*numpy.random.randn(*b.shape))
scipy_c = scipy.signal.fftconvolve(a, b)
if new_style_scipy_fftn:
scipy.fftpack.fftn = scipy_fftpack.fftn
scipy.fftpack.ifftn = scipy_fftpack.ifftn
else:
scipy.signal.signaltools.fftn = scipy_fftpack.fftn
scipy.signal.signaltools.ifftn = scipy_fftpack.ifftn
scipy_replaced_c = scipy.signal.fftconvolve(a, b)
self.assertTrue(numpy.allclose(scipy_c, scipy_replaced_c))
if new_style_scipy_fftn:
scipy.fftpack.fftn = scipy_fftn
scipy.fftpack.ifftn = scipy_ifftn
else:
scipy.signal.signaltools.fftn = scipy_fftn
scipy.signal.signaltools.ifftn = scipy_ifftn
def setUp(self):
self.input_array_slicer = [slice(None), slice(256)]
self.FFTW_array_slicer = [slice(128), slice(None)]
self.input_array = empty_aligned((128, 512), dtype='complex128')
self.output_array = empty_aligned((256, 256), dtype='complex128')
self.internal_array = empty_aligned((256, 256), dtype='complex128')
self.fft = utils._FFTWWrapper(self.internal_array,
self.output_array,
input_array_slicer=self.input_array_slicer,
FFTW_array_slicer=self.FFTW_array_slicer)
self.input_array[:] = (numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape))
self.internal_array[:] = 0
self.internal_array[self.FFTW_array_slicer] = (
self.input_array[self.input_array_slicer])
def __init__(self, mask, tcc):
self.tcc = tcc # TCC
self.mask = mask # Mask
self.order = tcc.order # TCC Order
self.resist_a = 80 # Resist model parameter: sharpness
self.resist_t = 0.6 # Resist model parameter: threshold
self.kernels = tcc.kernels # Kernels
self.coefs = tcc.coefs # Coefs
self.norm = self.mask.y_gridnum*self.mask.x_gridnum
self.x1 = np.floor(self.mask.x_gridnum/2) - self.tcc.s.fnum
self.x2 = np.floor(self.mask.x_gridnum/2) + self.tcc.s.fnum + 1
self.y1 = np.floor(self.mask.y_gridnum/2) - self.tcc.s.gnum
self.y2 = np.floor(self.mask.y_gridnum/2) + self.tcc.s.gnum + 1
self.spat_part = pyfftw.empty_aligned((self.mask.y_gridnum,self.mask.x_gridnum),\
dtype='complex128')
self.freq_part = pyfftw.empty_aligned((self.mask.y_gridnum,self.mask.x_gridnum),\
dtype='complex128')
self.ifft_image = pyfftw.FFTW(self.freq_part,self.spat_part,axes=(0,1),\
direction='FFTW_BACKWARD')
def compute_pow2_real_wisdom(path,
pow2=range(20),
rigor='measure',
threads=16):
"""
???
If you plan with FFTW_PATIENT, it will automatically disable
threads for sizes that don't benefit from parallelization.
"""
flags = [RIGOR_MAP[rigor],
'FFTW_DESTROY_INPUT']
wisdom_fnames = []
for pow2_i in pow2:
N = 2**pow2_i
for direction in ['forward', 'backward']:
logger.info('building wisdom for real 2**{} {}'.format(pow2_i,
direction))
if direction == 'forward':
x_input = pyfftw.empty_aligned(N, dtype='float64')
x_output = pyfftw.empty_aligned(int(N // 2) + 1, dtype='complex128')
else:
x_output = pyfftw.empty_aligned(N, dtype='float64')
x_input = pyfftw.empty_aligned(int(N // 2) + 1, dtype='complex128')
plan = pyfftw.FFTW(x_input,
x_output,
direction=DIRECTION_MAP[direction],
flags=flags,
threads=threads)
wisdom_fnames.append(wisdom_fname(path,
'real',
pow2_i,
threads,
direction,
rigor))
logger.info('writing to {}'.format(wisdom_fnames[-1]))
with open(wisdom_fnames[-1], 'w') as fid:
dump(pyfftw.export_wisdom(), fid, -1)
pyfftw.forget_wisdom()
return wisdom_fnames
