用于科学计算的 NumPy Einsum

简介

在科学计算中，高效地执行线性代数运算是至关重要的。NumPy 是 Python 中用于数值计算的基础库，它为这些运算提供了众多工具。在这些工具中，einsum（爱因斯坦求和约定）作为一个特别强大的函数脱颖而出，它能够简洁地表达复杂的数组运算。

本教程将引导你了解什么是 einsum，以及如何在你的 Python 代码中有效地使用它。在本实验结束时，你将能够使用 einsum 函数执行各种数组运算，并理解它相对于传统 NumPy 函数的优势。

Skills Graph

理解 NumPy Einsum 基础

爱因斯坦求和约定 (einsum) 是一个强大的 NumPy 函数，它允许你使用简洁的符号来表达许多数组运算。它遵循爱因斯坦求和约定，这在物理学和数学中常用于简化复杂的方程。

打开 Python 交互式环境

让我们从打开 Python 交互式环境开始。在桌面打开一个终端并输入：

python3

你应该会看到 Python 提示符 (>>>)，这表明你现在已经进入了 Python 交互式环境。

导入 NumPy

首先，我们需要导入 NumPy 库：

import numpy as np

什么是 Einsum？

NumPy 中的 einsum 函数允许你使用字符串符号指定数组运算，该符号描述了应该对数组的哪些索引（维度）进行操作。

einsum 运算的基本格式为：

np.einsum('notation', array1, array2, ...)

其中，符号字符串描述了要执行的操作。

一个简单的例子：向量点积

让我们从一个简单的例子开始：计算两个向量的点积。用数学符号表示，两个向量 u 和 v 的点积为：

\sum_i u_i \times v_i

以下是如何使用 einsum 来计算它：

## 创建两个随机向量
u = np.random.rand(5)
v = np.random.rand(5)

## 打印向量以查看它们的值
print("Vector u:", u)
print("Vector v:", v)

## 使用 einsum 计算点积
dot_product = np.einsum('i,i->', u, v)
print("Dot product using einsum:", dot_product)

## 使用 NumPy 的 dot 函数进行验证
numpy_dot = np.dot(u, v)
print("Dot product using np.dot:", numpy_dot)

符号 'i,i->' 表示：

• i 表示第一个数组 (u) 的索引
• 第二个 i 表示第二个数组 (v) 的索引
• 箭头 -> 后面没有内容表示我们想要一个标量结果（对所有索引求和）

理解 Einsum 符号

einsum 符号遵循以下一般模式：

'index1,index2,...->output_indices'

• index1、index2：每个输入数组维度的标签
• output_indices：输出数组维度的标签
• 输入数组中重复出现的索引会被求和
• 出现在输出中的索引会保留在结果中

例如，在符号 'ij,jk->ik' 中：

• i、j 是第一个数组的维度
• j、k 是第二个数组的维度
• j 同时出现在两个输入数组中，因此我们对这个维度求和
• i、k 出现在输出中，因此这些维度会被保留

这正是矩阵乘法的公式！

常见的 Einsum 运算

既然我们已经了解了 einsum 的基础知识，那么让我们来探索一些可以使用这个强大函数执行的常见运算。

矩阵转置

矩阵转置意味着交换其行和列。用数学符号表示，如果 A 是一个矩阵，其转置 A^T 定义为：

A^T_{ij} = A_{ji}

让我们看看如何使用 einsum 执行矩阵转置：

## 创建一个随机矩阵
A = np.random.rand(3, 4)
print("Original matrix A:")
print(A)
print("Shape of A:", A.shape)  ## Should be (3, 4)

## 使用 einsum 进行转置
A_transpose = np.einsum('ij->ji', A)
print("\nTransposed matrix using einsum:")
print(A_transpose)
print("Shape of transposed A:", A_transpose.shape)  ## Should be (4, 3)

## 使用 NumPy 的转置函数进行验证
numpy_transpose = A.T
print("\nTransposed matrix using A.T:")
print(numpy_transpose)

