4. MP2 核坐标梯度的重新推导

这一篇文档中，我们要讨论如何从比较简单的 MP2 核坐标梯度，反推出可以相对来说节省内存的算法。最后我们再总结通过 PySCF 所改编的实施方法。

由于角标开始复杂化，我们这里使用支持希腊字母、并且实现方式与 numpy.einsum 相近的 opt_einsum.contract 作为张量求和引擎。

4.1. 初始化

from pyscf import gto, scf, mp, grad, hessian, ao2mo, lib
import numpy as np
from functools import partial
from pyscf.scf import cphf
from pyscf.ao2mo import _ao2mo
import warnings
import opt_einsum

from pyxdh.DerivOnce import GradMP2

np.set_printoptions(4, suppress=True, linewidth=180)
warnings.filterwarnings("ignore")
einsum = opt_einsum.contract

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()

<pyscf.gto.mole.Mole at 0x7f53d062d640>

mf_scf = scf.RHF(mol).run()
mf_mp2 = mp.MP2(mol).run()
mf_grad = grad.mp2.Gradients(mf_mp2).run()
mf_grad.de

array([[-0.1109, -0.0858,  0.0086],
       [ 0.0768,  0.0067,  0.023 ],
       [ 0.0133,  0.059 ,  0.0179],
       [ 0.0209,  0.0201, -0.0494]])

gradh = GradMP2({"scf_eng": mf_scf})
gradh.E_1

array([[-0.1109, -0.0858,  0.0086],
       [ 0.0768,  0.0067,  0.023 ],
       [ 0.0133,  0.059 ,  0.0179],
       [ 0.0209,  0.0201, -0.0494]])

nocc, nmo, nao = mol.nelec[0], mol.nao, mol.nao
natm, nbas = mol.natm, mol.nbas
nvir = nmo - nocc
so, sv, sa = slice(0, nocc), slice(nocc, nmo), slice(0, nmo)

C, e = mf_scf.mo_coeff, mf_scf.mo_energy
Co, Cv = C[:, so], C[:, sv]
eo, ev = e[so], e[sv]
t2 = mf_mp2.t2
D = mf_scf.make_rdm1()

4.2. MP2 核坐标梯度

4.2.1. pyxdh 的做法

\[\begin{split}\begin{align} \partial_{A_t} E_\mathrm{MP2, c} &= D_{pq}^\mathrm{MP2} B_{pq}^{A_t} + W_{pq}^\mathrm{MP2} [\mathrm{I}] S_{pq}^{A_t} + 2 T_{ij}^{ab} (ia|jb)^{A_t} \\ \partial_{A_t} E_\mathrm{HF, tot} &= h_{\mu \nu}^{A_t} D_{\mu \nu} + \frac{1}{2} (\mu \nu | \kappa \lambda)^{A_t} D_{\mu \nu} D_{\kappa \lambda} - \frac{1}{4} (\mu \kappa | \nu \lambda)^{A_t} D_{\mu \nu} D_{\kappa \lambda} - 2 F_{ij} S_{ij}^{A_t} + \partial_{A_t} E_\mathrm{nuc} \end{align}\end{split}\]

(   # MP2 Contribution Derivative
    + einsum("pq, Apq -> A", gradh.D_r, gradh.B_1)
    + einsum("pq, Apq -> A", gradh.W_I, gradh.S_1_mo)
    + 2 * einsum("iajb, Aiajb -> A", gradh.T_iajb, gradh.eri1_mo[:, so, sv, so, sv])
).reshape(-1, 3)

array([[ 0.0298,  0.0308,  0.0364],
       [-0.0179, -0.0035, -0.0059],
       [-0.0062, -0.0225, -0.0046],
       [-0.0057, -0.0049, -0.0259]])

(   # Total Derivative
    + einsum("pq, Apq -> A", gradh.D_r, gradh.B_1)
    + einsum("pq, Apq -> A", gradh.W_I, gradh.S_1_mo)
    + 2 * einsum("iajb, Aiajb -> A", gradh.T_iajb, gradh.eri1_mo[:, so, sv, so, sv])
    + einsum("pq, Apq -> A", gradh.D, gradh.H_1_ao)
    + 0.5 * einsum("Auvkl, uv, kl -> A", gradh.eri1_ao, D, D)
    - 0.25 * einsum("Aukvl, uv, kl -> A", gradh.eri1_ao, D, D)
    - 2 * einsum("ij, Aij -> A", gradh.F_0_mo[so, so], gradh.S_1_mo[:, so, so])
    + mf_grad.grad_nuc().flatten()
).reshape(-1, 3)

array([[-0.1109, -0.0858,  0.0086],
       [ 0.0768,  0.0067,  0.023 ],
       [ 0.0133,  0.059 ,  0.0179],
       [ 0.0209,  0.0201, -0.0494]])

对于当前的问题，我们需要作至少两部分考虑：对梯度的导数的分离、与避免大于 \(n_\mathrm{occ}^2 n_\mathrm{AO}^2\) 大小的内存储存。如果需要硬盘空间，应避免大于 \(n_\mathrm{occ} n_\mathrm{AO}^3\)。

之所以不提出更高的要求，也是因为目前 PySCF 的代码就是如此实现的。在内存中没有 \(t_{ij}^{ab}\) 的情况下，程序无法处理 MP2 的核坐标梯度。一个额外允许的条件是，\(n_\mathrm{AO}^3\) 的内存量是可以接受的。尽管说对于小分子大基组的情况，这种三次方的内存实际上划不来；但当情况到苯分子的 aug-cc-pVQZ 时，\(n_\mathrm{occ}^2 n_\mathrm{AO}^2 \simeq n_\mathrm{AO}^3\)。我们姑且接受这种做法。

4.2.2. 一般程序的实现过程

实际上，一般的程序都不是用下述方式生成 MP2 相关能所贡献的梯度：

\[\partial_{A_t} E_\mathrm{MP2, c} = D_{pq}^\mathrm{MP2} B_{pq}^{A_t} + W_{pq}^\mathrm{MP2} [\mathrm{I}] S_{pq}^{A_t} + 2 T_{ij}^{ab} (ia|jb)^{A_t}\]

相信这是因为电子积分通常都是以原子轨道形式存在。程序无法一次性地处理张量，因此必须要作某种分割；最直观的分割方式是对性质 (原子坐标) \(A_t\) 作分割。

