I found scipy.signal.fftconvolve
, as also pointed out by magnus, but didn't realize at the time that it's n-dimensional. Since it's built-in and produces the right values, it seems like the ideal solution.
From Example of 2D Convolution:
In [1]: a = asarray([[ 1, 2, 3],
...: [ 4, 5, 6],
...: [ 7, 8, 9]])
In [2]: b = asarray([[-1,-2,-1],
...: [ 0, 0, 0],
...: [ 1, 2, 1]])
In [3]: scipy.signal.fftconvolve(a, b, mode = 'same')
Out[3]:
array([[-13., -20., -17.],
[-18., -24., -18.],
[ 13., 20., 17.]])
Correct! The STSCI version, on the other hand, requires some extra work to make the boundaries correct?
In [4]: stsci.convolve2d(a, b, fft = True)
Out[4]:
array([[-12., -12., -12.],
[-24., -24., -24.],
[-12., -12., -12.]])
(The STSCI method also requires compiling, which I was unsuccessful with (I just commented out the non-python parts), has some bugs like this and modifying the inputs ([1, 2] becomes [[1, 2]]), etc. So I changed my accepted answer to the built-in fftconvolve()
function.)
Correlation, of course, is the same thing as convolution, but with one input reversed:
In [5]: a
Out[5]:
array([[3, 0, 0],
[2, 0, 0],
[1, 0, 0]])
In [6]: b
Out[6]:
array([[3, 2, 1],
[0, 0, 0],
[0, 0, 0]])
In [7]: scipy.signal.fftconvolve(a, b[::-1, ::-1])
Out[7]:
array([[ 0., -0., 0., 0., 0.],
[ 0., -0., 0., 0., 0.],
[ 3., 6., 9., 0., 0.],
[ 2., 4., 6., 0., 0.],
[ 1., 2., 3., 0., 0.]])
In [8]: scipy.signal.correlate2d(a, b)
Out[8]:
array([[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[3, 6, 9, 0, 0],
[2, 4, 6, 0, 0],
[1, 2, 3, 0, 0]])
and the latest revision has been sped up by using power-of-two sizes internally (and then I sped it up more by using real FFT for real input and using 5-smooth lengths instead of powers of 2 :D ).