def test_update_unaligned_data_with_FFTW_UNALIGNED(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 + 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
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 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_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 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 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 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 test_call_with_keyword_output_update(self):
"""Test the class call with a keyword output update.
"""
output_array = byte_align(
(numpy.random.randn(*self.output_array.shape) + 1j * numpy.random.randn(*self.output_array.shape)), n=16
)
returned_output_array = self.fft(output_array=byte_align(output_array.copy(), n=16)).copy()
self.fft.update_arrays(self.input_array, output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(returned_output_array == output_array))
def test_call_with_keyword_input_update(self):
"""Test the class call with a keyword input update.
"""
input_array = byte_align(
numpy.random.randn(*self.input_array.shape) + 1j * numpy.random.randn(*self.input_array.shape), n=16
)
output_array = self.fft(input_array=byte_align(input_array.copy(), n=16)).copy()
self.fft.update_arrays(input_array, self.output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(output_array == self.output_array))
def inner_product(template, signal, snf, freq_scale, template_pre_fft=False):
if template.size > signal.size:
size = template.size
else:
size = signal.size
if not template_pre_fft:
template = np.flip(template)
#template_f = fft.fft(template, n=size)
template = pyfftw.byte_align(template, n=pyfftw.simd_alignment)
template_f = pyfftw.interfaces.numpy_fft.fft(template)
else:
template_f = template
# signal_f = np.fft.fft(signal, n=size)
signal = signal + 1000.0
signal = pyfftw.byte_align(signal, n=pyfftw.simd_alignment)
signal_f = pyfftw.interfaces.numpy_fft.fft(signal)
signal_f = np.pad(signal_f, (0, template_f.size - signal_f.size),
mode='constant')
freqs = freq_scale * pyfftw.interfaces.numpy_fft.fftfreq(size)
#@jit(nopython=True, cache=True)
def blah(template_f, signal_f, snf, freqs):
num = template_f * signal_f
den = snf
eq = num / den
cond = freqs >= 0
eq2 = eq[:eq.size // 2]
eq2 = np.extract(cond, eq)
freqs = np.extract(cond, freqs)
#ip = 4*np.real(np.sum(eq2)*freq_scale/eq2.size)
ip = 4 * np.real(ing.trapz(eq2, x=freqs))
#ip = 4*np.real(np.trapz(eq2, x=freqs))
return ip, eq
ip, eq = blah(template_f, signal_f, snf, freqs)
#ip, eq = blah(template_f, signal_f, snf, None)
if DEBUG:
print('inner_product(): debug output')
print('Freq. Scale: {}'.format(freq_scale))
print('Inner Product: {}'.format(ip))
print('Max over filter: {}'.format(np.max(eq)))
return ip, eq
def test_call_with_keyword_input_update(self):
'''Test the class call with a keyword input update.
'''
input_array = byte_align(
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())).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_positional_input_update(self):
'''Test the class call with a positional input update.
'''
input_array = byte_align(
(numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape)), n=16)
output_array = self.fft(byte_align(input_array.copy(), n=16)).copy()
self.fft.update_arrays(input_array, self.output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(output_array == self.output_array))
def test_call_with_positional_input_update(self):
'''Test the class call with a positional input update.
'''
input_array = byte_align(
(numpy.random.randn(*self.input_array.shape) +
1j * numpy.random.randn(*self.input_array.shape)),
n=16)
output_array = self.fft(byte_align(input_array.copy(), n=16)).copy()
self.fft.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 = byte_align(
(numpy.random.randn(*self.output_array.shape) +
1j * numpy.random.randn(*self.output_array.shape)),
n=16)
returned_output_array = self.fft(
output_array=byte_align(output_array.copy(), n=16)).copy()
self.fft.update_arrays(self.input_array, output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(returned_output_array == output_array))
def __init__(self, data, axes, threads=0):
self.axes = axes
self.threads = threads
self.use_pyfftw = False
if threads > 0 and pyfftw_available:
self.use_pyfftw = True
self.data = data
self.data_temp = pyfftw.byte_align(0 * data)
self.fftw = pyfftw.FFTW(
data,
self.data_temp,
axes=axes,
threads=threads,
direction="FFTW_FORWARD",
flags=("FFTW_MEASURE",),
)
self.ifftw = pyfftw.FFTW(
data,
self.data_temp,
axes=axes,
threads=threads,
direction="FFTW_BACKWARD",
flags=("FFTW_MEASURE",),
)
print(f"fftw simd_aligned={self.fftw.simd_aligned}")
print(f"ifftw simd_aligned={self.fftw.simd_aligned}")
else:
# I perform one fft to have numpy cache the plan
_ = np.fft.ifftn(np.fft.fftn(data, axes=axes), axes=axes)
def test_call_with_different_input_dtype(self):
'''Test the class call with an array with a different input dtype
'''
input_array = byte_align(numpy.complex64(
numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape)))
output_array = self.fft(byte_align(input_array.copy())).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_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 [None, 3, 7, 9, 16, 24, 23, 63, 64]:
d = byte_align(a, n=n)
self.assertTrue(numpy.array_equal(a, d))
d = byte_align(b, n=n)
self.assertTrue(numpy.array_equal(b, d))
d = byte_align(c, n=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 = 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 test_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 [None, 3, 7, 9, 16, 24, 23, 63, 64]:
d = byte_align(a, n=n)
self.assertTrue(numpy.array_equal(a, d))
d = byte_align(b, n=n)
self.assertTrue(numpy.array_equal(b, d))
d = byte_align(c, n=n)
self.assertTrue(numpy.array_equal(c, d))
def fastifftn(input_array, **kwargs):
"""
Auxiliary function to use pyFFTW. It does the align, planning and
apply inverse FFTW transform
Parameters
---------
input_array : array_like
Array to be FFTWed
Returns
-------
ifftw_array : array_like
Inverse Fourier transformed array
"""
# number of cores available
ncores = kwargs["ncores"] # multiprocessing.cpu_count()
# stating the precision.
cprecision = kwargs["cprecision"] # np.complex64 # single precision
planner_type = kwargs["planner_type"] # 'FFTW_MEASURE'
# align array
ifftw_array = pyfftw.byte_align(input_array, dtype=cprecision, n=16)
ifftw_array = pyfftw.interfaces.numpy_fft.ifftn(
ifftw_array,
overwrite_input=True,
planner_effort=planner_type,
threads=ncores)
return ifftw_array
def test_byte_align_integer_shape(self):
shape = 100
a = numpy.random.randn(shape)
# Test a few alignments
for n in [None, 3, 7, 9, 16, 24, 23, 63, 64]:
expected_alignment = get_expected_alignment(n)
b = byte_align(a, n=n)
self.assertTrue(b.ctypes.data % expected_alignment == 0)
def test_byte_align_integer_shape(self):
shape = 100
a = numpy.random.randn(shape)
# Test a few alignments
for n in [None, 3, 7, 9, 16, 24, 23, 63, 64]:
expected_alignment = get_expected_alignment(n)
b = byte_align(a, n=n)
self.assertTrue(b.ctypes.data % expected_alignment == 0)
def calc_bessel(x, y, radius, gradr=False):
rho = (x**2 + y**2)**0.5
zero = special.jn_zeros(0, 1)
# See RT2 p75 for derivation
norm_fac = np.pi**0.5*special.jv(1, zero)*radius
z = rho*zero/radius
in_beam = rho <= radius
E = pyfftw.byte_align(special.jv(0, z)*in_beam/norm_fac)
if gradr:
gradrhoE = special.jvp(0, z, 1)*in_beam/norm_fac*zero/radius
gradxE = pyfftw.byte_align(mathx.divide0(gradrhoE*x, rho))
gradyE = pyfftw.byte_align(mathx.divide0(gradrhoE*y, rho))
return E, (gradxE, gradyE)
else:
return E
def test_call_with_positional_updates(self):
'''Test the class call with a positional array updates.
'''
input_array = byte_align((numpy.random.randn(*self.input_array.shape)
+ 1j*numpy.random.randn(*self.input_array.shape)))
output_array = byte_align((numpy.random.randn(*self.output_array.shape)
+ 1j*numpy.random.randn(*self.output_array.shape)))
returned_output_array = self.fft(
byte_align(input_array.copy()),
byte_align(output_array.copy())).copy()
self.update_arrays(input_array, output_array)
self.fft.execute()
self.assertTrue(numpy.alltrue(returned_output_array == output_array))
def run(self):
"""
Propagate field through each module, with the resulting field at the
exit of each module stored in a dictionary, with module name as key.
"""
self.field = byte_align(self.field)
for module in self.modules:
self.field = module(self.domain, self.field)
self.fields[module.name] = self.field
def create_fftw_objects(array, allow_new_plan=True):
"""
Creates FFTW object for forward and backward Fourier transforms. The input array will be
transformed in place. The function tries to retrieve FFTW plans from wisdom only.
If no plan exists for the input array, a new plan is cached and then retrieved.
:param array: Numpy array to be transformed. 2 dimensions or more.
:param allow_new_plan: If false raise an exception instead of caching a new plan.
:return:
"""
try:
# try using cached FFTW plan
fftw_forward = pyfftw.FFTW(array,
array,
axes=(-1, -2),
threads=FFTW_THREADS,
flags=(FFTW_EFFORT, 'FFTW_WISDOM_ONLY',
'FFTW_DESTROY_INPUT'))
fftw_backward = pyfftw.FFTW(array,
array,
axes=(-1, -2),
direction='FFTW_BACKWARD',
threads=FFTW_THREADS,
flags=(FFTW_EFFORT, 'FFTW_WISDOM_ONLY',
'FFTW_DESTROY_INPUT'))
return fftw_forward, fftw_backward
except RuntimeError as e:
if (not allow_new_plan):
fftw_forward = pyfftw.builders.fft2(array)
fftw_backward = pyfftw.builders.ifft2(array)
return fftw_forward, fftw_backward
# if ('No FFTW wisdom is known for this plan.' != str(e)) or (not allow_new_plan):
# raise
# create new FFTW object, not to be used, but the plan will remain in the cache
dummy = pyfftw.byte_align(np.zeros_like(array))
# this is necessary because FFTW overwrites the array during planning
pyfftw.FFTW(dummy,
dummy,
axes=(-1, -2),
threads=FFTW_THREADS,
flags=(FFTW_EFFORT, 'FFTW_DESTROY_INPUT'),
planning_timelimit=FFTW_TIMELIMIT)
pyfftw.FFTW(dummy,
dummy,
axes=(-1, -2),
direction='FFTW_BACKWARD',
threads=FFTW_THREADS,
flags=(FFTW_EFFORT, 'FFTW_DESTROY_INPUT'),
planning_timelimit=FFTW_TIMELIMIT)
return create_fftw_objects(array, False)
def map2rfft(self, _map):
oupt = pyfftw.empty_aligned(self.ell_mat.rshape, dtype='complex128')
fft = pyfftw.FFTW(pyfftw.byte_align(_map, dtype='float64'),
oupt,
axes=(0, 1),
direction='FFTW_FORWARD',
flags=self.flags,
threads=self.threads)
fft()
return oupt
def test_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 [None, 3, 7, 9, 16, 24, 23, 63, 64]:
expected_alignment = get_expected_alignment(n)
d = byte_align(a, dtype='float32', n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.dtype == 'float32')
d = byte_align(b, dtype='float32', n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.dtype == 'float32')
d = byte_align(c, dtype='float64', n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.dtype == 'float64')
def test_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 [None, 3, 7, 9, 16, 24, 23, 63, 64]:
expected_alignment = get_expected_alignment(n)
d = byte_align(a, n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.__class__ == a.__class__)
d = byte_align(b, n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.__class__ == b.__class__)
d = byte_align(c, n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.__class__ == c.__class__)
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_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 [None, 3, 7, 9, 16, 24, 23, 63, 64]:
expected_alignment = get_expected_alignment(n)
d = byte_align(a, dtype="float32", n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.dtype == "float32")
d = byte_align(b, dtype="float32", n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.dtype == "float32")
d = byte_align(c, dtype="float64", n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.dtype == "float64")
def test_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 [None, 3, 7, 9, 16, 24, 23, 63, 64]:
expected_alignment = get_expected_alignment(n)
d = byte_align(a, n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.__class__ == a.__class__)
d = byte_align(b, n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.__class__ == b.__class__)
d = byte_align(c, n=n)
self.assertTrue(d.ctypes.data % expected_alignment == 0)
self.assertTrue(d.__class__ == c.__class__)
def test_call_with_different_striding(self):
"""Test the input update with different strides to internal array.
"""
shape = self.input_array.shape + (2,)
input_array = byte_align(numpy.random.randn(*shape) + 1j * numpy.random.randn(*shape), n=16)
fft = FFTW(input_array[:, :, 0], self.output_array)
test_output_array = fft().copy()
new_input_array = byte_align(input_array[:, :, 0].copy(), n=16)
new_output = fft(new_input_array).copy()
# Test the test!
self.assertTrue(new_input_array.strides != input_array[:, :, 0].strides)
self.assertTrue(numpy.alltrue(test_output_array == new_output))
def test_call_with_invalid_output_striding(self):
'''Test the class call with an invalid strided output update.
'''
# Add an extra dimension to bugger up the striding
new_shape = self.output_array.shape + (2,)
output_array = byte_align(numpy.random.randn(*new_shape)
+ 1j*numpy.random.randn(*new_shape))
self.assertRaisesRegex(ValueError, 'Invalid output striding',
self.fft, **{'output_array': output_array[:,:,1]})
def test_update_unaligned_data_with_FFTW_UNALIGNED(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 +
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
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 fft1d(s, fs=1, N=None, axis=0, WRAP=True):
"""
Perform 1D FFT transform (using hanning window).
@param: s - signal (1-D array (or python list)),
@param: fs - sampling frequency (defalt is 1),
@param: N - FFT point number (int),
@param: axis - the axis along which to perform fft
(0: along column; 1: along row),
@param: WRAP - whether to unwrap the phase angle.
----
@return: f - frequency,
@return: mag - magnitude of each frequency,
@return: phase - phase angle of each frequency.
"""
s = pyfftw.byte_align(s, dtype='float')
if N is None:
# N = 2**math.floor(math.log2(s.shape[0]))
N = s.shape[axis]
if s.shape[axis] < N:
warnings.warn('WARNING: N > array length!!!')
tmp = np.zeros(N)
tmp[:s.shape[axis]] = s
s = tmp
df = fs / N
if N % 2 == 0:
f = np.arange(N // 2 + 1) * df
else:
f = np.arange((N + 1) / 2) * df
w = np.hanning(N)
if len(s.shape) > 1:
if axis == 0:
for i in range(s.shape[1]):
s[:N, i] = w * spsig.detrend(s[:N, i], type='constant')
fft = 2 * pyfftw.interfaces.numpy_fft.rfft(s[:N, :], N, axis=axis)
else:
for i in range(s.shape[0]):
s[i, :N] = w * spsig.detrend(s[i, :N], type='constant')
fft = 2 * pyfftw.interfaces.numpy_fft.rfft(s[:, :N], N, axis=axis)
else:
s[:N] = w * spsig.detrend(s[:N], type='constant')
fft = 2 * pyfftw.interfaces.numpy_fft.rfft(s[:N], N, axis=axis)
mag = 2 * np.abs(fft) / N
mag[0] /= 2 # amplitude of the constant component
phase = np.angle(fft)
if not WRAP:
phase = np.unwrap(phase)
return f, mag, phase
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 = byte_align(numpy.random.randn(*shape) + 1j * numpy.random.randn(*shape), n=16)
fft = FFTW(self.input_array, self.output_array)
self.assertRaisesRegex(ValueError, "Invalid input shape", self.fft, **{"input_array": input_array[:, :, 0]})
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 = byte_align(numpy.random.randn(*shape)
+ 1j*numpy.random.randn(*shape))
self.assertRaisesRegex(ValueError, 'Invalid input shape',
self.fft, **{'input_array': input_array[:,:,0]})
def test_call_with_different_striding(self):
'''Test the input update with different strides to internal array.
'''
shape = self.input_array.shape + (2, )
input_array = byte_align(numpy.random.randn(*shape) +
1j * numpy.random.randn(*shape),
n=16)
fft = FFTW(input_array[:, :, 0], self.output_array)
test_output_array = fft().copy()
new_input_array = byte_align(input_array[:, :, 0].copy(), n=16)
new_output = fft(new_input_array).copy()
# Test the test!
self.assertTrue(
new_input_array.strides != input_array[:, :, 0].strides)
self.assertTrue(numpy.alltrue(test_output_array == new_output))
def _fft(self, a):
if self._fftw is None:
a = pyfftw.byte_align(a, n=self._n_simd)
self._setup_fftw(a)
# Save a bit of useless checking in FFTW if possible.
if a is self._fftw.input_array:
a = None
if self._inverse is None:
b = None
else:
b = self._inverse._fftw.input_array
if b is self._fftw.output_array:
b = None
return self._fftw(a, b)
def calc_gaussian(k, x, y, waist0s, z=0, r0s=(0, 0), q0s=(0, 0), gradr=False):
"""Sample 2D paraxial Gaussian beam.
The normalization is such that the analytic (true, not sampled) field has unity integrated |E(r)|^2 across a transverse
plane.
Args:
k (scalar): Wavenumber.
x (N*1 array): X values to sample.
y (Length M array): Y values to sample.
waist0s (scalar or pair of): Waist in x and y planes.
z (numeric that broadcasts with x and y): Z values relative to the waist to sample.
r0s (pair of scalars): Coordinates of real-space center of Gaussian.
q0s (pair of scalars): Coordinates of angular-space center of Gaussian.
gradr (bool): If True, return transverse partial derivatives.
Returns:
If gradr: Er, (gradxEr, gradyEr). Otherwise just Er.
"""
waist0s = sa.to_scalar_pair(waist0s)
# 1D functions are normalized to unity amplitude on the waist. We want result to have normalized power.
factor = (np.pi*np.prod(waist0s)/2)**0.5
Ex = calc_gaussian_1d(k, x, waist0s[0], z, r0s[0], q0s[0], False)
Ey = calc_gaussian_1d(k, y, waist0s[1], z, r0s[1], q0s[1], True)
Er = pyfftw.byte_align(Ex*Ey)/factor
if gradr:
gradxEx = calc_gaussian_1d(k, x, waist0s[0], z, r0s[0], q0s[0], False, True)
gradyEy = calc_gaussian_1d(k, y, waist0s[1], z, r0s[1], q0s[1], True, True)
gradxEr = pyfftw.byte_align(gradxEx*Ey)/factor
gradyEr = pyfftw.byte_align(Ex*gradyEy)/factor
return Er, (gradxEr, gradyEr)
else:
return Er
def test_call_with_auto_input_alignment(self):
'''Test the class call with a keyword input update.
'''
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()
# Offset by one from 16 byte aligned to guarantee it's not
# 16 byte aligned
a = input_array
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
# Just confirm that a usual update will fail
self.assertRaisesRegex(ValueError, 'Invalid input alignment',
self.fft.update_arrays,
*(a_, self.output_array))
self.fft(a_, self.output_array)
self.assertTrue(numpy.alltrue(output_array == self.output_array))
# now try with a single byte offset and SIMD off
ar, ai = numpy.float32(numpy.random.randn(2, 257))
a = ar[1:] + 1j * ai[1:]
b = a.copy()
a_size = len(a.ravel()) * a.itemsize
update_array = numpy.frombuffer(numpy.zeros(a_size + 1,
dtype='int8')[1:].data,
dtype=a.dtype).reshape(a.shape)
fft = FFTW(a, b, flags=('FFTW_UNALIGNED', ))
# Confirm that a usual update will fail (it's not on the
# byte boundary)
self.assertRaisesRegex(ValueError, 'Invalid input alignment',
fft.update_arrays, *(update_array, b))
fft(update_array, b)
def alm2map(self, alm, lib_almout=None):
assert alm.size == self.alm_size, (alm.size, self.alm_size)
if lib_almout is None:
oupt = pyfftw.empty_aligned(self.ell_mat.shape, dtype='float64')
inpt = pyfftw.empty_aligned(self.ell_mat.rshape,
dtype='complex128')
ifft = pyfftw.FFTW(inpt,
oupt,
axes=(0, 1),
direction='FFTW_BACKWARD',
flags=self.flags,
threads=self.threads)
return ifft(
pyfftw.byte_align(self.alm2rfft(alm), dtype='complex128'))
else:
return lib_almout.alm2map(lib_almout.udgrade(self, alm))
def test_call_with_auto_input_alignment(self):
'''Test the class call with a keyword input update.
'''
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()
# Offset by one from 16 byte aligned to guarantee it's not
# 16 byte aligned
a = input_array
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
# Just confirm that a usual update will fail
self.assertRaisesRegex(ValueError, 'Invalid input alignment',
self.fft.update_arrays, *(a_, self.output_array))
self.fft(a_, self.output_array)
self.assertTrue(numpy.alltrue(output_array == self.output_array))
# now try with a single byte offset and SIMD off
ar, ai = numpy.float32(numpy.random.randn(2, 257))
a = ar[1:] + 1j*ai[1:]
b = a.copy()
a_size = len(a.ravel())*a.itemsize
update_array = numpy.frombuffer(
numpy.zeros(a_size + 1, dtype='int8')[1:].data,
dtype=a.dtype).reshape(a.shape)
fft = FFTW(a, b, flags=('FFTW_UNALIGNED',))
# Confirm that a usual update will fail (it's not on the
# byte boundary)
self.assertRaisesRegex(ValueError, 'Invalid input alignment',
fft.update_arrays, *(update_array, b))
fft(update_array, b)
def pyfftw_byte_aligned(array, dtype=None, n=None):
"""
Construct a byte-aligned array for efficient use by :mod:`pyfftw`.
This function is a wrapper for :func:`pyfftw.byte_align`
Parameters
----------
array : ndarray
Input array
dtype : dtype, optional (default None)
Output array dtype
n : int, optional (default None)
Output array should be aligned to n-byte boundary
Returns
-------
a : ndarray
Array with required byte-alignment
"""
return pyfftw.byte_align(array, n=n, dtype=dtype)
def empty(N, dtype=np.float, bytes=16):
return pyfftw.byte_align(Empty(N, dtype=dtype), n=bytes)
def test_alignment(self):
'''Test to see if the alignment is returned correctly
'''
in_shape = self.input_shapes['2d']
out_shape = self.output_shapes['2d']
input_dtype_alignment = self.get_input_dtype_alignment()
output_dtype_alignment = self.get_output_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)
a = byte_align(a, n=16)
b = byte_align(b, n=16)
a_orig = a.copy()
b_orig = b.copy()
# 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
a[:] = a_orig
fft, ifft = self.run_validate_fft(a, b, axes,
create_array_copies=False)
self.assertTrue(fft.input_alignment == 16)
self.assertTrue(fft.output_alignment == 16)
a[:] = a_orig
fft, ifft = self.run_validate_fft(a, b_, axes,
create_array_copies=False)
self.assertTrue(fft.input_alignment == input_dtype_alignment)
self.assertTrue(fft.output_alignment == output_dtype_alignment)
a_[:] = a_orig
fft, ifft = self.run_validate_fft(a_, b, axes,
create_array_copies=False)
self.assertTrue(fft.input_alignment == input_dtype_alignment)
self.assertTrue(fft.output_alignment == output_dtype_alignment)
a_[:] = a_orig
fft, ifft = self.run_validate_fft(a_, b_, axes,
create_array_copies=False)
self.assertTrue(fft.input_alignment == input_dtype_alignment)
self.assertTrue(fft.output_alignment == output_dtype_alignment)
a[:] = a_orig
fft, ifft = self.run_validate_fft(a, b, axes,
create_array_copies=False, force_unaligned_data=True)
self.assertTrue(fft.input_alignment == input_dtype_alignment)
self.assertTrue(fft.output_alignment == output_dtype_alignment)
def zeros(N, dtype=np.float, bytes=16):
return pyfftw.byte_align(Zeros(N, dtype=dtype), n=bytes)
def _Xfftn(a, s, axes, overwrite_input,
planner_effort, threads, auto_align_input, auto_contiguous,
avoid_copy, inverse, real, normalise_idft=True, ortho=False):
'''Generic transform interface for all the transforms. No
defaults exist. The transform must be specified exactly.
'''
a_orig = a
invreal = inverse and real
if inverse:
direction = 'FFTW_BACKWARD'
else:
direction = 'FFTW_FORWARD'
if planner_effort not in _valid_efforts:
raise ValueError('Invalid planner effort: ', planner_effort)
s, axes = _cook_nd_args(a, s, axes, invreal)
input_shape, output_shape = _compute_array_shapes(
a, s, axes, inverse, real)
a_is_complex = numpy.iscomplexobj(a)
# Make the input dtype correct by transforming to an available type
if a.dtype.char not in _rc_dtype_pairs:
dtype = _default_dtype
if a.dtype == numpy.dtype('float16') and '32' in pyfftw._supported_types:
# convert half-precision to single precision, if available
dtype = numpy.dtype('float32')
# warn when losing precision but not when using a higher precision
if dtype.itemsize < a.dtype.itemsize:
warnings.warn("Narrowing conversion from %s to %s precision" % (a.dtype, dtype))
if not real or inverse:
# It's going to be complex
dtype = numpy.dtype(_rc_dtype_pairs[dtype.char])
# finally convert the input array
a = numpy.asarray(a, dtype=dtype)
elif not (real and not inverse) and not a_is_complex:
# We need to make it a complex dtype
a = numpy.asarray(a, dtype=_rc_dtype_pairs[a.dtype.char])
elif (real and not inverse) and a_is_complex:
# It should be real
a = numpy.asarray(a, dtype=_rc_dtype_pairs[a.dtype.char])
# Make the output dtype correct
if not real:
output_dtype = a.dtype
else:
output_dtype = _rc_dtype_pairs[a.dtype.char]
if not avoid_copy:
a_copy = a.copy()
output_array = pyfftw.empty_aligned(output_shape, output_dtype)
flags = [planner_effort]
if not auto_align_input:
flags.append('FFTW_UNALIGNED')
if overwrite_input:
flags.append('FFTW_DESTROY_INPUT')
if not a.shape == input_shape:
if avoid_copy:
raise ValueError('Cannot avoid copy: '
'The transform shape is not the same as the array size. '
'(from avoid_copy flag)')
# This means we need to use an _FFTWWrapper object
# and so need to create slicers.
update_input_array_slicer, FFTW_array_slicer = (
_setup_input_slicers(a.shape, input_shape))
# Also, the input array will be a different shape to the shape of
# `a`, so we need to create a new array.
input_array = pyfftw.empty_aligned(input_shape, a.dtype)
FFTW_object = _FFTWWrapper(input_array, output_array, axes, direction,
flags, threads, input_array_slicer=update_input_array_slicer,
FFTW_array_slicer=FFTW_array_slicer,
normalise_idft=normalise_idft, ortho=ortho)
# We copy the data back into the internal FFTW object array
internal_array = FFTW_object.input_array
internal_array[:] = 0
internal_array[FFTW_array_slicer] = (
a_copy[update_input_array_slicer])
else:
# Otherwise we can use `a` as-is
input_array = a
if auto_contiguous:
# We only need to create a new array if it's not already
# contiguous
if not (a.flags['C_CONTIGUOUS'] or a.flags['F_CONTIGUOUS']):
if avoid_copy:
raise ValueError('Cannot avoid copy: '
'The input array is not contiguous and '
'auto_contiguous is set. (from avoid_copy flag)')
input_array = pyfftw.empty_aligned(a.shape, a.dtype)
if (auto_align_input and not pyfftw.is_byte_aligned(input_array)):
if avoid_copy:
raise ValueError('Cannot avoid copy: '
'The input array is not aligned and '
'auto_align is set. (from avoid_copy flag)')
input_array = pyfftw.byte_align(input_array)
FFTW_object = pyfftw.FFTW(input_array, output_array, axes, direction,
flags, threads, normalise_idft=normalise_idft, ortho=ortho)
if not avoid_copy:
# Copy the data back into the (likely) destroyed array
FFTW_object.input_array[:] = a_copy
return FFTW_object