但这种做法实际上并不合适。举例来说，在生成 \((\partial_t \mu | \nu)\) 时，程序的 API 只提供了生成维度 \((t, \mu_{s_1}, \nu_{s_2})\) 的电子积分。其中的原子轨道 \(\mu_{s_1}, \nu_{s_2}\) 是可以被分割的；但 \(t\) 并不能分割。这或许是出于积分效率上的考量。

因此，对 \(A_t\) 直接的分割并不合适。当然，对于核坐标梯度，对原子本身 \(A\) 的分割是可以接受的；这会在后面具体实现时遇到。

退而求其次的方法，是对原子轨道本身作分割。这就要求处理性质的时候，必须要写成原子轨道张量的分割之间的乘积。尽管公式推导时，使用分子轨道角标会非常便利；但在编写实际程序时，必须要考虑将思路转换一下，用原子轨道来实现大部分功能。

因此，我们还是必须要回到原子轨道矩阵的运算。我们的目标是回到 Aikens TCA (24)：

\[\partial_{A_t} E_\mathrm{MP2, c} = D_{\mu \nu}^\mathrm{MP2} F_{\mu \nu}^{A_t} + W_{\mu \nu}^\mathrm{MP2} S_{\mu \nu}^{A_t} + \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \lambda)^{A_t}\]

记号不同

这里使用 \(\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2}\) 表示 Non-separable 部分的二阶密度，不包含 Separable 部分。这种定义方式是为了能将定义比较方便地拓展到 DFT 代码中。因此，这里与 Aikens TCA (24) 式存在差别。那边的公式第一项是 \(D_{\mu \nu}^\mathrm{MP2} h_{\mu \nu}^{A_t}\)。

这里与分子轨道的角标有很多的不同。

• 首先，这里没有出现 \(B_{pq}^{A_t}\)，转而使用了 \(h_{\mu \nu}^{A_t}\)。对于一阶梯度而言，这是很重要的；因为 \(h_{\mu \nu}^{A_t}\) 是可以直接求得积分，而 \(B_{pq}^{A_t}\) 则是需要经过复杂积分求取的。

• 作为代价，\(W_{\mu \nu}^\mathrm{MP2}\) 就不是简单的 \(W_{pq}^\mathrm{MP2} [\mathrm{I}]\) 的原子轨道表示了。

• \(\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2}\) 与 \(2 T_{ij}^{ab}\) 有关。尽管用原子轨道看起来会把问题弄得更复杂，张量大小更多；但实际上对 FLOP 没有很大的影响 (FLOP 的增加不在 \(O(N^5)\) 的算法上)。

我们总结一下原子轨道基的计算过程。随后再对此作展开描述。

• D_r_ao

\[D_{\mu \nu}^\mathrm{MP2} = C_{\mu p} D_{pq}^\mathrm{MP2} C_{\nu q}\]

D_r = gradh.D_r
D_r_ao = einsum("μp, pq, νq -> μν", C, D_r, C)

• W_I

\[\begin{split}\begin{aligned} W_{ij}^\mathrm{MP2} [\mathrm{I}] &= - 2 T_{ik}^{ab} (ja|kb) \\ W_{ab}^\mathrm{MP2} [\mathrm{I}] &= - 2 T_{ij}^{ac} (ib|jc) \\ W_{ai}^\mathrm{MP2} [\mathrm{I}] &= - 4 T_{jk}^{ab} (ij|bk) \\ W_{ia}^\mathrm{MP2} [\mathrm{I}] &= 0 \end{aligned}\end{split}\]

• W_II

\[W_{pq}^\mathrm{MP2} [\mathrm{II}] = - D_{pq}^\mathrm{MP2} \varepsilon_q\]

• W_III

\[W_{ij}^\mathrm{MP2} [\mathrm{III}] = - \frac{1}{2} A_{ij, pq} D_{pq}^\mathrm{MP2}\]

• W, W_ao

\[\begin{split}\begin{align} W_{pq}^\mathrm{MP2} &= W_{pq}^\mathrm{MP2} [\mathrm{I}] + W_{pq}^\mathrm{MP2} [\mathrm{II}] + W_{pq}^\mathrm{MP2} [\mathrm{III}] \\ W_{\mu \nu}^\mathrm{MP2} &= C_{\mu p} W_{pq}^\mathrm{MP2} C_{\nu q} \end{align}\end{split}\]

Ax0_Core = gradh.Ax0_Core
W_I = gradh.W_I
W_II = - einsum("pq, q -> pq", D_r, e)
W_III = np.zeros((nmo, nmo))
W_III[so, so] = - 0.5 * Ax0_Core(so, so, sa, sa)(D_r)
W = W_I + W_II + W_III
W_ao = einsum("μp, pq, νq -> μν", C, W, C)

• D2_r_ao

\[\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} = 2 T_{ij}^{ab} C_{\mu i} C_{\nu a} C_{\kappa j} C_{\lambda b}\]

D2_r_ao = 2 * einsum("iajb, μi, νa, κj, λb -> μνκλ", gradh.T_iajb, Co, Cv, Co, Cv)

最终的梯度结算：

\[\partial_{A_t} E_\mathrm{MP2, c} = D_{\mu \nu}^\mathrm{MP2} h_{\mu \nu}^{A_t} + W_{\mu \nu}^\mathrm{MP2} S_{\mu \nu}^{A_t} + \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \lambda)^{A_t}\]

(   # MP2 Contribution Derivative
    + einsum("μv, Aμv -> A", D_r_ao, gradh.F_1_ao)
    + einsum("μv, Aμv -> A", W_ao, gradh.S_1_ao)
    + einsum("μvkl, Aμvkl -> A", D2_r_ao, gradh.eri1_ao)
).reshape(-1, 3)

array([[ 0.0298,  0.0308,  0.0364],
       [-0.0179, -0.0035, -0.0059],
       [-0.0062, -0.0225, -0.0046],
       [-0.0057, -0.0049, -0.0259]])

4.3. 各张量的生成

4.3.1. \(D_{pq}^\mathrm{MP2, oo-vv}\) 的生成

我们首先会考虑比较简单的张量乘积：

