10. 简单理解真正 \(O(N^4)\) 的 MP2 方法：LS-THC-MP2

创建时间：2021-09-18

在这份文档中，我们会简单地了解一种真正 \(O(N^4)\) 的 MP2 实现方法。它称为 LS-THC-MP2 (Least-Squares Tensor HyperContraction Second-order Møller-Plesset perturbation)。尽管这个概念应该早些时候 (自 2012 年起[1]) 有人提过，但这篇文档较多地使用 Devin A. Matthews[2] 的思路。

尽管很想说不要被这么长的名称吓到，但它一般来说还要求 LT-DF (Laplace-Transformation Density-Fitting)，DF 也经常称为 RI (Resolution-Identity)。因此，本文的方法具体来说应当称为 LS-THC-LT-DF-MP2。事实上，这些麻烦的名称总地来说，都是矩阵或张量分解的思路。若对于矩阵或张量分解 (譬如 Cholesky、SVD 分解等) 和 MP2 方法本身有基本的了解，这篇文档应当是可以接受的。

计算化学方法的加速通常依赖两种策略：利用零值稀疏性本身 (剑宗？见招拆招)，与利用稠密矩阵的低维分解 (气宗？万物皆准)。这篇文档明确地是使用后者的立场。该方法本身对于 MP2 的意义相对来说比较微妙，因为 LS-THC-MP2 恐怕只有在很庞大的体系下才比 DF-MP2 更快。但对于 CC2 或 MP3 等传统上需要 \(O(N^5)\) 甚至 \(O(N^6)\) 的方法 (或许包括最近兴起的 ADC(2) 方法)，在数百个电子的体系下，从高标度有效地加速到 \(O(N^4)\) 确实是可能且实用的。若不考虑问题的稀疏性，这意味着我们曾经一直视为代价巨大的一些后自洽场 (Post-SCF) 方法，其计算复杂度甚至内存复杂度与没有任何积分加速的自洽场是对等的。到 \(O(N^4)\) 为止，基本可以认为基于稠密矩阵的一部分 Post-SCF 方法的计算复杂度已经优化到极限了；剩下的问题是如何利用稀疏性、以及是否可以在相同复杂度下作计算简化。

本文档始终使用经过变通的 Einstein Summation Convention。这份文档一定程度上包含了对 DF-MP2 与 LT-DF-SOS-MP2 的简单介绍，但稍详细的讨论可以参考这两份文档。我们会利用许多 PySCF 所提供的便利接口。

警告

这份文档的实现不一定等同于文献的实现方式，也不能保证其在任意体系的正确性。这份文档仅仅是叙述 LS-THC 方法，并尝试给出作者自己所理解的模型实现而已。

%matplotlib notebook
from pyscf import gto, scf, df, mp, dft, ao2mo
from pyscf.dft.numint import NumInt
from pyscf.dft import Grids
from pyscf.ao2mo import _ao2mo
import numpy as np
import scipy
from opt_einsum import contract as einsum
import matplotlib.pyplot as plt

np.set_printoptions(6, suppress=True, linewidth=150)

10.1. 准备工作

10.1.1. 分子体系定义

我们在这篇文档使用闭壳层 cc-pVQZ 的水分子。使用这么小的分子显然是大材小用了，因为 LS-THC 即使是对于 MP3 也要求在数百个电子的体系才能观察到效率飞跃，对于 MP2 更是如此；但足够说明问题就行。

mol = gto.Mole(atom="O; H 1 0.94; H 1 0.94 2 104.5", basis="cc-pVQZ", verbose=0).build()
nocc, nao, nmo = mol.nelec[0], mol.nao, mol.nao
nvir = nmo - nocc
so, sv, sa = slice(0, nocc), slice(nocc, nmo), slice(0, nmo)

下面的代码给出自洽场实例、分子轨道与轨道能量：

mf = scf.RHF(mol).run()
C, e = mf.mo_coeff, mf.mo_energy
Co, Cv = C[:, so], C[:, sv]
eo, ev = e[so], e[sv]

下面的代码给出辅助基 (auxiliary basis) 的定义：

with_df = df.DF(mol, df.make_auxbasis(mol, mp2fit=True))
naux = with_df.get_naoaux()

当前的问题下，一些重要的体系大小表征值是

• \(n_\mathrm{occ} = 5\) 占据轨道数

• \(n_\mathrm{vir} = 110\) 非占轨道 (虚轨道) 数

• \(n_\mathrm{AO} = n_\mathrm{MO} = 115\) 分子轨道数

• \(n_\mathrm{aux} = 242\) 辅助基数

• \(n_\mathrm{occ} n_\mathrm{vir} = 550\)

nocc, nvir, nao, naux, nocc*nvir

(5, 110, 115, 242, 550)

10.1.2. 未加速与近似的 MP2 相关能

我们再次回顾 MP2 的相关能计算。MP2 相关能可以表示为

\[ E_\mathsf{corr} = 2 g_{ij}^{ab} g_{ij}^{ab} / D_{ij}^{ab} - g_{ij}^{ab} g_{ij}^{ba} / D_{ij}^{ab} \]

特别地，相关能还可以拆分为自旋相反部分 (OS, Opposite-Spin) 与相同部分 (SS, Same-Spin)：

