Search code examples
pythonscipyconvolution

Need a circular FFT convolution in Python


I need a faster analog of

scipy.signal.convolve2d(data, filter, boundary="wrap", mode="same")

Cannot you advice me how to replace it?

P.S. scipy.signal.fftconvolve is fast enough, but it does not have boundary option and I cannot make it work in circular convolution mode.


Solution

  • If you compute the following:

    from scipy.fftpack import fft2, ifft2
    
    f2 = ifft2(fft2(data, shape=data.shape) * fft2(filter, shape=data.shape)).real
    

    then f2 contains the same values as convolve2d(data, filt, boundary='wrap', mode='same'), but the values are shifted ("rolled", in numpy terminology) in each axis. (This is an application of the convolution theorem.)

    Here's a short function that rolls the result to the give same result as the convolve2d function call:

    def fftconvolve2d(x, y):
        # This assumes y is "smaller" than x.
        f2 = ifft2(fft2(x, shape=x.shape) * fft2(y, shape=x.shape)).real
        f2 = np.roll(f2, (-((y.shape[0] - 1)//2), -((y.shape[1] - 1)//2)), axis=(0, 1))
        return f2
    

    For example,

    In [91]: data = np.random.rand(256, 256)
    
    In [92]: filt = np.random.rand(16, 16)
    
    In [93]: c2d = convolve2d(data, filt, boundary='wrap', mode='same')
    
    In [94]: f2 = fftconvolve2d(data, filt)
    

    Verify that the results are the same:

    In [95]: np.allclose(c2d, f2)
    Out[95]: True
    

    Check the performance:

    In [96]: %timeit c2d = convolve2d(data, filt, boundary='wrap', mode='same')
    44.9 ms ± 77.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
    
    In [97]: %timeit f2 = fftconvolve2d(data, filt)
    5.23 ms ± 11.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    

    The FFT version is much faster (but note that I chose the dimensions of data to be a power of 2).