Efficiently computing all pairwise products of a given vector's elements in NumPy

Question

I am looking for an "optimal" way to compute all pairwise products of a given vector's elements. If the vector is of size N, the output will be a vector of size N * (N + 1) // 2 and contain x[i] * x[j] values for all (i, j) pairs with i <= j. The naive way to compute this is as follows:

import numpy as np

def get_pairwise_products_naive(vec: np.ndarray):
    k, size = 0, vec.size
    output = np.empty(size * (size + 1) // 2)
    for i in range(size):
        for j in range(i, size):
            output[k] = vec[i] * vec[j]
            k += 1
    return output

Desiderata:

Minimize extra memory allocations/usage: Directly write to the output buffer if possible.
Use vectorized NumPy routines instead of explicit loops.
Avoid extra (unnecessary) calculations.

I have been playing with routines such as outer, triu_indices and einsum as well as some indexing/view tricks, but haven't been able to find a solution that fits the above desiderata.

user_565635 · Answer

I would probably compute M = v^Tv and then flatten the lower or higher triangular portion of this matrix.

def pairwise_products(v: np.ndarray):
    assert len(v.shape) == 1
    n = v.shape[0]
    m = v.reshape(n, 1) @ v.reshape(1, n)
    return m[np.tril_indices_from(m)].ravel()

I would also like to mention numba, which would make your 'naive' approach most likely faster than this one.

import numba

@numba.njit
def pairwise_products_numba(vec: np.ndarray):
    k, size = 0, vec.size
    output = np.empty(size * (size + 1) // 2)
    for i in range(size):
        for j in range(i, size):
            output[k] = vec[i] * vec[j]
            k += 1
    return output

Just testing the above pairwise_products(np.arange(5000)) takes ~0.3 sec whereas the numba version takes ~0.05 sec (ignoring the first run which is used to just-in-time compile the function).

user_3293881 · Answer

Approach #1

For a vectorized one with NumPy, you can use a masking one after getting all the pairwise multiplications with outer-multiplication, like so -

def pairwise_multiply_masking(a):
    return (a[:,None]*a)[~np.tri(len(a),k=-1,dtype=bool)]

Approach #2

For really big input 1D arrays, we might want to resort to iterative slicing method that uses one-loop -

def pairwise_multiply_iterative_slicing(a):
    n = len(a)
    N = (n*(n+1))//2
    out = np.empty(N, dtype=a.dtype)
    c = np.r_[0,np.arange(n,0,-1)].cumsum()
    for ii,(i,j) in enumerate(zip(c[:-1],c[1:])):
        out[i:j] = a[ii:]*a[ii]
    return out

Benchmarking

We will include pairwise_products and pairwise_products_numba from @orlp's solution in the setup.

Using benchit package (few benchmarking tools packaged together; disclaimer: I am its author) to benchmark proposed solutions.

import benchit
funcs = [pairwise_multiply_masking, pairwise_multiply_iterative_slicing, pairwise_products_numba, pairwise_products]
in_ = [np.random.rand(n) for n in [10,50,100,200,500,1000,5000]]
t = benchit.timings(funcs, in_)
t.plot(logx=True, save='timings.png')
t.speedups(-1).plot(logx=True, logy=False, save='speedups.png')

Results (timings and speedups over pairwise_products) -

As can be seen with the plot trends, for really large arrays, the slicing based one will start winning, otherwise vectorized one does a good job.

Suggestions

We can also look into numexpr for performing the outer multiplications more efficienctly for large arrays.

user_4045774 · Answer

You could also parallelize this algorithm. If it would be possible to allocate a large enough array (a smaller view on this array almost costs nothing) only once and overwrite it afterwards larger speedups can be can be achieved.

Example

@numba.njit(parallel=True)
def pairwise_products_numba_2_with_allocation(vec):
    k, size = 0, vec.size
    k_vec=np.empty(vec.size,dtype=np.int64)
    output = np.empty(size * (size + 1) // 2)

    #precalculate the indices
    for i in range(size):
        k_vec[i] = k
        k+=(size-i)

    for i in numba.prange(size):
        k=k_vec[i]
        for j in range(size-i):
            output[k+j] = vec[i] * vec[j+i]

    return output

@numba.njit(parallel=True)
def pairwise_products_numba_2_without_allocation(vec,output):
    k, size = 0, vec.size
    k_vec=np.empty(vec.size,dtype=np.int64)

    #precalculate the indices
    for i in range(size):
        k_vec[i] = k
        k+=(size-i)

    for i in numba.prange(size):
        k=k_vec[i]
        for j in range(size-i):
            output[k+j] = vec[i] * vec[j+i]

    return output

Timings

A=np.arange(5000)
k, size = 0, A.size
output = np.empty(size * (size + 1) // 2)

%timeit res_1=pairwise_products_numba_2_without_allocation(A,output)
#7.84 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_2=pairwise_products_numba_2_with_allocation(A)
#16.9 ms ± 325 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_3=pairwise_products_numba(A) #@orlp
#43.3 ms ± 134 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Efficiently computing all pairwise products of a given vector's elements in NumPy

3 Answers