Perform matrix multiplication between two arrays and get result only on masked places

Question

I have two dense matrices, A [200000,10], B [10,100000]. I need to multiply them to get matrix C. I can\'t do that directly, s

Answer 1

Here's one approach using np.einsum for a vectorized solution -

from scipy import sparse
from scipy.sparse import coo_matrix

# Get row, col for the output array
r,c,_= sparse.find(W)

# Get the sum-reduction using valid rows and corresponding cols from A, B
out = np.einsum('ij,ji->i',A[r],B[:,c])

# Store as sparse matrix
out_sparse = coo_matrix((out, (r, c)), shape=W.shape)

Sample run -

1) Inputs :

In [168]: A
Out[168]: 
array([[4, 6, 1, 1, 1],
       [0, 8, 1, 3, 7],
       [2, 8, 3, 2, 2],
       [3, 4, 1, 6, 3]])

In [169]: B
Out[169]: 
array([[5, 2, 4],
       [2, 1, 3],
       [7, 7, 2],
       [5, 7, 5],
       [8, 5, 0]])

In [176]: W
Out[176]: 
<4x3 sparse matrix of type ''
    with 5 stored elements in Compressed Sparse Row format>

In [177]: W.toarray()
Out[177]: 
array([[ True, False, False],
       [False, False, False],
       [ True,  True, False],
       [ True, False,  True]], dtype=bool)

2) Using dense array to perform direct calculations and verify results later on :

In [171]: (A.dot(B))*W.toarray()
Out[171]: 
array([[52,  0,  0],
       [ 0,  0,  0],
       [73, 57,  0],
       [84,  0, 56]])

3) Use the proposed codes and get sparse matrix output :

In [172]: # Using proposed codes
     ...: r,c,_= sparse.find(W)
     ...: out = np.einsum('ij,ji->i',A[r],B[:,c])
     ...: out_sparse = coo_matrix((out, (r, c)), shape=W.shape)
     ...: 

4) Finally verify results by converting to dense/array version and checking against direct version -

In [173]: out_sparse.toarray()
Out[173]: 
array([[52,  0,  0],
       [ 0,  0,  0],
       [73, 57,  0],
       [84,  0, 56]])