4. 需要的 NumPy 技巧：初步

这一节我们会回顾一些 numpy 的相关问题。整个 pyxdh 的电子积分、基函数等问题都建立在 PySCF 库，但其余的量化方法都建立在 numpy 中；对 numpy 的熟悉将是至关重要的。

一般认为，numpy 具有正常的效率与并行能力；但处理特别的计算问题时，numpy 并不很鲁棒。对于多节点计算与 GPU 计算，需要 numpy 以外的环境。

在以后，我们会使用修改版的 Einstein Summation Convention，即无上下标区分的张量角标，以简化公式表达，与 np.einsum 函数形成稳定的关联。

numpy 具有非常详尽的 API 文档；在 PyCharm 中，numpy 的 python 部分函数源码也可以通过 Ctrl + 鼠标左键得到，其文档 (docstring) 可以通过 Ctrl + Q 获得。

[1]:

import numpy as np
import scipy
import scipy.linalg

4.1. 矩阵与张量定义

numpy 的最基本单元是 ndarray 类；一般来说，我们的矩阵与张量都是这个类的实例 (instance)。从列表 lst 通过 np.array 可以生成一个矩阵 mat：

[2]:

lst = [
    [ 0,  1,  2],
    [10, 11, 12],
    [20, 21, 22],
]
mat = np.array(lst)
mat

[2]:

array([[ 0,  1,  2],
       [10, 11, 12],
       [20, 21, 22]])

矩阵的最基本信息包括矩阵维度、元素类型、占用内存空间大小等。这些可以从成员直接查看：

[3]:

print(mat.dtype)
print(mat.shape)
print(mat.size)
print(mat.nbytes)

int64
(3, 3)
9
72

与 Fortran 不同，numpy 沿用 Python (包括绝大多数其它语言) 的 0 索引方式：

[4]:

mat[1]

[4]:

array([10, 11, 12])

但与 Fortran 相同的是，通常可以使用冒号取其中的某一列 (这是 C 程序所不便利之处)：

[5]:

mat[:, 1]

[5]:

array([ 1, 11, 21])

如果要生成随机但维度确定的矩阵，则可以使用 np.random 中的函数。这里不再详述。

任务 (1)

生成一个不等维度的三维列表，并将其化为三维张量．观察三维张量的属性、并了解它是如何索引并输出的．

在 Hessian 任务中，我们需要了解最高八维张量的计算，以及最高四维张量的输出．对任意维度张量的操控与调试是后面笔记中所经常使用的能力．

4.2. 元素操作

numpy 中的许多操作，包括张量与数 (scaler)、算符 (operation) 的操作，是元素操作 (elementwise manuplation)。举例来说，对于矩阵与数的加法

\[\mathbf{A} = \mathbf{M} + 5\]

这就等价于

\[A_{ij} = M_{ij} + 5\]

[6]:

A = mat + 5
A

[6]:

array([[ 5,  6,  7],
       [15, 16, 17],
       [25, 26, 27]])

幂次计算则可以表示为

\[B_{ij} = M_{ij}^3\]

[7]:

B = mat ** 3
B

[7]:

array([[    0,     1,     8],
       [ 1000,  1331,  1728],
       [ 8000,  9261, 10648]])

对于指数计算也类似。譬如下述的的过程表示

\[C_{ij} = \exp (- M_{ij} / 5)\]

[8]:

C = np.exp(- mat / 5)
C

[8]:

array([[1.        , 0.81873075, 0.67032005],
       [0.13533528, 0.11080316, 0.09071795],
       [0.01831564, 0.01499558, 0.01227734]])

任务 (2)

(可选) 上述 \(M_{ij}^3\) 并不是通常矩阵运算中的 \(\mathbf{M}^3\)。请指出如何计算 \(\mathbf{M}^3\)。
(可选) 上述 \(\exp(- M_{ij} / 5)\) 并不是通常矩阵运算中的 \(\exp(- \mathbf{M} / 5)\)。请指出如何计算 \(\exp(- \mathbf{M} / 5)\)。

4.3. 矩阵乘积

在量化计算中，最基本的操作是矩阵乘积。矩阵乘积在 numpy 中的实现方法至少有三种；在这里，我们会拿一个简单的例子说明。

[9]:

A = np.array([[1, 2, 3], [11, 12, 13]])                           # A.shape = (2, 3)
B = np.array([[1, 2, 3, 4], [11, 12, 13, 14], [21, 22, 23, 24]])  # B.shape = (3, 4)

