Speed up for loop in convolution for numpy 3D array?

Question

Performing convolution along Z vector of a 3d numpy array, then other operations on the results, but it is slow as it is implemented now. Is the for loop what is slowing me down here or is is the convolution? I tried reshaping to a 1d vector and perform the convolution in 1 pass (as I did in Matlab), without the for loop, but it doesn't improve performance. My Matlab version is about 50% faster than anything I can come up with in Python. Relevant section of code:

convolved=np.zeros((y_lines,x_lines,z_depth))
for i in range(0, y_lines):
    for j in range(0, x_lines):
        convolved[i,j,:]= fftconvolve(data[i,j,:], Gauss) #80% of time here
        result[i,j,:]= other_calculations(convolved[i,j,:]) #20% of time here

Is there a better way to do this than a for loop? Heard of Cython but I have limited experience in Python as of now, would aim for the simplest solution.

Answer 1

The fftconvolve function you are using is presumably from SciPy. If so, be aware that it takes N-dimensional arrays. So a faster way to do your convolution would be to generate the 3d kernel that corresponds to doing nothing in the x and y dimensions and doing a 1d gaussian convolution in z.

Some code and timing results are below. On my machine and with some toy data, this led to a 10× speedup, as you can see:

import numpy as np
from scipy.signal import fftconvolve
from scipy.ndimage.filters import gaussian_filter

# use scipy filtering functions designed to apply kernels to isolate a 1d gaussian kernel
kernel_base = np.ones(shape=(5))
kernel_1d = gaussian_filter(kernel_base, sigma=1, mode='constant')
kernel_1d = kernel_1d / np.sum(kernel_1d)

# make the 3d kernel that does gaussian convolution in z axis only
kernel_3d = np.zeros(shape=(1, 1, 5,))
kernel_3d[0, 0, :] = kernel_1d

# generate random data
data = np.random.random(size=(50, 50, 50))

# define a function for loop based convolution for easy timeit invocation
def convolve_with_loops(data):
    nx, ny, nz = data.shape
    convolved=np.zeros((nx, ny, nz))
    for i in range(0, nx):
        for j in range(0, ny):
            convolved[i,j,:]= fftconvolve(data[i, j, :], kernel_1d, mode='same') 
    return convolved

# compute the convolution two diff. ways: with loops (first) or as a 3d convolution (2nd)
convolved = convolve_with_loops(data)
convolved_2 = fftconvolve(data, kernel_3d, mode='same')

# raise an error unless the two computations return equivalent results
assert np.all(np.isclose(convolved, convolved_2))

# time the two routes of the computation
%timeit convolved = convolve_with_loops(data)
%timeit convolved_2 = fftconvolve(data, kernel_3d, mode='same')

---

timeit results:

10 loops, best of 3: 198 ms per loop
100 loops, best of 3: 18.1 ms per loop

Answer 2

I think you have already found the source code of the fftconvolve function. Normally, for real inputs, it uses the numpy.fft.rfftn and .irfftn functions, which compute N-dimensional transforms. Since the goal is to do multiple 1-D transformations, you can basically rewrite fftconvolve like this (simplified):

from scipy.signal.signaltools import _next_regular

def fftconvolve_1d(in1, in2):

    outlen = in1.shape[-1] + in2.shape[-1] - 1
    n = _next_regular(outlen)

    tr1 = np.fft.rfft(in1, n)
    tr2 = np.fft.rfft(in2, n)
    out = np.fft.irfft(tr1 * tr2, n)

    return out[..., :outlen].copy()

And calculate the desired result:

result = fftconvolve_1d(data, Gauss)

This works because numpy.fft.rfft and .irfft (notice the lack of n in the name) transform over a single axis of the input array (the last axis by default). This is roughly 40% faster than the OP code on my system.

---

Further speedup can be achieved by using a different FFT back-end.

For one, the functions in scipy.fftpack appear to be somewhat faster than their Numpy equivalents. However, the output format of the Scipy variants is pretty awkward (see docs) and this makes it hard to do the multiplication.

Another possible backend is FFTW through the pyFFTW wrapper. Downsides are that transforms are preceded by a slow "planning phase" and that inputs have to be 16-byte aligned to achieve best performance. This is explained pretty well in the pyFFTW tutorial. The resulting code could be for example:

from scipy.signal.signaltools import _next_regular
import pyfftw
pyfftw.interfaces.cache.enable()  # Cache for the "planning"
pyfftw.interfaces.cache.set_keepalive_time(1.0)

def fftwconvolve_1d(in1, in2):

    outlen = in1.shape[-1] + in2.shape[-1] - 1
    n = _next_regular(outlen)

    tr1 = pyfftw.interfaces.numpy_fft.rfft(in1, n)
    tr2 = pyfftw.interfaces.numpy_fft.rfft(in2, n)

    sh = np.broadcast(tr1, tr2).shape
    dt = np.common_type(tr1, tr2)
    pr = pyfftw.n_byte_align_empty(sh, 16, dt)
    np.multiply(tr1, tr2, out=pr)
    out = pyfftw.interfaces.numpy_fft.irfft(pr, n)

    return out[..., :outlen].copy()

With aligned inputs and cached "planning" I saw a speedup of almost 3x over the code in the OP. Memory alignment can be easily checked by looking at the memory address in the ctypes.data attribute of a Numpy array.