符号 'ij->ji' 表示：

• ij 表示输入矩阵的索引（i 表示行，j 表示列）
• ji 表示输出矩阵的索引（j 表示行，i 表示列）
• 我们本质上是交换了索引的位置

矩阵乘法

矩阵乘法是线性代数中的基本运算。对于两个矩阵 A 和 B，它们的乘积 C 定义为：

C_{ik} = \sum_j A_{ij} \times B_{jk}

以下是如何使用 einsum 执行矩阵乘法：

## 创建两个随机矩阵
A = np.random.rand(3, 4)  ## 3x4 矩阵
B = np.random.rand(4, 2)  ## 4x2 矩阵

print("Matrix A shape:", A.shape)
print("Matrix B shape:", B.shape)

## 使用 einsum 进行矩阵乘法
C = np.einsum('ij,jk->ik', A, B)
print("\nResult matrix C using einsum:")
print(C)
print("Shape of C:", C.shape)  ## Should be (3, 2)

## 使用 NumPy 的 matmul 函数进行验证
numpy_matmul = np.matmul(A, B)
print("\nResult matrix using np.matmul:")
print(numpy_matmul)

符号 'ij,jk->ik' 表示：

• ij 表示矩阵 A 的索引（i 表示行，j 表示列）
• jk 表示矩阵 B 的索引（j 表示行，k 表示列）
• ik 表示输出矩阵 C 的索引（i 表示行，k 表示列）
• 重复的索引 j 会被求和（矩阵乘法）

逐元素乘法

逐元素乘法是指将两个数组的对应元素相乘。对于两个形状相同的矩阵 A 和 B，它们的逐元素乘积 C 为：

C_{ij} = A_{ij} \times B_{ij}

以下是如何使用 einsum 执行逐元素乘法：

## 创建两个形状相同的随机矩阵
A = np.random.rand(3, 3)
B = np.random.rand(3, 3)

print("Matrix A:")
print(A)
print("\nMatrix B:")
print(B)

## 使用 einsum 进行逐元素乘法
C = np.einsum('ij,ij->ij', A, B)
print("\nElement-wise product using einsum:")
print(C)

## 使用 NumPy 的 multiply 函数进行验证
numpy_multiply = A * B
print("\nElement-wise product using A * B:")
print(numpy_multiply)

符号 'ij,ij->ij' 表示：

• ij 表示矩阵 A 的索引
• ij 表示矩阵 B 的索引
• ij 表示输出矩阵 C 的索引
• 没有索引被求和，这意味着我们只是将对应元素相乘

高级 Einsum 运算

既然我们已经熟悉了基本的 einsum 运算，那么让我们来探索一些更高级的应用。这些运算展示了 einsum 函数真正的强大之处和灵活性。

提取对角线元素

提取矩阵的对角线元素是线性代数中的常见操作。对于矩阵 A，其对角线元素构成向量 d，其中：

d_i = A_{ii}

以下是如何使用 einsum 提取对角线元素：

## 创建一个随机方阵
A = np.random.rand(4, 4)
print("Matrix A:")
print(A)

## 使用 einsum 提取对角线元素
diagonal = np.einsum('ii->i', A)
print("\nDiagonal elements using einsum:")
print(diagonal)

## 使用 NumPy 的 diagonal 函数进行验证
numpy_diagonal = np.diagonal(A)
print("\nDiagonal elements using np.diagonal():")
print(numpy_diagonal)

符号 'ii->i' 表示：

• ii 表示矩阵 A 对角线元素的重复索引
• i 表示我们将这些元素提取到一个一维数组中

矩阵的迹

矩阵的迹是其对角线元素之和。对于矩阵 A，其迹为：

\text{trace}(A) = \sum_i A_{ii}

以下是如何使用 einsum 计算矩阵的迹：

