9. 简单理解 SCF 中的 DIIS#
创建时间：2019-10-23
这篇文档将会简单地叙述 GGA 为代表的 SCF DIIS。
DIIS 是一种 (几乎) 专门用于自洽场过程加速的算法。关于 DIIS 的算法与数学论述，这里并不作展开。这里推荐 C. David Sherrill 的笔记 [1] 与 Psi4NumPy 的 Jupyter Notebook [2] 作为拓展阅读。
这篇笔记会借助 PySCF 的 DIIS 程序，对 Fock 矩阵进行外推。我们将描述在第 \(t\) 步 DIIS 过程之后，如何更新第 \(t+1\) 步的 Fock 矩阵。我们 并不会 从头写一个 DIIS 程序；这一方面是出于程序复杂性考虑，因为一般的 DIIS 程序应当要允许便利地增添、删除迭代过程中间的向量，能够处理解矩阵方程时会出现的线性依赖问题，并且要能保证一定的程序效率。另一方面，若对 DIIS 的更新过程有所了解，那么原则上我们已经理解了 DIIS 程序了，剩下的细节将只是时间与耐心上的问题。
import numpy as np
import scipy
from pyscf import gto, dft, lib
from functools import partial
np.einsum = partial(np.einsum, optimize=["greedy", 1024 ** 3 * 2 / 8])
np.set_printoptions(5, linewidth=150, suppress=True)
9.1. PySCF 的 DIIS 使用#
9.1.1. 分子体系定义与 DIIS 程序#
首先，我们的分子体系是不对称的双氧水。
mol = gto.Mole()
mol.atom = """
O 0.0 0.0 0.0
O 0.0 0.0 1.5
H 1.0 0.0 0.0
H 0.0 0.7 1.0
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
nao = nmo = mol.nao
nocc = mol.nelec[0]
nvir = nmo - nocc
so, sv, sa = slice(0, nocc), slice(nocc, nmo), slice(0, nmo)
为了简化程序，我们借助 PySCF 的 DFT 自洽场类 scf_eng，用来生成 Fock 矩阵和密度矩阵。
scf_eng = dft.RKS(mol)
scf_eng.xc = "B3LYPg"
S = mol.intor("int1e_ovlp")
mo_occ = np.zeros(nmo)
mo_occ[:nocc] = 2
我们先简单地用下述程序展示 PySCF 中的 DIIS 是如何使用的。我们在后面会介绍 DIIS 类的具体的一些函数；这个用作演示用途的 DIIS 也时通过下述程序生成的。
下面的自洽场程序取自 pyxdh 文档 [3]，但作了一些修改与简化。这个自洽场程序大体思路与 Szabo [4] 第三章的叙述契合，也与 Psi4NumPy 和 PySCF 的不少演示程序比较相似。
与 Szabo 第三章叙述不太一样的地方有两处。其中一行代码是
D = coef * D + (1 - coef) * D_old
这行代码仅仅是用来对 Naive SCF 作的修改。Szabo 的第三章可以称作是 Naive SCF，即单纯地将 Fock 矩阵对角化生成分子轨道，再得到密度代入 Fock 矩阵中。但这一行会将上一次迭代的密度 \(D_{\mu \nu}^{t-1}\) 与这一次迭代的密度 \(D_{\mu \nu}^{t}\) 混合，产生新的密度代入 Fock 矩阵中。这仅仅是为了防止 Naive SCF 振荡收敛，不会用于 DIIS 加速的算法。
另一行代码是
F = func(diis=diis, F=F, C=C, mo_occ=mo_occ)
这行代码是用于指定 DIIS 的更新。方式在这份文档中，DIIS 的更新方式有
func_no_special：Naive SCF，不引入 DIISdiis_err_deviation：通过迭代过程中的 Fock 矩阵差值 \(\Delta F_{\mu \nu}^{t} = F_{\mu \nu}^{t} - F_{\mu \nu}^{t - 1}\) 更新 DIIS 状态diis_err_gradient：通过占据-非占 Fock 矩阵 \(F_{ai}^{t}\) 更新 DIIS 状态
之所以用这种不太常见也不太直观的代码方式指定 DIIS 更新方式，单纯地是因为节省文档中的代码空间，避免代码冗余。
def scf_process(func, coef=1.0, maxcycle=128):
diis = lib.diis.DIIS()
C = e = NotImplemented # Orbital (canonical) coefficient
D = np.zeros((nao, nao)) # Density in this iteration
D_old = np.zeros((nao, nao)) + 1e-4 # Density in last iteration
count = 0 # Iteration count (1, 2, ...)
while (not np.allclose(D, D_old)): # atol=1e-8, rtol=1e-5
if count > maxcycle:
raise ValueError("SCF not converged!")
count += 1
D_old = D
F = scf_eng.get_fock(dm=D) # Generate initial Fock matrix from Density
if count > 1: # avoid the case: C = NotImplemented
F = func(diis=diis, F=F, C=C, mo_occ=mo_occ) # Different DIIS approaches
# func_no_special : nothing happens
# diis_err_deviation : F = diis.update(F)
# diis_err_gradient : F = diis.update(F, scf_eng.get_grad(C, mo_occ))
e, C = scipy.linalg.eigh(F, S) # Solve FC = SCe
D = scf_eng.make_rdm1(mo_coeff=C, mo_occ=mo_occ) # D = 2 * C(occ).T @ C(occ)
D = coef * D + (1 - coef) * D_old # For convergence of original SCF
# func_no_special: D = 0.3 * D + 0.7 * D_old
# other cases : nothing happens
E_tot = scf_eng.energy_tot(dm=D)
print("SCF Converged in ", count, " loops")
print("Total energy (B3LYP) ", E_tot, " a.u.")
9.1.2. DIIS 方法加速效果#
现在我们可以来看每种 DIIS 更新方式会产生的效果了。
Naive SCF
def func_no_special(*args, **kwargs):
return kwargs["F"]
scf_process(func_no_special, coef=0.3)
SCF Converged in 51 loops
Total energy (B3LYP) -151.37754201912352 a.u.
DIIS：Fock 矩阵差值
def diis_err_deviation(*args, **kwargs):
diis, F = kwargs["diis"], kwargs["F"]
return diis.update(F)
scf_process(diis_err_deviation)
SCF Converged in 15 loops
Total energy (B3LYP) -151.37754323564036 a.u.
DIIS：占据-非占 Fock 矩阵
def diis_err_gradient(*args, **kwargs):
diis, F, C, mo_occ = kwargs["diis"], kwargs["F"], kwargs["C"], kwargs["mo_occ"]
return diis.update(F, scf_eng.get_grad(C, mo_occ))
scf_process(diis_err_gradient)
SCF Converged in 22 loops
Total energy (B3LYP) -151.37754323564045 a.u.
尽管从原理上，使用占据-非占 Fock 矩阵的方法应当更好；但在当前的体系下 Fock 矩阵差值的方法能更快地收敛。
我们能发现，若使用 PySCF 的 DIIS 类 (Psi4NumPy 的 helper_HF.py 的 DIIS_helper 类也是相似的)，我们在实际进行 DIIS 时只需要相对于 Naive SCF 增加一行，非常方便地就可以加速收敛：
F = diis.update(F) # diis_err_deviation
或者
F = diis.update(F, scf_eng.get_grad(C, mo_occ)) # diis_err_gradient
简单地说，这就是将 Fock 矩阵在每次迭代过程中，利用以前储存下来的 Fock 矩阵的信息进行再次更新。
9.2. DIIS 细节#
在这一段中，我们主要通过占据-非占 Fock 矩阵更新法的程序，对迭代 6 次时的 DIIS 状态进行较为细致的分析，并以此推导出第 6 次 DIIS 更新后的 Fock 矩阵。
第 6 次迭代时的 DIIS 类 diis 与更新前的 Fock 矩阵 F_old \(F_{\mu \nu}^{t=6}\)、更新后的 Fock 矩阵 F \(\mathscr{F}_{\mu \nu}\) 的获得方式是通过限制迭代次数而给出的：
diis = lib.diis.DIIS()
C = e = NotImplemented # Orbital (canonical) coefficient
D = np.zeros((nao, nao)) # Density in this iteration
D_old = np.zeros((nao, nao)) + 1e-4 # Density in last iteration
count = 0 # Iteration count (1, 2, ...)
F_old = NotImplemented # Variable in last iteration
while (not np.allclose(D, D_old)): # atol=1e-8, rtol=1e-5
count += 1
D_old = D
F = scf_eng.get_fock(dm=D) # Generate initial Fock matrix from Density
if count == 6:
F_old = F.copy()
F = diis.update(F, scf_eng.get_grad(C, mo_occ))
break
elif count > 1: # avoid the case: C = NotImplemented
F = diis.update(F, scf_eng.get_grad(C, mo_occ))
e, C = scipy.linalg.eigh(F, S) # Solve FC = SCe
D = scf_eng.make_rdm1(mo_coeff=C, mo_occ=mo_occ) # D = 2 * C(occ).T @ C(occ)
这里补充一个程序细节，更新后的 Fock 矩阵 \(\mathscr{F}_{\mu \nu}\) 也可以通过 diis.extrapolate 获得：
np.allclose(F, diis.extrapolate().reshape(nmo, nmo))
True
而 diis.update 除了给出更新后的 Fock 矩阵外，还将当前的 Fock 矩阵与误差信息更新入 diis 中，为下一次迭代的 DIIS 更新作准备。
9.2.1. DIIS 储存内容#
一般来说，DIIS 储存两部分内容：待外推信息 \(p_I^t\) 与误差信息 \(e_J^t\)。
我们在 SCF 过程中使用 DIIS 的目的是借助以前若干步迭代过程中的信息，对 Fock 矩阵作外推，得到在当前迭代步 \(t\) 下更好的 Fock 矩阵。因此，待外推信息 \(p_I^t\) 是第 \(t\) 次迭代过程计算得到的 Fock 矩阵 \(F_{\mu \nu}^t\)。
这里会有些诡异的地方是，待外推信息 \(p_I^t\) 是单下标 \(I\) 的向量，但原子轨道基组下的 Fock 矩阵 \(F_{\mu \nu}^t\) 是双下标的矩阵。事实上，\(p_I^t\) 在实践过程中就是将 \(F_{\mu \nu}^t\) 压平成一维向量。待外推信息 \(p_I^t\) 可以通过 diis.get_vec 给出，我们将这些向量整合为 vecs 变量中，其角标储存是 \((t, I)\)：
vecs = np.array([diis.get_vec(i) for i in range(diis.get_num_vec())])
我们记每次迭代的误差信息为 \(e_J^t\)。对于占据-非占 Fock 矩阵更新法而言，\(e_J^t\) 即是分子轨道基组下的非占-占据 Fock 矩阵 \(F_{ai}^t\)。我们知道，对于 SCF 收敛之后的状态下，\(F_{ai} = 0\)；但这在自洽场迭代过程中，该量一般地并不是零，甚至说自洽场过程就是为了达成 \(F_{ai} = 0\) 的目的也不为过。因此，\(F_{ai}^t\) 的状况可以看作是自洽场是否收敛得较好的判标；于是我们定义 \(e_J^t\) 为压平之后的 \(F_{ai}^t\)。
误差信息 \(e_J^t\) 可以通过 diis.get_err_vec 给出，我们将这些向量整合为 err_vecs 变量中，其角标储存是 \((t, J)\)：
err_vecs = np.array([diis.get_err_vec(i) for i in range(diis.get_num_vec())])
我们指出，\(p_I^t\) 与 \(e_J^t\) 下标所指代的维度未必要是一样的。
print(vecs.shape)
print(err_vecs.shape)
(5, 484)
(5, 117)
备注
从上述的叙述与代码中，能看到我们只进行了 6 次迭代，其中迭代过程的误差信息与待外推信息只储存了 5 次 (\(t\) 从 1 计数，外推信息的 \(t \in [2, 6]\))。我们定义当前作为迭代次数的上标 \(t\) 的集合是 \(\mathscr{T} = \{2, 3, 4, 5, 6\}\)；但在 PySCF 的 DIIS 类 diis 中，通过程序取出这些向量的指标时则应当使用 0, 1, 2, 3, 4。
我们知道，DIIS 为了进行外推，会储存许多待外推信息与误差信息；但对于大分子而言，这会占用许多内存空间。出于这个目的 (以及出于收敛性的考量)，DIIS 通常只会存比较少量地待外推信息与误差信息。PySCF 的 DIIS 一般只储存 6 次迭代过程的信息。这意味着，若我们进行了 15 次迭代，待外推矩阵至多也只会储存 6 个，其余的待外推信息或误差信息都会舍去。
为简化讨论，我们在这篇文档中不讨论如何舍去已经储存的待外推信息与误差信息。
9.2.2. DIIS 外推：理论#
有了所有的待外推信息 \(p_I^t\) 与误差信息 \(e_J^t\) 后，我们可以作出外推结果 \(\mathscr{p}_I = \mathscr{F}_{\mu \nu}\)。上述公式中看似有问题的单下标转双下标可以通过互阵的 reshape 实现。
外推的含义是
\(\mathscr{T}\) 代表 DIIS 当前储存的每个外推信息对应的被迭代次数的集合，在这里恰好是从 2 开始的所有的被迭代次数。如果我们现在的迭代次数非常大，但只允许 DIIS 储存不多于 6 个待外推信息，那么求和指标 \(t\) 的取值范围 \(\mathscr{T}\) 将会舍去这些迭代次数，从而保持其集合元素数量 \(|\mathscr{T}|\) 不超过 6。
我们人为地引入权重 \(w_t\) 以归一条件：
如果我们假定待外推的信息 \(p_I^t\) 与对应的误差信息 \(e_J^t\) 呈线性关系，那么被外推的信息 \(\mathscr{p}_I\) 的误差 \(\mathscr{e}_J\) 应当满足
我们希望误差 \(\Vert \mathscr{e}_J \Vert_2^2\) 最小化，但同时又满足 \(w_t\) 的归一化条件；那么我们通过 Lagrange 乘子法，构造以下损失函数
我们现在定义
那么损失函数可以写为
对上述损失函数求取关于 \(w_t\) 的偏导数，则得到
我们显然是希望让损失函数对关于 \(w_t\) 的偏导数为零；那么联立归一化条件 \(\sum_{t \in \mathscr{T}} w_t = 1\)，我们应当得到以下矩阵方程：
其中，\(t_0, t_1, \cdots \in \mathscr{T}\) 是互不相同的指标。求解上述方程，就可以获得权重 \(w_t\)，进而给出 \(\mathscr{F}_{\mu \nu} = \mathscr{p}_I = \sum_{t \in \mathscr{T}} w_t p_I^t\)，达成目标。
9.2.3. DIIS 外推：实现#
首先，我们出，diis 的一个隐含变量 diis._H 储存的就是矩阵方程 LHS 的矩阵部分：
A = diis._H[:diis.get_num_vec()+1, :diis.get_num_vec()+1]
A
array([[ 0. , 1. , 1. , 1. , 1. , 1. ],
[ 1. , 22.07074, -0.48488, 0.19048, -1.05937, -0.92243],
[ 1. , -0.48488, 21.30345, 0.23478, 1.63515, -0.41527],
[ 1. , 0.19048, 0.23478, 1.35725, -0.18202, 0.38307],
[ 1. , -1.05937, 1.63515, -0.18202, 5.33326, -0.14979],
[ 1. , -0.92243, -0.41527, 0.38307, -0.14979, 1.18526]])
我们能很方便地构建上述矩阵的第 1 行以下、第 1 列以右的子矩阵 \(B_{ts} = \sum_{J} e_J^t e_J^s\)：
np.einsum("tI, sI -> ts", err_vecs, err_vecs)
array([[22.07074, -0.48488, 0.19048, -1.05937, -0.92243],
[-0.48488, 21.30345, 0.23478, 1.63515, -0.41527],
[ 0.19048, 0.23478, 1.35725, -0.18202, 0.38307],
[-1.05937, 1.63515, -0.18202, 5.33326, -0.14979],
[-0.92243, -0.41527, 0.38307, -0.14979, 1.18526]])
我们可以直接解上述的矩阵方程：
b = np.zeros(diis.get_num_vec() + 1)
b[0] = 1
w = np.linalg.solve(A, b)
w = w[1:]
w
array([0.05123, 0.02428, 0.31624, 0.13904, 0.46921])
那么我们可以通过 \(\mathscr{F}_{\mu \nu} = \mathscr{p}_I = \sum_{t \in \mathscr{T}} w_t p_I^t\) 给出外推 Fock 矩阵 F_ex，并且与 diis 给出的外推的 Fock 矩阵 F 进行比较：
F_ex = np.einsum("t, tI -> I", w, vecs).reshape(nmo, nmo)
np.allclose(F_ex, F)
True
小技巧
在求解 DIIS 所给出的权重 \(w_t\) 向量的过程中，会遇到线性依赖关系，或者说会遇到 \(B_{ts}\) 数值上不满秩的情况。在这种情况下，求解矩阵方程可能会失败。
一种解决方案是，干脆关闭 DIIS，使用 Naive SCF 作最后的收尾工作。由于 DIIS 已经将电子态密度收敛到相当不错的状态了，因此应当能预期这种情况下 Naive SCF 可以正常地进行收敛。
另一种解决方式是对矩阵方程 \(\mathbf{A} \boldsymbol{x} = \boldsymbol{b}\) 的矩阵 \(\mathbf{A}\) 作对角化，并舍去其中绝对值极小的本征值与本征向量，求解一个子空间的线性方程组问题。这种解决方案应用在 PySCF 的 DIIS 程序中。
9.2.4. 对 Fock 矩阵差值方法的补充说明#
Fock 矩阵差值方法的计算过程与占据-非占 Fock 矩阵方法的实现过程几乎是相同的。唯一的区别是：
占据-非占 Fock 矩阵方法 \(e_J^t = F_{ai}^t\)
Fock 矩阵差值方法 \(e_J^t = \Delta F_{\mu \nu}^{t} = F_{\mu \nu}^{t} - F_{\mu \nu}^{t - 1}\)