def test_direction_property(self):
'''Test to see if the direction property returns the correct thing
'''
self.assertEqual(self.fft.direction, 'FFTW_FORWARD')
new_fft = FFTW(self.input_array,
self.output_array,
direction='FFTW_BACKWARD')
self.assertEqual(new_fft.direction, 'FFTW_BACKWARD')
def test_ortho_property(self):
'''ortho property defaults to False
'''
self.assertEqual(self.fft.ortho, False)
newfft = FFTW(self.input_array,
self.output_array,
ortho=True,
normalise_idft=False)
self.assertEqual(newfft.ortho, True)
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_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 test_extra_dimension_fail(self):
in_shape = self.input_shapes['2d']
_out_shape = self.output_shapes['2d']
out_shape = (2, _out_shape[0], _out_shape[1])
axes = (1, )
a, b = self.create_test_arrays(in_shape, out_shape)
with self.assertRaisesRegex(ValueError, 'Invalid shapes'):
FFTW(a, b, direction=self.direction)
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_missized_nonfft_axes_fail(self):
in_shape = self.input_shapes['3d']
_out_shape = self.output_shapes['3d']
out_shape = (_out_shape[0], _out_shape[1] + 1, _out_shape[2])
axes = (2, )
a, b = self.create_test_arrays(in_shape, out_shape)
with self.assertRaisesRegex(ValueError, 'Invalid shapes'):
FFTW(a, b, direction=self.direction)
def test_default_args(self):
in_shape = self.input_shapes['2d']
out_shape = self.output_shapes['2d']
a, b = self.create_test_arrays(in_shape, out_shape)
# default args should fail for backwards transforms
# (as the default is FFTW_FORWARD)
with self.assertRaisesRegex(ValueError, 'Invalid direction'):
FFTW(a, b)
def test_default_args(self):
in_shape = self.input_shapes['2d']
out_shape = self.output_shapes['2d']
a, b = self.create_test_arrays(in_shape, out_shape)
fft = FFTW(a, b)
fft.execute()
ref_b = self.reference_fftn(a, axes=(-1, ))
self.assertTrue(numpy.allclose(b, ref_b, rtol=1e-2, atol=1e-3))
def test_axes_property(self):
'''Test to see if the axes property returns the correct thing
'''
self.assertEqual(self.fft.axes, (1,))
new_fft = FFTW(self.input_array, self.output_array, axes=(-1, -2))
self.assertEqual(new_fft.axes, (1, 0))
new_fft = FFTW(self.input_array, self.output_array, axes=(-2, -1))
self.assertEqual(new_fft.axes, (0, 1))
new_fft = FFTW(self.input_array, self.output_array, axes=(1, 0))
self.assertEqual(new_fft.axes, (1, 0))
new_fft = FFTW(self.input_array, self.output_array, axes=(1,))
self.assertEqual(new_fft.axes, (1,))
new_fft = FFTW(self.input_array, self.output_array, axes=(0,))
self.assertEqual(new_fft.axes, (0,))
def test_call_with_normalisation_off(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=False)
_input_array /= ifft.N
self.assertTrue(numpy.allclose(self.input_array, _input_array))
def test_output_dtype(self):
'''Test to see if the output_dtype property returns the correct thing
'''
self.assertEqual(self.fft.output_dtype, self.output_array.dtype)
new_input_array = numpy.complex64(self.input_array)
new_output_array = numpy.complex64(self.output_array)
new_fft = FFTW(new_input_array, new_output_array)
self.assertEqual(new_fft.output_dtype, new_output_array.dtype)
def test_output_strides(self):
'''Test to see if the output_shape property returns the correct thing
'''
self.assertEqual(self.fft.output_shape, self.output_array.shape)
new_input_array = self.output_array[::2, ::4]
new_output_array = self.output_array[::2, ::4]
new_fft = FFTW(new_input_array, new_output_array)
self.assertEqual(new_fft.output_shape, new_output_array.shape)
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_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_planning_time_limit(self):
in_shape = self.input_shapes['1d']
out_shape = self.output_shapes['1d']
axes = (0, )
a, b = self.create_test_arrays(in_shape, out_shape)
# run this a few times
runs = 10
t1 = time.time()
for n in range(runs):
forget_wisdom()
fft = FFTW(a, b, axes=axes)
unlimited_time = (time.time() - t1) / runs
time_limit = (unlimited_time) / 8
# Now do it again but with an upper limit on the time
t1 = time.time()
for n in range(runs):
forget_wisdom()
fft = FFTW(a, b, axes=axes, planning_timelimit=time_limit)
limited_time = (time.time() - t1) / runs
import sys
if sys.platform == 'win32':
# Give a 6x margin on windows. The timers are low
# precision and FFTW seems to take longer anyway.
# Also, we need to allow for processor contention which
# Appveyor seems prone to.
self.assertTrue(limited_time < time_limit * 6)
else:
# Otherwise have a 2x margin
self.assertTrue(limited_time < time_limit * 2)
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)
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_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 psd(buf_in, buf_out):
"""
Perform discrete fourier transforms using the FFTW library and use it to
get the power spectral density. FFTW optimizes
the fft algorithm based on the size of the arrays, with SIMD parallelized
commands. This optimization requires initialization, so this is a factory
function that returns a numba gufunc that performs the FFT. FFTW works on
fixed memory buffers, so you must tell it what memory to use ahead of time.
When using this with ProcessingChain, to ensure the correct buffers are used
call ProcessingChain.get_variable('var_name') to give it the internal memory
buffer directly (with raw_to_dsp, you can just give it the name and it will
automatically happen!). The possible dtypes for the input/outputs are:
- complex64 (size n) -> float32/float (size n)
- complex128 (size n) -> float64/double (size n)
- complex256/clongdouble (size n) -> float128/longdouble (size n)
- float32/float (size n) -> float32/float (size n/2+1)
- float64/double (size n) -> float64/double (size n/2+1)
- float128/longdouble (size n) -> float128/longdouble (size n/2+1)
"""
# build intermediate array for the dft, which will be abs'd to get the PSD
buf_dft = np.ndarray(
buf_out.shape, np.dtype('complex' + str(buf_out.dtype.itemsize * 16)))
try:
dft_fun = FFTW(buf_in, buf_dft, axes=(-1, ), direction='FFTW_FORWARD')
except ValueError:
raise ValueError("""Incompatible array types/shapes. Allowed:
- complex64 (size n) -> float32/float (size n)
- complex128 (size n) -> float64/double (size n)
- complex256/clongdouble (size n) -> float128/longdouble (size n)
- float32/float (size n) -> float32/float (size n/2+1)
- float64/double (size n) -> float64/double (size n/2+1)
- float128/longdouble (size n) -> float128/longdouble (size n/2+1)""")
typesig = 'void(' + str(buf_in.dtype) + '[:, :], ' + str(
buf_out.dtype) + '[:, :])'
sizesig = '(m, n)->(m, n)' if buf_in.shape == buf_out.shape else '(m, n),(m, l)'
@guvectorize([typesig], sizesig, forceobj=True)
def psd(wf_in, psd_out):
dft_fun(wf_in, buf_dft)
np.abs(buf_dft, psd_out)
return psd
def inv_dft(buf_in, buf_out):
"""
Perform inverse discrete fourier transforms using FFTW. FFTW optimizes
the fft algorithm based on the size of the arrays, with SIMD parallelized
commands. This optimization requires initialization, so this is a factory
function that returns a numba gufunc that performs the FFT. FFTW works on
fixed memory buffers, so you must tell it what memory to use ahead of time.
When using this with ProcessingChain, to ensure the correct buffers are used
call ProcessingChain.get_variable('var_name') to give it the internal memory
buffer directly (with raw_to_dsp, you can just give it the name and it will
automatically happen!). The possible dtypes for the input/outputs are:
- complex64 (size n/2+1) -> float32/float (size n)
- complex128 (size n/2+1) -> float64/double (size n)
- complex256/clongdouble (size n/2+1) -> float128/longdouble (size n)
- complex64 (size n) -> complex64 (size n)
- complex128 (size n) -> complex128 (size n)
- complex256/clongdouble (size n) -> complex256/clongdouble (size n)
"""
try:
idft_fun = FFTW(buf_in,
buf_out,
axes=(-1, ),
direction='FFTW_BACKWARD')
except ValueError:
raise ValueError("""Incompatible array types/shapes. Allowed:
- complex64 (size n/2+1) -> float32/float (size n)
- complex128 (size n/2+1) -> float64/double (size n)
- complex256/clongdouble (size n/2+1) -> float128/longdouble (size n)
- complex64 (size n) -> complex64 (size n)
- complex128 (size n) -> complex128 (size n)
- complex256/clongdouble (size n) -> complex256/clongdouble (size n)"""
)
typesig = 'void(' + str(buf_in.dtype) + '[:, :], ' + str(
buf_out.dtype) + '[:, :])'
sizesig = '(m, n)->(m, n)' if buf_in.shape == buf_out.shape else '(m, n),(m, l)'
@guvectorize([typesig], sizesig, forceobj=True)
def inv_dft(wf_in, dft_out):
idft_fun(wf_in, dft_out)
return inv_dft
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 cryo_epsdr(vol, samples_idx, max_d):
p = vol.shape[0]
k = vol.shape[2]
i, j = np.meshgrid(np.arange(max_d + 1), np.arange(max_d + 1))
dists = np.square(i) + np.square(j)
dsquare = np.sort(np.unique(dists[np.where(dists <= max_d**2)]))
corrs = np.zeros(len(dsquare))
corr_count = np.zeros(len(dsquare))
x = np.sqrt(dsquare)
dist_map = np.zeros(dists.shape)
for i in range(max_d + 1):
for j in range(max_d + 1):
d = i**2 + j**2
if d <= max_d**2:
idx, _ = bsearch(dsquare, d - 1e-13, d + 1e-13)
dist_map[i, j] = idx
dist_map = dist_map.astype('int') - 1
valid_dists = np.where(dist_map != -1)
mask = np.zeros((p, p))
mask[samples_idx] = 1
tmp = np.zeros((2 * p + 1, 2 * p + 1))
tmp[:p, :p] = mask
ftmp = fft2(tmp)
c = ifft2(ftmp * np.conj(ftmp))
c = c[:max_d + 1, :max_d + 1]
c = np.round(c.real).astype('int')
r = np.zeros(len(corrs))
# optimized version
vol = vol.transpose((2, 0, 1)).copy()
input_fft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
output_fft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
input_ifft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
output_ifft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
flags = ('FFTW_MEASURE', 'FFTW_UNALIGNED')
a_fft2 = FFTW(input_fft2,
output_fft2,
axes=(0, 1),
direction='FFTW_FORWARD',
flags=flags)
a_ifft2 = FFTW(input_ifft2,
output_ifft2,
axes=(0, 1),
direction='FFTW_BACKWARD',
flags=flags)
sum_s = np.zeros(output_ifft2.shape, output_ifft2.dtype)
sum_c = c * vol.shape[0]
for i in range(k):
proj = vol[i]
input_fft2[samples_idx] = proj[samples_idx]
a_fft2()
np.multiply(output_fft2, np.conj(output_fft2), out=input_ifft2)
a_ifft2()
sum_s += output_ifft2
for curr_dist in zip(valid_dists[0], valid_dists[1]):
dmidx = dist_map[curr_dist]
corrs[dmidx] += sum_s[curr_dist].real
corr_count[dmidx] += sum_c[curr_dist]
idx = np.where(corr_count != 0)[0]
r[idx] += corrs[idx] / corr_count[idx]
cnt = corr_count[idx]
idx = np.where(corr_count == 0)[0]
r[idx] = 0
x[idx] = 0
return r, x, cnt
def test_differing_aligned_arrays_update(self):
'''Test to see if the alignment code is working as expected
'''
# Start by creating arrays that are only on various byte
# alignments (4, 16 and 32)
_input_array = empty_aligned(len(self.input_array.ravel())*2+5,
dtype='float32', n=32)
_output_array = empty_aligned(len(self.output_array.ravel())*2+5,
dtype='float32', n=32)
_input_array[:] = 0
_output_array[:] = 0
input_array_4 = (
numpy.frombuffer(_input_array[1:-4].data, dtype='complex64')
.reshape(self.input_array.shape))
output_array_4 = (
numpy.frombuffer(_output_array[1:-4].data, dtype='complex64')
.reshape(self.output_array.shape))
input_array_16 = (
numpy.frombuffer(_input_array[4:-1].data, dtype='complex64')
.reshape(self.input_array.shape))
output_array_16 = (
numpy.frombuffer(_output_array[4:-1].data, dtype='complex64')
.reshape(self.output_array.shape))
input_array_32 = (
numpy.frombuffer(_input_array[:-5].data, dtype='complex64')
.reshape(self.input_array.shape))
output_array_32 = (
numpy.frombuffer(_output_array[:-5].data, dtype='complex64')
.reshape(self.output_array.shape))
input_arrays = {4: input_array_4,
16: input_array_16,
32: input_array_32}
output_arrays = {4: output_array_4,
16: output_array_16,
32: output_array_32}
alignments = (4, 16, 32)
# Test the arrays are aligned on 4 bytes...
self.assertTrue(is_byte_aligned(input_arrays[4], n=4))
self.assertTrue(is_byte_aligned(output_arrays[4], n=4))
# ...and on 16...
self.assertFalse(is_byte_aligned(input_arrays[4], n=16))
self.assertFalse(is_byte_aligned(output_arrays[4], n=16))
self.assertTrue(is_byte_aligned(input_arrays[16], n=16))
self.assertTrue(is_byte_aligned(output_arrays[16], n=16))
# ...and on 32...
self.assertFalse(is_byte_aligned(input_arrays[16], n=32))
self.assertFalse(is_byte_aligned(output_arrays[16], n=32))
self.assertTrue(is_byte_aligned(input_arrays[32], n=32))
self.assertTrue(is_byte_aligned(output_arrays[32], n=32))
if len(pyfftw.pyfftw._valid_simd_alignments) > 0:
max_align = pyfftw.pyfftw._valid_simd_alignments[0]
else:
max_align = simd_alignment
for in_align in alignments:
for out_align in alignments:
expected_align = min(in_align, out_align, max_align)
fft = FFTW(input_arrays[in_align], output_arrays[out_align])
self.assertTrue(fft.input_alignment == expected_align)
self.assertTrue(fft.output_alignment == expected_align)
for update_align in alignments:
if update_align < expected_align:
# This should fail (not aligned properly)
self.assertRaisesRegex(ValueError,
'Invalid input alignment',
fft.update_arrays,
input_arrays[update_align],
output_arrays[out_align])
self.assertRaisesRegex(ValueError,
'Invalid output alignment',
fft.update_arrays,
input_arrays[in_align],
output_arrays[update_align])
else:
# This should work (and not segfault!)
fft.update_arrays(input_arrays[update_align],
output_arrays[out_align])
fft.update_arrays(input_arrays[in_align],
output_arrays[update_align])
fft.execute()
def cryo_epsdr(image, samples_idx, max_d):
"""
Estimate the 1D isotropic autocorrelation of an image. The samples to use
are given in samples_idx. The correlation is computed up to a maximal
distance of max_d.
Parameters
----------
image : numpy.ndarray
square pxp image.
samples_idx : tuple
pixel indices to use for autocorrelation estimation.
max_d : int
Correlations are computed up to a maximal distance of max_d pixels.
Default p-1.
Returns
-------
r : numpy.ndarray
1D vector with samples of the isotropic autocorrelation function.
x : numpy.ndarray
Distaces at which the samples of the autocorrelation function are
given. A vector of the same length as R.
cnt : numpy.ndarray
Number of autocorrelation samples available for each distance.
"""
p = image.shape[0]
# Generate all possible squared distances. For a vertical shift i and
# horizontal shift j, dists(i,j) contains the corresponding isotropic
# correlation distance i^2+j^2. dsquare is then the set of all possible
# squared distances, that is, all distances that can be generated by
# integer steps in the horizontal and vertical directions.
i = np.array([[x for x in range(max_d + 1)] for x in range(max_d + 1)])
dists = i ** 2 + i.transpose() ** 2
dsquare = np.sort(np.unique(dists[np.where(dists <= max_d ** 2)]))
x = dsquare ** 0.5 # Distances at which the correlations are computed.
# Create a distance map whose value at the index (i,j) is the index in the
# array dsquare whose value is i^2+j^2. Pairs (i,j) that correspond to
# distances that are larger than max_d are indicated by (-1).
dist_map = distmap(max_d, dsquare, dists.shape)
valid_dists = np.where(dist_map != -1)
# Compute the number of terms in the expression sum_{j}x(j)x(j+d) for each
# distance d. As the correlation is two-dimensioanl, we compute for each
# sum of the form R(k1,k2)=sum_{i,j} X_{i,j} X_{i+k1,j+k2}, how many
# summands are in the in it for each (k1,k2). This is done by setting the
# participating image samples to 1 and computing autocorrelation again.
mask = np.zeros((p, p))
mask[samples_idx] = 1
tmp = np.zeros((2 * p + 1, 2 * p + 1))
tmp[:p, :p] = mask
ftmp = fft2(tmp)
c = ifft2(ftmp * np.conj(ftmp))
c = c[:max_d + 1, :max_d + 1]
c = np.round(c.real).astype('int')
r = np.zeros(len(dsquare)) # r(i) is the value of the ACF at distance x(i)
# Compute non-periodic autocorrelation of masked image with itself (mask
# all pixels that are not used to autocorrelation estimation).
input_fft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
output_fft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
input_ifft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
output_ifft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
flags = ('FFTW_MEASURE', 'FFTW_UNALIGNED')
a_fft2 = FFTW(input_fft2, output_fft2, axes=(0, 1), direction='FFTW_FORWARD', flags=flags)
a_ifft2 = FFTW(input_ifft2, output_ifft2, axes=(0, 1), direction='FFTW_BACKWARD', flags=flags)
sum_c = c
input_fft2[samples_idx] = image[samples_idx]
a_fft2()
np.multiply(output_fft2, np.conj(output_fft2), out=input_ifft2)
a_ifft2()
sum_s = output_ifft2
# Accumulate all autocorrelation values R(k1,k2) such that k1^2+k2^2=const
# (all autocorrelations of a certain distance).
# corrs(i) contains the sum of all products of the form x(j)x(j+d), where
# d=sqrt(dsquare(i)).
# corr_count is the number of pairs x(j)x(j+d) for each d.
corr_count, corrs = accumelate_corrs(len(dsquare), valid_dists, dist_map, sum_c, sum_s)
# Remove zero correlation sums (distances for which we had no samples at
# that distance)
idx = np.where(corr_count != 0)[0]
r[idx] += corrs[idx] / corr_count[idx]
cnt = corr_count[idx]
idx = np.where(corr_count == 0)[0]
r[idx] = 0
x[idx] = 0
return r, x, cnt
def __init__(self,
nfft,
D=1,
analysis_window=None,
synthesis_window=None,
transform='numpy',
axis=0,
bits=32):
self.nfft = nfft
self.D = D
self.axis = axis
if bits == 32:
time_dtype = 'float32'
freq_dtype = 'complex64'
elif bits == 64:
time_dtype = 'float64'
freq_dtype = 'complex128'
else:
raise ValueError("Invalid number of bits. Must be 32 or 64.")
if axis == 0:
self.x = np.squeeze(np.zeros((self.nfft, self.D),
dtype=time_dtype))
self.X = np.squeeze(
np.zeros((self.nfft // 2 + 1, self.D), dtype=freq_dtype))
elif axis == 1:
self.x = np.squeeze(np.zeros((self.D, self.nfft),
dtype=time_dtype))
self.X = np.squeeze(
np.zeros((self.D, self.nfft // 2 + 1), dtype=freq_dtype))
else:
raise ValueError("Invalid 'axis' option. Must be 0 or 1.")
if analysis_window is not None:
if axis == 0 and D > 1:
self.analysis_window = analysis_window[:, np.newaxis].astype(
time_dtype)
else:
self.analysis_window = analysis_window.astype(time_dtype)
else:
self.analysis_window = None
if synthesis_window is not None:
if axis == 0 and D > 1:
self.synthesis_window = synthesis_window[:, np.newaxis].astype(
time_dtype)
else:
self.synthesis_window = synthesis_window.astype(time_dtype)
else:
self.synthesis_window = None
if transform == 'fftw':
if pyfftw_available:
from pyfftw import empty_aligned, FFTW
self.transform = transform
# allocate input (real) and output for pyfftw
if self.D == 1:
self.a = empty_aligned(self.nfft, dtype=time_dtype)
self.b = empty_aligned(self.nfft // 2 + 1,
dtype=freq_dtype)
self._forward = FFTW(self.a, self.b)
self._backward = FFTW(self.b,
self.a,
direction='FFTW_BACKWARD')
else:
if axis == 0:
self.a = empty_aligned([self.nfft, self.D],
dtype=time_dtype)
self.b = empty_aligned([self.nfft // 2 + 1, self.D],
dtype=freq_dtype)
elif axis == 1:
self.a = empty_aligned([self.D, self.nfft],
dtype=time_dtype)
self.b = empty_aligned([self.D, self.nfft // 2 + 1],
dtype=freq_dtype)
self._forward = FFTW(self.a, self.b, axes=(self.axis, ))
self._backward = FFTW(self.b,
self.a,
axes=(self.axis, ),
direction='FFTW_BACKWARD')
else:
warnings.warn(
"Could not import pyfftw wrapper for fftw functions. Using numpy's rfft instead."
)
self.transform = 'numpy'
elif transform == 'mkl':
if mkl_available:
import mkl_fft
self.transform = 'mkl'
else:
warnings.warn(
"Could not import mkl wrapper. Using numpy's rfft instead."
)
self.transform = 'numpy'
else:
self.transform = 'numpy'
def generate_wisdom(self):
for each_dtype in _supported_nptypes_complex:
n = 16
a = empty_aligned((1, 1024), each_dtype, n=n)
b = empty_aligned(a.shape, dtype=a.dtype, n=n)
fft = FFTW(a, b)
def do_fft(self, inputa): outputa = np.zeros(self.fft_size, dtype=complex) fft_plan = FFTW(inputa, outputa, direction='FFTW_FORWARD') #fft_plan = FFTW.Plan(inputa, outputa, direction='forward', flags=['estimate']) fft_plan.execute() return (np.log10(np.abs(outputa)) * 20)[:self.fft_size/2]
def run_validate_fft(self,
a,
b,
axes,
fft=None,
ifft=None,
force_unaligned_data=False,
create_array_copies=True,
threads=1,
flags=('FFTW_ESTIMATE', )):
''' Run a validation of the FFTW routines for the passed pair
of arrays, a and b, and the axes argument.
a and b are assumed to be the same shape (but not necessarily
the same layout in memory).
fft and ifft, if passed, should be instantiated FFTW objects.
If force_unaligned_data is True, the flag FFTW_UNALIGNED
will be passed to the fftw routines.
The threads argument runs the validation with multiple threads.
flags is passed to the creation of the FFTW object.
'''
if create_array_copies:
# Don't corrupt the original mutable arrays
a = a.copy()
b = b.copy()
a_orig = a.copy()
flags = list(flags)
if force_unaligned_data:
flags.append('FFTW_UNALIGNED')
if fft == None:
fft = FFTW(a,
b,
axes=axes,
direction='FFTW_FORWARD',
flags=flags,
threads=threads)
else:
fft.update_arrays(a, b)
if ifft == None:
ifft = FFTW(b,
a,
axes=axes,
direction='FFTW_BACKWARD',
flags=flags,
threads=threads)
else:
ifft.update_arrays(b, a)
a[:] = a_orig
# Test the forward FFT by comparing it to the result from numpy.fft
fft.execute()
ref_b = self.reference_fftn(a, axes=axes)
# This is actually quite a poor relative error, but it still
# sometimes fails. I assume that numpy.fft has different internals
# to fftw.
self.assertTrue(numpy.allclose(b, ref_b, rtol=1e-2, atol=1e-3))
# Test the inverse FFT by comparing the result to the starting
# value (which is scaled as per FFTW being unnormalised).
ifft.execute()
# The scaling is the product of the lengths of the fft along
# the axes along which the fft is taken.
scaling = numpy.prod(numpy.array(a.shape)[list(axes)])
self.assertEqual(ifft.N, scaling)
self.assertEqual(fft.N, scaling)
self.assertTrue(
numpy.allclose(a / scaling, a_orig, rtol=1e-2, atol=1e-3))
return fft, ifft
def __init__(
self,
nfft,
D=1,
analysis_window=None,
synthesis_window=None,
transform="numpy",
axis=0,
precision="double",
bits=None,
):
self.nfft = nfft
self.D = D
self.axis = axis
if bits is not None and precision is not None:
warnings.warn(
'Deprecated keyword "bits" ignored in favor of new keyword "precision"',
DeprecationWarning,
)
elif bits is not None and precision is None:
warnings.warn(
'Keyword "bits" is deprecated and has been replaced by "precision"'
)
if bits == 32:
precision = "single"
elif bits == 64:
precision = "double"
if (
precision == np.float32
or precision == np.complex64
or precision == "single"
):
time_dtype = np.float32
freq_dtype = np.complex64
else:
time_dtype = np.float64
freq_dtype = np.complex128
if axis == 0:
self.x = np.squeeze(np.zeros((self.nfft, self.D), dtype=time_dtype))
self.X = np.squeeze(
np.zeros((self.nfft // 2 + 1, self.D), dtype=freq_dtype)
)
elif axis == 1:
self.x = np.squeeze(np.zeros((self.D, self.nfft), dtype=time_dtype))
self.X = np.squeeze(
np.zeros((self.D, self.nfft // 2 + 1), dtype=freq_dtype)
)
else:
raise ValueError("Invalid 'axis' option. Must be 0 or 1.")
if analysis_window is not None:
if axis == 0 and D > 1:
self.analysis_window = analysis_window[:, np.newaxis].astype(time_dtype)
else:
self.analysis_window = analysis_window.astype(time_dtype)
else:
self.analysis_window = None
if synthesis_window is not None:
if axis == 0 and D > 1:
self.synthesis_window = synthesis_window[:, np.newaxis].astype(
time_dtype
)
else:
self.synthesis_window = synthesis_window.astype(time_dtype)
else:
self.synthesis_window = None
if transform == "fftw":
if pyfftw_available:
from pyfftw import empty_aligned, FFTW
self.transform = transform
# allocate input (real) and output for pyfftw
if self.D == 1:
self.a = empty_aligned(self.nfft, dtype=time_dtype)
self.b = empty_aligned(self.nfft // 2 + 1, dtype=freq_dtype)
self._forward = FFTW(self.a, self.b)
self._backward = FFTW(self.b, self.a, direction="FFTW_BACKWARD")
else:
if axis == 0:
self.a = empty_aligned([self.nfft, self.D], dtype=time_dtype)
self.b = empty_aligned(
[self.nfft // 2 + 1, self.D], dtype=freq_dtype
)
elif axis == 1:
self.a = empty_aligned([self.D, self.nfft], dtype=time_dtype)
self.b = empty_aligned(
[self.D, self.nfft // 2 + 1], dtype=freq_dtype
)
self._forward = FFTW(self.a, self.b, axes=(self.axis,))
self._backward = FFTW(
self.b, self.a, axes=(self.axis,), direction="FFTW_BACKWARD"
)
else:
warnings.warn(
"Could not import pyfftw wrapper for fftw functions. Using numpy's rfft instead."
)
self.transform = "numpy"
elif transform == "mkl":
if mkl_available:
import mkl_fft
self.transform = "mkl"
else:
warnings.warn(
"Could not import mkl wrapper. Using numpy's rfft instead."
)
self.transform = "numpy"
else:
self.transform = "numpy"
