3. Autograd (3)：梯度下降法解量子化学 RHF 自洽场能量#
创建时间：2019-12-17
这一篇文档将会简单地介绍量子化学计算中经常使用的 RHF 方法，在自动求导下的一种应用。
RHF (Restricted Hartree-Fock) 方法是量化计算的基本方法；对于更高精度的分子体系或晶体等周期性体系的计算，譬如 DFT、Post-HF 方法如 CC 和 CI，通常都是以 HF 方法为原型开发。但 RHF 方法的求解方程通常都是通过自洽场方法 (self-consistent field, SCF) 迭代计算；一种经典的迭代计算算法是 DIIS (见 简单理解 SCF 中的 DIIS)。关于 SCF 求解 RHF 的入门文档可以参考 Szabo, Ostlund [1]。关于自洽场求解 RHF 的程序相关问题，可以参考 pyxdh 文档。
而事实上，更经典但现在不太用于 RHF 方程求解的方法是将其化为局域凸优化问题。尽管通常这类凸优化问题属于类牛顿法的范畴，但如果将 Hessian 矩阵始终视为单位矩阵的常数倍，这就化为了普通的梯度下降问题，也恰好是自动求导可以解决的问题范畴。这篇文档就用 PyTorch 简单地实现这样一个功能。
在继续这篇文档之前，我们先引入程序库：
import torch
import math
import numpy as np
import scipy
from pyscf import gto, scf
import warnings
torch.set_printoptions(precision=5, sci_mode=False, linewidth=120)
np.set_printoptions(precision=5, suppress=True, linewidth=120)
warnings.filterwarnings("ignore")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
注意
这篇文档是找了一个作者稍微熟悉的优化问题，用自动求导来解决。但这不是通常处理这类问题的做法，况且并不高效。
同时需要指出的是，一般来说 RHF 并不是一个凸优化问题；但在比较合理的初猜下，可以将问题 近似地 看成凸优化问题。后面尽管默认了凸优化这个条件，但许多现实的 RHF 收敛问题实际上来源于非凸优化的复杂性。作者也不是搞优化问题的，因此很多术语可能用得不对。
因此，这篇文档是 娱乐向 的文档。
3.1. 分子体系的定义#
3.1.1. 分子体系能量#
我们使用的分子体系是不对称的双氧水分子，基组为 6-31G，用 PySCF 进行计算。分子体系的总能量为 energy_tot \(E_\mathrm{tot}\) -150.5850338 Hartree。这篇文档的目标就是计算得到这个数值。
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()
<pyscf.gto.mole.Mole at 0x7fa30ea25160>
scf_eng = scf.RHF(mol).run()
energy_tot = scf_eng.e_tot
energy_tot
-150.5850337808369
3.1.2. 恒定物理量与矩阵#
这里列举的是随分子不同而不同的量。这些量可能计算耗时小，也可能是计算过程中的重要构成部分；但它们与凸优化问题的变量无关，在这篇文档中姑且称为恒定量。
分子体系总能量中，一部分能量是无需通过量化方法，而只需要经典物理的库伦排斥力计算得到。这部分能量是原子核排斥能 energy_nuc \(E_\mathrm{nuc}\) 37.8846744 Hartree。为了简化叙述，我们就直接借用了 PySCF 的程序的结果。
energy_nuc = scf_eng.energy_nuc()
energy_nuc
37.88467440864127
下面是原子轨道角标的矩阵或张量。无电子相互作用的 Hamiltonian 积分 (Hamiltonian Core 积分) \(h_{\mu \nu}\) 记为 H，重叠积分 \(S_{\mu \nu}\) 记为 S，双电子积分 \((\mu \nu | \kappa \lambda)\) 记为 eri_ao。
T = mol.intor("int1e_kin")
Vnuc = mol.intor("int1e_nuc")
H = T + Vnuc
S = mol.intor("int1e_ovlp")
eri_ao = mol.intor("int2e")
上述矩阵或张量都是通过 PySCF 给出的；PySCF 使用的默认引擎是 numpy。我们需要将它们转为 PyTorch 引擎的矩阵或张量。
H = torch.tensor(H, device=device)
S = torch.tensor(S, device=device)
eri_ao = torch.tensor(eri_ao, device=device)
3.1.3. 凸优化自变量与密度矩阵#
在凸优化问题中，RHF 问题可以看作是对下述原子轨道空间下 \(\mathbf{X}\) 矩阵 (或写作矩阵元的形式，\(X_{\mu \nu}\)) 作为自变量的函数优化问题 (Helgaker et al. [2], eq 10.7.69，但这篇文档对密度矩阵 \(R_{\mu \nu}\) 的定义恰好是 Helgaker 课本中的两倍)：
其中，根据 Helgaker, eq 10.7.32，有
关于这些公式的意义与计算过程，我们后面还会有更详细的描述。这里需要留意的是，\(\mathbf{G}\) 是一个四维张量；其它地方的记号与矩阵乘法没有区别。
我们现在需要知道的是，其一，\(E_\mathrm{tot} [\mathbf{X}]\) 是我们的凸优化目标；对应于统计或机器学习中的概念，就类似于损失函数 (loss function)。其二，尽管 \(\mathbf{X}\) 是我们的学习目标，但真正有物理意义的、可以被实验观测的量是 \(\mathbf{R}[\mathbf{X}]\)；该量的意义是电子云密度 (简称密度)。换一种说法的话，RHF 方法作为量子化学方法的核心目标是给出基态的电子云密度；有了电子云密度之后，许多其它的分子性质是可以从中导出的。因此，对于 RHF 方法而言，其电子云密度的价值近乎等价于作为 Schrodinger 方程的求解目的的波函数。
方才我们已经用 PySCF 计算了分子的 RHF 的能量。PySCF 也可以给出分子基态下的 RHF 密度 D \(D_{\mu \nu}\)：
D = scf_eng.make_rdm1()
D = torch.tensor(D, device=device)
对于 RHF 问题而言，我们会说，\(\mathbf{X}\) 取到函数 \(E_\mathrm{tot} [\mathbf{X}]\) 处于最小值的极小值几乎等价于问题
我们等下会验证这个结论。
需要补充的是，PySCF 给出电子密度的方式与上述过程完全不相同。在一般的自洽场迭代过程中，电子态密度是由分子轨道系数得到的：
或用 Einstein Summation Convention 的语言，
这与 \(\mathbf{D} = \mathbf{R} [\mathbf{X}]\) 看似是完全无关的公式，但至少从结果上，两者确实相同。
3.2. 损失函数的定义#
3.2.1. PyTorch 矩阵幂运算#
这一段只是讨论一个技术细节。尽管 PyTorch 功能确实很强大，但一个比较尴尬的地方是 PyTorch 不支持不少需要一些数学技巧才能保证效率与稳定性的算法。其中一个算法是矩阵幂 (作者 Bing 了一刻钟发现实在找不到，于是作出了断言，但高 zi 贵 bi 的作者是不会提 issue 或者 stackoverflow 的 >.>)。
为了解决这样一个量化的问题，我们无奈需要自己造轮子，目标是只要能用就行。
下面的程序 calc_exp 是根据矩阵幂的定义给出的：
其中用到了矩阵指数运算 \(\mathbf{X}^k\) 的函数 mat_power；该函数通过递归定义。
我们知道，这种函数的求和上界应当是 \(K \rightarrow \infty\)；但实际的程序不能允许这种操作，必须让 \(K\) 作有界的截断。这种有界的截断的判据是，若通过截断计算得到的 \(\exp(\mathbf{X})\) 和 numpy 给出的不截断的 \(\exp(\mathbf{X})\) 近乎相等 (np.allclose 的判标下)，那么就允许截断。
def mat_power(X, order):
if order > 0:
return X @ mat_power(X, order - 1)
return torch.eye(X.shape[0], dtype=X.dtype, device=X.device)
def calc_exp(X, debug=False):
target = scipy.linalg.expm(X.cpu().clone().detach().numpy())
result = mat_power(X, 0) + mat_power(X, 1)
order = 1
while not np.allclose(result.cpu().clone().detach().numpy(), target):
order += 1
result += mat_power(X, order) / math.factorial(order)
if debug: print(np.linalg.norm(result.cpu().clone().detach().numpy() - target))
if debug: print("Current order is", order)
return result
读者可以尝试几个矩阵幂运算的例子。但需要留意，这个程序的鲁棒性显然是不高的；遇到上述定义下 \(\mathbf{X}\) 超出收敛域的情形下，这个程序就失效了。
3.2.2. 分子能量函数 \(E_\mathrm{tot} [\mathbf{X}]\)#
下面我们回到分子能量函数 (损失函数) \(E_\mathrm{tot} [\mathbf{X}]\) 的定义中。事实上让程序来做这件事其实是相当容易的，只需要几行代码就能搞定：
def calc_RX(X, S, R0):
return calc_exp(-X @ S) @ R0 @ calc_exp(S @ X)
def calc_EX(X, S, R0, H, eri_ao, energy_nuc):
RX = calc_RX(X, S, R0)
return (
torch.einsum("uv, uv ->", RX, H)
+ 0.5 * torch.einsum("uv, uvkl, kl ->", RX, eri_ao, RX)
- 0.25 * torch.einsum("uv, ukvl, kl ->", RX, eri_ao, RX)
+ energy_nuc
)
但我们需要对其中的过程作一些说明。首先是关于密度的导出式 calc_RX：
这里的 \(\mathbf{S}\) 是重叠积分。\(\mathbf{R}_0\) 原则上是任意的、满足下述三式的矩阵 (Helgaker, eq 10.7.25-27)：
上述三式对应的性质分别称为对称性、电子数守恒、基于重叠积分 \(\mathbf{S}\) 的幂等性。其中，\(N\) 代表体系的电子数。
但实际上，从数值稳定性的角度上讲，我们不希望让 \(\mathbf{X}\) 太大以至于超出收敛域；因此 \(\mathbf{R}_0\) 总是需要尽可能接近真实密度 \(\mathbf{D}\)。在这篇文档中，我们称 \(\mathbf{R}_0\) 为密度初猜。
同时，\(\mathbf{X}\) 是一种旋转矩阵；它必须要满足反对称性，即 \(\mathbf{X}^\mathrm{T} = - \mathbf{X}\)。这是相当重要且有用的性质，但我们这里不作更多分析。
而对于能量函数 calc_EX：
由于其中的四脚标张量的存在稍有麻烦，在程序编写便利、以及借用强大的 einsum 功能的角度上，我们采用 Einstein Summation Convention 可以将上式写作
各个矩阵元都已经在上文中被定义了，因此程序的书写也并不困难。
3.2.3. 最小值的极小点分析#
这一段我们分析极小点附近的性质。我们之前提到过，\(\mathbf{R}_0\) 总是需要尽可能接近真实密度 \(\mathbf{D}\)。如果这两者相等是什么情形？
现在我们假定 \(E_\mathrm{tot} [\mathbf{X}]\) 已经取到了处于最小值的极小点。我们说起过这种情况下与 \(\mathbf{D} = \mathbf{R} [\mathbf{X}]\) 几乎等价。但我们并没有验证过这个结论。下面我们就来验证之。
若同时有 \(\mathbf{R}_0 = \mathbf{D}\) 与 \(\mathbf{D} = \mathbf{R} [\mathbf{X}]\)，那么
一个显然满足上式的 \(\mathbf{X}\) 是零矩阵。在这种情况下，我们应当会期待分子体系总能量 energy_autograd 就是使用现成量化软件给出的总能量 energy_tot：
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = D.clone()
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
float(energy_autograd)
-150.58503378083833
并且，我们会期待 \(\mathbf{X}\) 确实取到了极小值，通过自动求导得到的导数是不会发生任何变化的零矩阵：
energy_autograd.backward()
gX = X.grad
torch.norm(gX)
tensor( 0.00000, device='cuda:0', dtype=torch.float64)
从这个意义上，我们可以说，这一套工作流程是基本合理的。
3.2.4. 真实密度初猜的分析#
但上面一段是假设我们已经知道体系密度的理想情况。现实情况是，我们并不知道体系的真正密度，从而需要求自变量 \(\mathbf{X}\) 的结果。
现实中有许多给出密度初猜 \(\mathbf{R}_0\) 的方式。其中最简单的密度初猜是将 Fock 矩阵当作 Hamiltonian Core 矩阵，进行对角化后得到的密度矩阵。在 PySCF 中，它通过 init_guess_by_1e 方法给出。同时，初始的自变量 (或称旋转矩阵) \(\mathbf{X}\) 定为零矩阵。
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
以这个初始密度代入能量的计算公式中，会得到一个比刚才的 -150 Hartree 高出不少的能量：
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
energy_autograd
tensor(-140.54132, device='cuda:0', dtype=torch.float64, grad_fn=<AddBackward0>)
得到这样的结果也是情理之中：这完全满足变分原理的要求，即任意非基态的波函数 (在 RHF 中，也可以称为其对应的基态密度 \(\mathbf{D}\))，那么其给出的能量会高于最低能量。这也满足优化问题的要求，即任何偏离作为最小值的极小点附近的变量 \(\mathbf{X}\)，一定给出高于最小值的结果来。
下面我们来考虑关于 \(\mathbf{X}\) 的梯度。首先，由于 \(\mathbf{X}\) 在最小值附近，我们假设其附近是凸性的，因此其梯度必然不为零：
energy_autograd.backward()
gX = X.grad
torch.norm(gX)
tensor(6.70115, device='cuda:0', dtype=torch.float64)
其次，我们说起过，\(\mathbf{X}\) 具有反对称的性质。如果让能量函数 \(E_\mathrm{tot} [\mathbf{X}]\) 作反向传播，那么 \(\mathbf{X}\) 的梯度也应当是满足这种反对称性质。
torch.allclose(-gX.T, gX)
True
3.3. RHF 的梯度下降求取#
3.3.1. 学习率递降 class Scheduler#
这其实是一个非常细节的问题，但却需要不少代码来解决。
我们都知道在深度学习中，梯度下降算法与学习率递降都是非常有必要的。但由于我们需要在训练的过程中始终保证作为自变量的 \(\mathbf{X}\) 具有反对称性，因此它不能像 PyTorch 深度网络中参数一样，使用较为自由的优化方式；而是需要在优化 (机器学习中类似于训练) 过程中始终重新将 \(\mathbf{X}\) 反对称化。
因此，这里暂时不使用 PyTorch 提供的 optim 类的 Adam 或 ReduceLROnPlateau 类来解决问题，而需要手动造与学习率有关的轮子；至于梯度下降法只能用 Naive 的版本，而不使用 Adam 等更高级的版本。至于是不是一定不能用 PyTorch 提供的自动化工具，作者暂时还不清楚。
class Sheduler:
def __init__(self, init_lr=0.05, min_lr=1e-6, factor=0.95, patience=3, stop_train=30, stop_thresh=1e-6, debug=True):
self.lr = init_lr
self.min_lr = min_lr
self.factor = factor
self.patience = patience
self.stop_train = stop_train
self.stop_thresh = stop_thresh
self.debug = debug
self.loss = 1e10
self.epoch = 0
self.flag_min = False
self.end_train = False
self.loss_list = []
def step(self, loss, epoch_id=-1):
self.epoch += 1
self.loss_list.append(float(loss))
if loss < self.loss:
self.loss = loss
self.epoch = 0
if not self.flag_min:
if self.epoch >= self.patience:
self.lr *= self.factor
self.epoch = 0
print("learning rate decrease to {:7.4f} on epoch {:5d}.".format(self.lr, epoch)) if self.debug else None
if self.lr < self.min_lr:
self.flag_min = True
print("Hit minimum learning rate.") if self.debug else None
if self.flag_min:
if self.epoch >= self.stop_train:
print("End learning at epoch {:5d} since {:5d} epoch loss is larger than lowest loss.".format(epoch, self.stop_train)) if self.debug else None
self.end_train = True
if np.all(np.abs(np.array(self.loss_list[-self.stop_train:]) - float(self.loss)) < self.stop_thresh) and len(self.loss_list) > 30:
print("End learning at epoch {:5d} since last {:5d} epoch loss is allclose to lowest loss".format(epoch, self.stop_train) +
" within threshold {:7.2e}.".format(self.stop_thresh)) if self.debug else None
self.end_train = True
这个类的优化过程稍微有些一点点复杂 (但比起现成的程序已经简单太多)，一些简单的说明是：
优化过程的学习率从
init_lr(默认 0.05) 渐渐降低为min_lr(默认 1e-6)，每次降低factor(默认 0.95) 倍，但降到最低后就不再下降；学习率降低的判据是，如果连续
patience(默认 3) 次的损失函数大于最低的损失函数值，就降低学习率上述过程比通常深度学习中的要求高很多；这是因为通常的深度学习不是凸优化问题，在损失函数势能面上游走不见得是坏事；但 RHF 问题我们假设在合理的密度初猜下是凸优化，损失函数不应该经常上升；
停止优化的第一判据是，若学习率
self.lr降到最低后连续stop_train(默认 30) 次损失函数的值大于最低的那一次，那么停止优化；停止优化的第二判据是，若最后连续
stop_train(默认 30) 次损失函数都与最低的那一次之间差不超过stop_thresh(默认 1e-6) Hartree 大小，那么停止优化。
3.3.2. 优化过程#
有了上面文字的铺垫，后面的优化过程若读者有任何一点深度学习的脚本经验，应该是非常容易理解的。这里用到的初始条件是
密度初猜 \(\mathbf{R}_0\)：采用 Hamiltonian Core 作为 Fock 矩阵对角化所得密度
初始自变量 \(\mathbf{X}\)：采用零矩阵
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
随后学习率从 0.04 开始递降，进行损失函数 (即能量函数 \(E_\mathrm{tot} [\mathbf{X}]\)) 的优化。
sheduler = Sheduler(init_lr=0.04)
for epoch in range(0, 5000):
# Calculate loss function
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
energy_autograd.backward()
# Make optimization step
t = X - X.grad * sheduler.lr
# Force anti-symmetrize X
X = ((t - t.T) / 2).clone().detach().requires_grad_(True)
# Update learning rate or stop training
sheduler.step(energy_autograd, epoch)
if sheduler.end_train:
print("End at epoch {:5d}.".format(epoch))
break
print(float(energy_autograd))
learning rate decrease to 0.0380 on epoch 7.
learning rate decrease to 0.0361 on epoch 10.
learning rate decrease to 0.0343 on epoch 182.
learning rate decrease to 0.0326 on epoch 196.
End learning at epoch 859 since last 30 epoch loss is allclose to lowest loss within threshold 1.00e-06.
End at epoch 859.
-150.5850332855298
这就完成了借助于自动求导的 RHF 的能量计算问题！我们可以拿它与 PySCF 作为量化程序计算得到的能量作比较：
energy_tot
-150.5850337808369
3.4. 小结#
我们在这份文档中简单地讨论了在给定电子积分的情况下，使用 PyTorch 的自动求导功能，以 Naive 梯度下降法给出 RHF 能量。这样的程序既可以在 CPU 下运行，也能在 GPU 下运行。
简单的代码总结可以是，在定义了与量子化学问题无关的矩阵指数 mat_power、矩阵幂 calc_exp、学习率递降管理器 Sheduler 后，我们可以进行 RHF 的能量计算，其代码也非常简单。
分子定义
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()
<pyscf.gto.mole.Mole at 0x7fa2800cb908>
电子积分定义
H = torch.tensor(mol.intor("int1e_kin") + mol.intor("int1e_nuc"), device=device)
S = torch.tensor(mol.intor("int1e_ovlp"), device=device)
eri_ao = torch.tensor(mol.intor("int2e"), device=device)
核库伦排斥能定义
natm = mol.natm
Z_A, A_t = mol.atom_charges(), mol.atom_coords()
r_AB = np.linalg.norm(A_t[:, None, :] - A_t[None, :, :], axis=-1)
r_AB += np.diag(np.ones(natm) * np.inf)
energy_nuc = 0.5 * (Z_A[None, :] * Z_A[:, None] / r_AB).sum()
密度矩阵与能量损失函数定义
def calc_RX(X, S, R0):
return calc_exp(-X @ S) @ R0 @ calc_exp(S @ X)
def calc_EX(X, S, R0, H, eri_ao, energy_nuc):
RX = calc_RX(X, S, R0)
return (
torch.einsum("uv, uv ->", RX, H)
+ 0.5 * torch.einsum("uv, uvkl, kl ->", RX, eri_ao, RX)
- 0.25 * torch.einsum("uv, ukvl, kl ->", RX, eri_ao, RX)
+ energy_nuc
)
密度初猜与自变量定义
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
梯度下降得到 RHF 能量
sheduler = Sheduler(init_lr=0.04, debug=False)
for epoch in range(0, 5000):
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
energy_autograd.backward()
t = X - X.grad * sheduler.lr
X = ((t - t.T) / 2).clone().detach().requires_grad_(True)
sheduler.step(energy_autograd, epoch)
if sheduler.end_train: break
最终能量
float(energy_autograd)
-150.5850332855298