\[\begin{split}\begin{aligned} D_{ij}^\text{MP2} &= - 2 T_{ik}^{ab} t_{jk}^{ab} \\ D_{ab}^\text{MP2} &= 2 T_{ij}^{ac} t_{ij}^{bc} \end{aligned}\end{split}\]

这个张量乘积可以很容易地实现。

# Code block 1
D_r_oovv = np.zeros((nmo, nmo))
D_r_oovv[so, so] = - 2 * einsum("ikab, jkab -> ij", 2 * t2 - t2.swapaxes(-1, -2), t2)
D_r_oovv[sv, sv] = 2 * einsum("ijac, ijbc -> ab", 2 * t2 - t2.swapaxes(-1, -2), t2)
np.allclose(gradh.D_r[so, so], D_r_oovv[so, so]), np.allclose(gradh.D_r[sv, sv], D_r_oovv[sv, sv])

(True, True)

上述过程可以很容易地通过拆分角标直接实现。譬如下述代码简单地对 \(i\) 做拆分，就很容易地在比较小的内存限制下工作，额外内存消耗 \(n_\mathrm{occ} n_\mathrm{vir}^2\)：

# Code block 2
D_r_oovv = np.zeros((nmo, nmo))
for i in range(nocc):
    D_r_oovv[so, so] += - 2 * einsum("iab, jab -> ij", 2 * t2[:, i] - t2[:, i].swapaxes(-1, -2), t2[:, i])
    D_r_oovv[sv, sv] += 2 * einsum("jac, jbc -> ab", 2 * t2[i] - t2[i].swapaxes(-1, -2), t2[i])
np.allclose(gradh.D_r[so, so], D_r_oovv[so, so]), np.allclose(gradh.D_r[sv, sv], D_r_oovv[sv, sv])

(True, True)

但需要注意，说到底上面并不是经过优化了的代码。特别是针对这类问题，einsum 函数在转置 (Code block 1) 或在重复遍历 (Code block 2) 上，会花费大量的时间。我们在文档最后会特别提到这两个函数的优化问题。

这份文档的主要目标暂时不是作程序的最优化，而是首先对内存占用进行优化，指明一条可以程序实现的通路；后面的程序优化问题就希望是非常细节的优化了。

4.3.2. \(\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2}\) 的生成：初步方案

MP2 梯度没有简单的代码。所有剩余的代码都要经过 \(\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2}\) 的生成。

\[\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} = 2 T_{ij}^{ab} C_{\mu i} C_{\nu a} C_{\kappa j} C_{\lambda b}\]

我们首先不考虑内存的消耗 (因为储存这个二阶密度矩阵需要 \(n_\mathrm{AO}^4\) 的内存，显然这是不能接受的)，看看通过 PySCF 代码改编而来的生成方式。

\[\begin{split}\begin{align} \mathtt{D2t1}_{i j \nu \lambda} &= t_{ij}^{ab} C_{\nu a} C_{\lambda b} \\ \mathtt{D2t2}_{i \nu \lambda j} &= 4 \times \mathtt{D2t1}_{i j \nu \lambda} - 2 \times \mathtt{D2t1}_{i j \lambda \nu} \\ \mathtt{D2t3}_{\mu \nu \lambda j} &= C_{\mu i} \mathtt{D2t2}_{i \nu \lambda j} \\ \mathtt{D2t4}_{\mu \nu \lambda \kappa} &= \mathtt{D2t3}_{\mu \nu \lambda j} C_{\kappa j} \end{align}\end{split}\]

我们能注意到，\(\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} = \mathtt{D2t4}_{\mu \nu \lambda \kappa}\)。

# D2t1 = einsum("ijab, νa, λb -> ijνλ", t2, Cv, Cv)
D2t1 = _ao2mo.nr_e2(
    t2.reshape(nocc**2, nvir**2), Cv.T,
    (0, nao, 0, nao), "s1", "s1").reshape((nocc, nocc, nao, nao))
# D2t2 = 4 * einsum("ijνλ -> iνλj", D2t1) - 2 * einsum("ijλν -> iνλj", D2t1)
D2t2 = 4 * D2t1.transpose((0, 2, 3, 1)) - 2 * D2t1.transpose((0, 3, 2, 1))
D2t3 = einsum("μi, iνλj -> μνλj", Co, D2t2)
D2t4 = einsum("μνλj, κj -> μνλκ", D2t3, Co)

np.allclose(D2t4.swapaxes(-1, -2), D2_r_ao)

True

其中，D2t1 与 D2t2 的生成对内存的消耗是我们可以接受的；在 PySCF 中，D2t2 用 part_dm2 变量表示。这两步若不使用 einsum，效率会比较高。但 D2t3 与 D2t4 两步的消耗分别是 \(n_\mathrm{occ} n_\mathrm{AO}^3\) 与 \(n_\mathrm{AO}^4\)，我们无法承受。在任何实际的程序中，这两步都必须要进行分割。

在 PySCF 中，被分割的对象是 \(\mu\) 角标，从而将 D2t4 的内存消耗降到 \(n_\mathrm{blk} n_\mathrm{AO}^3\)；但代价是无法生成完整的 D2t4 张量，必须要将其立即应用于中间矩阵或最终梯度性质的计算中。

4.3.3. 二阶约化密度对梯度的贡献

我们先考察下述计算过程：

\[\partial_{A_t} E_\mathrm{MP2, c} \leftarrow \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \lambda)^{A_t}\]

einsum("μvkl, Aμvkl -> A", D2_r_ao, gradh.eri1_ao).reshape(-1, 3)

array([[ 0.0182,  0.0146,  0.0167],
       [-0.0142, -0.0026, -0.0048],
       [-0.003 , -0.0113, -0.0028],
       [-0.001 , -0.0007, -0.009 ]])

推导的目标是将 \(\mu\) 的角标成功分离。最好可以一定程度上利用对称性。

首先，我们要注意到 \(\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} = \Gamma_{\kappa \lambda \mu \nu}^\mathrm{MP2}\)。这个性质与 \(T_{ij}^{ab} = T_{ji}^{ba}\) 的对称性有关。

np.allclose(D2_r_ao, D2_r_ao.transpose((2, 3, 0, 1)))

True

随后，我们希望将 \((\mu \nu | \kappa \lambda)^{A_t}\) 作展开：

