How to matrix multiply 3 matrices in numpy?-CodePudding

Given an F matrix which is 3x3 and some point matrices which are both N x 3 I would like to compute E = b.T @ F @ a efficiently.

F = np.arange(9).reshape(3, 3)
>>> [[0 1 2]
     [3 4 5]
     [6 7 8]]

a = np.array([[1, 2, 1], [3, 4, 1], [5, 6, 1], [7, 8, 1]])
>>>[[1 2 1]
    [3 4 1]
    [5 6 1]
    [7 8 1]]

b = np.array([[10, 20, 1],[30, 40, 1],[50, 60, 1],[70, 80, 1]])
>>>[[10 20  1]
    [30 40  1]
    [50 60  1]
    [70 80  1]]

The expected output is: E = [388 1434 3120 5446]

I can get it with a simple for loop but I would like to do this with all numpy. I tried reshaping the a and b matrices but that didn't quite work.

N = b.shape[0] #4
a_reshaped = a.reshape(N, 3, 1)
b_reshaped = b.reshape(1, N, 3)

F_repeated = np.repeat(F[None,:], N, axis=0)
E = b_reshaped @ F_repeated @ a_reshaped 
>>>
[[[ 388]
  [ 788]
  [1188]
  [1588]]

 [[ 714]
  [1434]
  [2154]
  [2874]]

 [[1040]
  [2080]
  [3120]
  [4160]]

 [[1366]
  [2726]
  [4086]
  [5446]]]

If I then take the diagonal values I get the expected result however, this is very inefficient.

Any suggestions?

Edit: Here is the for loop version of what I am describing:

E = []
for k in range(N):
    error = b[k].T @ F @ a[k]
    E.append(error)
E = np.array(E)
>>>[388 1434 3120 5446]

CodePudding user response：

Following on the comment of @havingaball, you cannot get a 1D array from a series of 2D matrix multiplications.

It looks like what you want is:

(b@[email protected]).diagonal()

(or (a@(b@F).T).diagonal())

Output: array([ 388, 1434, 3120, 5446])

CodePudding user response：

I found a solution that avoids making an NxN matrix:

partial_mult_b = b @ F
E_prime = partial_mult_b * a
E = np.sum(E_prime, axis=1)

CodePudding user response：

First, the 'naive' iterative version - to be clear of the action:

In [67]: c = np.zeros(4,int)
In [68]: for i in range(4):
    ...:     c[i] = b[i]@F@a[i]
    ...: 
In [69]: c
Out[69]: array([ 388, 1434, 3120, 5446])

the easy einsum version

In [70]: np.einsum('ij,jk,ik->i',b,F,a)
Out[70]: array([ 388, 1434, 3120, 5446])

batched matmul version

In [71]: b[:,None,:]@F@a[:,:,None]
Out[71]: 
array([[[ 388]],

       [[1434]],

       [[3120]],

       [[5446]]])
In [72]: (b[:,None,:]@F@a[:,:,None]).squeeze()
Out[72]: array([ 388, 1434, 3120, 5446])

this:

In [73]: ((b@F)*a).sum(axis=1)
Out[73]: array([ 388, 1434, 3120, 5446])

is the same as expanding the einsum to two steps:

In [74]: np.einsum('ik,ik->i',np.einsum('ij,jk->ik',b,F),a)
Out[74]: array([ 388, 1434, 3120, 5446])