Vectorize this convolution type loop more efficiently in numpy

Question

I have to do many loops of the following type

for i in range(len(a)):
    for j in range(i+1):
        c[i] += a[j]*b[i-j]

where a and b are short arrays (of the same size, which is between about 10 and 50). This can be done efficiently using a convolution:

import numpy as np
np.convolve(a, b) 

However, this gives me the full convolution (i.e. the vector is too long, compared to the for loop above). If I use the 'same' option in convolve, I get the central part, but what I want is the first part. Of course, I can chop off what I don't need from the full vector, but I would like to get rid of the unnecessary computation time if possible. Can someone suggest a better vectorization of the loops?

Answer 1

You could write a small C extension in Cython:

# cython: boundscheck=False
cimport numpy as np
import numpy as np  # zeros_like

ctypedef np.float64_t np_t
def convolve_cy_np(np.ndarray[np_t] a not None,
                   np.ndarray[np_t] b not None,
                   np.ndarray[np_t] c=None):
    if c is None:
       c = np.zeros_like(a)
    cdef Py_ssize_t i, j, n = c.shape[0]
    with nogil:
        for i in range(n):
            for j in range(i + 1):
                c[i] += a[j] * b[i - j]
    return c

It performs well for n=10..50 compared to np.convolve(a,b)[:len(a)] on my machine.

Also it seems like a job for numba.

Answer 2

There is no way to do a convolution with vectorized array manipulations in numpy. Your best bet is to use np.convolve(a, b, mode='same') and trim off what you don't need. Thats going to probably be 10x faster than the double loop in pure python that you have above. You could also roll your own with Cython if you are really concerned about the speed, but it likely won't be much if any faster than np.convolve().