\[\partial_{A_t} E_\mathrm{MP2, c} \leftarrow \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \lambda)^{A_t} = \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} \big[ (\partial_{A_t} \mu \nu | \kappa \lambda) + (\mu \partial_{A_t} \nu | \kappa \lambda) + (\mu \nu | \partial_{A_t} \kappa \lambda) + (\mu \nu | \kappa \partial_{A_t} \lambda) \big]\]

我们的分割对象是 \(\mu\) 原子轨道，因此将所有出现 \(\partial_{A_t}\) 的记号全都转换到 \(\mu\) 的轨道上。利用下面的角标对换：

\[\begin{split}\begin{align} \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \partial_{A_t} \nu | \kappa \lambda) \rightarrow \Gamma_{\nu \mu \kappa \lambda}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) \quad & (\mu, \nu) \rightarrow (\nu, \mu) \\ \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \partial_{A_t} \kappa \lambda) \rightarrow \Gamma_{\kappa \lambda \mu \nu}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) \quad & (\kappa \lambda, \mu \nu) \rightarrow (\mu \nu, \kappa \lambda) \\ \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \partial_{A_t} \lambda) \rightarrow \Gamma_{\lambda \kappa \nu \mu}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) \quad & (\lambda \kappa, \mu \nu) \rightarrow (\mu \nu, \kappa \lambda) \\ \end{align}\end{split}\]

利用到 \(\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} = \Gamma_{\kappa \lambda \mu \nu}^\mathrm{MP2}\) 的性质，下述关系成立 (在对 \(\mu \nu \kappa \lambda\) 求和的情况下)：

\[\begin{split}\begin{align} \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \partial_{A_t} \kappa \lambda) &= \Gamma_{\kappa \lambda \mu \nu}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) = \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) \\ \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \partial_{A_t} \lambda) &= \Gamma_{\lambda \kappa \nu \mu}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) = \Gamma_{\nu \mu \lambda \kappa}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) \end{align}\end{split}\]

这次利用到 \((\partial_{A_t} \mu \nu | \kappa \lambda) = (\partial_{A_t} \mu \nu | \lambda \kappa)\) 的性质，因此 (在对 \(\mu \nu \kappa \lambda\) 求和的情况下)：

\[\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \partial_{A_t} \lambda) = \Gamma_{\nu \mu \lambda \kappa}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda) = \Gamma_{\nu \mu \kappa \lambda}^\mathrm{MP2}(\partial_{A_t} \mu \nu | \kappa \lambda)\]

因此，最终的表达式是

\[\partial_{A_t} E_\mathrm{MP2, c} \leftarrow \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} (\mu \nu | \kappa \lambda)^{A_t} = 2 \big( \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} + \Gamma_{\nu \mu \kappa \lambda}^\mathrm{MP2} \big) (\partial_{A_t} \mu \nu | \kappa \lambda)\]

这是可以通过对 \(\mu\) 进行分割，进而求和得到的结果。这里恰好是因为必须要对核坐标分割 (得到 \(\mu_A\))，我们就不进行更细致的分割了。实际程序中，要进行更小单位的分割，但过程是类似的。

回顾到 D2t2 是可以完整放在内存中的量，因此推导从 D2t2 开始：

\[\Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} = C_{\mu i} \mathtt{D2t2}_{i \nu \lambda j} C_{\kappa j}\]

\[\begin{split}\begin{align} \partial_{A_t} E_\mathrm{MP2, c} &\leftarrow 2 \big( \Gamma_{\mu \nu \kappa \lambda}^\mathrm{MP2} + \Gamma_{\nu \mu \kappa \lambda}^\mathrm{MP2} \big) (\partial_{A_t} \mu \nu | \kappa \lambda) \\ &= - 2 \big( \Gamma_{\mu \nu_A \kappa \lambda}^\mathrm{MP2} + \Gamma_{\nu \mu_A \kappa \lambda}^\mathrm{MP2} \big) (\partial_t \mu_A \nu | \kappa \lambda) \\ &= - 2 \big( C_{\mu_A i} \mathtt{D2t2}_{i \nu \lambda j} C_{\kappa j} + C_{\nu i} \mathtt{D2t2}_{i \mu_A \lambda j} C_{\kappa j} \big) (\partial_t \mu_A \nu | \kappa \lambda) \end{align}\end{split}\]

我们会在程序中，声明

\[\begin{split}\begin{align} \mathtt{D2t5}_{\mu_A \nu \kappa \lambda} &= C_{\mu_A i} \mathtt{D2t2}_{i \nu \lambda j} C_{\kappa j} + C_{\nu i} \mathtt{D2t2}_{i \mu_A \lambda j} C_{\kappa j} \\ \mathtt{eri\_ip1t1}_{t \mu_A \nu \kappa \lambda} &= (\partial_t \mu_A \nu | \kappa \lambda) \end{align}\end{split}\]

de_contrib_D2_r = np.zeros((natm, 3))
for A in range(natm):
    Ash0, Ash1, A0, A1 = mol.aoslice_by_atom()[A]
    # Evaluate D2t5
    D2t5 = einsum("μi, iνλj -> μνλj", Co[A0:A1], D2t2)
    D2t5 += einsum("νi, iμλj -> μνλj", Co, D2t2[:, A0:A1])
    D2t5 = einsum("μνλj, κj -> μνλκ", D2t5, Co)
    # Evaluate eri_ip1t1
    shls_slice = (Ash0, Ash1, 0, nbas, 0, nbas, 0, nbas)
    eri_ip1t1 = mol.intor("int2e_ip1", shls_slice=shls_slice)
    # Count contribution to derivative
    de_contrib_D2_r[A] = - 2 * einsum("μνλκ, tμνλκ -> t", D2t5, eri_ip1t1)
de_contrib_D2_r

array([[ 0.0182,  0.0146,  0.0167],
       [-0.0142, -0.0026, -0.0048],
       [-0.003 , -0.0113, -0.0028],
       [-0.001 , -0.0007, -0.009 ]])

最后指出，在 PySCF 中，会对 D2t5 的 \((\kappa, \lambda)\) 角标作对称化，使得积分 eri_ip1t1 可以利用到 s2kl 的对称性，从而将积分时间降低一半。

整个过程中，最耗时的步骤是 D2t5 的生成，需要 \(n_\mathrm{AO}^4 n_\mathrm{occ}\)；需要的内存空间是 \(n_\mathrm{blk} n_\mathrm{AO}^3\)。

