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)

  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)

  1. (可选) 上述 \(M_{ij}^3\) 并不是通常矩阵运算中的 \(\mathbf{M}^3\)。请指出如何计算 \(\mathbf{M}^3\)
  2. (可选) 上述 \(\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)

  1. 使用 np.dot 的写法,写一段三矩阵连续相乘的代码;并指出在三矩阵相乘时,使用 A.dot(B)np.dot(A, B) 的便利之处.在量化计算中,我们经常会遇到原子轨道 (AO) 矩阵转换到分子轨道 (MO) 矩阵的操作,这是一个三矩阵的计算.

  2. 使用 @np.einsum 写一段三矩阵连续相乘的代码,思考哪种方式更适合自己的公式代码化的思路.我相信不同的人有不同的见解.注意 @ 需要用到两次,但 np.einsum 应该只能用到一次,即不应该出现如下代码:

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

    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)

  1. (可选) 对三个 (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)

  1. 我们现在已经了解矩阵之间、矩阵与向量的乘法关系的代码编写了;对于三维张量与一维向量或二维矩阵之间的乘法运算,能否给出类似的代码?
  2. (可选) 若现在是三维张量与三维张量之间的乘法运算,np.dotnp.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)

  1. 现在考虑涉及三角标的轮换的转置.如果现在的转置目标是 \(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_primeC_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.einsumoptimize 选项需要设定为 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