我们知道，(2, 3) 维矩阵与 (3, 4) 维矩阵的乘积是 (2, 4) 维。这两个矩阵的乘积可以用 Einstein Convention 表示为

\[C_{ij} = A_{ik} B_{kj}\]

矩阵乘积的最标准写法是使用 @ 或等价地 np.matmul；另两种写法是使用 np.dot，以及使用 np.einsum。后者则更直观地体现角标变化。

[10]:

C_byat = A @ B
C_bydot = A.dot(B)  # equivalent to np.dot(A, B)
C_byein = np.einsum("ik, kj -> ij", A, B)
C_bydot

[10]:

array([[ 86,  92,  98, 104],
       [416, 452, 488, 524]])

我们可以使用 np.allclose 来查看矩阵之间是否近乎相等：

[11]:

print(np.allclose(C_byat, C_bydot))
print(np.allclose(C_byein, C_bydot))

True
True

任务 (3)

使用 np.dot 的写法，写一段三矩阵连续相乘的代码；并指出在三矩阵相乘时，使用 A.dot(B) 比 np.dot(A, B) 的便利之处．在量化计算中，我们经常会遇到原子轨道 (AO) 矩阵转换到分子轨道 (MO) 矩阵的操作，这是一个三矩阵的计算．
使用 @ 与 np.einsum 写一段三矩阵连续相乘的代码，思考哪种方式更适合自己的公式代码化的思路．我相信不同的人有不同的见解．注意 @ 需要用到两次，但 np.einsum 应该只能用到一次，即不应该出现如下代码：
```
np.einsum("ik, kj -> ij", A, np.einsum("kl, lj -> kj", B, C))
```
(可选) 将
```
np.einsum("ik, kj -> ij", A, B)
```
更改为
```
np.einsum("ik, kj", A, B)
```
并查看效果．尽管我自己不太推荐这种做法，但这是 Einstein Convention 中“角标出现两次或以上时就对角标求和”的具体例子．

4.4. 矩阵乘积效率评估

下面我们简单地讨论上述三种矩阵乘法的效率。在 Jupyter Notebook 中，可以使用 timeit magic command (API 文档) 进行大致的效率评估。

[12]:

%%timeit -r 7 -n 10000
A @ B

1.26 µs ± 90 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

[13]:

%%timeit -r 7 -n 10000
A.dot(B)

1.12 µs ± 62.6 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

[14]:

%%timeit -r 7 -n 10000
np.einsum("ik, kj -> ij", A, B)

1.91 µs ± 160 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

我们能发现，在效率上，np.dot 通常是最快的，而 np.einsum 通常是最慢的。但这未必是确定的；在矩阵较大时，np.matmul 会稍快一些，这可以留待读者测试。

尽管一般来说，np.einsum 的效率较低，但在处理巨大的或量化张量时，np.einsum 通常不会给出太差的效率；甚至有时自己使用 np.matmul 的效率不见得比 np.einsum 高。

任务 (4)

(可选) 对三个 (1000, 1000) 矩阵相乘的情况作效率测评。

4.5. 矩阵与向量运算：内积运算

类似于普通的矩阵乘法，矩阵与向量也有普通的矩阵乘法或称内积运算。该过程即

\[b_i = A_{ij} x_j\]

[15]:

A = np.array([[1, 2, 3], [11, 12, 13]])  # A.shape = (2, 3)
x = np.array([5, 6, 7])                  # x.shape = (3, )

上述的三种矩阵乘法操作在矩阵与向量计算中仍然成立：

[16]:

b_byat = A @ x
b_bydot = A.dot(x)  # equivalent to np.dot(A, x)
b_byein = np.einsum("ij, j -> i", A, x)
b_bydot

[16]:

array([ 38, 218])

[17]:

print(np.allclose(b_byat, b_bydot))
print(np.allclose(b_byein, b_bydot))

True
True

任务 (5)

我们现在已经了解矩阵之间、矩阵与向量的乘法关系的代码编写了；对于三维张量与一维向量或二维矩阵之间的乘法运算，能否给出类似的代码？
(可选) 若现在是三维张量与三维张量之间的乘法运算，np.dot 与 np.matmul 是如何工作的？

4.6. 张量转置

