[Numpy-discussion] reformulating a series of correlations as one fft, ifft pair

[Paul Northug]

I am having trouble reformulating a series of correlations as a single
fft, ifft pair.

I have a set of kernels phi : (M = channel, N = kernel, T = time)
correlated with signal a : (N, P+T-1) yielding x : (M, T).

The correlation, for now, is only in the last dimension, with the two
other dimensions degenerate. Below are 4 implementations (x == y == z
!= w). I would like to reformulate the operation as a single pair of
3d ffts to compare how it fares for different problem dimensions but I
am not able to do it. Implementation #4 is incorrect. Please note the
.conj() to get correlations rather than convolutions.

Any tips would be appreciated. Optimally, I would not have to pad with
too many zeros as the matrices can be quite large. The challenge is to
beat implementation 2 as P grows. M, N, P are typically on the order
of 10^2 and T, 10^3.

#
import numpy as np
from scipy.signal import fftn, ifftn, correlate

M, N, P, T = 2, 3, 4, 5

phi = np.random.randn(M, N, P)
a = np.random.randn(N, P+T-1)
x = np.zeros((M, T))

# 1. correlate

x = np.array([correlate(phi[i], a, 'valid').squeeze() for i in range(M)])

# 2. dot products

y = np.zeros_like(x)
for b in range(P): y += np.dot(phi[:,:,b], a[:,b:b+T])

# 3. a series of ffts (code adapted from scipy convolve)

s1 = np.array(phi.shape[1:])
s2 = np.array(a.shape)
size = s1+s2-1
z = np.empty_like(y)
fsize = (2**np.ceil(np.log2(size))).astype(np.int) # power of 2 pad
af = fftn(a, fsize)
fslice = tuple([slice(0, sz) for sz in size])
for i in range(M):
    phif = fftn(phi[i], fsize).conj()
    z[i] = np.real(ifftn(phif * af)[fslice])[0,:T]

# 4. <BROKEN> single fft, ifft pair - some rows are correct

s1 = np.array(phi.shape)
s2 = np.array(a[np.newaxis].shape)
size = s1+s2-1
fsize = (2**np.ceil(np.log2(size))).astype(np.int)
af = fftn(a[np.newaxis], fsize)
phif = fftn(phi, fsize).conj()
fslice = tuple([slice(0, sz) for sz in size])
w = np.real(ifftn(phif * af)[fslice])[:,:,:T]