## 使用上面的矩阵 A
trace = np.einsum('ii->', A)
print("Trace of matrix A using einsum:", trace)

## 使用 NumPy 的 trace 函数进行验证
numpy_trace = np.trace(A)
print("Trace of matrix A using np.trace():", numpy_trace)

符号 'ii->' 表示：

• ii 表示对角线元素的重复索引
• 空的输出索引表示我们对所有对角线元素求和以得到一个标量

批量矩阵乘法

在对高维数组执行运算时，einsum 真正展现出了它的优势。例如，批量矩阵乘法涉及将两个批次中的矩阵对进行相乘。

如果我们有一批形状为 (n, m, p) 的矩阵 A 和一批形状为 (n, p, q) 的矩阵 B，批量矩阵乘法将得到一个形状为 (n, m, q) 的结果 C：

C_{ijk} = \sum_l A_{ijl} \times B_{ilk}

以下是如何使用 einsum 执行批量矩阵乘法：

## 创建矩阵批次
n, m, p, q = 5, 3, 4, 2  ## 批次大小和矩阵维度
A = np.random.rand(n, m, p)  ## 5 个矩阵的批次，每个矩阵为 3x4
B = np.random.rand(n, p, q)  ## 5 个矩阵的批次，每个矩阵为 4x2

print("Shape of batch A:", A.shape)
print("Shape of batch B:", B.shape)

## 使用 einsum 进行批量矩阵乘法
C = np.einsum('nmp,npq->nmq', A, B)
print("\nShape of result batch C:", C.shape)  ## 应为 (5, 3, 2)

## 让我们检查批次中的第一个矩阵乘法结果
print("\nFirst result matrix from batch using einsum:")
print(C[0])

## 使用 NumPy 的 matmul 函数进行验证
numpy_batch_matmul = np.matmul(A, B)
print("\nFirst result matrix from batch using np.matmul:")
print(numpy_batch_matmul[0])

符号 'nmp,npq->nmq' 表示：

• nmp 表示批次 A 中矩阵的索引（n 表示批次，m 表示行，p 表示列）
• npq 表示批次 B 中矩阵的索引（n 表示批次，p 表示行，q 表示列）
• nmq 表示输出批次 C 中矩阵的索引（n 表示批次，m 表示行，q 表示列）
• 重复的索引 p 会被求和（矩阵乘法）

为什么使用 Einsum？

你可能会疑惑，既然 NumPy 已经为这些操作提供了专门的函数，为什么我们还要使用 einsum 呢？以下是一些优点：

1. 统一接口：einsum 为许多数组操作提供了一个单一的函数。
2. 灵活性：它可以表达那些原本需要多个步骤才能完成的操作。
3. 可读性：一旦你理解了符号，代码会变得更加简洁。
4. 性能：在很多情况下，einsum 操作经过了优化，效率很高。

对于复杂的张量操作，einsum 通常能提供最清晰、最直接的实现方式。

总结

在本次实验中，你探索了 NumPy 中强大的 einsum 函数，该函数实现了用于数组操作的爱因斯坦求和约定。让我们回顾一下你所学的内容：

1. Einsum 基本概念：你学习了如何使用 einsum 符号来表示数组操作，其中索引代表数组的维度，重复的索引表示求和。

2. 常见操作：你使用 einsum 实现了几个基本操作：

   • 向量点积
   • 矩阵转置
   • 矩阵乘法
   • 逐元素乘法

3. 高级应用：你探索了更复杂的操作：

   • 提取对角线元素
   • 矩阵的迹
   • 批量矩阵乘法

einsum 函数为 NumPy 中的数组操作提供了一种统一且灵活的方法。虽然像 np.dot、np.matmul 和 np.transpose 这样的专用函数可用于特定操作，但 einsum 为各种操作提供了一致的接口，这在处理高维数组时尤为有用。

在你继续科学计算和数据科学之旅时，einsum 将成为你执行复杂数组操作的强大工具，让你的代码简洁易读。