处理量化问题时经常会遇到张量转置的问题。张量转置的弱化问题是矩阵转置；这几乎是显然的。在 numpy 中，通常可以用 np.transpose 或者直接在矩阵后使用 T 方法获得矩阵：

[18]:

lst = [
    [0, 1, 2],
    [10, 11, 12]
]
mat = np.array(lst)
mat

[18]:

array([[ 0,  1,  2],
       [10, 11, 12]])

[19]:

mat.T

[19]:

array([[ 0, 10],
       [ 1, 11],
       [ 2, 12]])

[20]:

print(np.allclose(np.transpose(mat), mat.T))
print(np.allclose(mat.transpose(), mat.T))

True
True

但对于一般的张量转置问题，就不是那么显然了。除了 np.transpose 或者 T 方法外， np.einsum 也能进行转置，其优势仍然是直观，但效率偏低，比较适合作为验证工具；np.swapaxes 也可以解决张量转置工具，但只适合于对换两个角标。

我们拿一个不等维度的三维张量进行说明。转置的目标是

\[R_{ijk} \rightarrow R_{ikj}\]

待转置张量是 R：

[21]:

R = np.arange(24).reshape(2, 3, 4)
R

[21]:

array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])

成功转置的张量是 R_T：

[22]:

R_T = np.einsum("ijk -> ikj", R)
R_T

[22]:

array([[[ 0,  4,  8],
        [ 1,  5,  9],
        [ 2,  6, 10],
        [ 3,  7, 11]],

       [[12, 16, 20],
        [13, 17, 21],
        [14, 18, 22],
        [15, 19, 23]]])

[23]:

print(np.allclose(R.transpose(0, 2, 1), R_T))
print(np.allclose(R.swapaxes(1, 2), R_T))

True
True

任务 (6)

现在考虑涉及三角标的轮换的转置．如果现在的转置目标是 \(R_{ijk} \rightarrow R_{jki}\)，是应该使用 R.transpose(1, 2, 0)，还是 R.transpose(2, 0, 1)？请用 np.einsum("ijk -> jki", R)，或者用 shape 查看转之后张量的形状信息辅助验证．

4.7. 参考任务解答

4.7.1. 任务 (1)

4.7.1.1. 任务 (1.1) 可选

尽管原题目是希望通过 python 三维列表生成 numpy 三维张量，但我们也可以使用偷懒一些的方法。使用 np.arange 可以生成类似于 python 的 range 迭代器，但 np.arange 最终返回的是以 ndarray 形式的迭代器，即一维向量：

[24]:

tensor = np.arange(24)
tensor

[24]:

array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,
       17, 18, 19, 20, 21, 22, 23])

随后我们将上述向量重塑为 (2, 3, 4) 的三维张量的形状：

[25]:

tensor.shape = (2, 3, 4)
tensor

[25]:

array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])

以此为出发，可以了解 ndarray 三维张量的索引方式。应当能发现 numpy 的索引基本与 C 的索引方式等同。

4.7.2. 任务 (2)

4.7.2.1. 任务 (2.1) 可选