4.3.4. \(W_{pq}^\mathrm{MP2} [\mathrm{I}]\) 的生成与重要中间矩阵 \(I_{\mu \nu}\)

对于这个问题，尽管可以像 \(D_{ij}^\mathrm{MP2}\) 或 \(D_{ab}^\mathrm{MP2}\) 一样直接对分子轨道作张量缩并 (事实上这样也确实更快)；但这里依照 PySCF 的流程，先导出其中一个重要的中间矩阵，称为 Imat \(I_{\mu \nu}\)。这个中间矩阵在后面会经常用到。这里的 \(I_{\mu \nu}\) 不代表恒等矩阵；我们用 \(\delta_{\mu \nu}\) 表示恒等矩阵。

我们先用最简单的代码给出 \(I_{\mu \nu}\) 的定义。

\[\begin{split}\begin{align} I_{pi} &= 2 T_{ik}^{ab} (pa|kb) \\ I_{pa} &= 2 T_{ik}^{ab} (ip|kb) \\ I_{\mu \nu} &= S_{\mu \eta} C_{\eta p} I_{pq} C_{\nu q} \end{align}\end{split}\]

Imat_mo = np.zeros((nmo, nmo))
Imat_mo[:, so] = 2 * einsum("iakb, pakb -> pi", gradh.T_iajb, gradh.eri0_mo[sa, sv, so, sv])
Imat_mo[:, sv] = 2 * einsum("iakb, ipkb -> pa", gradh.T_iajb, gradh.eri0_mo[so, sa, so, sv])
Imat = gradh.S_0_ao @ C @ Imat_mo @ C.T

实际上分子轨道的 Imat_mo 已经能得到与 \(W_{pq}^\mathrm{MP2} [\mathrm{I}]\) 有关的张量了。

\[W_{ij}^\mathrm{MP2} [\mathrm{I}] = - I_{ji}, \; W_{ab}^\mathrm{MP2} [\mathrm{I}] = - I_{ab}, \; W_{ai}^\mathrm{MP2} [\mathrm{I}] = - 2 * I_{ia}\]

np.allclose(- Imat_mo[so, so].T, gradh.W_I[so, so]), \
np.allclose(- Imat_mo[sv, sv].T, gradh.W_I[sv, sv]), \
np.allclose(- 2 * Imat_mo[so, sv].T, gradh.W_I[sv, so])

(True, True, True)

随后我们就考察 Imat \(I_{\mu \nu}\) 的生成。先讨论生成目的。尽管 \(T_{ik}^{ab} (ja|kb)\) 的直接缩并的耗时会更低，但我们以后碰到的缩并不止有这种形式；还包括 Lagrangian 量的 \(T_{ik}^{cb} (ac|kb)\) 等等形式。因此，处理这类型的张量缩并总计算量估计为 \(n_\mathrm{AO}^5\)，而非 \(n_\mathrm{occ} n_\mathrm{AO}^4\)。我们认为这种情况下，在原子轨道预先处理会比较好；在原子轨道下还可以进行有效的分割。

那么我们就将分子轨道的计算式用原子轨道表示：

\[I_{pi} = 2 T_{ik}^{ab} (pa|kb) = 2 T_{ik}^{ab} C_{\mu p} C_{\xi a} C_{\kappa k} C_{\lambda b} (\mu \xi | \kappa \lambda) = 2 T_{ik}^{ab} C_{\mu p} C_{\xi a} C_{\kappa k} C_{\lambda b} (\xi \mu | \kappa \lambda)\]

注意到 \(\Gamma_{\nu \xi \kappa \lambda}^\mathrm{MP2} = 2 T_{ij}^{ab} C_{\nu i} C_{\xi a} C_{\kappa j} C_{\lambda b}\)，我们发现上式几乎可以构成 \(\Gamma_{\nu \xi \kappa \lambda}^\mathrm{MP2}\)，但还缺少了 \(C_{\nu i}\)。补上这一项的方式是插入恒等式 \((\mathbf{C}^{-1})_{\nu i} C_{\nu i} = 1\)。

\[I_{pi} = 2 T_{ik}^{ab} (pa|kb) = C_{\mu p} \Gamma_{\nu \xi \kappa \lambda}^\mathrm{MP2} (\xi \mu | \kappa \lambda) (\mathbf{C}^{-1})_{\nu i}\]

这里假设了 \(\mathbf{C}\) 作为系数矩阵是满秩的；但不满秩的情况下上式也能成立。回顾到 RHF 的条件之一，即轨道系数正交 \(C_{\nu i} S_{\nu \eta} C_{\eta j} = \delta_{ij}\)，那么可以认为 \((\mathbf{C}^{-1})_{\nu i} = S_{\nu \eta} C_{\eta i}\)。从而

\[I_{pi} = 2 T_{ik}^{ab} (pa|kb) = C_{\mu p} \Gamma_{\nu \xi \kappa \lambda}^\mathrm{MP2} (\xi \mu | \kappa \lambda) S_{\nu \eta} C_{\eta i}\]

从而这里就可以定义原子轨道基下的矩阵

\[\begin{split}\begin{align} I_{\mu \nu}^\mathrm{occ} &= (\xi \mu | \kappa \lambda) \Gamma_{\nu \xi \kappa \lambda}^\mathrm{MP2} \\ I_{pi} &= C_{\mu p} I_{\mu \nu}^\mathrm{occ} S_{\nu \eta} C_{\eta i} \end{align}\end{split}\]

基于同样的方法，定义 (注意 \(\xi\) 与 \(\nu\) 角标的顺序)

\[\begin{split}\begin{align} I_{\mu \nu}^\mathrm{vir} &= (\xi \mu | \kappa \lambda) \Gamma_{\xi \nu \kappa \lambda}^\mathrm{MP2} \\ I_{pa} &= C_{\mu p} I_{\mu \nu}^\mathrm{vir} S_{\nu \eta} C_{\eta a} \end{align}\end{split}\]

可以验证，对于占据部分的 \(I_{\mu \nu}^\mathrm{occ}\)，非占的轨道下标缩并结果 \(C_{\mu p} I_{\mu \nu}^\mathrm{occ} S_{\nu \eta} C_{\eta a} = 0\)；反之亦然。因此，我们规定