\[\begin{split} \begin{align*} E_\textsf{OS} &= g_{ij}^{ab} g_{ij}^{ab} / D_{ij}^{ab} \\ E_\textsf{SS} &= g_{ij}^{ab} g_{ij}^{ab} / D_{ij}^{ab} - g_{ij}^{ab} g_{ij}^{ba} / D_{ij}^{ab} = \frac{1}{2} (g_{ij}^{ab} - g_{ij}^{ba})^2 / D_{ij}^{ab} \end{align*} \end{split}\]

其中，

• g \(g_{ij}^{ab} = C_{\mu i} C_{\nu a} (\mu \nu | \kappa \lambda) C_{\kappa j} C_{\lambda b}\) (dim: \(i, a, j, b\)) 分子轨道下的双电子积分

• D \(D_{ij}^{ab} = \varepsilon_i - \varepsilon_a + \varepsilon_j - \varepsilon_b\) (dim: \(i, a, j, b\)) MP2 分母

g = ao2mo.incore.general(mol.intor("int2e"), (Co, Cv, Co, Cv), compact=False)
D = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]

corr_simple_os = einsum("iajb, iajb, iajb ->", g, g, 1/D)
corr_simple_ss = corr_simple_os - einsum("iajb, ibja, iajb ->", g, g, 1/D)
print("==> [Simple MP2 OS corr energy]         : {:20.12f}".format(corr_simple_os))
print("==> [Simple MP2 SS corr energy]         : {:20.12f}".format(corr_simple_ss))
print("==> [Simple MP2 corr energy]            : {:20.12f}".format(corr_simple_os + corr_simple_ss))
print("==> [Simple MP2 corr energy from PySCF] : {:20.12f}".format(mp.MP2(mf).run().e_corr))

==> [Simple MP2 OS corr energy]         :      -0.240427624844
==> [Simple MP2 SS corr energy]         :      -0.071748736656
==> [Simple MP2 corr energy]            :      -0.312176361500
==> [Simple MP2 corr energy from PySCF] :      -0.312176361500

光从能量表达式来看，好像只涉及到 \(i, j, a, b\) 四个分子轨道角标的运算与求和；但导致 MP2 能量计算耗时的主要因素在于 \(g_{ij}^{ab}\) 的构造过程，它是五次方计算量。一般来说，计算量最大的一步出现在

\[ (i \nu | \kappa \lambda) = C_{\mu i} (\mu \nu | \kappa \lambda) \]

这个单步运算涉及到 \(i, \mu, \nu, \kappa, \lambda\) 五个角标，因此是 \(O(n_\mathrm{occ} n_\mathrm{AO}^4)\) 复杂度。而且由于硬盘通常无法储存所有的原子轨道 ERI \((\mu \nu | \kappa \lambda)\)，因此该计算过程对大体系而言还会出现 \(O(N^4) \sim O(N^5)\) 的多次电子积分计算过程。

10.2. 一次加速：Density Fitting

10.2.1. 背景：对 ERI 积分张量的本征值分解近似

我们回顾 Density Fitting 的提出动因。一般来说，原子轨道的双电子积分 (ERI, Electron-Repulsion Integral)

\[ (\mu \nu | \kappa \lambda) = \iint \phi_\mu (\boldsymbol{r}_1) \phi_\nu (\boldsymbol{r}_1) \frac{1}{|\boldsymbol{r}_1 - \boldsymbol{r}_2|} \phi_\kappa (\boldsymbol{r}_2) \phi_\lambda (\boldsymbol{r}_2) \, \mathrm{d} \boldsymbol{r}_1 \, \mathrm{d} \boldsymbol{r}_2 \]

如果作为 \(n_\mathrm{AO}^2 \times n_\mathrm{AO}^2 = 13225 \times 13225\) 的矩阵，是非常不满秩的。如果我们对该矩阵作本征值分解，会发现大于 \(10^{-5} \, \text{a.u.}\) 的本征值只有 782 个。如果只使用者 782 个本征值与本征向量重构近似的原子轨道的双电子积分 ERI，我们可以得到与精确值误差只有 \(10^{-6} \, \text{a.u.}\) 的 MP2 相关能。事实上，我们还能承受更加激进得多的近似。

eri = mol.intor("int2e").reshape(nao**2, nao**2)

eig_eri, coef_eri = scipy.linalg.eigh(eri)
(eig_eri > 1e-5).sum()

782

Y_ao_approx = coef_eri[:, -782:] * eig_eri[-782:]**0.5
eri_approx = (Y_ao_approx @ Y_ao_approx.T).reshape(nao, nao, nao, nao)
g_approx = ao2mo.incore.general(eri_approx, (Co, Cv, Co, Cv), compact=False)
corr_approx = 2 * einsum("iajb, iajb, iajb ->", g_approx, g_approx, 1/D) - einsum("iajb, ibja, iajb ->", g_approx, g_approx, 1/D)
print("==> [MP2 corr energy by approx eri]     : {:20.12f}".format(corr_approx))
print("==> [MP2 corr error from approx eri]    : {:20.12f}".format(corr_approx - (corr_simple_os + corr_simple_ss)))

==> [MP2 corr energy by approx eri]     :      -0.312175302330
==> [MP2 corr error from approx eri]    :       0.000001059170

10.2.2. Density Fitting 公式、程序与结果

基于 ERI 张量严重不满秩的特性 (一定程度上表明张量数学结构上的稀疏性，而非因为张量有很多零值)，我们有可能大量地简化计算。将 ERI 拆分成两个张量的乘积：

\[ (\mu \nu | \kappa \lambda) = Y_{\mu \nu, P} Y_{\kappa \lambda, P} \]

