Compute inner dot products of multidimensional numpy array elements

user_11779593 0
Viewed 87

What is the most elegant way to compute "inner" dot products of a multidimensional numpy array?

Let's assume I have 2 arrays a and b both of shape (2, 2, 2) (could be (n, n, 2) with n>= 2) and I want to compute the inner_dot(a, b) with the following definition: 

np.array([[np.dot(a[0, 0, :], b[0, 0, :]), np.dot(a[1, 0, :], b[1, 0, :])],
          [np.dot(a[0, 1, :], b[0, 1, :]), np.dot(a[1, 1, :], b[1, 1, :])]])

Here is an example:

a = np.arange(8).reshape(2, 2, 2)
b = np.arange(8).reshape(2, 2, 2)

Expected result:

array([[ 1, 41],
       [13, 85]])
scipy numpy python
3 Answers

You can directly multiply both arrays and sum along the last axis. Also since you want the top right corner of the output as np.dot(a[1, 0, :], b[1, 0, :]), rather than 0,1 and presumably in the same fashion for the off-diagonal elements, you can transpose the result to get as expected:

(a*b).sum(-1).T
array([[ 1, 41],
       [13, 85]])
user_96986840

Use np.einsum -

In [54]: np.einsum('ijk,ijk->ji',a,b)
Out[54]: 
array([[ 1, 41],
       [13, 85]])

With @ opetator/ np.matmul -

In [16]: (a[:,:,None,:]@b[:,:,:,None])[...,0,0].T
Out[16]: 
array([[ 1, 41],
       [13, 85]])
user_32938810

Just for the sake of completeness, I run a quick comparison of the run time:

import numpy as np
import timeit

a = np.arange(1024).reshape(64, 8, 2)
b = np.arange(1024).reshape(64, 8, 2)

example_1 = lambda: (a*b).sum(-1).T
print(timeit.timeit(example_1, number=10000))

example_2 = lambda: np.einsum('ijk,ijk->ji',a,b)
print(timeit.timeit(example_2, number=10000))

example_3 = lambda: (a[:,:,None,:]@b[:,:,:,None])[...,0,0].T
print(timeit.timeit(example_1, number=10000))

On my comupter, the results are (in seconds):

example 1: 0.10384939999999787
example 2: 0.05444100000000063
example 3: 0.097186200000003
user_92260110
Related