2.3. 需要的 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
2.3.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 任务中,我们需要了解最高八维张量的计算,以及最高四维张量的输出.对任意维度张量的操控与调试是后面笔记中所经常使用的能力.
2.3.2. 元素操作¶
numpy 中的许多操作,包括张量与数 (scaler)、算符 (operation) 的操作,是元素操作 (elementwise manuplation)。举例来说,对于矩阵与数的加法
这就等价于
[6]:
A = mat + 5
A
[6]:
array([[ 5, 6, 7],
[15, 16, 17],
[25, 26, 27]])
幂次计算则可以表示为
[7]:
B = mat ** 3
B
[7]:
array([[ 0, 1, 8],
[ 1000, 1331, 1728],
[ 8000, 9261, 10648]])
对于指数计算也类似。譬如下述的的过程表示
[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)\)。
2.3.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 表示为
矩阵乘积的最标准写法是使用 @
或等价地 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 中“角标出现两次或以上时就对角标求和”的具体例子.
2.3.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) 矩阵相乘的情况作效率测评。
2.3.5. 矩阵与向量运算:内积运算¶
类似于普通的矩阵乘法,矩阵与向量也有普通的矩阵乘法或称内积运算。该过程即
[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
是如何工作的?
2.3.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
:
[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
查看转之后张量的形状信息辅助验证.
2.3.7. 参考任务解答¶
2.3.7.1. 任务 (1)¶
2.3.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 的索引方式等同。
2.3.7.2. 任务 (2)¶
2.3.7.2.1. 任务 (2.1) 可选¶
矩阵计算中,\(\mathbf{B}' = \mathbf{M}^3\) 事实上用 Einstein Notation,应当表示为
即三次连续的矩阵乘法。因此,
[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]])
2.3.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 展开;我们这里也使用这一事实:
下面的代码就是将上述 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\) 附近时,还无法看出收敛的迹象。
2.3.7.3. 任务 (3)¶
这一部分任务中,我们指定 A
, B
, C
如下。矩阵计算的目标是
[32]:
np.random.seed(0)
A = np.random.randn(2, 3)
B = np.random.randn(3, 4)
C = np.random.randn(4, 5)
2.3.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]])
2.3.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
的使用方法反而会将问题复杂化。
2.3.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]])
同时,如果现在的问题是
或者使用普通的求和记号,即
那么这种情况用完整的张量缩并字符串会比较方便:
[40]:
np.einsum("ik, kl, lj -> ik", A, B, C)
[40]:
array([[-2.91662457, 0.68839088, 3.11223391],
[-3.70501713, 3.21276204, -3.10759101]])
2.3.7.4. 任务 (4)¶
2.3.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
或其它优化方式。关于这部分,我们将在下一节文档中作更详细的描述。
2.3.7.5. 任务 (5)¶
2.3.7.5.1. 任务 (5.1)¶
现在假定 A
为维度 (2, 3, 5) 的张量,B
为 (5, 4) 的张量,那么对于下述问题,
可以清楚验证的方式是 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
2.3.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
,其形式为
结果是 (2, 3, 3) 维度的张量:
[49]:
np.allclose(A @ B, np.einsum("pik, pkj -> pij", A, B))
[49]:
True
而对于 A.dot(B)
,其形式为
结果是 (2, 3, 2, 3) 维度的张量:
[50]:
np.allclose(A.dot(B), np.einsum("pik, qkj -> piqj", A, B))
[50]:
True
2.3.7.6. 任务 (6)¶
2.3.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