这种拆分方式并非是唯一的。我们方才在程序中定义的 Y_ao_approx 就是一种做法，其中角标 \(P\) 所代表的维度是 782 维度。但这种拆解需要求解本征问题 (求解本征问题实际上是 \(O(n_\mathrm{AO}^6)\) 计算复杂度)，反而增加了计算耗时；并且 782 维度对于水分子的 cc-pVQZ 实在太大，并不经济。

为了解决这些问题，从 1990 年代发展出一套称为 Density Fitting 的方法，能有效地在 \(O(n_\mathrm{AO}^4)\) 或更低标度，给出对 ERI 的张量分解。这种分解并不类似于数学中的 Cholesky 或 SVD 分解；数学所讨论的分解需要提前计算矩阵并将矩阵储存在内存，但比较流行的 DF 方法通过引入辅助基 (Auxiliary Basis Set) 的方式，不需要了解分子形状或预先计算 ERI 就可以进行矩阵分解。

我们已经知道对于当前的水分子 cc-pVQZ 而言，辅助基大小是 \(n_\mathrm{aux} = 242\)。DF 的方式不是唯一的；我们这里使用 PySCF 默认的方式 (RI-V)。我们通过下述代码生成 Y_ov \(Y_{ia, P} = Y_{\mu \nu, P} C_{\mu i} C_{\nu a}\) (dim: \(i, a, P\)) ERI 在辅助基下的拆解张量 (由于 PySCF 的实现方式，通常也称为 Cholesky 3c-2e (3-center 2-electron) 积分)。

Y_ov = _ao2mo.nr_e2(next(with_df.loop(naux)), C, (0, nocc, nocc, nmo), aosym="s2", mosym="s1") \
             .reshape(naux, nocc, nvir).transpose(1, 2, 0)

那么分子轨道下的双电子积分可以近似为

\[ g_{ij}^{ab} \simeq Y_{ia, P} Y_{jb, P} \]

上述的单步计算过程由于设计角标 \(i, j, a, b, P\)，因此计算复杂度是 \(O(n_\mathrm{nocc}^2 n_\mathrm{vir}^2 n_\mathrm{aux})\)。尽管也是 \(O(N^5)\) 的复杂度，但若基组较大或者硬盘可以存下所有 \(Y_{\mu \nu, P}\)，那么 DF-MP2 方法相比于没有任何数值近似的 MP2 方法的效率完全就是划得来甚至优势巨大的。这个真的是谁用谁知道啊，我以前算过的大多数不算太大的分子都是 DF-MP2 比 \(O(n_\mathrm{AO}^4)\) 的自洽场快多了。而 DF-MP2 所引入的数值误差只有 \(2 \times 10^{-5} \, \mathrm{a.u.} \sim 0.01 \, \mathrm{kcal/mol}\)，是可以接受的。

g_df = einsum("iaP, jbP -> iajb", Y_ov, Y_ov)
corr_df_os = einsum("iajb, iajb, iajb ->", g_df, g_df, 1/D)
corr_df_ss = corr_simple_os - einsum("iajb, ibja, iajb ->", g_df, g_df, 1/D)
print("==> [DF-MP2 OS corr energy]             : {:20.12f}".format(corr_df_os))
print("==> [DF-MP2 SS corr energy]             : {:20.12f}".format(corr_df_ss))
print("==> [DF-MP2 corr energy]                : {:20.12f}".format(corr_df_os + corr_df_ss))
print("==> [DF-MP2 corr error]                 : {:20.12f}".format((corr_df_os + corr_df_ss) - (corr_simple_os + corr_simple_ss)))

==> [DF-MP2 OS corr energy]             :      -0.240393392264
==> [DF-MP2 SS corr energy]             :      -0.071794050091
==> [DF-MP2 corr energy]                :      -0.312187442355
==> [DF-MP2 corr error]                 :      -0.000011080854

10.3. 对 OS 部分的二次加速：LT 方法

LT (Laplace-Transform) 方法尽管套了 Laplace 的大名，但和 Fourier 变换或者数理方程所学的 Laplace 变换其实不太一样。真正需要我们利用的性质其实是下述非常简单的积分式与其数值格点积分的近似：

\[ x^{-1} = \int_0^{+ \infty} e^{- x t} \, \mathrm{d} t \simeq x^{-1} \simeq w_l e^{- x t_l} \]

上式的最右边是要对数值积分格点角标 \(l\) 进行求和的；lt_w \(w_l\) 是积分权重、lt_x \(t_l\) 表示格点位置。我们在这份文档中取下述粗糙的代码所给出的 \(n_\textrm{LT} = 18\) 个格点；但需要指出，只要 MP2 分母 \(D_{ij}^{ab}\) 的绝对值在 \(0.01 \, \mathrm{a.u.}\) 到 \(400 \, \mathrm{a.u.}\) 之间，那么所需的格点数量至少可以缩减到 10 个。尽管我们在分析复杂度时会考虑 LT 格点数 \(n_\textrm{LT}\)，但需要知道它在不随体系增大而增大。

lt_x = 2.5 ** np.arange(-12, 6)
lt_w = np.log(2.5) * lt_x

因此，很有意思地，\((D_{ij}^{ab})^{-1}\) 的计算变成了张量运算：

\[ (D_{ij}^{ab})^{-1} = - (- D_{ij}^{ab})^{-1} = w_l e^{D_{ij}^{ab} t_l} = - w_l e^{(\varepsilon_i - \varepsilon_a + \varepsilon_j - \varepsilon_b) t_l} \]

