3. 重新学习 MP2 能量笔记
在 pyxdh 中，所有张量全部都储存在内存中，包括原子轨道基组的双电子积分。但在现实问题中，这对内存量与磁盘量上，是不允许的。
所有的 MP2 算法，都多少会使用内存降低方法。我们需要逐个地考虑它们。
3.1. 初始化
[1]:
import numpy as np
from functools import partial
np.einsum = partial(np.einsum, optimize=["greedy", 1024 ** 3 * 2 / 8])
[2]:
import numpy as np
from pyscf import gto, scf, mp, ao2mo, lib
from pyscf.ao2mo import _ao2mo
from functools import partial
import h5py
np.set_printoptions(4, suppress=True, linewidth=180)
np.einsum = partial(np.einsum, optimize=["greedy", 1024 ** 3 * 2 / 8])
[3]:
mol = gto.Mole()
mol.atom = """
N 0. 0. 0.
H 1.5 0. 0.2
H 0.1 1.2 0.
H 0. 0. 1.
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
[3]:
<pyscf.gto.mole.Mole at 0x7f8865ac2be0>
[4]:
nocc, nmo, nao, nbas = mol.nelec[0], mol.nao, mol.nao, mol.nbas
nvir = nmo - nocc
so, sv = slice(0, nocc), slice(nocc, nmo)
3.2. 最简单方法回顾
最简单的方法是将原子轨道基的双电子积分完全储存到内存中，并作双电子积分转换。
[5]:
mf_scf = scf.RHF(mol).run()
C, e = mf_scf.mo_coeff, mf_scf.mo_energy
Co, Cv = C[:, so], C[:, sv]
eo, ev = e[so], e[sv]
[6]:
eri0_ao = mol.intor("int2e")
eri0_iajb = np.einsum("ui, va, uvkl, kj, lb -> iajb", Co, Cv, eri0_ao, Co, Cv)
D_iajb = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]
t_iajb = eri0_iajb / D_iajb
T_iajb = 2 * t_iajb - t_iajb.swapaxes(-1, -3)
np.einsum("iajb, iajb, iajb ->", T_iajb, t_iajb, D_iajb)
[6]:
-0.14554742350036615
整个过程最耗时部分，对于当前的问题反而是原子基组双电子积分 \((\mu \nu | \kappa \lambda)\) 的计算。这一步其实是 \(O(N^4)\) 量级的。对于更大的分子，显然是原子轨道到分子轨道基转换的 \(O(N^5)\) 量级更加耗时。因此不能说这里的的效率测评有意义。
[7]:
%%timeit -n 100
eri0_ao = mol.intor("int2e")
10.4 ms ± 1.11 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
[8]:
%%timeit -n 100
eri0_iajb = np.einsum("ui, va, uvkl, kj, lb -> iajb", Co, Cv, eri0_ao, Co, Cv)
D_iajb = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]
t_iajb = eri0_iajb / D_iajb
T_iajb = 2 * t_iajb - t_iajb.swapaxes(-1, -3)
np.einsum("iajb, iajb, iajb ->", T_iajb, t_iajb, D_iajb)
1.11 ms ± 232 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
通过 PySCF 默认路径的结果如下，可以验证上面结果的正确性。
[9]:
mf_mp2 = mp.MP2(mf_scf).run()
mf_mp2.e_corr
[9]:
-0.1455474235003661
事实上，上面的双电子转换效率并不快。出于双电子积分本身的对称性、与 MP2 激发系数自身的对称性，PySCF 的默认程序 (incore 转换) 本身会更快一些。
[10]:
%%timeit -n 100
mf_mp2 = mp.MP2(mf_scf).run()
3.02 ms ± 821 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.3. Incore 转换
PySCF 默认使用 incore 转换。对于小分子体系而言，实际上原子基组的双电子积分 \((\mu \nu | \kappa \lambda)\) 是被储存下来的：
[11]:
mf_scf._eri.shape
[11]:
(7260,)
这个数值大小其实是对称性简化之后的双电子积分大小 (我们暂时规定了 \(n_\mathrm{AO} = n_\mathrm{MO}\))
[12]:
(nmo * (nmo + 1) // 2 * (nmo * (nmo + 1) // 2 + 1)) // 2
[12]:
7260
因此，这里的 incore 转换实际上要求内存占用至少包括原子基组双电子积分 \(\sim n_\mathrm{AO}^4/8\)。在此基础上，所有 \((ia|jb)\) 积分也全部储存，即额外内存 \(n_\mathrm{occ}^2 n_\mathrm{vir}^2\)。同时，转换过程还需要 \(n_\mathrm{occ} n_\mathrm{vir} n_\mathrm{AO}^2\) 大小。因此可以说是开销不计成本。
PySCF 中对 incore 的处理是非常简单粗暴的。直接提供简化或未简化的积分、以及用于缩并的四个系数矩阵就可以了。
[13]:
np.allclose(ao2mo.general(mf_scf._eri, (Co, Cv, Co, Cv)).reshape((nocc, nvir, nocc, nvir)), eri0_iajb)
[13]:
True
[14]:
np.allclose(ao2mo.general(eri0_ao, (Co, Cv, Co, Cv)).reshape((nocc, nvir, nocc, nvir)), eri0_iajb)
[14]:
True
进行过转换的积分会用于给出 MP2 激发系数 \(t_{ij}^{ab}\)。但与 pyxdh 不同地是，pyxdh 使用维度 \((i, a, j, b)\) 储存激发系数，而 PySCF 使用维度 \((i, j, a, b)\)。
[15]:
np.allclose(mf_mp2.t2.swapaxes(-2, -3), t_iajb)
[15]:
True
3.4. Outcore 转换
在 PySCF 中，outcore 转换必须要求 \((ia|jb)\) 大小的张量可以储存在内存中，因此需要比 \(n_\mathrm{occ}^2 n_\mathrm{vir}^2\) 再大一些。它的实现可以非常简单，也可以非常复杂。
但如果只求取 MP2 的能量，那么我们不一定真的需要这么大的张量。
3.4.1. 简单粗暴的 outcore
如果替换 ao2mo.general 的第一个参量为 gto.Mole 实例，那么会自动地使用 outcore 算法。
[16]:
np.allclose(ao2mo.general(mol, (Co, Cv, Co, Cv)).reshape((nocc, nvir, nocc, nvir)), eri0_iajb)
[16]:
True
[17]:
%%timeit -n 100
ao2mo.general(eri0_ao, (Co, Cv, Co, Cv))
1.21 ms ± 489 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
[18]:
%%timeit -n 10
ao2mo.general(mol, (Co, Cv, Co, Cv))
40.7 ms ± 4.35 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
3.4.2. 实际的 outcore (1)：壳层与原子对原子轨道的分割
之后几小节的讨论对象是 mp.mp2._ao2mo_ovov。首先，我们要学习或回顾原子轨道和轨道壳层的分割。
不论任何量化程序，原子轨道的数量总是恒定的。它会用 \(\mu, \nu, \cdots\) 角标表示。但在实际的积分中，不是任何原子轨道的分割都可以带入计算。在 PySCF 中，积分的最小单元是壳层 (shell)。
譬如，一个 \(s\) 壳层对应 1 根原子轨道，一个 \(p\) 壳层对应 3 根原子轨道。对于 \(d\) 壳层，根据球谐 (sph) 或笛卡尔 (cart) 的不同，会分别取 5 或 6 根原子轨道。现在一般的基组都使用球谐轨道；笛卡尔基组已经很少使用了 (特殊的 6-31G 系列基组会使用)。
这里顺便原子的分割。尽管 MP2 能量不需要用到，但核坐标梯度的程序会经常使用到。原子轨道之所以称为“原子轨道”，是因为它是以特定原子为中心的三维 Gaussian 函数。因此，原子轨道是可以对应到特定原子的。
现在我们考虑程序的问题。壳层的分割可以用下面的函数表示：
[19]:
ao_loc = mol.ao_loc_nr()
print(ao_loc)
print(len(ao_loc))
[ 0 1 2 3 6 9 10 11 12 13 14 15]
12
我们知道，氨分子的 6-31G 基组有 15 根原子轨道。上面的结果意味着该分子有 12 个壳层。这 12 个壳层的起点体现在上面的列表了。
第 3, 4 壳层 (起点分别为 3, 6) 不同于其他壳层；这一个壳层中包含 3 个原子轨道，因此可以推断为 p 轨道。
我们这里回顾原子对原子轨道与轨道壳层的分割。PySCF 的原子分割是：
[20]:
mol.aoslice_by_atom()
[20]:
array([[ 0, 5, 0, 9],
[ 5, 7, 9, 11],
[ 7, 9, 11, 13],
[ 9, 11, 13, 15]])
上面的每一行代表每一个原子；第一个原子的壳层为 0-4，原子轨道为 0-8；第二个原子的壳层为 5-6，原子轨道 9-10；以此类推。
现在遇到这样的现实的问题：假设氨分子非常大，我们的内存无法储存这种级别的分子。以什么标准对积分过程进行分割？
在 outcore 过程中，第一步是先计算占据轨道的积分：
第二部才是计算最终的分子轨道积分：
由于第一步缩并过程需要完整的 \(\mu, \lambda\)，因此这两个原子轨道角标不能进行分割。待分割的角标是 \(\nu, \kappa\)。
以多小的单位进行分割？在 PySCF 中，推荐的大小是 4 根原子轨道一组进行分割。由于电子积分的最小单位是壳层，因此如果基组包含 \(d\) 轨道 (5 根原子轨道)，因此一般来说实际消耗会比 4 根轨道还要再大一些。
之所以不使用更大的分割单位，我想是因为小一些的分割可以保证利用好原子轨道 ERI 积分的对称性；但不使用更小的分割，是为了减小积分函数的调用次数，避免不必要的调用时间。
PySCF 中的辅助函数 ao2mo.outcore.balance_partition 可以用于处理壳层的分割：
[21]:
sh_ranges = ao2mo.outcore.balance_partition(ao_loc, 4)
dmax = max(x[2] for x in sh_ranges)
print(sh_ranges)
print(dmax)
[(0, 3, 3), (3, 4, 3), (4, 6, 4), (6, 10, 4), (10, 11, 1)]
4
上述参量 sh_ranges 表示被分割的壳层。譬如以 (4, 6, 4) 为例，表示了第 2 个分割包含壳层 4 与壳层 5，该分割包含了 4 根轨道。
3.4.3. 实际的 outcore (2)：分割的原子轨道积分
随后的任务是分步地处理下述计算：
其中，下标 \(s\) 表示分割。随后将 \((i \nu_s | \kappa_s j)\) 以维度 \((i, j, \nu_s, \kappa_s)\) 的形式储存到外部硬盘，并对该张量赋予名称 \(s\)。
由于分割出来的原子轨道大小不超过壳层数，因此如果不考虑 \(d, f\) 等高角动量轨道，那么内存消耗大约是 \(16 n_\mathrm{AO}^2\)。
一个额外的工作是，为了避免内存空间的重复分配耗时，我们先预置一块内存 buf_eri，专门用于储存临时积分 \((\mu \nu_s | \kappa_s \lambda)\)。
[22]:
buf_eri = np.empty((nao, dmax, dmax, nao))
循环过程中，我们所需要使用的壳层分割和原子轨道分割列表分别用 list_shell_slice 与 list_ao_slice 储存：
[23]:
list_shell_slice, list_ao_slice = [], []
for ip, (ish0, ish1, _) in enumerate(sh_ranges):
for jsh0, jsh1, nj in sh_ranges[:ip+1]:
list_shell_slice.append((ish0, ish1, jsh0, jsh1))
i0, i1 = ao_loc[ish0], ao_loc[ish1]
j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
list_ao_slice.append((i0, i1, j0, j1))
[24]:
for slice_shell, slice_ao in zip(list_shell_slice, list_ao_slice):
print(slice_shell, slice_ao)
(0, 3, 0, 3) (0, 3, 0, 3)
(3, 4, 0, 3) (3, 6, 0, 3)
(3, 4, 3, 4) (3, 6, 3, 6)
(4, 6, 0, 3) (6, 10, 0, 3)
(4, 6, 3, 4) (6, 10, 3, 6)
(4, 6, 4, 6) (6, 10, 6, 10)
(6, 10, 0, 3) (10, 14, 0, 3)
(6, 10, 3, 4) (10, 14, 3, 6)
(6, 10, 4, 6) (10, 14, 6, 10)
(6, 10, 6, 10) (10, 14, 10, 14)
(10, 11, 0, 3) (14, 15, 0, 3)
(10, 11, 3, 4) (14, 15, 3, 6)
(10, 11, 4, 6) (14, 15, 6, 10)
(10, 11, 6, 10) (14, 15, 10, 14)
(10, 11, 10, 11) (14, 15, 14, 15)
第一步完整的积分过程可以用下面的程序表示：
[25]:
ftmp = lib.H5TmpFile()
count = 0
for (ish0, ish1, jsh0, jsh1), (i0, i1, j0, j1) in zip(list_shell_slice, list_ao_slice):
# slice mu nu kappa lambda
shls_slice = (0, nbas) + (ish0, ish1) + (jsh0, jsh1) + (0, nbas)
# calculate ( mu nu | kappa lambda ). note that kappa is sliced, so aosym is none, i.e. ( mu nu | kappa lambda ) != ( mu nu | lambda kappa )
eri = mol.intor("int2e", shls_slice=shls_slice, aosym="s1", out=buf_eri)
# reshape to tensor (mu, nu, kappa, lambda) from one-dim array
eri.shape = (nao, (i1-i0), (j1-j0), nao)
# ( mu nu | kappa lambda ) -> ( i nu | kappa j )
tensor_ijvk = np.einsum("uvkl, ui, lj -> ijvk", eri, Co, Co)
# dump result to h5py file
ftmp.create_dataset(str(count), data=tensor_ijvk)
count += 1
ftmp.keys()
[25]:
<KeysViewHDF5 ['0', '1', '10', '11', '12', '13', '14', '2', '3', '4', '5', '6', '7', '8', '9']>
现在我们验证其中一个积分。我们能看到第 4 个分割是 \(\nu_4\) 包含原子轨道 6-9，\(\kappa_4\) 包含原子轨道 3-5。
[26]:
list_ao_slice[4]
[26]:
(6, 10, 3, 6)
不妨直接地验证一下
[27]:
np.allclose(np.einsum("uvkl, ui, kj -> ijvl", mol.intor("int2e"), Co, Co)[:, :, 6:10, 3:6], ftmp["4"][()])
[27]:
True
3.4.4. 实际的 outcore (3)：分子轨道的 ERI 积分导出
最终的激发系数导出方式是
上面的步骤中，由于缩并时角标 \(\nu, \kappa\) 的存在，因此这两个变量必须是连续的。可以通过内存操作被分割的角标是 \(i, j\)。
上一步中，分割的大小被限制到 4 个原子轨道，是比较激进的限制方式；但这里，程序倾向于使用宽松的限制方式，因此 \(i, j\) 的分割越小越好。这是因为上一步的分割限制涉及到因对称性所导致的浪费；如果分割太大，那么浪费也会越多。但这一步中，分割太小会导致大量硬盘 I/O，从而限制效率。
在 PySCF 的实现中，假设内存有至少 \(8 n_\mathrm{occ} n_\mathrm{AO}^2\)。这要求了其中一个角标 \(j\) 在内存中是连续储存的。不连续的部分就是 \(i\)，且最小的分割是 \(n_\mathrm{occblk} = 4\)。
在演示的实例中，我们将分割的大小稍微缩小到 3，这是因为氨分子一共就 5 个占据轨道。因此，对于氨分子，两次分割分别是
[28]:
occblk = 3
list_occ_slice = []
for i in range(0, nocc, occblk):
num_blk = min(occblk, nocc-i)
list_occ_slice.append((i, i + num_blk, num_blk))
list_occ_slice
[28]:
[(0, 3, 3), (3, 5, 2)]
在实现积分转换之前，我们要能先把 \((i_s \nu | \kappa j)\) 储存到内存。为此，我们声明一块内存，用于储存这部分转换到一半的积分：(维度定为 \((i_s, j, \nu, \kappa)\))
[29]:
blk_eri = np.empty((occblk, nocc, nao, nao))
我们需要写一个函数，以加载这部分积分。这一步是 I/O 关键步骤，因为尽管我们要求的内存量是 \(n_\mathrm{occblk} n_\mathrm{occ} n_\mathrm{AO}^2\)，但要求单次分割的硬盘 I/O 量是 \(n_\mathrm{occ}^2 n_\mathrm{AO}^2\)，因此总共要求的硬盘 I/O 量是 \(n_\mathrm{occ}^3 n_\mathrm{AO}^2 / n_\mathrm{occblk}\)。如果 \(n_\mathrm{occblk}\) 设置得太小，那么在 I/O 上就变相地变成五次方时间消耗，非常划不来。
[30]:
def load(i0, i1, eri):
for idx, (v0, v1, k0, k1) in enumerate(list_ao_slice):
eri[:i1-i0, :, v0:v1, k0:k1] = ftmp[str(idx)][i0:i1]
if v0 != k0:
dat = ftmp[str(idx)][:, i0:i1]
eri[:i1-i0, :, k0:k1, v0:v1] = dat.transpose((1, 0, 3, 2))
我们验证一下这个程序的正确性。如果现在取第一个分割，即对 \(i\) 取 0-2 号轨道，那么 i0 = 0, i1 = 3。我们需要事先准备好转换到一半的 eri \((i_s \nu | \kappa j)\) 的空间 (维度 \((i_s, j, \nu, \kappa)\))。这个函数会原地对 eri 内部的值直接改动，不返回结果。
[31]:
eri = np.ndarray((3, nocc, nao, nao), buffer=blk_eri)
load(0, 3, eri)
np.allclose(np.einsum("uvkl, ui, kj -> ijvl", mol.intor("int2e"), Co, Co)[0:3], eri)
[31]:
True
随后，我们就可以考虑积分的过程了。尽管 \((i_s \nu | \kappa j)\) 的维度是 \((i_s, j, \nu, \kappa)\)，但最终生成的 \((ia|jb)\) 的维度是 \((i, a, j, b)\)。其过程如下。
[32]:
feri = lib.H5TmpFile()
h5_oovv = feri.create_dataset("ovov", (nocc, nvir, nocc, nvir), "f8")
for i0, i1, inum in list_occ_slice:
# Not using `blk_eri` directly, since i1-i0 = inum could be different in iterations
eri = np.ndarray((inum, nocc, nao, nao), buffer=blk_eri)
# load (i_s nu | kappa j) to `eri`
load(i0, i1, eri)
# calculate (i_s a | j b) tensor contraction
eri_mo = np.einsum("ijvl, va, lb -> iajb", eri, Cv, Cv)
# dump to h5py file (hard-disk)
h5_oovv[i0:i1] = eri_mo
最后，我们验证一下计算得到的 \((ia|jb)\) h5_oovv 是否确实是之前我们给出过的 eri_iajb：
[33]:
np.allclose(eri0_iajb, h5_oovv)
[33]:
True
3.4.5. 浮点计算关键步骤：_ao2mo.nr_e2 函数
刚才一直强调硬盘 I/O 的消耗，但整个程序最耗时的步骤是浮点计算的张量缩并：
# calculate (i_s a | j b) tensor contraction
eri_mo = np.einsum("ijvl, va, lb -> iajb", eri, Cv, Cv)
在 PySCF 中，它并非使用 Python 程序进行处理，而是用 C 程序进行计算。
相同的工作可以使用 np.einsum 或 lib.einsum 完成，但对于这个特定的任务，_ao2mo.nr_e2 的处理效率非常高。文档的最后就讨论这个函数。
我们现在假设遇到的问题是 \(n_\mathrm{occ} = 20\), \(n_\mathrm{vir}=100\), \(n_\mathrm{MO} = n_\mathrm{AO} = 120\)。不考虑对称性，我们随机一个 eri 与 Cv 矩阵，作为当前模型问题的矩阵。
[34]:
nao = nmo = 120
nocc, nvir = 20, 100
[35]:
eri = np.random.randn(nocc, nocc, nao, nao)
Cv = np.random.randn(nao, nvir)
我们首先可以验证 _ao2mo.nr_e2 函数与 np.einsum 的作用效果相同。
[36]:
np.allclose(
_ao2mo.nr_e2(eri.reshape(nocc**2, nao**2), Cv, (0, nvir, 0, nvir), "s1", "s1").reshape(nocc, nocc, nvir, nvir),
np.einsum("ijvl, va, lb -> ijab", eri, Cv, Cv)
)
[36]:
True
但从耗时上，_ao2mo.nr_e2 的消耗远远比 np.einsum 低：
[37]:
%%timeit -n 10
_ao2mo.nr_e2(eri.reshape(nocc**2, nao**2), Cv, (0, nvir, 0, nvir), "s1", "s1").reshape(nocc, nocc, nvir, nvir)
22.4 ms ± 4 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
[38]:
%%timeit -n 10
np.einsum("ijvl, va, lb -> ijab", eri, Cv, Cv)
47.1 ms ± 1.46 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
在对程序没有更深了解的情况下，只能说 _ao2mo.nr_e2 是一个优秀的特化程序；可以的情况下就尽量使用。