Vectorize NumPY Dot Product Between a 2D array and Another 2D Array per Row-CodePudding

Hello Stack Overflow,

I am able to get desired results through a loop but wondering how I can apply a dot product per row of a 2D array. I'm trying to take each row from A and apply a series of dot products with B. It's a bit wonky right now with my intermediate results but I do get my desired results. Looking to vectorize and eliminate the loop.

A = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]])
B = np.array([[9,-1,8],[-12,20,-5],[6,1,4]])
int_result = np.ones((len(A),3))
for r in np.arange(0,len(A)):
    int_result[r] = A[r].T.dot(B).dot(A[r])
print(int_result)
desired_result = int_result[:,0]
print(desired_result)

Intermediate Results

[[ 117.  117.  117.]  
[ 744.  744.  744.]  
[1911. 1911. 1911.]  
[3618. 3618. 3618.]]

Desired Results

[ 117.  744. 1911. 3618.]

CodePudding user response：

This is 2 np.matmul multiplications, using the first dimension of A as the batch dimension:

In [55]: A = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]])
    ...: B = np.array([[9,-1,8],[-12,20,-5],[6,1,4]])
In [56]: A
Out[56]: 
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])
In [57]: B
Out[57]: 
array([[  9,  -1,   8],
       [-12,  20,  -5],
       [  6,   1,   4]])
In [58]: A.shape
Out[58]: (4, 3)
In [59]: B.shape
Out[59]: (3, 3)

Adding dimensions to A to make it (4,1,3) and (4,3,1), we can do:

In [60]: (A[:,None,:]@B@A[:,:,None]).shape
Out[60]: (4, 1, 1)

and removing the size 1 dimensions:

In [61]: (A[:,None,:]@B@A[:,:,None]).squeeze()
Out[61]: array([ 117,  744, 1911, 3618])

The einsum equivalent is equally good:

In [62]: np.einsum('ij,kj,ik->i',A,B,A)
Out[62]: array([ 117,  744, 1911, 3618])

In your loop:

In [65]: r=3; A[r]@B@A[r]    # A[r] is 1d, so T isn't needed
Out[65]: 3618