如果我们进而定义

• go \(g_i^l = w_l^{1/4} e^{\varepsilon_i t_l}\) (dim: \(i, l\))

• gv \(g_a^l = w_l^{1/4} e^{- \varepsilon_a t_l}\) (dim: \(a, l\))

go = lt_w**0.25 * np.exp( eo[:, None] * lt_x)
gv = lt_w**0.25 * np.exp(-ev[:, None] * lt_x)

那么

\[ (D_{ij}^{ab})^{-1} = - g_i^l g_a^l g_j^l g_b^l \]

进一步地，

\[ E_\textsf{OS} = g_{ij}^{ab} g_{ij}^{ab} / D_{ij}^{ab} \simeq - Y_{ia,P} Y_{jb,P} Y_{ia,Q} Y_{jb,Q} g_i^l g_a^l g_j^l g_b^l \]

看起来这个表达式很复杂，但其计算复杂度确实是 \(O(n_\mathrm{LT} n_\mathrm{occ} n_\mathrm{vir} n_\mathrm{aux}^2)\) 即 \(O(N^4)\) 的。

#                        g_ovov    , g_ovov    , 1/D_ovov
corr_ltdf_os = - einsum("iaP , jbP , iaQ , jbQ , il, al, jl, bl ->",
                         Y_ov, Y_ov, Y_ov, Y_ov, go, gv, go, gv)
print("==> [LT-DF-MP2 OS corr energy]          : {:20.12f}".format(corr_ltdf_os))
print("==> [LT-DF-MP2 OS corr error from LT]   : {:20.12f}".format(corr_ltdf_os - corr_df_os))
print("==> [LT-DF-MP2 OS corr error from DF]   : {:20.12f}".format(corr_df_os - corr_simple_os))

==> [LT-DF-MP2 OS corr energy]          :      -0.240362101359
==> [LT-DF-MP2 OS corr error from LT]   :       0.000031290905
==> [LT-DF-MP2 OS corr error from DF]   :       0.000034232581

在当前的 LT 格点设置下，因 LT 所带来的误差接近于因 DF 所带来的误差。

需要说明的是，只有 MP2 的 OS 部分可以由 LT 加速；SS 部分在 LT 下一般只会花更多时间。

10.4. 三次加速：LS-THC

10.4.1. 背景

但我们不止步于 DF (Density-Fitting) 的拆解 \(g_{ij}^{ab} \simeq Y_{ia, P} Y_{jb, P}\)；我们希望作更进一步的拆解。其原因是，一方面来说，即使是 DF 也无法将 MP3, CC2 等方法降低到 \(O(N^4)\) 标度；另一方面，\(Y_{ia, P}\) 仍然是 \(O(N^3)\) 的张量，一旦我们需要经常性复用张量 \(Y_{ia, P}\)，那么就意味着必须要准备一块 \(O(N^3)\) 的硬盘或内存。这对于较大的体系而言，是不太容易承受的。

因此就需要寻找一种将 \(g_{ij}^{ab}\) 拆分为 \(O(N^2)\) 张量乘积的方法。拆分的大体思路是

\[ g_{ij}^{ab} \simeq X_i^R X_a^R V_{RS} X_j^S X_b^S \]

需要注意的是，这里的角标 \(R, S\) 并非表示辅助基，且一般来说其索引范围要比辅助基大一些。这种张量拆分在化学中称为 THC (Tensor HyperContraction)。

当然，知道拆分的思路不意味着计算量就真的得以简化了，而且拆分必然伴随着一定程度的精度损失；为此至少要回答两个问题。第一是，张量拆分的思路不只一种，完全也可能是 \(g_{ij}^{ab} \simeq X_i^R X_a^R X_j^R X_b^R\)；所以为何我们不选用后者的拆分而选用前者。第二是，拆分张量的计算量不应超过 \(O(N^4)\)，且应当有小至少一个数量级的内存消耗；我们要如何保证张量拆分的效率，同时又有可以控制的拆分精度损失？

对于这两个问题，我们的文档中不会作说明。这些恐怕都是前人的经验教训了。我们后文只对第二个问题作实现上的说明。

这里补充说明一些，尽管现在有自动化的张量拆分库，譬如 TensorLy 可以对张量作 CANDECOMP/PARAFAC 或 Tucker 拆分，但这些算法为了达到所需精度需要大量的迭代次数，不一定能保证部分对称张量在拆分时所需要满足的对称性要求，拆分有很强的随机性，而且张量经常需要全部存到内存中。这种自动化的、不需要对张量本身性质作了解的拆分方式，或许会在机器学习的神经网络层或图像识别与压缩中发挥巨大优势，但在化学庞大的张量下使用可能是不合适的。

10.4.2. 轨道 (母) 格点拆分

我们重新回顾分子轨道 ERI 积分的定义：

\[ g_{ij}^{ab} = \iint \phi_i (\boldsymbol{r}_1) \phi_a (\boldsymbol{r}_1) \frac{1}{|\boldsymbol{r}_1 - \boldsymbol{r}_2|} \phi_j (\boldsymbol{r}_2) \phi_b (\boldsymbol{r}_2) \, \mathrm{d} \boldsymbol{r}_1 \, \mathrm{d} \boldsymbol{r}_2 \]

