5. 需要的 NumPy 技巧:进阶
上一节我们回顾了了一部分 numpy 技巧,它们可以解决通常的量化问题或者其它应用问题。这一节我们回顾一些 numpy 中不太常用的,但在 pyxdh
包中使用到的一些技巧。
[1]:
import numpy as np
5.1. ndarray 数据储存
在进行下面的话题前,我们有必要简单了解 ndarray 储存方式,以及使用 np.base
判断 ndarray 之间是否共享相同的数据内存空间。这有助于我们判断是否修改了张量的数值,或者判断是否复制了多余且无用的内存数据。这部分的描述也可以参考 一份公开书籍。
ndarray 类型可以看成由底层数据与修饰信息构成;底层数据就是一个一个数值,而修饰信息则是由维度、拆分方式等构成。一般来说,修饰信息的内存占用总远小于底层数据。
在 numpy 中,除非必要,否则大多数时候不会真正地复制底层数据,只复制修饰信息。我们可以使用下面的 get_memloc
函数来获取 ndarray 底层数据的头指针在内存的位置信息:
[2]:
def get_base(array):
while isinstance(array.base, np.ndarray):
array = array.base
return array
def get_memloc(array):
return get_base(array).__array_interface__['data'][0]
[3]:
R = np.arange(24).reshape((2, 3, 4))
get_base(R)
[3]:
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])
[4]:
get_memloc(R)
R_NoCopy = R # Nothing changed
R_ShallowCopy = R.view()
R_DeepCopy = R.copy()
print(R is R_NoCopy) # Equivalent to `id(R) == id(R_NoCopy)`
print(R is R_ShallowCopy) # Different id means objects are actually different
print(get_memloc(R) == get_memloc(R_ShallowCopy)) # ... but a shallow copy means that base data is not actually copied
print(get_memloc(R) == get_memloc(R_DeepCopy)) # A deep copy means all things are cloned to another block of memory
print(np.all(R == R_DeepCopy)) # ... however, the value of copy is the exactly the same to origin
True
False
True
False
True
is
作为 python 命令是判别两个对象是否完全相同;但 get_memloc
所给出的是作为 ndarray 底层数据的指针;因此,如果对于两个 ndarray 对象,is
或者 get_memloc
给出相等的判断,那么底层数据就没有被复制。
上述代码中,直接赋值与浅复制就没有复制底层的数据,因此代码的运行几乎没有事件与内存的消耗,但修改其中的数据的过程是不可逆的。而深复制则会复制整个底层的数据,但反过来修改深复制后的数据不会对原始的数据造成影响。
任务 (1)
(可选) 这里的
get_base
函数使用了 Python 的isinstance
函数.该函数在 PySCF 的源代码中涉及到临时信息如何储存的部分中经常出现.但是我们也可以不使用循环与isinstance
函数,而使用递归方法来写get_base
.试写出这样一个函数.提示:你应当可以发现
type(R.base) == np.ndarray
但R.base.base is None == True
.
一般来说,单纯的张量转置与重塑不会进行深复制:
[5]:
print(get_memloc(R) == get_memloc(R.reshape(-1, 6)))
print(get_memloc(R) == get_memloc(R.reshape(8, 3)))
print(get_memloc(R) == get_memloc(R.transpose(1, 0, 2).reshape(3, 2, 2, 2)))
True
True
True
但下述的代码的最后行则会给出深复制的结果:
[6]:
print(R.transpose(0, 2, 1).shape)
print(get_memloc(R) == get_memloc(R.transpose(0, 2, 1).reshape(2, 4, 3)))
print(get_memloc(R) == get_memloc(R.transpose(0, 2, 1).reshape(8, 3)))
(2, 4, 3)
True
False
原因应该是,对于一个 ndarray 对象,其储存需要使用连续的内存;如果转置与重塑矩阵使得内存的储存无法连续,那么 numpy 会进行深复制。
同时,需要指出,即使转置后的矩阵与转置前的矩阵值不相等,但底层的数据仍然是相等的:
[7]:
A = np.random.random((10, 10))
print(np.allclose(A, A.T))
print(get_memloc(A) == get_memloc(A.T))
False
True
最后我们讨论运算下的情况。尽管对于 python 中简单普通的数值运算,下述关系 (有时) 是成立的:
[8]:
(1.5 + 1.5) is (9. / 3.)
[8]:
True
但对于 numpy 的矩阵操作,这通常不成立:
[9]:
A = B = np.random.random((10, 10))
print(A @ B is A @ B)
print(np.all(A @ B == A @ B))
False
True
有趣的是,使用 np.einsum
进行矩阵转置仍然是浅复制:
[10]:
print(get_memloc(R) == get_memloc(np.einsum("ijk -> jki", R)))
True
5.2. 向量截取
附近几小节 (到向量直和为止) 的内容可以参考 numpy 的 Boardcasting 文档。
我们曾经说过,numpy 可以使用类似于 Fortran 的截取方式:
[11]:
R = np.arange(15)
R[1:12:3]
[11]:
array([ 1, 4, 7, 10])
我们现在就可以检查截取 R[:, 1, 3]
与 R
是否共享相同的内存空间:
[12]:
get_memloc(R) == get_memloc(R[1:12:3])
[12]:
True
但是需要留意,因为这种截取共享了相同的内存空间,因此对其更改会直接影响原始数据:
[13]:
R[1:12:3] = [-1, -2, -3, -4]
R
[13]:
array([ 0, -1, 2, 3, -2, 5, 6, -3, 8, 9, -4, 11, 12, 13, 14])
这种截取方式称为 slice,与下述的代码完全等价:
[14]:
R[slice(1, 12, 3)]
[14]:
array([-1, -2, -3, -4])
如果对这种截取进行重新赋值,也会直接地改变原始向量的数据:
[15]:
R[slice(1, 12, 3)] = [-5, -6, -7, -8]
R
[15]:
array([ 0, -5, 2, 3, -6, 5, 6, -7, 8, 9, -8, 11, 12, 13, 14])
使用 slice(start, stop, step)
的方便之处是可以将分割作为一种对象来赋值,这在以后处理原子轨道的分割时会经常使用。
除了 slice 分割之外,还可以使用列表或者 range 截取向量中的元素:
[16]:
R[[1, 4, 7, 10]]
[16]:
array([-5, -6, -7, -8])
[17]:
R[range(1, 12, 3)]
[17]:
array([-5, -6, -7, -8])
尽管列表的截取相对来说更为自由 (可以截取任意位置的元素,不受制于 slice 的 start, stop, step 模式),但它们实际上很可能作了一层深复制:
[18]:
print(get_memloc(R) == get_memloc(R[range(1, 12, 3)]))
print(get_memloc(R) == get_memloc(R[[1, 4, 7, 10]]))
False
False
但这仍然可以更改底层的数据:
[19]:
R[[1, 4, 7, 10]] = [0, 0, 0, 0]
R
[19]:
array([ 0, 0, 2, 3, 0, 5, 6, 0, 8, 9, 0, 11, 12, 13, 14])
对于这种更改底层数据的特性,需要补充的例子是,如果重复更改底层数据,那么一般来说会用最后一个值覆盖:
[20]:
R[[1, 2, 2, 1]] = [-1, -2, -3, -4]
R
[20]:
array([ 0, -4, -3, 3, 0, 5, 6, 0, 8, 9, 0, 11, 12, 13, 14])
5.3. 矩阵截取与对角元
下面我们讨论矩阵的截取。
[21]:
A = np.arange(25).reshape(5, 5)
A
[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, 24]])
平时最为常用的截取方式仍然是使用 slice 截取子矩阵:
[22]:
A[slice(3), slice(3)]
[22]:
array([[ 0, 1, 2],
[ 5, 6, 7],
[10, 11, 12]])
但按相同的方式,在列表和 range 截取的结果会完全不同:
[23]:
A[range(3), range(3)]
[23]:
array([ 0, 6, 12])
我们仍然应当能发现这几种截取使用不同的底层数据复制方式:
[24]:
print(get_memloc(A) == get_memloc(A[slice(3), slice(3)]))
print(get_memloc(A) == get_memloc(A[range(3), range(3)]))
True
False
但这两种方法都能对底层数据作修改:
[25]:
A[slice(3), slice(3)] = np.arange(9).reshape(3, 3)
A[range(3), range(3)] = - np.arange(3)
A
[25]:
array([[ 0, 1, 2, 3, 4],
[ 3, -1, 5, 8, 9],
[ 6, 7, -2, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]])
依据列表或 range 的截取方式,我们可以获得矩阵的对角元。截取对角元在处理分子轨道基组下的矩阵、或处理 U 矩阵的对角元时会有所帮助:
[26]:
A[range(5), range(5)]
[26]:
array([ 0, -1, -2, 18, 24])
在 numpy 中,函数 np.diagonal
也能截取对角元:
[27]:
A.diagonal()
[27]:
array([ 0, -1, -2, 18, 24])
尽管这种截取是浅复制:
[28]:
get_memloc(A.diagonal()) == get_memloc(A)
[28]:
True
但这种截取不可以进行赋值:
[29]:
A.diagonal() = - np.arange(5)
File "<ipython-input-29-c113c2cc9ca7>", line 1
A.diagonal() = - np.arange(5)
^
SyntaxError: can't assign to function call
如果希望对矩阵求迹,还可以使用 np.trace
:
[30]:
A.trace()
[30]:
39
5.4. 张量截取与对角元
下面我们考虑张量的截取。在梯度求取过程中,我们会经常遇到最后两维度的大小等于原子轨道数的张量;这里用 (2, 3, 3) 维张量来描述这一问题。
[31]:
R = np.arange(18).reshape(2, 3, 3)
R
[31]:
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8]],
[[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17]]])
如果原子轨道数为 3,其中有两个原子,第一个原子占用轨道 1,第二个原子占用轨道 2, 3;那么截取其中第二个原子的部分就可以写为
[32]:
R[:, slice(1, 3), slice(1, 3)]
[32]:
array([[[ 4, 5],
[ 7, 8]],
[[13, 14],
[16, 17]]])
如果要取该张量最后两维度的对角元,则写为
[33]:
R.diagonal(axis1=-2, axis2=-1)
[33]:
array([[ 0, 4, 8],
[ 9, 13, 17]])
5.5. 张量与向量和
向量的直和会在出现 \(\varepsilon_i - \varepsilon_a\) 的情形中出现.在不考虑内存占用效率的情况下,我们可以定义 \(D_i^a = \varepsilon_i - \varepsilon_a\) 矩阵;这个矩阵的生成方式即是向量 \(\varepsilon_i\) 与 \(- \varepsilon_a\) 的直和.
求向量直和的情形可以出现在 CP-HF 方程的等式左,也可以出现在 MP2 方法的分母的求轨道能差中.在其它语言中,这些计算完全可以通过二重或四重循环解决;但一般认为 Python 自身的的循环效率较低,因此若单个循环体的计算量不大,可以不利用循环之处要尽量不用.在这种情况下,牺牲内存而保证一定效率的构建 \(D_i^a\) 矩阵或 \(D_{ij}^{ab}\) 张量的做法通常会提高计算效率.
在 numpy 中没有直接配备向量直和的函数.通常,我们会采用下述方式生成直和:
[34]:
eo = - np.arange(1, 4)
ev = np.arange(1, 6)
D_ia = eo[:, None] - ev[None, :]
D_ia
[34]:
array([[-2, -3, -4, -5, -6],
[-3, -4, -5, -6, -7],
[-4, -5, -6, -7, -8]])
这个过程可以看作是将 eo
与 ev
分别展开为两个矩阵从而相减;但两者分别展开为 (3, 1) 与 (1, 5) 维矩阵。随后,维度不为 1 的部分与维度为 1 的部分之间作类似于向量与数的元素操作,从而这两个矩阵相减后成为 (3, 5) 的矩阵。
在 numpy 中,维度前置为 1 的部分在直和中是可有可无的:
[35]:
np.allclose(eo[:, None] - ev, D_ia)
[35]:
True
5.6. np.allclose
的补充说明
我们已经使用很多次 np.allclose
了,它是判断两个张量之间是否几乎相等的函数。之所以并未必完全相等,是因为 np.allclose
实际上允许一些微小的数值偏差。事实上,如果要判断两个矩阵是否 完全 相等,应当是通过 np.all
来判断。譬如对于下述矩阵:
[36]:
np.random.seed(2)
A = np.random.random((3, 3))
B = A.copy()
这两个矩阵应当是完全相等的:
[37]:
A == B
[37]:
array([[ True, True, True],
[ True, True, True],
[ True, True, True]])
np.all
会判断布尔张量的元素是否全部为 True:
[38]:
np.all(A == B)
[38]:
True
显然,如果用 np.allclose
也会给出 True:
[39]:
np.allclose(A, B)
[39]:
True
但如果引入非常微小的偏差,np.all
会给出这两个矩阵不相等,但 np.allclose
则会给出大致相等的结论:
[40]:
B = A + 1e-10
print(A == B)
print(np.all(A == B))
print(np.allclose(A, B))
[[False False False]
[False False False]
[False False False]]
False
True
np.all
与 np.allclose
都是判断一个 ndarray 对象元素是否全部为 True。np.all
是对 A == B
的元素作判断,np.allclose
则是对 np.isclose(A, B)
的元素作判断:
[41]:
np.isclose(A, B)
[41]:
array([[ True, True, True],
[ True, True, True],
[ True, True, True]])
但如果微小的偏差渐渐地不再微小,那么 np.isclose
就不会给出全部 True 的 ndarray 对象;并且即使只有一个值为 False,那么 np.allclose
就会给出 False:
[42]:
B = A + 1e-6
print(np.isclose(A, B))
np.allclose(A, B)
[[ True False True]
[ True True True]
[ True True True]]
[42]:
False
但如果降低 np.allclose
或 np.isclose
的判断阈值,就可以给出 True 的结果了。atol
代表绝对值阈值,默认为 1e-8;rtol
代表相对值阈值,默认为 1e-5,这些可以参考 API 文档。
[43]:
print(np.allclose(A, B, atol=8e-7))
print(np.allclose(A, B, rtol=5e-5))
True
True
也可以混合两种判断阈值:
[44]:
print(np.allclose(A, B, atol=5e-7))
print(np.allclose(A, B, rtol=3e-5))
print(np.allclose(A, B, atol=5e-7, rtol=3e-5))
False
False
True
以后我们会非常经常地使用 np.allclose
来判断我们的计算结果与 PySCF 程序的,或与数值结果的是否相似。
5.7. np.einsum
的使用:布尔值
在整个笔记中,我们会大量地使用 np.einsum
函数。这个函数可以非常直观地将公式与程序对应起来。从作者的角度,np.einsum
可以确实地颠覆了对量化程序或者普遍的张量缩并程序编写的认知。
以前代码中,我们已经对 np.einsum
有基本的认识,一般来说足够应对以后的张量乘积的处理;而后面三小节我们会稍稍了解一下这个函数的进阶使用方式。
optimize
选项是优化算法的选项。optimize
可以设定为布尔值,以判断是否自动优化算法。所谓自动优化,是指根据张量的维度与乘积的方式自动优化计算过程。我们模拟双电子积分转换的情形,来了解自动优化的意义。双电子积分的计算公式是
我们用 C
代表系数矩阵 \(C_{\mu p}\),而用 eri
代表 AO 基组下的双电子积分 (ERI) 张量;原子轨道基组与分子轨道数都定为 5。下面的程序并不处理真正的化学问题,但可以评价处理化学问题过程中算法的效率。
[45]:
nmo = nao = 5
C = np.random.random((nmo, nmo))
eri = np.random.random((nmo, nmo, nmo, nmo))
[46]:
%%timeit -r 7 -n 10
np.einsum("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=True)
412 µs ± 129 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
[47]:
%%timeit -r 7 -n 10
np.einsum("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=False)
9.05 ms ± 185 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
我们会发现,如果使用打开 optimize=True
的代码,运行效率会快很多,并且时间复杂度是 \(O(N^5)\);否则时间复杂度就会成为 \(O(N^8)\)。事实上,如果基组数再大一些,optimize=False
的代码就会基本卡死。
但是必须指出,默认情况下,np.einsum
不会打开 自动优化的代码。我们从下面模拟 AO 基组密度生成的例子 (也是我们以前使用过的普通矩阵乘法的例子) 就能说明原因:
[48]:
%%timeit -r 7 -n 1000
np.einsum("ui, vi -> uv", C, C) # Equivalent to `optimize=False` added
2.06 µs ± 315 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
[49]:
%%timeit -r 7 -n 1000
np.einsum("ui, vi -> uv", C, C, optimize=True)
70.4 µs ± 3.11 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
有两个缘由导致上述的情况。其一是,两个系数矩阵相乘的可能方法 (几乎) 只有一种,时间复杂度是 \(O(N^3)\) 并且无法优化。其二是,打开 optimize=True
后,程序会使用与 optimize=False
完全不同的方式构建张量,并且还会使用 Python 而非 C 的代码进行最优时间复杂度的算法搜索。因此,在处理小矩阵计算,或者无法优化代码的计算中,optimize=True
具有非常大的劣势。反之,一旦计算过程可以被优化,optimize=True
则会体现出极大的优势。大多数情况下,手写张量缩并程序会比 np.einsum
的效率慢一些。
5.8. np.einsum
的使用:np.einsum_path
np.einsum_path
则是给出当前计算的算法,它不实际执行计算;它将输出二维元组。我们仍然拿双电子积分转换为例:
[50]:
path, info = np.einsum_path("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=True)
我们先来看相对友好的路径信息:
[51]:
print(info)
Complete contraction: up,vq,uvkl,kr,ls->pqrs
Naive scaling: 8
Optimized scaling: 5
Naive FLOP count: 1.953e+06
Optimized FLOP count: 2.500e+04
Theoretical speedup: 78.122
Largest intermediate: 6.250e+02 elements
--------------------------------------------------------------------------
scaling current remaining
--------------------------------------------------------------------------
5 uvkl,up->klpv vq,kr,ls,klpv->pqrs
5 klpv,vq->klpq kr,ls,klpq->pqrs
5 klpq,kr->lpqr ls,lpqr->pqrs
5 lpqr,ls->pqrs pqrs->pqrs
我们知道,双电子积分从 AO 到 MO 的转化若未经优化是 \(O(N^8)\) 的时间复杂度,而优化后是 \(O(N^5)\) 的时间复杂度。
第 2, 3 行:优化前后的时间复杂度;
第 4, 5 行:优化前后的代码的计算次数;
第 6 行:代表优化前后的速度比值;
第 7 行:算法计算过程中,中间张量所占用的空间 (以元素数记而非实际比特大小)
第 8 行及以后:算法的实际计算过程.
每个张量计算后,生成的临时矩阵会放在下一次张量缩并列表的最后。
随后我们来看实际交给程序阅读的路径信息:
[52]:
path
[52]:
['einsum_path', (0, 2), (0, 3), (0, 2), (0, 1)]
其第一个元素是 'einsum_path'
,之后的信息是每个最基本张量乘积缩并的操作对象的编号。具体来说,我们一开始代入程序的张量顺序是
张量缩并总是两两进行的。第一次缩并是计算 \((p \nu | \kappa \lambda) = C_{\mu p} (\mu \nu | \kappa \lambda)\),这两者分别对应到长度为 5 的张量缩并列表中的第 (0, 2) 个元素 (从 0 计数)。计算得到的 \((p \nu | \kappa \lambda)v\) 会放到列表的最后。第一次缩并结束得到的张量列表则是
而第二次缩并是计算 \((pq | \kappa \lambda) = C_{\nu q} (p \nu | \kappa \lambda)\),被缩并的两者分别是列表中第 (0, 3) 个元素。依次类推,就可以得到上述的张量缩并列表。
交由程序阅读的路径信息实际上可以代回 np.einsum
的 optimize
选项:
[53]:
%%timeit -r 7 -n 10
np.einsum("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=['einsum_path', (0, 2), (0, 3), (0, 2), (0, 1)])
323 µs ± 114 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
我们也许能发现,其计算速度还比 optimize=True
快一些:
[54]:
%%timeit -r 7 -n 10
np.einsum("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=True)
453 µs ± 194 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
这是因为 optimize=True
还需要通过一小段 python 程序,通过张量的维度信息推测最优的缩并方式,但直接向 np.einsum
提交计算路径则省去了这一步。我们也许会感到,np.einsum
判断程序路径确实会花不少时间;但这部分的计算量不会随着张量的增大而增大,因此大多数情况下是可以忽略的。只是如果打算在使用 np.einsum
的情况下将代码效率优化到极限,那么可以手动地设置张量计算与缩并路径。
下面两段程序可以判断 optimize=True
在推测最优缩并过程时所多花费的时间;对于作者的电脑而言,大约是 100 微秒:
[55]:
%%timeit -r 7 -n 10
np.einsum_path("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=True)
229 µs ± 92.9 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
[56]:
%%timeit -r 7 -n 10
np.einsum_path("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=['einsum_path', (0, 2), (0, 3), (0, 2), (0, 1)])
125 µs ± 22.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
我们也可以对 optimize=False
的情形给出缩并列表;当然,它会给出 \(O(N^8)\) 的计算量级:
[57]:
print(np.einsum_path("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C, optimize=False)[1])
Complete contraction: up,vq,uvkl,kr,ls->pqrs
Naive scaling: 8
Optimized scaling: 8
Naive FLOP count: 1.953e+06
Optimized FLOP count: 1.953e+06
Theoretical speedup: 1.000
Largest intermediate: 6.250e+02 elements
--------------------------------------------------------------------------
scaling current remaining
--------------------------------------------------------------------------
8 ls,kr,uvkl,vq,up->pqrs pqrs->pqrs
5.9. np.einsum
的使用:中转内存设置
下面是以前我在实际工作中遇到过的情况。一般来说,optimize=True
已经能很好地处理张量乘积的算法问题了;但在处理某些涉及格点积分问题时,其计算效率低得令人难以忍受.这是因为 np.einsum
默认的临时张量内存大小设定为所有张量列表中,或张量结果中最大的内存占用值.这也就意味着,默认情况下中间张量的大小其实不是很大。
在计算 GGA 对 \(F_{\mu \nu}^{A_t}\) 即 Fock 矩阵对原子坐标的偏导数的贡献时,默认中间张量内存大小限制就成为了计算耗时的瓶颈。我们使用 np.einsum
默认使用的贪心算法 greedy
来给出设定不同内存大小时,计算路径与理论提速之间的差别。
我们首先定义一些张量。我们在这里不对它们赋予实际意义。
[58]:
grid_fg = np.random.random((3000))
grid_rho_1 = np.random.random((3, 3000))
grid_A_rho_2 = np.random.random((4, 3, 3, 3000))
grid_ao_1 = np.random.random((3, 3000, 22))
grid_ao_0 = np.random.random((3000, 22))
下述计算的时间复杂度在默认内存限制下,仍然为 \(O(N^6)\):
[59]:
%%timeit -r 7 -n 7
np.einsum(
"g, Atrg, rgu, gv -> Atuv",
grid_fg, grid_A_rho_2, grid_ao_1, grid_ao_0,
optimize="greedy" # Equivalent to `optimize=True`
)
149 ms ± 2.15 ms per loop (mean ± std. dev. of 7 runs, 7 loops each)
[60]:
print(np.einsum_path(
"g, Atrg, rgu, gv -> Atuv",
grid_fg, grid_A_rho_2, grid_ao_1, grid_ao_0,
optimize="greedy" # Equivalent to `optimize=True`
)[1])
Complete contraction: g,Atrg,rgu,gv->Atuv
Naive scaling: 6
Optimized scaling: 6
Naive FLOP count: 2.091e+08
Optimized FLOP count: 1.569e+08
Theoretical speedup: 1.333
Largest intermediate: 6.600e+04 elements
--------------------------------------------------------------------------
scaling current remaining
--------------------------------------------------------------------------
2 gv,g->vg Atrg,rgu,vg->Atuv
6 vg,rgu,Atrg->Atuv Atuv->Atuv
但如果我们放开内存限制到 2 GB 的大小 (之所以要除以 8 是因为一般我们使用 64 位浮点数进行矩阵计算,而每 1 Byte 指代 8 位,因此每个浮点数占用 8 Byte 空间),那么代码变成 \(O(N^5)\) 复杂度,提速约 4 倍 (实际提速与理论提速似乎还有不少区别,我还解释不了):
[61]:
%%timeit -r 7 -n 7
np.einsum(
"g, Atrg, rgu, gv -> Atuv",
grid_fg, grid_A_rho_2, grid_ao_1, grid_ao_0,
optimize=["greedy", 1024 ** 3 * 2 / 8] # 1024 ** 3 * 2 -> 2GB; 8 -> float64
)
29.7 ms ± 925 µs per loop (mean ± std. dev. of 7 runs, 7 loops each)
[62]:
print(np.einsum_path(
"g, Atrg, rgu, gv -> Atuv",
grid_fg, grid_A_rho_2, grid_ao_1, grid_ao_0,
optimize=["greedy", 1024 ** 3 * 2 / 8] # 1024 ** 3 * 2 -> 2GB; 8 -> float64
)[1])
Complete contraction: g,Atrg,rgu,gv->Atuv
Naive scaling: 6
Optimized scaling: 5
Naive FLOP count: 2.091e+08
Optimized FLOP count: 3.967e+07
Theoretical speedup: 5.271
Largest intermediate: 7.920e+05 elements
--------------------------------------------------------------------------
scaling current remaining
--------------------------------------------------------------------------
2 gv,g->vg Atrg,rgu,vg->Atuv
5 rgu,Atrg->tAug vg,tAug->Atuv
5 tAug,vg->Atuv Atuv->Atuv
但我们注意到,其代价是生成了非常大的中间张量 (以 tAug
指代);其大小是
而在张量列表与张量结果中,最大的张量也只有
因此,我们实际上是以牺牲内存占用为代价,换取计算效率的提升.
最后指出,np.einsum
中,除了 greedy
贪心算法寻找张量缩并过程之外,还有 optimal
方式,即遍历所有可能的张量缩并过程并寻找最优解。但遍历所有可能缩并的时间复杂度是 \(O(N!)\);若 \(N\) 代表被缩并张量数,因此一旦涉及比较多的张量,optimal
可能会出现非常糟糕的性能表现。通常来说,greedy
已经能给出不错的缩并,因此使用 greedy
已经足够了。
技巧
以后,我们会默认使用 ["greedy", 1024 ** 3 * 2 / 8]
作为 np.einsum
中的 optimize
选项,即使用贪心算法,在 2 GB 大小下作算法优化.
同时,以后在程序运行之前会使用下述代码,使得当前程序的 np.einsum
的 optimize
默认选项会由 ["greedy", 1024 ** 3 * 2 / 8]
所覆盖.
from functools import partial
np.einsum = partial(np.einsum, optimize=["greedy", 1024 ** 3 * 2 / 8])
5.10. 参考任务解答
5.10.1. 任务 (1)
5.10.1.1. 任务 (1.1) 可选
[63]:
R = np.arange(24).reshape((2, 3, 4))
def get_base(array):
if isinstance(array.base, np.ndarray):
return array.base
else:
return array
[64]:
get_base(R)
[64]:
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])
[65]:
print(get_base(R).base)
None