Python-了解NumPy的einsum

Python-了解NumPy的einsum

einsum是做什么的？

假设我们有两个多维数组，A和B。现在假设我们要…

• 乘 A用B一种特殊的方式来创造新的产品阵列; 然后也许
• 沿特定轴求和该新数组；然后也许
• 以特定顺序转置新数组的轴。

有一个很好的机会，

einsum

可以帮助我们做到这一点更快，内存更是有效的NumPy的功能组合，喜欢

multiply，sum

和

transpose

允许。

einsum是如何工作的？

这是一个简单（但并非完全无关紧要）的示例。取以下两个数组：

A = np.array([0, 1, 2])B = np.array([[ 0,  1,  2,  3],   [ 4,  5,  6,  7],   [ 8,  9, 10, 11]])

我们将逐个元素相乘A，B然后沿着新数组的行求和。在“普通” NumPy中，我们将编写：

>>> (A[:, np.newaxis] * B).sum(axis=1)array([ 0, 22, 76])

因此，这里的索引 *** 作A将两个数组的第一个轴对齐，以便可以广播乘法。然后将乘积数组中的行相加以返回答案。

现在，如果我们想使用它

einsum

，我们可以这样写：

>>> np.einsum('i,ij->i', A, B)array([ 0, 22, 76])

该签名字符串’i,ij->i’是这里的关键，需要解释的一点。你可以将其分为两半。在左侧（的左侧->），我们标记了两个输入数组。在的右侧->，我们标记了要结束的数组。

接下来是什么：

• A

  有一个轴；我们已经标记了它

  i

  。并且

  B

  有两个轴；我们将轴

  0

  标记为

  i

  ，将轴

  1

  标记为

  j

  。

• 通过在两个输入数组中重复标签i，我们告诉我们einsum这两个轴应该相乘。换句话说，就像A数组B一样，我们将

  array

  与

  array

  的每一列相乘

  A[:, np.newaxis] * B

  。

• 注意，j在我们期望的输出中，它没有显示为标签；我们刚刚使用过i（我们想以一维数组结尾）。通过省略标签，我们告诉einsum来总结沿着这条轴线。换句话说，我们就像对行进行求和

  .sum(axis=1)

  。

基本上，这是你需要了解的所有信息

einsum

。玩一会会有所帮助；如果我们将两个标签都留在输出中，则会

'i,ij->ij'

返回2D产品数组（与相同

A[:, np.newaxis] * B）

。如果我们说没有输出标签，

'i,ij->

我们将返回一个数字（与相同(

A[:, np.newaxis] * B).sum()

）。

einsum但是，最重要的是，它不会首先构建临时的产品阵列；它只是对产品进行累加。这样可以节省大量内存。

一个更大的例子

为了解释点积，这里有两个新数组：

A = array([[1, 1, 1],[2, 2, 2],[5, 5, 5]])B = array([[0, 1, 0],[1, 1, 0],[1, 1, 1]])

我们将使用计算点积

np.einsum('ij,jk->ik', A, B)

。这是一张图片，显示了从函数获得的A和B和输出数组的标签：

你会看到j重复的标签-这意味着我们会将的行A与的列相乘B。此外，j输出中不包含标签-我们对这些产品求和。标签i和k被保留用于输出，因此我们得到一个2D数组。

这一结果与其中标签阵列比较可能是更加明显j的不求和。在下面的左侧，你可以看到写入产生的3D数组

np.einsum('ij,jk->ijk', A, B)

（即，我们保留了

label j

）：

在此处输入图片说明

求和轴j给出了预期的点积，如右图所示。

一些练习

为了获得更多的感觉

einsum

，使用下标符号实现熟悉的

NumPy

数组 *** 作可能会很有用。任何涉及乘法和求和轴组合的内容都可以使用编写

einsum

。

令A和B为两个具有相同长度的一维数组。例如

A = np.arange(10)和B = np.arange(5, 15)

。

• 总和A可以写成：

np.einsum('i->', A)

• 可以将按元素相乘A * B编写为：

np.einsum('i,i->i', A, B)

• 内积或点积np.inner(A, B)或np.dot(A, B)可以写成：

np.einsum('i,i->', A, B) # or just use 'i,i'

• 外部乘积np.outer(A, B)可以写成：

np.einsum('i,j->ij', A, B)

对于2D数组，

C

和

D

，只要轴是兼容的长度（相同长度或其中一个长度为1），下面是一些示例：

• C（主对角线总和）的轨迹

  np.trace(C)

  可以写成：

np.einsum('ii', C)

• 元素方式乘法

  C

  和转置

  D，C * D.T

  可以写成：

np.einsum('ij,ji->ij', C, D)

• 可以将每个元素乘以C该数组D（以构成4D数组）

  C[:, :, None, None] * D

  ，可以写成：

np.einsum('ij,kl->ijkl', C, D)