What's a quick way to calculate the dot product of arbitrary rows from two sparse matrices

Question

For example...

import numpy as np
from scipy.sparse import csr_matrix

X = csr_matrix([[1,2,3], [4,5,6], [7,8,9]])
Y = csr_matrix([[1,2,3], [4,5,6], [7,8,9], [11,12,13]])

# Print matrices
X.toarray()
[[1, 2, 3],
 [4, 5, 6],
 [7, 8, 9]]

Y.toarray()
[[ 1,  2,  3],
 [ 4,  5,  6],
 [ 7,  8,  9],
 [11, 12, 13]]

I have a set of pairs of indices (x,y) representing a row from X and a row from Y. I'd like to take the dot product of the corresponding rows, but I can't figure out how to do this efficiently.

Here's what I've tried

# build arbitrary combinations of row from X and row from Y. Need to calculate dot product of each pair
x_idxs = np.array([2,2,1,0])
y_idxs = np.arange(Y.shape[0])

# current method (slow)
def get_dot_product(x_idx, y_idx):
    return np.dot(X[x_idx].toarray()[0], Y[y_idx].toarray()[0])

func_args = np.transpose(np.array([x_idxs, y_idxs]))
np.apply_along_axis(func1d=lambda x: get_dot_product(x[0], x[1]), axis=1, arr=func_args)

which works but is slow as X and Y get large. Is there a more efficient way?

Update

Following Warren's elegant but slow solution, here's a better example for testing (along with a benchmark)

X = csr_matrix(np.tile(np.repeat(1, 50000),(10000,1)))
Y = X
y_idxs = np.arange(Y.shape[0])
x_idxs = y_idxs

import time
start_time = time.time()
func_args = np.transpose(np.array([x_idxs, y_idxs]))
bg = np.apply_along_axis(func1d=lambda x: get_dot_product(x[0], x[1]), axis=1, arr=func_args)
print("--- %s seconds ---" % (time.time() - start_time)) # 15.48 seconds

start_time = time.time()
ww = X[x_idxs].multiply(Y[y_idxs]).sum(axis=1)
print("--- %s seconds ---" % (time.time() - start_time)) # 38.29 seconds

Answer 1

With your X, Y, x_idxs and y_idxs, you can do:

In [160]: X[x_idxs].multiply(Y[y_idxs]).sum(axis=1)
Out[160]: 
matrix([[ 50],
        [122],
        [122],
        [ 74]])

That uses "fancy" indexing (i.e. indexing with an arbitrary sequence to pull out the desired set of rows), followed by pointwise multiplication and a sum along axis 1 to compute the dot products.

The result is in a numpy matrix, which you can convert to a regular numpy array and flatten as needed. You could even use the somewhat cryptic A1 attribute (a shortcut for the getA1 method):

In [178]: p = X[x_idxs].multiply(Y[y_idxs]).sum(axis=1)

In [179]: p
Out[179]: 
matrix([[ 50],
        [122],
        [122],
        [ 74]])

In [180]: p.A1
Out[180]: array([ 50, 122, 122,  74])

Update, with timing...

Here's a complete script to compare the performance of my version with the original, using arrays X and Y that are actually sparse (with density approximately 0.001, i.e. about 0.1% nonzero elements).

import numpy as np
from scipy import sparse


def get_dot_product(x_idx, y_idx):
    return np.dot(X[x_idx].toarray()[0], Y[y_idx].toarray()[0])

print("Generating random sparse integer matrix X...")
X = (100000*sparse.rand(50000, 120000, density=0.001, format='csr')).astype(np.int64)
X.eliminate_zeros()
print("X has shape %s with %s nonzero elements." % (X.shape, X.nnz))
Y = X
y_idxs = np.arange(Y.shape[0])
x_idxs = y_idxs

import time
start_time = time.time()
func_args = np.transpose(np.array([x_idxs, y_idxs]))
bg = np.apply_along_axis(func1d=lambda x: get_dot_product(x[0], x[1]), axis=1, arr=func_args)
print("--- %8.5f seconds ---" % (time.time() - start_time))

start_time = time.time()
ww = X[x_idxs].multiply(Y[y_idxs]).sum(axis=1)
print("--- %8.5f seconds ---" % (time.time() - start_time))

Output:

Generating random sparse integer matrix X...
X has shape (50000, 120000) with 5999934 nonzero elements.
--- 18.29916 seconds ---
---  0.32749 seconds ---

For less sparse matrices, the speed difference is not so large, and for sufficiently dense matrices, the original version is faster.