矩阵计算中，\(\mathbf{B}' = \mathbf{M}^3\) 事实上用 Einstein Notation，应当表示为

\[B'_{ij} = M_{ik} M_{kl} M_{lj}\]

即三次连续的矩阵乘法。因此，

[26]:

lst = [
    [ 0,  1,  2],
    [10, 11, 12],
    [20, 21, 22],
]
mat = np.array(lst)

[27]:

B_prime = mat @ mat @ mat
B_prime

[27]:

array([[ 1650,  1809,  1968],
       [12150, 13299, 14448],
       [22650, 24789, 26928]])

np.linalg.matrix_power (API 文档) 可以实现整数的矩阵幂次的运算：

[28]:

np.linalg.matrix_power(mat, 3)

[28]:

array([[ 1650,  1809,  1968],
       [12150, 13299, 14448],
       [22650, 24789, 26928]])

scipy.linalg.fractional_matrix_power (API 文档) 功能更强大，可以实现非整数的矩阵幂次运算 (但这里作为范例仍然使用整数)：

[29]:

scipy.linalg.fractional_matrix_power(mat, 3.)

[29]:

array([[ 1650,  1809,  1968],
       [12150, 13299, 14448],
       [22650, 24789, 26928]])

4.7.2.2. 任务 (2.2) 可选

我们先给出矩阵计算下 \(\mathbf{C}' = \exp(- \mathbf{M} / 5)\) 的结果。这通过 `scipy.linalg.fractional_matrix_power <https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.expm.html>`__ 实现：

[30]:

C_prime = scipy.linalg.expm(- mat / 5)
C_prime

[30]:

array([[ 1.28820913,  0.07655702, -0.1350951 ],
       [-0.27198272,  0.68929065, -0.34943599],
       [-0.83217457, -0.69797572,  0.43622312]])

下面我们使用更容易理解的语言解释矩阵计算的指数运算。指数运算从定义上是 Taylor 展开；我们这里也使用这一事实：

\[\mathbf{C}' = \sum_{n = 1}^{\infty} \frac{(- \mathbf{M} / 5)^n}{n!}\]

下面的代码就是将上述 Taylor 展开截断到 \(n = 50\) 所给出；我们应当能看到 C_prime 与 C_prime_truncated 几乎相同：

[31]:

C_prime_truncated = np.zeros_like(mat, dtype=np.float64)
for n in range(0, 50):
    C_prime_truncated += np.linalg.matrix_power(- mat / 5., n) * (1. / np.math.factorial(n))
C_prime_truncated

[31]:

array([[ 1.28820913,  0.07655702, -0.1350951 ],
       [-0.27198272,  0.68929065, -0.34943599],
       [-0.83217457, -0.69797572,  0.43622312]])

不过对于这个例子，在 Taylor 展开截断较少，譬如 \(n = 10\) 附近时，还无法看出收敛的迹象。

4.7.3. 任务 (3)

这一部分任务中，我们指定 A, B, C 如下。矩阵计算的目标是

\[D_{ij} = A_{ik} B_{kl} C_{lj}\]

[32]:

np.random.seed(0)
A = np.random.randn(2, 3)
B = np.random.randn(3, 4)
C = np.random.randn(4, 5)

4.7.3.1. 任务 (3.1)

\(D_{ij}\) 的表达式为 (维度为 (2, 5))

[33]:

D = A.dot(B).dot(C)
D

[33]:

array([[ 1.49830283,  1.73384943, -6.13304366,  2.74972964,  1.03516199],
       [-3.99086154,  2.1109627 , -7.82833666,  2.93379787,  3.17459153]])

如果要使用 np.dot，代码会变得不太容易理解，但仍然能得到正确结果：

[34]:

np.dot(A, np.dot(B, C))

[34]:

array([[ 1.49830283,  1.73384943, -6.13304366,  2.74972964,  1.03516199],
       [-3.99086154,  2.1109627 , -7.82833666,  2.93379787,  3.17459153]])

4.7.3.2. 任务 (3.2)

使用 @ 的代码如下：

[35]:

A @ B @ C

[35]:

array([[ 1.49830283,  1.73384943, -6.13304366,  2.74972964,  1.03516199],
       [-3.99086154,  2.1109627 , -7.82833666,  2.93379787,  3.17459153]])

使用 @ 的代码显然是所有矩阵乘法中最短小的写法；而 np.einsum 可以认为是最为省事的做法，因为它能完整地还原公式中张量的乘法过程：

[36]:

np.einsum("ik, kl, lj -> ij", A, B, C)

[36]:

array([[ 1.49830283,  1.73384943, -6.13304366,  2.74972964,  1.03516199],
       [-3.99086154,  2.1109627 , -7.82833666,  2.93379787,  3.17459153]])

尽管显然下述的代码也能给出正确的结果：

[37]:

np.einsum("ik, kj -> ij", A, np.einsum("kl, lj -> kj", B, C))

[37]:

array([[ 1.49830283,  1.73384943, -6.13304366,  2.74972964,  1.03516199],
       [-3.99086154,  2.1109627 , -7.82833666,  2.93379787,  3.17459153]])

但上面这种对 np.einsum 的使用方法反而会将问题复杂化。

4.7.3.3. 任务 (3.3)

我们不妨尝试 \(D_{ij} = A_{ik} B_{kl} C_{lj}\)：

[38]:

np.einsum("ik, kl, lj", A, B, C)

[38]:

array([[ 1.49830283,  1.73384943, -6.13304366,  2.74972964,  1.03516199],
       [-3.99086154,  2.1109627 , -7.82833666,  2.93379787,  3.17459153]])

但假设，现在遇到的张量乘法的情形是 \(D_{ji} = A_{ik} B_{kl} C_{lj}\)，那么上述的代码不能一次性地给出正确的结果，而一定要写为

[39]:

np.einsum("ik, kl, lj -> ji", A, B, C)

[39]:

array([[ 1.49830283, -3.99086154],
       [ 1.73384943,  2.1109627 ],
       [-6.13304366, -7.82833666],
       [ 2.74972964,  2.93379787],
       [ 1.03516199,  3.17459153]])

同时，如果现在的问题是

\[D_{ik} = A_{ik} B_{kl} C_{lj}\]

或者使用普通的求和记号，即

\[D_{ik} = \sum_{jl} A_{ik} B_{kl} C_{lj}\]

那么这种情况用完整的张量缩并字符串会比较方便：

[40]:

np.einsum("ik, kl, lj -> ik", A, B, C)

[40]:

array([[-2.91662457,  0.68839088,  3.11223391],
       [-3.70501713,  3.21276204, -3.10759101]])

4.7.4. 任务 (4)

4.7.4.1. 任务 (4.1) 可选

[41]:

A = np.random.randn(1000, 1000)
B = np.random.randn(1000, 1000)
C = np.random.randn(1000, 1000)

[42]:

%%timeit -r 7 -n 10
A @ B @ C

54.2 ms ± 3.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

[43]:

%%timeit -r 7 -n 10
A.dot(B).dot(C)

57.2 ms ± 2.99 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

[44]:

%%timeit -r 7 -n 10
np.einsum("ik, kl, lj -> ij", A, B, C, optimize=True)

60.8 ms ± 4.25 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

可以看出，对于较大的矩阵操作，np.einsum 的效率并没有比其它两种乘法更差。因此，在实际实践中，即使微秒级的运算中 np.einsum 的效率非常低，但真正影响量化计算的微秒、秒级运算中，np.einsum 完全是可以使用的。

需要注意，np.einsum 中 optimize 选项需要设定为 True 或其它优化方式。关于这部分，我们将在下一节文档中作更详细的描述。

4.7.5. 任务 (5)

4.7.5.1. 任务 (5.1)

现在假定 A 为维度 (2, 3, 5) 的张量，B 为 (5, 4) 的张量，那么对于下述问题，

\[C_{ijk} = A_{ijl} B_{lk}\]

可以清楚验证的方式是 np.einsum：

[45]:

np.random.seed(0)
A = np.random.randint(50, size=(2, 3, 5))
B = np.random.randint(50, size=(5, 4))

[46]:

C = np.einsum("ijl, lk -> ijk", A, B)
C

[46]:

array([[[2494, 3229, 2399, 2356],
        [3049, 2791, 1724, 3636],
        [2151, 1825, 1122, 2712]],

       [[4248, 2183, 3673, 4987],
        [2861, 2394, 1875, 3623],
        [1482, 1038,  980, 1872]]])

[47]:

print(np.allclose(A @ B, C))
print(np.allclose(A.dot(B), C))

True
True

4.7.5.2. 任务 (5.2) 可选

现在假定 A 为维度 (2, 3, 5) 的张量，B 为维度 (2, 5, 3) 的张量：

[48]:

np.random.seed(0)
A = np.random.randint(50, size=(2, 3, 5))
B = np.random.randint(50, size=(2, 5, 3))

对于 A @ B，其形式为

\[C_{pij} = A_{pik} B_{pkj}\]

结果是 (2, 3, 3) 维度的张量：

[49]:

np.allclose(A @ B, np.einsum("pik, pkj -> pij", A, B))

[49]:

True

而对于 A.dot(B)，其形式为

\[C_{piqj} = A_{pik} B_{qkj}\]

结果是 (2, 3, 2, 3) 维度的张量：

[50]:

np.allclose(A.dot(B), np.einsum("pik, qkj -> piqj", A, B))

[50]:

True

4.7.6. 任务 (6)

4.7.6.1. 任务 (6.1)

待转置张量是 R，转置张量则是 R_T：

[51]:

R = np.arange(24).reshape(2, 3, 4)
R_T = np.einsum("ijk -> jki", R)
R_T.shape

[51]:

(3, 4, 2)

对于题目中给出的两种转置，

[52]:

print(R.transpose(1, 2, 0).shape)
print(R.transpose(2, 0, 1).shape)

(3, 4, 2)
(4, 2, 3)

我们能注意到，R.transpose(1, 2, 0) 应当是正确的转置。

[53]:

np.allclose(R.transpose(1, 2, 0), R_T)

[53]:

True