0

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]
Crimp City
  • 31
  • 4
  • `b` here is [4 x 3], so `b.T` would be [3x4] which obviously won't work with `F`. If we take what you have and instead run simply `b @ F @ a.T` then you will get a [4x4] matrix with your desired values along the diagonal. Is that helpful? – havingaball Oct 26 '21 at 21:26
  • 1
    Thanks for the response! Is there a way to avoid the [4x4] matrix? In general this matrix will be size [N, N]. This is a minimal working example but the size of my a and b matrices will be N = 4k to 10k ish. Do you think we can avoid the diagonal on NxN matrix and just get the values? – Crimp City Oct 26 '21 at 21:37
  • @hpaulj native question, but the einsum way is just an alternate notation, right? It doesn't change the computations? All combinations of dot product will still be performed? – mozway Oct 26 '21 at 22:00
  • @mozway, no, with a proper indexing, `einsum` does not do the full `ij,jk,lk->il' multiplication first. – hpaulj Oct 26 '21 at 22:20

3 Answers3

1

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])
hpaulj
  • 221,503
  • 14
  • 230
  • 353
0

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@F@a.T).diagonal()

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

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

mozway
  • 194,879
  • 13
  • 39
  • 75
  • Thanks for the response! Is there a way to avoid the [4x4] matrix? In general this matrix will be size [N, N]. This is a minimal working example but the size of my a and b matrices will be N = 4k to 10k ish. Do you think we can avoid the diagonal on NxN matrix and just get the values? – Crimp City Oct 26 '21 at 21:39
  • @Whenihittheground you should check [this question](https://stackoverflow.com/questions/14758283/is-there-a-numpy-scipy-dot-product-calculating-only-the-diagonal-entries-of-the), this is the same issue – mozway Oct 26 '21 at 21:51
  • thanks that is indeed the crux of my issue! I think my solution is very similar. – Crimp City Oct 26 '21 at 22:20
0

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)
Crimp City
  • 31
  • 4