\[\begin{split}\begin{align} I_{\mu \nu}^\mathrm{occ} &= I_{\mu \nu}^\mathrm{occ} + I_{\mu \nu}^\mathrm{vir} = (\xi \mu | \kappa \lambda) \big( \Gamma_{\nu \xi \kappa \lambda}^\mathrm{MP2} + \Gamma_{\xi \nu \kappa \lambda}^\mathrm{MP2} \big) \\ I_{pq} &= C_{\mu p} I_{\mu \nu} S_{\nu \eta} C_{\eta q} \end{align}\end{split}\]

生成 \(I_{\mu \nu}\) 的过程可以通过对原子轨道分批完成。所用到的重要中间量是 D2t5。这里要回顾二阶约化密度对梯度贡献的代码部分。

Imat = np.zeros((nao, nao))
for A in range(natm):
    Ash0, Ash1, A0, A1 = mol.aoslice_by_atom()[A]
    # Evaluate D2t5
    D2t5 = einsum("μi, iνλj -> μνλj", Co[A0:A1], D2t2)
    D2t5 += einsum("νi, iμλj -> μνλj", Co, D2t2[:, A0:A1])
    D2t5 = einsum("μνλj, κj -> μνλκ", D2t5, Co)
    # Evaluate eri_t1
    shls_slice = (Ash0, Ash1, 0, nbas, 0, nbas, 0, nbas)
    eri_t1 = mol.intor("int2e", shls_slice=shls_slice)
    # Count contribution to derivative
    Imat += einsum("ξμκλ, ξνκλ -> μν", eri_t1, D2t5)

Imat_mo = einsum("μp, μν, νη, ηq -> pq", C, Imat, mol.intor("int1e_ovlp"), C)
np.allclose(- Imat_mo[so, so].T, gradh.W_I[so, so]), \
np.allclose(- Imat_mo[sv, sv].T, gradh.W_I[sv, sv]), \
np.allclose(- 2 * Imat_mo[so, sv].T, gradh.W_I[sv, so])

(True, True, True)

4.3.5. Lagrangian 量 \(L_{ai}\)

有了 \(I_{pq}\) 之后，Lagrangian 量的求取就非常显然了。

\[\begin{split}\begin{align} L_{ai} &= A_{ai, pq} D_{pq}^\mathrm{MP2, oo-vv} - 4 T_{jk}^{ab} (ij|bk) + 4 T_{ij}^{bc} (ab|jc) \\ &= A_{ai, pq} D_{pq}^\mathrm{MP2, oo-vv} - 4 T_{jk}^{ab} (ji|kb) + 4 T_{ij}^{bc} (ab|jc) \\ &= A_{ai, pq} D_{pq}^\mathrm{MP2, oo-vv} - 2 I_{ia} + 2 I_{ai} \end{align}\end{split}\]

np.allclose(gradh.Ax0_Core(sv, so, sa, sa)(D_r_oovv) - 2 * Imat_mo[so, sv].T + 2 * Imat_mo[sv, so], gradh.L)

True

我们在这份文档中，不讨论 Ax0_Core 即 \(A_{ai, pq}\) 的实现。这会放在 DFT 代码实现中考虑。

4.3.6. MP2 弛豫密度 \(D_{ai}^\mathrm{MP2}\)

CP-HF 方程的求取在给出 \(L_{ai}\) 之后，就不是困难事了。

\[- (\varepsilon_a - \varepsilon_i) D_{ai}^\mathrm{MP2} - A_{ai, bj} D_{bj}^\mathrm{MP2} = L_{ai}\]

其计算复杂度是 \(O(T n_\mathrm{AO}^4)\)，即开销非常大的四次方算法。内存复杂度由 Ax0_Core 控制。

D_r_vo = cphf.solve(gradh.Ax0_Core(sv, so, sv, so), e, mf_scf.mo_occ, gradh.L)[0]

自此，弛豫密度部分就求取完毕了。这一项对于求取后续的各种梯度量非常重要。

D_r = D_r_oovv.copy()
D_r[sv, so] = D_r_vo

4.3.7. Hamiltonian Core 对梯度的贡献

\[\partial_{A_t} E_\mathrm{MP2, c} \leftarrow D_{\mu \nu}^\mathrm{MP2} F_{\mu \nu}^{A_t}\]

其中

\[D_{\mu \nu}^\mathrm{MP2} = C_{\mu p} D_{pq}^\mathrm{MP2} C_{\nu q}\]

D_r_ao = einsum("μp, pq, νq -> μν", C, D_r, C)

其中的 \(F_{\mu \nu}^{A_t}\) 我们暂不讨论其生成方式。我们这里就直接使用 PySCF 的结果。

mf_scf_hess = hessian.RHF(mf_scf)
F_1_ao = mf_scf_hess.make_h1(C, mf_scf.mo_occ)

de_contrib_hamilt = einsum("μν, Atμν -> At", D_r_ao, F_1_ao)
de_contrib_hamilt

array([[ 0.0256,  0.03  ,  0.0335],
       [-0.0145, -0.0026, -0.0045],
       [-0.005 , -0.02  , -0.0041],
       [-0.0062, -0.0074, -0.0249]])

4.3.8. 完整的 \(W_{\mu \nu}^\mathrm{MP2}\) 对梯度的贡献

由于我们已经完成了复杂的 \(W_{pq}^\mathrm{MP2} [\mathrm{I}]\) 的生成；其余步骤在给定 Ax0_Core 函数的情况下是显然容易实现的。我们假定 W_ao \(W_{\mu \nu}^\mathrm{MP2}\) 已经生成了。

稍微麻烦一些的问题是重叠矩阵的使用。PySCF 中不直接给出 \(S_{\mu \nu}^{A_t}\)，而是计算 \((\partial_t \mu | \nu)\)。因此，

\[\partial_{A_t} E_\mathrm{MP2, c} \leftarrow - \big( W_{\mu_A \nu}^\mathrm{MP2} + W_{\nu \mu_A}^\mathrm{MP2} \big) (\partial_t \mu_A | \nu)\]