我们从数值积分的角度对上式作重新解释。\(\boldsymbol{r}_1\) 数值积分坐标集合是 \(\{ \boldsymbol{r}_R \}\)，其对应的求和权重是 \(\{ w_R \}\)；\(\boldsymbol{r}_2\) 数值积分坐标集合是 \(\{ \boldsymbol{r}_S \}\)，其对应的求和权重是 \(\{ w_S \}\)。那么，

\[ g_{ij}^{ab} \simeq w_R w_S \times \phi_i (\boldsymbol{r}_R) \phi_a (\boldsymbol{r}_R) \frac{1}{|\boldsymbol{r}_R - \boldsymbol{r}_S|} \phi_j (\boldsymbol{r}_S) \phi_b (\boldsymbol{r}_S) \]

一种容易想到的思路是，我们定义

• Xo \(X_i^R = w_R^{1/2} \phi_i(\boldsymbol{r}_R)\) (dim: \(i, R\)) 加权的占据轨道 \(i\) 在坐标 \(\boldsymbol{r}_R\) 下的数值；

• Xv \(X_a^R = w_R^{1/2} \phi_a(\boldsymbol{r}_R)\) (dim: \(a, R\)) 加权的非占轨道 \(a\) 在坐标 \(\boldsymbol{r}_R\) 下的数值；

这种拆分思路也可以说是依格点拆分。由于后文才会讲述的原因，这边称该格点为母格点 (parent grid)。下述程序所给出的格点相当于 Matthews 文章[2] 所述的 \((L_\mathrm{max}, N_1, N_\textsf{H}) = (11, 19, 11)\) 格点，即角向采用 Lebedev 精确到 11 阶的格点 (共 50 个角向格点)、第二周期径向 19 格点、氢原子径向 11 格点。因此，\(n_\mathrm{gird}^\mathrm{parent} = 2250\)。

ni = NumInt()
grids = Grids(mol)
grids.atom_grid = {"O": (19, 50), "H": (11, 50)}
grids.prune = None
grids.build()
n_parent = grids.size
n_parent

2050

定义完格点后，我们自然地可以给出 Xo \(X_i^R\) 与 Xv \(X_a^R\)：

Xao = ni.eval_ao(mol, grids.coords)
Xo = einsum("Ru, ui, R -> iR", Xao, Co, grids.weights**0.5)
Xv = einsum("Ru, ua, R -> aR", Xao, Cv, grids.weights**0.5)

尽管说，我们如果定义 \(U_{RS} = |\boldsymbol{r}_R - \boldsymbol{r}_S|^{-1}\)，那么依据格点积分的定义，容易推知

\[ g_{ij}^{ab} \simeq X_i^R X_a^R U_{RS} X_j^S X_b^S \]

但这个格点可以说是非常小了，比 SG-1 恐怕还小一个数量级，完全无法用于 DFT 的格点积分。因此，如果真的用上面定义的 \(X_i^R, X_a^R, U_{RS}\) 来近似 ERI \(g_{ij}^{ab}\)，这就显得太粗糙了。

10.4.3. 最小二乘 (LS) 的应用

即使上面的表达式不能直接利用，我们仍然要说，上面的分析是很有帮助的。由于一般来说的张量拆解需要进行大量的迭代，才能将拆解矩阵稳定下来。但上述讨论中，我们认为 \(X_i^R, X_a^R\) 的选择是合理的，剩下的问题是 \(U_{RS} = |\boldsymbol{r}_R - \boldsymbol{r}_S|^{-1}\) 不适合用来近似 \(g_{ij}^{ab}\)。

那么我们需要选取另一个矩阵 \(V_{RS}\)，使得下述近似是合理的：

\[ g_{ij}^{ab} \simeq X_i^R X_a^R V_{RS} X_j^S X_b^S \]

这个近似将会基于最小二乘 (LS, Least-Square) 给出。这个思路是从 Devin A. Matthews[3] 的文章所来。相信该问题的最早前身是 T. J. Martínez 与 C. D. Sherrill[4] 的文章。

为了求使得 LS 误差最小的 \(V_{RS}\)，我们需要预先作一些准备。定义

• S \(S_{RS} = X_i^R X_a^R X_i^S X_a^S\) (dim: \(R, S\); shape: \(n_\mathrm{grid}^\mathrm{parent}, n_\mathrm{grid}^\mathrm{parent}\)) 格点重叠矩阵

S = einsum("iR, aR, iS, aS -> RS", Xo, Xv, Xo, Xv)

那么根据数值计算的原理，\(V_{RS}\) 的导出方式是

\[ V_{RS} := (\mathbf{S}^{-1})_{RT} X_i^T X_a^T g_{ij}^{ab} X_j^U X_b^U (\mathbf{S}^{-1})_{US} \]

尽管这是一个 \(O(n_\mathrm{occ}^2 n_\mathrm{vir}^2 n_\mathrm{grid}^\mathrm{parent})\) 即 \(O(N^5)\) 计算量，但若在上述表达式引入 DF，则计算量可以进一步降到 \(O(n_\mathrm{occ} n_\mathrm{vir} n_\mathrm{aux} n_\mathrm{grid}^\mathrm{parent})\) 即 \(O(N^4)\)：

\[ V_{RS} \simeq (\mathbf{S}^{-1})_{RT} X_i^T X_a^T Y_{ia,P} Y_{jb,P} X_j^U X_b^U (\mathbf{S}^{-1})_{US} \]

10.4.4. 方法 1：基于求本征问题的伪逆方法 (Pseudoinverse)

