""" Discrete Fourier Transforms - FFT.py The underlying code for these functions is an f2c translated and modified version of the FFTPACK routines. fft(a, n=None, axis=-1) ifft(a, n=None, axis=-1) rfft(a, n=None, axis=-1) irfft(a, n=None, axis=-1) hfft(a, n=None, axis=-1) ihfft(a, n=None, axis=-1) fftn(a, s=None, axes=None) ifftn(a, s=None, axes=None) rfftn(a, s=None, axes=None) irfftn(a, s=None, axes=None) fft2(a, s=None, axes=(-2,-1)) ifft2(a, s=None, axes=(-2, -1)) rfft2(a, s=None, axes=(-2,-1)) irfft2(a, s=None, axes=(-2, -1)) """ __all__ = ['fft','ifft', 'rfft', 'irfft', 'hfft', 'ihfft', 'rfftn', 'irfftn', 'rfft2', 'irfft2', 'fft2', 'ifft2', 'fftn', 'ifftn', 'refft', 'irefft','refftn','irefftn', 'refft2', 'irefft2'] from numpy.core import asarray, zeros, swapaxes, shape, conjugate, \ take import fftpack_lite as fftpack from helper import * _fft_cache = {} _real_fft_cache = {} def _raw_fft(a, n=None, axis=-1, init_function=fftpack.cffti, work_function=fftpack.cfftf, fft_cache = _fft_cache ): a = asarray(a) if n is None: n = a.shape[axis] if n < 1: raise ValueError("Invalid number of FFT data points (%d) specified." % n) try: wsave = fft_cache[n] except(KeyError): wsave = init_function(n) fft_cache[n] = wsave if a.shape[axis] != n: s = list(a.shape) if s[axis] > n: index = [slice(None)]*len(s) index[axis] = slice(0,n) a = a[index] else: index = [slice(None)]*len(s) index[axis] = slice(0,s[axis]) s[axis] = n z = zeros(s, a.dtype.char) z[index] = a a = z if axis != -1: a = swapaxes(a, axis, -1) r = work_function(a, wsave) if axis != -1: r = swapaxes(r, axis, -1) return r def fft(a, n=None, axis=-1): """ Compute the one dimensional fft on a given axis. Return the n point discrete Fourier transform of a. n defaults to the length of a. If n is larger than the length of a, then a will be zero-padded to make up the difference. If n is smaller than the length of a, only the first n items in a will be used. Parameters ---------- a : array input array n : int length of the fft axis : int axis over which to compute the fft Notes ----- The packing of the result is "standard": If A = fft(a, n), then A[0] contains the zero-frequency term, A[1:n/2+1] contains the positive-frequency terms, and A[n/2+1:] contains the negative-frequency terms, in order of decreasingly negative frequency. So for an 8-point transform, the frequencies of the result are [ 0, 1, 2, 3, 4, -3, -2, -1]. This is most efficient for n a power of two. This also stores a cache of working memory for different sizes of fft's, so you could theoretically run into memory problems if you call this too many times with too many different n's. """ return _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftf, _fft_cache) def ifft(a, n=None, axis=-1): """ Compute the one-dimensonal inverse fft on a given axis. Return the n point inverse discrete Fourier transform of a. n defaults to the length of a. If n is larger than the length of a, then a will be zero-padded to make up the difference. If n is smaller than the length of a, then a will be truncated to reduce its size. Parameters ---------- a : array input array n : int length of the fft axis : int axis over which to compute the inverse fft Notes ----- The input array is expected to be packed the same way as the output of fft, as discussed in the fft documentation. This is the inverse of fft: ifft(fft(a)) == a within numerical accuracy. This is most efficient for n a power of two. This also stores a cache of working memory for different sizes of fft's, so you could theoretically run into memory problems if you call this too many times with too many different n's. """ a = asarray(a).astype(complex) if n is None: n = shape(a)[axis] return _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftb, _fft_cache) / n def rfft(a, n=None, axis=-1): """ Compute the one-dimensional fft for real input. Return the n point discrete Fourier transform of the real valued array a. n defaults to the length of a. n is the length of the input, not the output. Parameters ---------- a : array input array with real data type n : int length of the fft axis : int axis over which to compute the fft Notes ----- The returned array will be the nonnegative frequency terms of the Hermite-symmetric, complex transform of the real array. So for an 8-point transform, the frequencies in the result are [ 0, 1, 2, 3, 4]. The first term will be real, as will the last if n is even. The negative frequency terms are not needed because they are the complex conjugates of the positive frequency terms. (This is what I mean when I say Hermite-symmetric.) This is most efficient for n a power of two. """ a = asarray(a).astype(float) return _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftf, _real_fft_cache) def irfft(a, n=None, axis=-1): """ Compute the one-dimensional inverse fft for real input. Notes ----- Return the real valued n point inverse discrete Fourier transform of a, where a contains the nonnegative frequency terms of a Hermite-symmetric sequence. n is the length of the result, not the input. If n is not supplied, the default is 2*(len(a)-1). If you want the length of the result to be odd, you have to say so. Parameters ---------- a : array input array with real data type n : int length of the fft axis : int axis over which to compute the fft Notes ----- If you specify an n such that a must be zero-padded or truncated, the extra/removed values will be added/removed at high frequencies. One can thus resample a series to m points via Fourier interpolation by: a_resamp = irfft(rfft(a), m). This is the inverse of rfft: irfft(rfft(a), len(a)) == a within numerical accuracy. """ a = asarray(a).astype(complex) if n is None: n = (shape(a)[axis] - 1) * 2 return _raw_fft(a, n, axis, fftpack.rffti, fftpack.rfftb, _real_fft_cache) / n def hfft(a, n=None, axis=-1): """ Compute the fft of a signal which spectrum has Hermitian symmetry. Parameters ---------- a : array input array n : int length of the hfft axis : int axis over which to compute the hfft Notes ----- These are a pair analogous to rfft/irfft, but for the opposite case: here the signal is real in the frequency domain and has Hermite symmetry in the time domain. So here it's hermite_fft for which you must supply the length of the result if it is to be odd. ihfft(hfft(a), len(a)) == a within numerical accuracy. See also -------- rfft ihfft """ a = asarray(a).astype(complex) if n is None: n = (shape(a)[axis] - 1) * 2 return irfft(conjugate(a), n, axis) * n def ihfft(a, n=None, axis=-1): """ Compute the inverse fft of a signal which spectrum has Hermitian symmetry. Parameters ---------- a : array input array n : int length of the ihfft axis : int axis over which to compute the ihfft Notes ----- These are a pair analogous to rfft/irfft, but for the opposite case: here the signal is real in the frequency domain and has Hermite symmetry in the time domain. So here it's hermite_fft for which you must supply the length of the result if it is to be odd. ihfft(hfft(a), len(a)) == a within numerical accuracy. See also -------- rfft hfft """ a = asarray(a).astype(float) if n is None: n = shape(a)[axis] return conjugate(rfft(a, n, axis))/n def _cook_nd_args(a, s=None, axes=None, invreal=0): if s is None: shapeless = 1 if axes is None: s = list(a.shape) else: s = take(a.shape, axes) else: shapeless = 0 s = list(s) if axes is None: axes = range(-len(s), 0) if len(s) != len(axes): raise ValueError, "Shape and axes have different lengths." if invreal and shapeless: s[axes[-1]] = (s[axes[-1]] - 1) * 2 return s, axes def _raw_fftnd(a, s=None, axes=None, function=fft): a = asarray(a) s, axes = _cook_nd_args(a, s, axes) itl = range(len(axes)) itl.reverse() for ii in itl: a = function(a, n=s[ii], axis=axes[ii]) return a def fftn(a, s=None, axes=None): """ Compute the N-dimensional Fast Fourier Transform. Parameters ---------- a : array_like Input array. s : sequence of ints Shape of each axis of the input (s[0] refers to axis 0, s[1] to axis 1, etc.). This corresponds to `n` for `fft(x, n)`. Along any axis, if the given shape is smaller than that of the input, the input is cropped. If it is larger, the input is padded with zeros. axes : tuple of int Axes over which to compute the FFT. Notes ----- Analogously to `fft`, the term for zero frequency in all axes is in the low-order corner, while the term for the Nyquist frequency in all axes is in the middle. If neither `s` nor `axes` is specified, the transform is taken along all axes. If `s` is specified and `axes` is not, the last ``len(s)`` axes are used. If `axes` is specified and `s` is not, the input shape along the specified axes is used. If `s` and `axes` are both specified and are not the same length, an exception is raised. """ return _raw_fftnd(a,s,axes,fft) def ifftn(a, s=None, axes=None): """ Compute the inverse of fftn. Parameters ---------- a : array input array s : sequence (int) shape of the ifft axis : int axis over which to compute the ifft Notes ----- The n-dimensional ifft of a. s is a sequence giving the shape of the input an result along the transformed axes, as n for fft. Results are packed analogously to fft: the term for zero frequency in all axes is in the low-order corner, while the term for the Nyquist frequency in all axes is in the middle. If neither s nor axes is specified, the transform is taken along all axes. If s is specified and axes is not, the last len(s) axes are used. If axes are specified and s is not, the input shape along the specified axes is used. If s and axes are both specified and are not the same length, an exception is raised. """ return _raw_fftnd(a, s, axes, ifft) def fft2(a, s=None, axes=(-2,-1)): """ Compute the 2d fft of an array. Parameters ---------- a : array input array s : sequence (int) shape of the fft axis : int axis over which to compute the fft Notes ----- This is really just fftn with different default behavior. """ return _raw_fftnd(a,s,axes,fft) def ifft2(a, s=None, axes=(-2,-1)): """ Compute the inverse 2d fft of an array. Parameters ---------- a : array input array s : sequence (int) shape of the ifft axis : int axis over which to compute the ifft Notes ----- This is really just ifftn with different default behavior. """ return _raw_fftnd(a, s, axes, ifft) def rfftn(a, s=None, axes=None): """ Compute the n-dimensional fft of a real array. Parameters ---------- a : array (real) input array s : sequence (int) shape of the fft axis : int axis over which to compute the fft Notes ----- A real transform as rfft is performed along the axis specified by the last element of axes, then complex transforms as fft are performed along the other axes. """ a = asarray(a).astype(float) s, axes = _cook_nd_args(a, s, axes) a = rfft(a, s[-1], axes[-1]) for ii in range(len(axes)-1): a = fft(a, s[ii], axes[ii]) return a def rfft2(a, s=None, axes=(-2,-1)): """ Compute the 2-dimensional fft of a real array. Parameters ---------- a : array (real) input array s : sequence (int) shape of the fft axis : int axis over which to compute the fft Notes ----- The 2-D fft of the real valued array a. This is really just rfftn with different default behavior. """ return rfftn(a, s, axes) def irfftn(a, s=None, axes=None): """ Compute the n-dimensional inverse fft of a real array. Parameters ---------- a : array (real) input array s : sequence (int) shape of the inverse fft axis : int axis over which to compute the inverse fft Notes ----- The transform implemented in ifftn is applied along all axes but the last, then the transform implemented in irfft is performed along the last axis. As with irfft, the length of the result along that axis must be specified if it is to be odd. """ a = asarray(a).astype(complex) s, axes = _cook_nd_args(a, s, axes, invreal=1) for ii in range(len(axes)-1): a = ifft(a, s[ii], axes[ii]) a = irfft(a, s[-1], axes[-1]) return a def irfft2(a, s=None, axes=(-2,-1)): """ Compute the 2-dimensional inverse fft of a real array. Parameters ---------- a : array (real) input array s : sequence (int) shape of the inverse fft axis : int axis over which to compute the inverse fft Notes ----- This is really irfftn with different default. """ return irfftn(a, s, axes) # Deprecated names from numpy import deprecate refft = deprecate(rfft, 'refft', 'rfft') irefft = deprecate(irfft, 'irefft', 'irfft') refft2 = deprecate(rfft2, 'refft2', 'rfft2') irefft2 = deprecate(irfft2, 'irefft2', 'irfft2') refftn = deprecate(rfftn, 'refftn', 'rfftn') irefftn = deprecate(irfftn, 'irefftn', 'irfftn')