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)

  1. (可选) 这里的 get_base 函数使用了 Python 的 isinstance 函数.该函数在 PySCF 的源代码中涉及到临时信息如何储存的部分中经常出现.但是我们也可以不使用循环与 isinstance 函数,而使用递归方法来写 get_base.试写出这样一个函数.

    提示:你应当可以发现 type(R.base) == np.ndarrayR.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]])

这个过程可以看作是将 eoev 分别展开为两个矩阵从而相减;但两者分别展开为 (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.allnp.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.allclosenp.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 可以设定为布尔值,以判断是否自动优化算法。所谓自动优化,是指根据张量的维度与乘积的方式自动优化计算过程。我们模拟双电子积分转换的情形,来了解自动优化的意义。双电子积分的计算公式是

\[(ij|kl) = C_{\mu p} C_{\nu q} (\mu \nu | \kappa \lambda) C_{\kappa r} C_{\lambda s}\]

我们用 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',之后的信息是每个最基本张量乘积缩并的操作对象的编号。具体来说,我们一开始代入程序的张量顺序是

\[[C_{\mu p}, C_{\nu q}, (\mu \nu | \kappa \lambda), C_{\kappa r}, C_{\lambda s}]\]

张量缩并总是两两进行的。第一次缩并是计算 \((p \nu | \kappa \lambda) = C_{\mu p} (\mu \nu | \kappa \lambda)\),这两者分别对应到长度为 5 的张量缩并列表中的第 (0, 2) 个元素 (从 0 计数)。计算得到的 \((p \nu | \kappa \lambda)v\) 会放到列表的最后。第一次缩并结束得到的张量列表则是

\[[C_{\nu q}, C_{\kappa r}, C_{\lambda s}, (p \nu | \kappa \lambda)]\]

而第二次缩并是计算 \((pq | \kappa \lambda) = C_{\nu q} (p \nu | \kappa \lambda)\),被缩并的两者分别是列表中第 (0, 3) 个元素。依次类推,就可以得到上述的张量缩并列表。

交由程序阅读的路径信息实际上可以代回 np.einsumoptimize 选项:

[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 指代);其大小是

\[3 \times 4 \times 22 \times 3000 = 792000 = 6187.5 \, \text{kB}\]

而在张量列表与张量结果中,最大的张量也只有

\[4 \times 3 \times 3 \times 3000 = 108000 = 843.8 \, \text{kB}\]

因此,我们实际上是以牺牲内存占用为代价,换取计算效率的提升.

最后指出,np.einsum 中,除了 greedy 贪心算法寻找张量缩并过程之外,还有 optimal 方式,即遍历所有可能的张量缩并过程并寻找最优解。但遍历所有可能缩并的时间复杂度是 \(O(N!)\);若 \(N\) 代表被缩并张量数,因此一旦涉及比较多的张量,optimal 可能会出现非常糟糕的性能表现。通常来说,greedy 已经能给出不错的缩并,因此使用 greedy 已经足够了。

技巧

以后,我们会默认使用 ["greedy", 1024 ** 3 * 2 / 8] 作为 np.einsum 中的 optimize 选项,即使用贪心算法,在 2 GB 大小下作算法优化.

同时,以后在程序运行之前会使用下述代码,使得当前程序的 np.einsumoptimize 默认选项会由 ["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