de_contrib_W_ao = np.zeros((natm, 3))
int1e_ipovlp = mol.intor("int1e_ipovlp")
for A in range(natm):
    _, _, A0, A1 = mol.aoslice_by_atom()[A]
    de_contrib_W_ao[A] -= einsum("μν, tμν -> t", W_ao[A0:A1], int1e_ipovlp[:, A0:A1])
    de_contrib_W_ao[A] -= einsum("νμ, tμν -> t", W_ao[:, A0:A1], int1e_ipovlp[:, A0:A1])
de_contrib_W_ao

array([[-0.0139, -0.0139, -0.0137],
       [ 0.0107,  0.0018,  0.0034],
       [ 0.0018,  0.0089,  0.0023],
       [ 0.0015,  0.0033,  0.008 ]])

4.4. MP2 二阶梯度实现总结

%reset -f

from pyscf import gto, scf, mp, grad, hessian, ao2mo
import numpy as np
from pyscf.scf import cphf
from pyscf.ao2mo import _ao2mo
import opt_einsum

from pyxdh.DerivOnce import GradMP2  # only to use it's Ax0_Core utility here

np.set_printoptions(4, suppress=True, linewidth=180)
einsum = opt_einsum.contract

# Preparation 1: pyscf objects
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()

mf_scf = scf.RHF(mol).run()
mf_mp2 = mp.MP2(mol).run()
mf_grad = grad.mp2.Gradients(mf_mp2)
mf_scf_hess = hessian.RHF(mf_scf)
gradh = GradMP2({"scf_eng": mf_scf})

# Preparation 2: dimension declaration
nocc, nmo, nao = mol.nelec[0], mol.nao, mol.nao
natm, nbas = mol.natm, mol.nbas
nvir = nmo - nocc
so, sv, sa = slice(0, nocc), slice(nocc, nmo), slice(0, nmo)

# Preparation 3: important intermediates
C, e, D = mf_scf.mo_coeff, mf_scf.mo_energy, mf_scf.make_rdm1()
mo_occ = mf_scf.mo_occ
Co, Cv = C[:, so], C[:, sv]
eo, ev = e[so], e[sv]
t2 = mf_mp2.t2
Ax0_Core = gradh.Ax0_Core

# Step 1: RDM1 (occ-occ block, vir-vir block)
D_r_oovv = np.zeros((nmo, nmo))
for i in range(nocc):
    D_r_oovv[so, so] += - 2 * einsum("iab, jab -> ij", 2 * t2[:, i] - t2[:, i].swapaxes(-1, -2), t2[:, i])
    D_r_oovv[sv, sv] += 2 * einsum("jac, jbc -> ab", 2 * t2[i] - t2[i].swapaxes(-1, -2), t2[i])

# Step 2: RDM2 by atomic slices, Add to gradient contribution
# Step 3: Imat

# D2t1 = einsum("ijab, νa, λb -> ijνλ", t2, Cv, Cv)
D2t1 = _ao2mo.nr_e2(
    t2.reshape(nocc**2, nvir**2), Cv.T,
    (0, nao, 0, nao), "s1", "s1").reshape((nocc, nocc, nao, nao))
# D2t2 = 4 * einsum("ijνλ -> iνλj", D2t1) - 2 * einsum("ijλν -> iνλj", D2t1)
D2t2 = 4 * D2t1.transpose((0, 2, 3, 1)) - 2 * D2t1.transpose((0, 3, 2, 1))
# Allocate gradient contribution
de_contrib_D2_r = np.zeros((natm, 3))
# Allocate Imat
Imat = np.zeros((nao, nao))
for A in range(natm):
    Ash0, Ash1, A0, A1 = mol.aoslice_by_atom()[A]
    # Evaluate D2t5
    D2t5 = einsum("μi, iνλj -> μνλj", Co[A0:A1], D2t2)
    D2t5 += einsum("νi, iμλj -> μνλj", Co, D2t2[:, A0:A1])
    D2t5 = einsum("μνλj, κj -> μνλκ", D2t5, Co)
    # Evaluate eri_ip1t1
    shls_slice = (Ash0, Ash1, 0, nbas, 0, nbas, 0, nbas)
    eri_ip1t1 = mol.intor("int2e_ip1", shls_slice=shls_slice)
    # Count contribution to derivative
    de_contrib_D2_r[A] = - 2 * einsum("μνλκ, tμνλκ -> t", D2t5, eri_ip1t1)
    # Evaluate eri_t1
    shls_slice = (Ash0, Ash1, 0, nbas, 0, nbas, 0, nbas)
    eri_t1 = mol.intor("int2e", shls_slice=shls_slice)
    # Count contribution to derivative
    Imat += einsum("ξμκλ, ξνκλ -> μν", eri_t1, D2t5)
Imat_mo = einsum("μp, μν, νη, ηq -> pq", C, Imat, mol.intor("int1e_ovlp"), C)

# Step 4: Lagrangian
L = Ax0_Core(sv, so, sa, sa)(D_r_oovv) - 2 * Imat_mo[so, sv].T + 2 * Imat_mo[sv, so]

# Step 5: full RDM1 and CP-HF
D_r_vo = cphf.solve(gradh.Ax0_Core(sv, so, sv, so), e, mf_scf.mo_occ, gradh.L)[0]
D_r = D_r_oovv.copy()
D_r[sv, so] = D_r_vo
D_r_ao = einsum("μp, pq, νq -> μν", C, D_r, C)

# Step 6: full W matrix
W_I = np.zeros((nmo, nmo))
W_I[so, so] = - Imat_mo[so, so].T
W_I[sv, sv] = - Imat_mo[sv, sv].T
W_I[sv, so] = - 2 * Imat_mo[so, sv].T
W_II = - einsum("pq, q -> pq", D_r, e)
W_III = np.zeros((nmo, nmo))
W_III[so, so] = - 0.5 * Ax0_Core(so, so, sa, sa)(D_r)
W = W_I + W_II + W_III
W_ao = einsum("μp, pq, νq -> μν", C, W, C)

# Step 7: gradient contribution of first derivative of fock matrix
F_1_ao = mf_scf_hess.make_h1(C, mf_scf.mo_occ)
de_contrib_hamilt = einsum("μν, Atμν -> At", D_r_ao, F_1_ao)