但上述计算存在两个比较严重的问题。首先是 \(\mathbf{S}\) 其实是很大的 \(n_\mathrm{grid}^\mathrm{parent} \times n_\mathrm{grid}^\mathrm{parent}\) 矩阵，其求逆的计算消耗尽管是 \(O(N^3)\) 但仍然应尽量避免。更严重的是，这个 \(\mathbf{S}\) 因为不满秩，其实压根不能求逆：

try: np.linalg.inv(S)
except Exception as e: print("!!! Error:", type(e), e)

!!! Error: <class 'numpy.linalg.LinAlgError'> Singular matrix

一种解决办法是利用求取伪逆的技术。首先是要对 \(S_{RS}\) 通过求本征问题，拆解为

\[ S_{RS} = R_{T'R} R_{T'S} \quad \text{or} \quad \mathbf{S} = \mathbf{R}^\mathrm{T} \mathbf{R} \]

其中，角标 \(T'\) 区别于母格点角标 \(R, S\)。我们称 \(T'\) 的角标为精简格点 (pruned grid) 角标。这个“精简”的格点数量原则上应是矩阵 \(S_{RS}\) 的秩；但现实中，\(S_{RS}\) 的秩是依据程序编写者心情而定的。我们这里使用对 \(S_{RS}\) 求本征问题来确定秩的大小；规定有效的本征向量对应的本征值与最大的本征值之比不小于 \(10^{-5}\)。从下面的程序可以看出，精简格点大小 \(n_\mathrm{grid}^\mathrm{pruned} = 321\)，比 DF 的辅助基大小 \(n_\mathrm{aux} = 242\) 要大一些。

eig, mat = np.linalg.eigh(S)
n_prune = ((eig / eig[-1]) > 1e-5).sum()
eig, mat = eig[:-n_prune-1:-1], mat[:, :-n_prune-1:-1]
n_prune

321

