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_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 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 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_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