# Step 8: gradient contribution of W matrix
de_contrib_W_ao = np.zeros((natm, 3))
int1e_ipovlp = mol.intor("int1e_ipovlp")
for A in range(natm):
    _, _, A0, A1 = mol.aoslice_by_atom()[A]
    de_contrib_W_ao[A] -= einsum("μν, tμν -> t", W_ao[A0:A1], int1e_ipovlp[:, A0:A1])
    de_contrib_W_ao[A] -= einsum("νμ, tμν -> t", W_ao[:, A0:A1], int1e_ipovlp[:, A0:A1])

# Step 9: sum up gradients
de_contrib_D2_r + de_contrib_hamilt + de_contrib_W_ao

array([[ 0.0298,  0.0308,  0.0364],
       [-0.0179, -0.0035, -0.0059],
       [-0.0062, -0.0225, -0.0046],
       [-0.0057, -0.0049, -0.0259]])

4.5. 程序的优化问题：以 \(D_{ij}^\mathrm{MP2}\) 为例

我们现在的目标是较大的分子，因此在作测试时，需要把体系扩大但不至于太大。我们令占据轨道数是 20，虚轨道数是 150。

NO, NV = 20, 150
T2 = np.random.randn(NO, NO, NV, NV)

现在如果用比较简单的两段代码 (用 V_1 与 V_2 分别表示) 完成下述工作：

\[V_{ij} = - 2 (2 t_{ik}^{ab} - t_{ik}^{ba}) t_{jk}^{ab}\]

V_1 = - 2 * einsum("ikab, jkab -> ij", 2 * T2 - T2.swapaxes(-1, -2), T2)
V_2 = np.zeros((NO, NO))
for k in range(NO):
    V_2 -= 2 * einsum("iab, jab -> ij", 2 * T2[:, k] - T2[:, k].swapaxes(-1, -2), T2[:, k])
np.allclose(V_1, V_2)

True

我们或许会认为，把所有的工作交给 einsum 完成会比较高效；但实际情况反而是有时显式地声明 for 循环效率更高。

%%timeit -n 5
V_1 = - 2 * einsum("ikab, jkab -> ij", 2 * T2 - T2.swapaxes(-1, -2), T2)

75.7 ms ± 13.7 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

%%timeit -n 5
V_2 = np.zeros((NO, NO))
for k in range(NO):
    V_2 -= 2 * einsum("iab, jab -> ij", 2 * T2[:, k] - T2[:, k].swapaxes(-1, -2), T2[:, k])

49.5 ms ± 7.48 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

当使用 profile 功能后，或许会发现其中的 reshape 相当耗时。实际上 einsum 或 tensordot 所要求的 reshape 是非常恐怖的。以 MOO_1 为例，在 50 次运行时的 %prun 下给出的结果是

 ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    150    1.996    0.013    1.996    0.013 {method 'reshape' of 'numpy.ndarray' objects}
      1    1.364    1.364    3.994    3.994 <ipython-input-275-26d8e19bfd5b>:1(func)
150/100    0.619    0.004    2.619    0.026 {built-in method numpy.core._multiarray_umath.implement_array_function}

按理最耗时的部分是张量乘积的 0.619 s，但实则是转置的 1.996 s。

如果我们手动作 reshape，反而效率会提升很多。

\[\begin{split}\begin{align} A_{i, kab} &= 2 t_{ik}^{ab} - t_{ik}^{ba} \\ B_{kab, j} &= t_{jk}^{ab} \\ V_{ij} &= - 2 A_{i, kab} B_{j, kab} \end{align}\end{split}\]

%%timeit -n 5
A = (2 * T2 - T2.swapaxes(-1, -2)).reshape(NO, -1)
B = T2.reshape(NO, -1).swapaxes(0, 1)
V_3 = - 2 * A.dot(B)

42.6 ms ± 6.11 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

这看起来作了很多手动的转置，效率非常低；但实际上是大大地提升了效率。但我们不能满足于此。对该过程重复 50 次，作 profile 得到的结果是

ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1    1.388    1.388    1.946    1.946 <ipython-input-284-9d7cb9c8a884>:1(func)
    50    0.557    0.011    0.557    0.011 {method 'dot' of 'numpy.ndarray' objects}

其中匿名函数 1 的耗时最长，有 1.388 秒。这一步应当归属于 \(2 t_{ik}^{ab} - t_{ik}^{ba}\) 的相减或标量乘的过程。为此，我们可以考虑下述思路。

\[\begin{split}\begin{align} A_{i, kab} &= t_{ik}^{ab} \\ B_{i, kab} &= t_{ik}^{ba} \\ V_{ij} &= - 4 A_{i, kab} A_{j, kab} + 2 A_{i, kab} B_{j, kab} \end{align}\end{split}\]

%%timeit -n 5
A = T2.reshape(NO, -1)
B = T2.swapaxes(-1, -2).reshape(NO, -1)
V_4 = - 4 * A.dot(A.T) + 2 * B.dot(A.T)

39.3 ms ± 11.1 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

该过程重复 50 次得到的 profile 结果是

ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   100    0.795    0.008    0.795    0.008 {method 'reshape' of 'numpy.ndarray' objects}
   100    0.738    0.007    0.738    0.007 {method 'dot' of 'numpy.ndarray' objects}
     1    0.012    0.012    1.545    1.545 <ipython-input-304-dde8e7fb18fe>:1(func)

似乎 reshape 又一次进入了函数，并且有较大的耗时 0.795 s；但实际上这个耗时包含在 MOD_3 的 1.388 s 中。该过程额外的损耗是进行了两次 dot 计算。对于当前的体系而言，有可能进行加减法运算的消耗反而比两次 dot 的消耗更大。因此，现在还无法判断两种算法孰优孰劣。目前看来，已经很难再减少 reshape 的耗时了；这部分的转置可能无法避免。

如果现在希望避免太大的内存消耗，我们可以取其中的 \(k\) 角标作显式 for 循环：

%%timeit -n 5
V_5 = np.zeros((NO, NO))
for k in range(NO):
    A = T2[:, k].reshape(NO, -1).copy()
    B = T2[:, k].swapaxes(-1, -2).reshape(NO, -1)
    V_5 += - 4 * A.dot(A.T) + 2 * B.dot(A.T)

44.9 ms ± 24.2 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

这就是以一定的时间换空间了。

最后指出，一般情况下 reshape 需要尽最大可能避免；因为 reshape 是单线程函数，但其它 numpy 函数往往是多线程的。