def _raw_fftnd(x, s, axes, direction, overwrite_x, work_function):
""" Internal auxiliary function for fftnd, ifftnd."""
if s is None:
s = x.shape
s = tuple(s)
if s!=x.shape:
assert len(s)<=len(x.shape)
for i in range(-len(s),0):
if x.shape[i]!=s[i]:
x = _fix_shape(x,s[i],i)
if axes is None:
return work_function(x,s,direction,overwrite_x=overwrite_x)
#XXX: should we allow/check for repeated indices in axes?
# If allowed then using it has an effect of reducing the shape
# implicitly.
s = [x.shape[i] for i in axes]
s = [1]*(len(x.shape)-len(s)) + s
swaps = []
state = range(-len(s),0)
for i in range(-1,-len(axes)-1,-1):
j = state[axes[i]]
if i==j: continue
state[i],state[j]=state[j],state[i]
swaps.append((j,i))
x = swapaxes(x, j,i)
r = work_function(x,s,direction,overwrite_x=overwrite_x)
swaps.reverse()
for i,j in swaps:
r = swapaxes(r,i,j)
return r
def fft(x, n=None, axis=-1, overwrite_x=0):
""" fft(x, n=None, axis=-1, overwrite_x=0) -> y
Return discrete Fourier transform of arbitrary type sequence x.
The returned complex array contains
[y(0),y(1),..,y(n/2-1),y(-n/2),...,y(-1)] if n is even
[y(0),y(1),..,y((n-1)/2),y(-(n-1)/2),...,y(-1)] if n is odd
where
y(j) = sum[k=0..n-1] x[k] * exp(-sqrt(-1)*j*k* 2*pi/n)
j = 0..n-1
Note that y(-j) = y(n-j).
Optional input:
n
Defines the length of the Fourier transform. If n is not
specified then n=x.shape[axis] is set. If n<x.shape[axis],
x is truncated. If n>x.shape[axis], x is zero-padded.
axis
The trasnform is applied along the given axis of the input
array (or the newly constructed array if n argument was used).
overwrite_x
If set to true, the contents of x can be destroyed.
Notes:
y == fft(ifft(y)) within numerical accuracy.
"""
tmp = asarray(x)
t = tmp.typecode()
if t=='D':
overwrite_x = overwrite_x or (tmp is not x and not \
hasattr(x,'__array__'))
work_function = fftpack.zfft
elif t=='F':
raise NotImplementedError
else:
overwrite_x = 1
work_function = fftpack.zrfft
#return _raw_fft(tmp,n,axis,1,overwrite_x,work_function)
if n is None:
n = tmp.shape[axis]
elif n != tmp.shape[axis]:
tmp = _fix_shape(tmp,n,axis)
overwrite_x = 1
if axis == -1 or axis == len(tmp.shape) - 1:
return work_function(tmp,n,1,0,overwrite_x)
tmp = swapaxes(tmp, axis, -1)
tmp = work_function(tmp,n,1,0,overwrite_x)
return swapaxes(tmp, axis, -1)
def _raw_fft(x, n, axis, direction, overwrite_x, work_function):
""" Internal auxiliary function for fft, ifft, rfft, irfft."""
if n is None:
n = x.shape[axis]
elif n != x.shape[axis]:
x = _fix_shape(x,n,axis)
overwrite_x = 1
if axis == -1 or axis == len(x.shape)-1:
r = work_function(x,n,direction,overwrite_x=overwrite_x)
else:
x = swapaxes(x, axis, -1)
r = work_function(x,n,direction,overwrite_x=overwrite_x)
r = swapaxes(r, axis, -1)
return r
def ifft(x, n=None, axis=-1, overwrite_x=0):
""" ifft(x, n=None, axis=-1, overwrite_x=0) -> y
Return inverse discrete Fourier transform of arbitrary type
sequence x.
The returned complex array contains
[y(0),y(1),...,y(n-1)]
where
y(j) = 1/n sum[k=0..n-1] x[k] * exp(sqrt(-1)*j*k* 2*pi/n)
Optional input: see fft.__doc__
"""
tmp = asarray(x)
t = tmp.typecode()
if t=='D':
overwrite_x = overwrite_x or (tmp is not x and not \
hasattr(x,'__array__'))
work_function = fftpack.zfft
elif t=='F':
raise NotImplementedError
else:
overwrite_x = 1
work_function = fftpack.zrfft
#return _raw_fft(tmp,n,axis,-1,overwrite_x,work_function)
if n is None:
n = tmp.shape[axis]
elif n != tmp.shape[axis]:
tmp = _fix_shape(tmp,n,axis)
overwrite_x = 1
if axis == -1 or axis == len(tmp.shape) - 1:
return work_function(tmp,n,-1,1,overwrite_x)
tmp = swapaxes(tmp, axis, -1)
tmp = work_function(tmp,n,-1,1,overwrite_x)
return swapaxes(tmp, axis, -1)
