计算多维numpy数组元素的内点积答案

【问题标题】：Compute inner dot products of multidimensional numpy array elements计算多维numpy数组元素的内点积
【发布时间】：2020-06-05 11:08:27
【问题描述】：

计算多维 numpy 数组的“内部”点积的最优雅方法是什么？

假设我有 2 个数组 a 和 b 形状都是 (2, 2, 2) （可能是 (n, n, 2) 和 n>= 2) 我想用以下定义计算 inner_dot(a, b)：

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, :])]])

这是一个例子：

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

预期结果：

array([[ 1, 41],
       [13, 85]])

【问题讨论】：

标签： python numpy scipy

【解决方案1】：

您可以直接将数组和sum 沿最后一个轴相乘。此外，由于您希望输出的右上角为np.dot(a[1, 0, :], b[1, 0, :])，而不是0,1，并且可能以相同的方式用于非对角元素，您可以transpose 得到预期的结果：

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

【讨论】：

【解决方案2】：

使用np.einsum -

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

与@ opetator/ np.matmul -

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

【讨论】：

【解决方案3】：

为了完整起见，我快速比较一下运行时间：

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))

在我的电脑上，结果是（以秒为单位）：

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

【讨论】：