Source code for mmf.math.fft

r"""Wrappers for the FFT.  Uses the :pkg:`anfft` package if available (which
wraps the FFTW3), otherwise uses :mod:`numpy.fft`."""

import functools
import warnings

__all__ = ['fft', 'ifft', 'fftn', 'ifftn', 'rfft', 'irfft', 'rfftn', 'irfftn',
           'fftshift', 'ifftshift', 'fftfreq']

__ANFFT = False
from numpy.fft import fftfreq, fftshift, ifftshift

try:
    import anfft
    __ANFFT = True
    __MEASURE = True
except ImportError, err:
    warnings.warn("Could not import `anfft`.  Using `numpy.fft`: " + str(err)) 
    import numpy.fft

# Wrap all __ANFFT functions to pass measure=True.
def _wrapfft(f):
    @functools.wraps(f)
    def fnew(a, *v, **kw):
        if __ANFFT:
            kw.setdefault('measure', __MEASURE)
            mod = anfft
        else:
            mod = numpy.fft
        return getattr(mod, f.__name__)(a, **kw)
    return fnew

@_wrapfft
def fft(): pass

@_wrapfft
def ifft(): pass

@_wrapfft
def rfft(): pass

@_wrapfft
def irfft(): pass

# Wrap all *fftn functions to allow `k` to be specified (as per anfft) or to
# fallback to numpy routines if `s` or `axes arguments are provided.
def _wrapfftn(f):
    @functools.wraps(f)
    def fnew(a, *v, **kw):
        if 'k' in kw and not __ANFFT:
            kw['axes'] = range(-kw.pop('k'), 0)
        if __ANFFT and not v and not 'axes' in kw and not 's' in kw:
            kw.setdefault('measure', __MEASURE)
            mod = anfft
        else:
            mod = numpy.fft
        return getattr(mod, f.__name__)(a, **kw)
    return fnew

@_wrapfftn
def fftn(): pass

@_wrapfftn
def ifftn(): pass

@_wrapfftn
def rfftn(): pass

@_wrapfftn
def irfftn(): pass

del warnings, _wrapfft, _wrapfftn