def test_get_fft_funcs():
m, n = [int(2 ** nextpow2(randint(5, 10))) for _ in xrange(2)]
x = randn(m, n)
xc = x + 1j
ifft, fft = get_fft_funcs(xc, xc)
assert fft is np.fft.fft, 'fft is not np.fft.fft'
assert ifft is np.fft.ifft, 'ifft is not np.fft.ifft'
ravelxc = xc.ravel()
assert_allclose(ifft(fft(ravelxc)), ravelxc)
ifft, fft = get_fft_funcs(x, x)
assert fft is np.fft.rfft
assert ifft is np.fft.irfft
ravelx = x.ravel()
assert_allclose(ifft(fft(ravelx)), ravelx)
def matrixcorr(x, nfft):
"""Cross-correlation of the columns of a matrix.
Parameters
----------
x : array_like
The matrix from which to compute the cross correlations of each column
with the others
nfft : int
The number of points used to compute the FFT (faster when this number
is a power of 2).
Returns
-------
c : array_like
The cross correlation of the columns `x`.
"""
_, n = x.shape
ifft, fft = get_fft_funcs(x)
X = fft(x.T, nfft)
Xc = X.conj()
mx, nx = X.shape
c = np.empty((mx ** 2, nx), dtype=X.dtype)
mult_mat_xcorr(X, Xc, c, n, nx)
return ifft(c, nfft).T
def crosscorr(x, y, nfft):
"""Compute the cross correlation of `x` and `y` using an FFT.
Parameters
----------
x, y : array_like
The arrays of which to compute the cross correlation.
nfft : int
The number of fft points.
Returns
-------
c : array_like
Cross correlation of `x` and `y`.
"""
ifft, fft = get_fft_funcs(x, y)
return ifft(fft(x, nfft) * fft(y, nfft).conj(), nfft)
def autocorr(x, nfft):
"""Compute the autocorrelation of `x` using a FFT.
Parameters
----------
x : array_like
Input array.
nfft : int
Number of FFT points.
Returns
-------
r : array_like
The autocorrelation of `x`.
"""
ifft, fft = get_fft_funcs(x)
a = np.abs(fft(x, nfft))
a *= a
return ifft(a, nfft)
def test_mult_mat_xcorr():
x = randn(randint(2, 4), randint(4, 5))
m, n = x.shape
ifft, fft = get_fft_funcs(x)
nfft = int(2 ** nextpow2(m))
X = fft(x.T, nfft)
Xc = X.conj()
mx, nx = X.shape
c = np.empty((mx ** 2, nx), dtype=X.dtype)
oc = c.copy()
mult_mat_xcorr(X, Xc, oc, n, nx)
for i in xrange(n):
c[i * n:(i + 1) * n] = X[i] * Xc
assert_allclose(c, oc)
cc, occ = map(lambda x: ifft(x, nfft).T, (c, oc))
assert_allclose(cc, occ)