• R \(R_{T' R}\) (dim: \(T, R\); shape: \(n_\mathrm{grid}^\mathrm{pruned}, n_\mathrm{grid}^\mathrm{parent}\))

R = (mat * eig**0.5).T
np.allclose(R.T @ R, S, atol=1e-6)

True

利用该矩阵，以及下述定义的

• E \(E_{TU} = X_i^T X_a^T Y_{ia,P} Y_{jb,P} X_j^U X_b^U\) (dim: \(T, S\); shape: \(n_\mathrm{grid}^\mathrm{parent}, n_\mathrm{grid}^\mathrm{parent}\)) 在母格点表示下的双电子积分，这一步的最大计算复杂度是 \(O(n_\mathrm{occ} n_\mathrm{vir} n_\mathrm{aux} n_\mathrm{grid}^\mathrm{parent})\)

E = einsum("iT, aT, iaP, jbP, jU, bU -> TU", Xo, Xv, Y_ov, Y_ov, Xo, Xv)

伪逆问题化为了求解线性方程组问题：

\[ \mathbf{R}^\mathrm{T} \mathbf{R} \mathbf{V} \mathbf{R}^\mathrm{T} \mathbf{R} = \mathbf{E} \]

• V \(V_{RS}\) (dim: \(R, S\); shape: \(n_\mathrm{grid}^\mathrm{parent}, n_\mathrm{grid}^\mathrm{parent}\)) 最小二乘下的母格点表示下的库仑势

Z = np.linalg.lstsq(R.T, E  , rcond=None)[0]
Y = np.linalg.lstsq(R.T, Z.T, rcond=None)[0].T
W = np.linalg.lstsq(R  , Y  , rcond=None)[0]
V = np.linalg.lstsq(R  , W.T, rcond=None)[0].T

我们需要说明，尽管我们确实成功地导出矩阵 \(R_{T'R}\)，表明母格点 \(n_\mathrm{gird}^\mathrm{parent} = 2250\) 可以修剪到 \(n_\mathrm{grid}^\mathrm{pruned} = 321\)；但由于张量拆解本身的定义原因，我们暂时无法将母格点下定义的 \(X_i^R\) 乘以一个转换矩阵到。因此 \(n_\mathrm{grid}^\mathrm{pruned}\) 在伪逆方法下，只能表征母格点的冗余程度，但并不能用修剪后的格点来加速计算。

最终，MP2 的能量写为

\[\begin{split} \begin{align*} E_\mathsf{corr} &= 2 g_{ij}^{ab} g_{ij}^{ab} / D_{ij}^{ab} - g_{ij}^{ab} g_{ij}^{ba} / D_{ij}^{ab} \\ &= -2 X_i^R X_a^R V_{RS} X_j^S X_b^S X_i^T X_a^T V_{TU} X_j^U X_b^U g_i^l g_a^l g_j^l g_b^l \\ &\quad + X_i^R X_a^R V_{RS} X_j^S X_b^S X_i^T X_b^T V_{TU} X_j^U X_a^U g_i^l g_a^l g_j^l g_b^l \end{align*} \end{split}\]

corr_thc_mp2 = (
    #             g_ovov            , g_ovov            , 1/D_ovov
    - 2 * einsum("iR, aR, RS, jS, bS, iT, aT, TU, jU, bU, il, al, jl, bl ->",
                  Xo, Xv,  V, Xo, Xv, Xo, Xv,  V, Xo, Xv, go, gv, go, gv)
    #             g_ovov            , g_ovov            , 1/D_ovov
    +     einsum("iR, aR, RS, jS, bS, iT, bT, TU, jU, aU, il, al, jl, bl ->",
                  Xo, Xv,  V, Xo, Xv, Xo, Xv,  V, Xo, Xv, go, gv, go, gv))
print("==> [LS-THC-MP2 corr energy pseudoinv]  : {:20.12f}".format(corr_thc_mp2))
print("==> [LS-THC-MP2 corr error pseudoinv]   : {:20.12f}".format(corr_thc_mp2 - (corr_simple_os + corr_simple_ss)))

==> [LS-THC-MP2 corr energy pseudoinv]  :      -0.312121723045
==> [LS-THC-MP2 corr error pseudoinv]   :       0.000054638455

由于所有的计算都表示为了矩阵乘法，因此尽管 MP2 能量表达式非常冗长，但最大计算量其实是 \(O(n_\mathrm{LT} n_\mathrm{occ} n_\mathrm{vir} (n_\mathrm{grid}^\mathrm{parent})^2)\) 即 \(O(N^4)\) 的。这是相对来说很大的四次方，因此只有当 \(n_\mathrm{LT} (n_\mathrm{grid}^\mathrm{parent})^2 < n_\mathrm{occ} n_\mathrm{vir} n_\mathrm{aux}\) 时，LS-THC-LT-DF-MP2 相比于 DF-MP2 才能发挥真正的优势。关于 MP2 是 \(O(N^4)\) 的证明，需要参阅 Martínez 的文章附录[1]。

10.4.5. 方法 2：基于列轮换 Cholesky 矩阵分解的直接格点削减方法

伪逆的做法是最为直观的；但其代价是，伪逆只能告诉我们母格点冗余程度是多大，但绝大多数的计算都无法利用修剪后的格点。为了利用修剪后的格点，另一种做法，也是 Devin. A. Matthews 文章[3] 的核心内容，是利用列轮换的 (Pivoting) Cholesky 矩阵分解。我们对这个技术不作太多介绍；其分解结果表述为

\[ S_{RS} = \Pi_{RR'} R_{TR'} R_{TS'} \Pi_{SS'} \quad \text{or} \quad \mathbf{S} = \mathbf{\Pi} \mathbf{R}^\mathrm{T} \mathbf{R} \mathbf{\Pi}^\mathrm{T} \]

其中，\(\Pi_{TR'}\) 相当于值非 1 即 0 的轮换矩阵，具有酉矩阵性质，作用等同于列轮换。\(R_{TR'}\) 是列轮换后实际的 Cholesky 分解上三角矩阵。

但还要知道，我们的目的是要通过 Cholesky 分解，削减一部分格点。因此，尽管程序所输出的角标 \(R'\) 的大小与母格点大小 \(n_\mathrm{grid}^\mathrm{parent}\) 相同，但通过设置 Cholesky 对角元阈值为最大对角元的 \(\epsilon=10^{-4}\) 倍，我们可以将母格点数量削减至 \(n_\mathrm{grid}^\mathrm{pruned} = 484\) 个。我们只取这 \(n_\mathrm{grid}^\mathrm{pruned}\) 个格点用以进行计算。这个格点数量比方才伪逆所求出的格点数量要多一些。

• R \(R_{R'S'}\) (shape: \(n_\mathrm{grid}^\mathrm{pruned}, n_\mathrm{grid}^\mathrm{pruned}\)) 一定阈值下的 Cholesky 分解矩阵

• P \(\Pi_{RR'}\) (shape: \(n_\mathrm{grid}^\mathrm{parent}, n_\mathrm{grid}^\mathrm{pruned}\)) 一定阈值下的列轮换矩阵

R, P_idxes, n_prune, info = scipy.linalg.lapack.dpstrf(S, tol=1e-8*np.abs(S).max())
assert info == 1  # Make sure pivoting Cholesky runs in success
# Set lower triangular as zero, since lapack does not zero tril if upper Cholesky is invoked
R[np.tril_indices(n_parent, k=-1)] = 0
# Generate permutation matrix (though permutation could be performed by array instead of matrix)
P = np.zeros((n_parent, n_parent))
P[P_idxes-1, np.arange(n_parent)] = 1
R, P = R[:n_prune, :n_prune], P[:, :n_prune]
n_prune

484

正因为列轮换矩阵非 1 即 0，因此 \(\Pi_{T R'}^2 = \Pi_{T R'}\)。又由于其实酉矩阵的性质。因此，

\[ X_i^R X_a^R \Pi_{R R'} = X_i^R X_a^S \delta_{RS} \Pi_{R R'} = X_i^R X_a^S \Pi_{R R'} \Pi_{S R'} \Pi_{R R'} = (X_i^R \Pi_{R R'}) (X_a^S \Pi_{S R'}) \]

这意味着我们可以通过定义

• xo \(X_i^{R'} = X_i^R \Pi_{R R'}\) (shape: \(n_\mathrm{occ}, n_\mathrm{grid}^\mathrm{pruned}\))

• xv \(X_a^{R'} = X_a^R \Pi_{R R'}\) (shape: \(n_\mathrm{vir}, n_\mathrm{grid}^\mathrm{pruned}\))

将所有的运算全部从母格点转移到修剪后的格点下。

xo = Xo @ P
xv = Xv @ P

• E \(E_{T'U'} = X_i^{T'} X_a^{T'} Y_{ia,P} Y_{jb,P} X_j^{U'} X_b^{U'}\) (dim: \(T', U'\); shape: \(n_\mathrm{grid}^\mathrm{pruned}, n_\mathrm{grid}^\mathrm{pruned}\)) 在修剪格点表示下的双电子积分，这一步的最大计算复杂度是 \(O(n_\mathrm{occ} n_\mathrm{vir} n_\mathrm{aux} n_\mathrm{grid}^\mathrm{pruned})\)

E = einsum("iT, aT, iaP, jbP, jU, bU -> TU", xo, xv, Y_ov, Y_ov, xo, xv)

尽管求逆问题在修剪过的格点下是正定、可以实现的；但既然我们已经获得了 Cholesky 分解矩阵，我们没有必要再花 \(O(N^3)\) 的精力求逆了，而是解四次线性方程。由于满秩性成立，这里的求解过程尽管与上一段一样，但并非是伪逆：

\[ \mathbf{R}^\mathrm{T} \mathbf{R} \mathbf{V} \mathbf{R}^\mathrm{T} \mathbf{R} = \mathbf{E} \]

• V \(V_{R'S'}\) (dim: \(R', S'\); shape: \(n_\mathrm{grid}^\mathrm{pruned}, n_\mathrm{grid}^\mathrm{pruned}\)) 最小二乘下修剪格点表示下的库仑势

Z = np.linalg.lstsq(R.T, E  , rcond=None)[0]
Y = np.linalg.lstsq(R.T, Z.T, rcond=None)[0].T
W = np.linalg.lstsq(R  , Y  , rcond=None)[0]
V = np.linalg.lstsq(R  , W.T, rcond=None)[0].T

最后我们可以求取 MP2 的总能量：

corr_thc_mp2 = (
    #             g_ovov            , g_ovov            , 1/D_ovov
    - 2 * einsum("iR, aR, RS, jS, bS, iT, aT, TU, jU, bU, il, al, jl, bl ->",
                  xo, xv,  V, xo, xv, xo, xv,  V, xo, xv, go, gv, go, gv)
    #             g_ovov            , g_ovov            , 1/D_ovov
    +     einsum("iR, aR, RS, jS, bS, iT, bT, TU, jU, aU, il, al, jl, bl ->",
                  xo, xv,  V, xo, xv, xo, xv,  V, xo, xv, go, gv, go, gv))
print("==> [LS-THC-MP2 corr energy pivotcho]   : {:20.12f}".format(corr_thc_mp2))
print("==> [LS-THC-MP2 corr error pivotcho]    : {:20.12f}".format(corr_thc_mp2 - (corr_simple_os + corr_simple_ss)))

==> [LS-THC-MP2 corr energy pivotcho]   :      -0.312116117678
==> [LS-THC-MP2 corr error pivotcho]    :       0.000060243822

10.4.6. 简单总结

对于 LS-THC-LT-DF-MP2，其一些重要的操作步骤是：

• 生成 Density-Fitting 的 3c-2e 积分 \(Y_{\mu \nu, P}\) 需要 \(O(n_\mathrm{AO}^2 n_\mathrm{aux}^2)\)，随后的分子轨道转化 \(Y_{ia, P}\) 需要 \(O(n_\mathrm{occ} n_\mathrm{AO}^2 n_\mathrm{aux})\) 的计算量；

• 生成母格点重叠矩阵 \(S_{RS}\) 的计算需要 \(O(n_\mathrm{occ} n_\mathrm{vir} (n_\mathrm{grid}^\mathrm{parent})^2)\) 计算量；随后的列轮换 Cholesky 分解需要 \(O((n_\mathrm{grid}^\mathrm{pruned})^2)\) 计算量；

• 生成修剪格点电子互斥积分 \(E_{T'U'}\) 需要 \(O(n_\mathrm{occ} n_\mathrm{vir} n_\mathrm{aux} n_\mathrm{grid}^\mathrm{pruned})\) 计算量；

• 求取修剪格点最小二乘库仑势 \(V_{R'S'}\) 需要 \(O((n_\mathrm{grid}^\mathrm{pruned})^2)\) 计算量；

• 计算 MP2 的交换部分 (对 SS 产生贡献的部分) 需要 \(O(n_\mathrm{LT} n_\mathrm{occ} n_\mathrm{vir} (n_\mathrm{grid}^\mathrm{pruned})^2)\) 计算量。

因此，总地来说，计算量是 \(O(N^4)\) 的。

事实上，MP3 能量也可以通过类似的方法给出，计算量同样是 \(O(N^4)\)。

---

[1] (1,2)

Hohenstein, E. G.; Parrish, R. M.; Martínez, T. J. Tensor Hypercontraction Density Fitting. I. Quartic Scaling Second- and Third-Order Møller-Plesset Perturbation Theory. J. Chem. Phys. 2012, 137 (4), 044103. doi: 10.1063/1.4732310.

[2] (1,2)

Matthews, D. A. A Critical Analysis of Least-Squares Tensor Hypercontraction Applied to MP3. J. Chem. Phys. 2021, 154 (13), 134102. doi: 10.1063/5.0038764.

[3] (1,2)

Matthews, D. A. Improved Grid Optimization and Fitting in Least Squares Tensor Hypercontraction. J. Chem. Theory Comput. 2020, 16 (3), 1382–1385. doi: 10.1021/acs.jctc.9b01205.

[4]

Parrish, R. M.; Hohenstein, E. G.; Martínez, T. J.; Sherrill, C. D. Tensor Hypercontraction. II. Least-Squares Renormalization. J. Chem. Phys. 2012, 137 (22), 224106. doi: 10.1063/1.4768233.