How to express this complicated expression using numpy slices

Question

I wish to implement the following expression in Python: $$ x_i = \sum_{j=1}^{i-1}k_{i-j,j}a_{i-j}a_j, $$ where $x$ and $y$ are numpy arrays of size $n$, and $k$ is a numpy array of size $n\times n$. The size $n$ might be up to about 10000, and the function is part of an inner loop that will be evaluated many times, so speed is important.

Ideally I'd like to avoid a for loop altogether, though I guess it's not the end of the world if there is one. The problem is that I'm having trouble seeing how to do it without having a couple of nested loops, and that's likely to make it rather slow.

Can anybody see how to express the above equation using numpy in a way that's efficient, and preferably also readable? More generally, what is the best way to approach this sort of thing?

Answer 1

Here is the Numba solution. On my machine the Numba version is >1000x faster than the python version without the decorator (for a 200x200 matrix, 'k' and 200-length vector 'a'). You can also use the @autojit decorator which adds about 10 microseconds per call so that the same code will work with multiple types.

from numba import jit, autojit

@jit('f8[:](f8[:,:],f8[:])')
#@autojit
def looped_ver(k, a):
    x = np.empty_like(a)
    for i in range(x.size):
        sm = 0.0
        for j in range(0, i+1):
            sm += k[i-j,j] * a[i-j] * a[j]
        x[i] = sm
    return x

Disclosure: I am one of the Numba developers.

Answer 2

Here's a start. First, my apologies for any mistakes.

I experimented with a couple of different approached. I was a bit confused by the limits on the summation - should the upper limit be $i$, rather than $i-1$?

Edit: No, the upper limit was correct as provided in the question. I have left it as is here because another answer now uses the same code, but the fix is simple.

First a looped version:

def looped_ver(k, a):
    x = np.empty_like(a)
    for i in range(x.size):
        sm = 0
        for j in range(0, i+1):
            sm += k[i-j,j] * a[i-j] * a[j]
        x[i] = sm
    return x

I made it a single loop with numpy slices:

def vectorized_ver(k, a):
    ktr = zeros_like(k)
    ar = zeros_like(k)
    sz = len(a)
    for i in range(sz):
        ktr[i,:i+1] = k[::-1].diagonal(-sz+i+1)
        a_ = a[:i+1]
        ar[i,:i+1] = a_[::-1] * a_
    return np.sum(ktr * ar, 1)

The numpy version with one explicit loop is about 25x faster on my computer when $n=5000$.

Then I wrote a Cython version of the (more-readable) looped code.

import numpy as np
import cython
cimport numpy as np

@cython.boundscheck(False)
@cython.wraparound(False)
def cyth_ver(double [:, ::1] k not None,
              double [:] a not None):
    cdef double[:] x = np.empty_like(a)
    cdef double sm
    cdef int i, j

    for i in range(len(a)):
        sm = 0.0
        for j in range(i+1):
            sm = sm + k[i-j,j] * a[i-j] * a[j]
        x[i] = sm
    return x

On my laptop, this one is about 200x faster than the looped version (and 8x faster than the 1-loop vectorized version). I'm sure others can do better.

I played with a Julia version, and it seemed (if I timed it properly) comparable to the Cython code.

Answer 3

What you want seems to be a convolution; I think the quickest way for achieving it would be the numpy.convolve function.

Your may have to fix the indices according to your exact needs but I think you would like to try something like:

import numpy as np
a = [1, 2, 3, 4, 5]
k = [2, 4, 6, 8, 10]

result = np.convolve(a, k*a[::-1])