'''
Wrapper for the MKL FFT routines.
Inspiration from:
http://stackoverflow.com/questions/11752898/threaded-fft-in-enthought-python
'''
import numpy as np
import ctypes as _ctypes
import os
from dftidefs import *
def load_libmkl():
if os.name == 'posix':
try:
lib_mkl = os.getenv('LIBMKL')
return _ctypes.cdll.LoadLibrary(lib_mkl)
except:
pass
try:
return _ctypes.cdll.LoadLibrary("libmkl_rt.dylib")
except:
raise ValueError('MKL Library not found')
else:
try:
return _ctypes.cdll.LoadLibrary("mk2_rt.dll")
except:
raise ValueError('MKL Library not found')
mkl = load_libmkl()
def mkl_rfft(a, n=None, axis=-1, norm=None, direction='forward', out=None):
'''
Forward one-dimensional double-precision real-complex FFT.
Uses the Intel MKL libraries distributed with Anaconda Python.
Normalisation is different from Numpy!
By default, allocates new memory like 'a' for output data.
Returns the array containing output data.
'''
if axis == -1:
axis = a.ndim-1
# This code only works for 1D and 2D arrays
assert a.ndim < 3
assert (axis < a.ndim and axis >= -1)
assert (direction == 'forward' or direction == 'backward')
if direction == 'forward':
assert a.dtype == np.float32 or a.dtype == np.float64
else:
assert a.dtype == np.complex64 or a.dtype == np.complex128
order = 'C'
if a.flags['F_CONTIGUOUS'] and not a.flags['C_CONTIGUOUS']:
order = 'F'
# Add zero padding or truncate if needed (incurs memory copy)
if n is not None:
m = n if direction == 'forward' else (n//2 + 1)
if a.shape[axis] < m:
# pad axis with zeros
pad_width = np.zeros((a.ndim, 2), dtype=np.int)
pad_width[axis,1] = m - a.shape[axis]
a = np.pad(a, pad_width, mode='constant')
elif a.shape[axis] > m:
# truncate along axis
b = np.swapaxes(a, axis, 0)[:m,]
a = np.swapaxes(b, 0, axis).copy()
elif direction == 'forward':
n = a.shape[axis]
elif direction == 'backward':
n = 2*(a.shape[axis]-1)
# determine output type
if direction == 'backward':
out_type = np.float64
if a.dtype == np.complex64:
out_type = np.float32
elif direction == 'forward':
out_type = np.complex128
if a.dtype == np.float32:
out_type = np.complex64
# Configure output array
assert a is not out
if out is not None:
assert out.dtype == out_type
for i in xrange(a.ndim):
if i != axis:
assert a.shape[i] == out.shape[i]
if direction == 'forward':
assert (n//2 + 1) == out.shape[axis]
else:
assert out.shape[axis] == n
assert not np.may_share_memory(a, out)
else:
size = list(a.shape)
size[axis] = n//2 + 1 if direction == 'forward' else n
out = np.empty(size, dtype=out_type, order=order)
# Define length, number of transforms strides
length = _ctypes.c_int(n)
n_transforms = _ctypes.c_int(np.prod(a.shape)/a.shape[axis])
# For strides, the C type used *must* be int64
strides = (_ctypes.c_int64*2)(0, a.strides[axis]/a.itemsize)
if a.ndim == 2:
if axis == 0:
distance = _ctypes.c_int(a.strides[1]/a.itemsize)
out_distance = _ctypes.c_int(out.strides[1]/out.itemsize)
else:
distance = _ctypes.c_int(a.strides[0]/a.itemsize)
out_distance = _ctypes.c_int(out.strides[0]/out.itemsize)
double_precision = True
if (direction == 'forward' and a.dtype == np.float32) or (direction == 'backward' and a.dtype == np.complex64):
double_precision = False
# Create the description handle
Desc_Handle = _ctypes.c_void_p(0)
if not double_precision:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_SINGLE, DFTI_REAL, _ctypes.c_int(1), length)
else:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_DOUBLE, DFTI_REAL, _ctypes.c_int(1), length)
# set the storage type
mkl.DftiSetValue(Desc_Handle, DFTI_CONJUGATE_EVEN_STORAGE, DFTI_COMPLEX_COMPLEX)
# set normalization factor
if norm == 'ortho':
if not double_precision:
scale = _ctypes.c_double(1/np.sqrt(n))
else:
scale = _ctypes.c_double(1/np.sqrt(n))
mkl.DftiSetValue(Desc_Handle, DFTI_FORWARD_SCALE, scale)
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
elif norm is None:
if not double_precision:
scale = _ctypes.c_double(1./n)
else:
scale = _ctypes.c_double(1./n)
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
# set all values if necessary
if a.ndim != 1:
mkl.DftiSetValue(Desc_Handle, DFTI_NUMBER_OF_TRANSFORMS, n_transforms)
mkl.DftiSetValue(Desc_Handle, DFTI_INPUT_DISTANCE, distance)
mkl.DftiSetValue(Desc_Handle, DFTI_OUTPUT_DISTANCE, out_distance)
mkl.DftiSetValue(Desc_Handle, DFTI_INPUT_STRIDES, _ctypes.byref(strides))
mkl.DftiSetValue(Desc_Handle, DFTI_OUTPUT_STRIDES, _ctypes.byref(strides))
if direction == 'forward':
fft_func = mkl.DftiComputeForward
elif direction == 'backward':
fft_func = mkl.DftiComputeBackward
else:
assert False
# Not-in-place FFT
mkl.DftiSetValue(Desc_Handle, DFTI_PLACEMENT, DFTI_NOT_INPLACE)
mkl.DftiCommitDescriptor(Desc_Handle)
fft_func(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p), out.ctypes.data_as(_ctypes.c_void_p) )
mkl.DftiFreeDescriptor(_ctypes.byref(Desc_Handle))
return out
def mkl_fft(a, n=None, axis=-1, norm=None, direction='forward', out=None):
'''
Forward/Backward one-dimensional single/double-precision complex-complex FFT.
Uses the Intel MKL libraries distributed with Anaconda Python.
Normalisation is different from Numpy!
By default, allocates new memory like 'a' for output data.
Returns the array containing output data.
'''
# This code only works for 1D and 2D arrays
assert a.ndim < 3
assert axis < a.ndim and axis >= -1
# Add zero padding if needed (incurs memory copy)
if n is not None and n != a.shape[axis]:
pad_width = np.zeros((a.ndim, 2), dtype=np.int)
pad_width[axis,1] = n - a.shape[axis]
a = np.pad(a, pad_width, mode='constant')
if n is not None:
if a.shape[axis] < n:
# pad axis with zeros
pad_width = np.zeros((a.ndim, 2))
pad_width[axis,1] = m - a.shape[axis]
a = np.pad(x, pad_width, mode='constant')
elif a.shape[axis] > n:
# truncate along axis
b = np.swapaxes(a, axis, 0)[:m,]
a = np.swapaxes(b, 0, axis).copy()
# Convert input to complex data type if real (also memory copy)
if a.dtype != np.complex128 and a.dtype != np.complex64:
if a.dtype == np.int64 or a.dtype == np.uint64 or a.dtype == np.float64:
a = np.array(a, dtype=np.complex128)
else:
a = np.array(a, dtype=np.complex64)
# Configure in-place vs out-of-place
inplace = False
if out is a:
inplace = True
elif out is not None:
assert out.dtype == a.dtype
assert a.shape == out.shape
assert not np.may_share_memory(a, out)
else:
out = np.empty_like(a)
# Define length, number of transforms strides
length = _ctypes.c_int(a.shape[axis])
n_transforms = _ctypes.c_int(np.prod(a.shape)/a.shape[axis])
# For strides, the C type used *must* be int64
strides = (_ctypes.c_int64*2)(0, a.strides[axis]/a.itemsize)
if a.ndim == 2:
if axis == 0:
distance = _ctypes.c_int(a.strides[1]/a.itemsize)
else:
distance = _ctypes.c_int(a.strides[0]/a.itemsize)
# Create the description handle
Desc_Handle = _ctypes.c_void_p(0)
if a.dtype == np.complex64:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_SINGLE, DFTI_COMPLEX, _ctypes.c_int(1), length)
elif a.dtype == np.complex128:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_DOUBLE, DFTI_COMPLEX, _ctypes.c_int(1), length)
# Set normalization factor
if norm == 'ortho':
if a.dtype == np.complex64:
scale = _ctypes.c_float(1/np.sqrt(a.shape[axis]))
else:
scale = _ctypes.c_double(1/np.sqrt(a.shape[axis]))
mkl.DftiSetValue(Desc_Handle, DFTI_FORWARD_SCALE, scale)
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
elif norm is None:
if a.dtype == np.complex64:
scale = _ctypes.c_float(1./a.shape[axis])
else:
scale = _ctypes.c_double(1./a.shape[axis])
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
# set all values if necessary
if a.ndim != 1:
mkl.DftiSetValue(Desc_Handle, DFTI_NUMBER_OF_TRANSFORMS, n_transforms)
mkl.DftiSetValue(Desc_Handle, DFTI_INPUT_DISTANCE, distance)
mkl.DftiSetValue(Desc_Handle, DFTI_OUTPUT_DISTANCE, distance)
mkl.DftiSetValue(Desc_Handle, DFTI_INPUT_STRIDES, _ctypes.byref(strides))
mkl.DftiSetValue(Desc_Handle, DFTI_OUTPUT_STRIDES, _ctypes.byref(strides))
if direction == 'forward':
fft_func = mkl.DftiComputeForward
elif direction == 'backward':
fft_func = mkl.DftiComputeBackward
else:
assert False
if inplace:
# In-place FFT
mkl.DftiCommitDescriptor(Desc_Handle)
fft_func(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p) )
else:
# Not-in-place FFT
mkl.DftiSetValue(Desc_Handle, DFTI_PLACEMENT, DFTI_NOT_INPLACE)
mkl.DftiCommitDescriptor(Desc_Handle)
fft_func(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p), out.ctypes.data_as(_ctypes.c_void_p) )
mkl.DftiFreeDescriptor(_ctypes.byref(Desc_Handle))
return out
def mkl_fft2(a, norm=None, direction='forward', out=None):
'''
Forward two-dimensional double-precision complex-complex FFT.
Uses the Intel MKL libraries distributed with Enthought Python.
Normalisation is different from Numpy!
By default, allocates new memory like 'a' for output data.
Returns the array containing output data.
'''
# convert input to complex data type if real (also memory copy)
if a.dtype != np.complex128 and a.dtype != np.complex64:
if a.dtype == np.int64 or a.dtype == np.uint64 or a.dtype == np.float64:
a = np.array(a, dtype=np.complex128)
else:
a = np.array(a, dtype=np.complex64)
# Configure in-place vs out-of-place
inplace = False
if out is a:
inplace = True
elif out is not None:
assert out.dtype == a.dtype
assert a.shape == out.shape
assert not np.may_share_memory(a, out)
else:
out = np.empty_like(a)
# Create the description handle
Desc_Handle = _ctypes.c_void_p(0)
dims = (_ctypes.c_int64*2)(*a.shape)
if a.dtype == np.complex64:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_SINGLE, DFTI_COMPLEX, _ctypes.c_int(2), dims)
elif a.dtype == np.complex128:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_DOUBLE, DFTI_COMPLEX, _ctypes.c_int(2), dims)
# Set normalization factor
if norm == 'ortho':
if a.dtype == np.complex64:
scale = _ctypes.c_double(1.0/np.sqrt(np.prod(a.shape)))
else:
scale = _ctypes.c_double(1.0/np.sqrt(np.prod(a.shape)))
mkl.DftiSetValue(Desc_Handle, DFTI_FORWARD_SCALE, scale)
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
elif norm is None:
if a.dtype == np.complex64:
scale = _ctypes.c_double(1.0/np.prod(a.shape))
else:
scale = _ctypes.c_double(1.0/np.prod(a.shape))
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
# Set input strides if necessary
if not a.flags['C_CONTIGUOUS']:
in_strides = (_ctypes.c_int*3)(0, a.strides[0]/a.itemsize, a.strides[1]/a.itemsize)
mkl.DftiSetValue(Desc_Handle, DFTI_INPUT_STRIDES, _ctypes.byref(in_strides))
if direction == 'forward':
fft_func = mkl.DftiComputeForward
elif direction == 'backward':
fft_func = mkl.DftiComputeBackward
else:
assert False
if inplace:
# In-place FFT
mkl.DftiCommitDescriptor(Desc_Handle)
fft_func(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p) )
else:
# Not-in-place FFT
mkl.DftiSetValue(Desc_Handle, DFTI_PLACEMENT, DFTI_NOT_INPLACE)
# Set output strides if necessary
if not out.flags['C_CONTIGUOUS']:
out_strides = (_ctypes.c_int*3)(0, out.strides[0]/out.itemsize, out.strides[1]/out.itemsize)
mkl.DftiSetValue(Desc_Handle, DFTI_OUTPUT_STRIDES, _ctypes.byref(out_strides))
mkl.DftiCommitDescriptor(Desc_Handle)
fft_func(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p), out.ctypes.data_as(_ctypes.c_void_p) )
mkl.DftiFreeDescriptor(_ctypes.byref(Desc_Handle))
return out
def rfft(a, n=None, axis=-1, norm=None, out=None):
return mkl_rfft(a, n=n, axis=axis, norm=norm, direction='forward', out=out)
def irfft(a, n=None, axis=-1, norm=None, out=None):
return mkl_rfft(a, n=n, axis=axis, norm=norm, direction='backward', out=out)
def fft(a, n=None, axis=-1, norm=None, out=None):
return mkl_fft(a, n=n, axis=axis, norm=norm, direction='forward', out=out)
def ifft(a, n=None, axis=-1, norm=None, out=None):
return mkl_fft(a, n=n, axis=axis, norm=norm, direction='backward', out=out)
def fft2(a, norm=None, out=None):
return mkl_fft2(a, norm=norm, direction='forward', out=out)
def ifft2(a, norm=None, out=None):
return mkl_fft2(a, norm=norm, direction='backward', out=out)
def cce2full(A):
# Assume all square for now
N = A.shape
N_half = N[0]//2 + 1
out = np.empty((A.shape[0], A.shape[0]), dtype=A.dtype)
out[:, :N_half] = A
out[1:, N_half:] = np.rot90(A[1:, 1:-1], 2).conj()
# Complete the first row
out[0, N_half:] = A[0, -2:0:-1].conj()
return out
def mkl_rfft2(a, norm=None, direction='forward', out=None):
'''
Forward two-dimensional double-precision complex-complex FFT.
Uses the Intel MKL libraries distributed with Enthought Python.
Normalisation is different from Numpy!
By default, allocates new memory like 'a' for output data.
Returns the array containing output data.
'''
assert (a.dtype == np.float32) or (a.dtype == np.float64)
out_type = np.complex128
if a.dtype == np.float32:
out_type = np.complex64
n = a.shape[1]
# Allocate memory if needed
if out is not None:
assert out.dtype == out_type
assert out.shape[1] == n//2 + 1
assert not np.may_share_memory(a, out)
else:
size = list(a.shape)
size[1] = n//2 + 1
out = np.empty(size, dtype=out_type)
# Create the description handle
Desc_Handle = _ctypes.c_void_p(0)
dims = (_ctypes.c_int64*2)(*a.shape)
if a.dtype == np.float32:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_SINGLE, DFTI_REAL, _ctypes.c_int(2), dims)
elif a.dtype == np.float64:
mkl.DftiCreateDescriptor(_ctypes.byref(Desc_Handle), DFTI_DOUBLE, DFTI_REAL, _ctypes.c_int(2), dims)
# Set the storage type
mkl.DftiSetValue(Desc_Handle, DFTI_CONJUGATE_EVEN_STORAGE, DFTI_COMPLEX_COMPLEX)
# Set normalization factor
if norm == 'ortho':
if a.dtype == np.float32:
scale = _ctypes.c_float(1.0/np.sqrt(np.prod(a.shape)))
else:
scale = _ctypes.c_double(1.0/np.sqrt(np.prod(a.shape)))
mkl.DftiSetValue(Desc_Handle, DFTI_FORWARD_SCALE, scale)
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
elif norm is None:
if a.dtype == np.float64:
scale = _ctypes.c_float(1.0/np.prod(a.shape))
else:
scale = _ctypes.c_double(1.0/np.prod(a.shape))
mkl.DftiSetValue(Desc_Handle, DFTI_BACKWARD_SCALE, scale)
# For strides, the C type used *must* be int64
in_strides = (_ctypes.c_int64*3)(0, a.strides[0]/a.itemsize, a.strides[1]/a.itemsize)
out_strides = (_ctypes.c_int64*3)(0, out.strides[0]/out.itemsize, out.strides[1]/out.itemsize)
# mkl.DftiSetValue(Desc_Handle, DFTI_INPUT_STRIDES, _ctypes.byref(in_strides))
mkl.DftiSetValue(Desc_Handle, DFTI_OUTPUT_STRIDES, _ctypes.byref(out_strides))
if direction == 'forward':
fft_func = mkl.DftiComputeForward
elif direction == 'backward':
fft_func = mkl.DftiComputeBackward
else:
assert False
# Not-in-place FFT
mkl.DftiSetValue(Desc_Handle, DFTI_PLACEMENT, DFTI_NOT_INPLACE)
# Set output strides if necessary
if not out.flags['C_CONTIGUOUS']:
out_strides = (_ctypes.c_int*3)(0, out.strides[0]/out.itemsize, out.strides[1]/out.itemsize)
mkl.DftiSetValue(Desc_Handle, DFTI_OUTPUT_STRIDES, _ctypes.byref(out_strides))
mkl.DftiCommitDescriptor(Desc_Handle)
fft_func(Desc_Handle, a.ctypes.data_as(_ctypes.c_void_p), out.ctypes.data_as(_ctypes.c_void_p) )
mkl.DftiFreeDescriptor(_ctypes.byref(Desc_Handle))
return out
if __name__ == "__main__":
import time
n_iter = 200
N = 256
np.seterr(all='raise')
A = np.complex64(np.random.randn(N, N))
C = np.zeros((N, N), dtype='complex64')
start_time = time.time()
for i in range(n_iter):
C += np.fft.fft(A)
print("--- %s seconds ---" % (time.time() - start_time))
A = np.complex64(np.random.randn(N, N))
C = np.zeros((N, N), dtype='complex64')
start_time = time.time()
for i in range(n_iter):
C += fft(A)
print("--- %s seconds ---" % (time.time() - start_time))
A = np.float32(np.random.randn(N, N))
C = np.zeros((N, N//2+1), dtype='complex64')
start_time = time.time()
for i in range(n_iter):
C += mkl_rfft(A)
print("--- %s seconds ---" % (time.time() - start_time))
