2. 重新学习 RHF 核坐标梯度笔记

2.1. 初始化

[1]:

from pyscf import gto, scf, grad, hessian, lib
import numpy as np
from functools import partial
import h5py
from pyscf.scf import _vhf

from pyxdh.DerivOnce import GradSCF

np.set_printoptions(4, suppress=True, linewidth=180)
np.einsum = partial(np.einsum, optimize=["greedy", 1024 ** 3 * 2 / 8])

[2]:

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

[2]:

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

[3]:

mf_scf = scf.RHF(mol).run()
mf_grad = grad.RHF(mf_scf).run()
mf_hess = hessian.RHF(mf_scf).run()
mf_grad.de

[3]:

array([[-0.1408, -0.1166, -0.0278],
       [ 0.0947,  0.0102,  0.0289],
       [ 0.0195,  0.0815,  0.0225],
       [ 0.0266,  0.025 , -0.0236]])

[4]:

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

[4]:

array([[-0.1408, -0.1166, -0.0278],
       [ 0.0947,  0.0102,  0.0289],
       [ 0.0195,  0.0815,  0.0225],
       [ 0.0266,  0.025 , -0.0236]])

2.2. PySCF 函数

2.2.1. `hcore_generator`

用于生成 \(h_{\mu \nu}^{A_t}\)，与 gradh.H_1_ao 对应。hcore_generator 是函数生成器。生成的函数的输入参量是原子序号 \(A\)，输出是 \(h_{\mu \nu}^{A_t}\)。

[5]:

hcore_deriv = mf_grad.hcore_generator(mol)

[6]:

np.allclose(hcore_deriv(0), gradh.H_1_ao[:3])

[6]:

True

2.2.2. `get_ovlp`

用于生成与 \(S_{\mu \nu}^{A_t}\) 有关的量 \(- (\partial_t \mu | \nu)\)。注意： 并非直接生成 \(S_{\mu \nu}^{A_t}\)。同时注意负号。

[7]:

np.allclose(mf_grad.get_ovlp(mol), - mol.intor("int1e_ipovlp"))

[7]:

True

2.2.3. `make_rdm1e`

用于生成轨道能加权密度 \(R_{\mu \nu} [\varepsilon_i]\)。

\[R_{\mu \nu} [\varepsilon_i] := 2 C_{\mu i} \varepsilon_i C_{\nu i}\]

[8]:

Co, eo = gradh.Co, gradh.eo

[9]:

np.allclose(
    2 * np.einsum("ui, i, vi -> uv", Co, eo, Co),
    mf_grad.make_rdm1e(mf_scf.mo_energy, mf_scf.mo_coeff, mf_scf.mo_occ),
)

[9]:

True

einsum 效率警告： 此处不适合用 einsum。

[10]:

np.allclose(
    2 * np.einsum("ui, i, vi -> uv", Co, eo, Co),
    2 * (Co * eo) @ Co.T,
)

[10]:

True

[11]:

%%timeit -r 7 -n 1000
2 * np.einsum("ui, i, vi -> uv", Co, eo, Co)

121 µs ± 11.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

[12]:

%%timeit -r 7 -n 1000
2 * (Co * eo) @ Co.T

5.36 µs ± 847 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

2.2.4. `get_j`, `get_k`

作张量缩并

\[\begin{split}\begin{align} J_{\mu \nu}^t [R_{\lambda \kappa}] &= - (\partial_t \mu \nu | \kappa \lambda) R_{\lambda \kappa} \\ K_{\mu \lambda}^t [R_{\nu \lambda}] &= - (\partial_t \mu \nu | \kappa \lambda) R_{\nu \kappa} \end{align}\end{split}\]

这个程序与 get_jk 有关。其底层调用是 scf._vhf.direct_mapdm。我们在后文讨论该函数。

[13]:

np.allclose(
    - np.einsum("tuvkl, kl -> tuv", mol.intor("int2e_ip1"), gradh.D),
    mf_grad.get_j(),
)

[13]:

True

[14]:

np.allclose(
    - np.einsum("tuvkl, vk -> tul", mol.intor("int2e_ip1"), gradh.D),
    mf_grad.get_k(),
)

[14]:

True

记号误认警告： 注意等式左右的下角标，这里刻意用了与 PySCF 本体程序类似的记号。但在理解程序上比较关键。如果采用 pyxdh 的习惯，公式应写为

\[\begin{split}\begin{align} J_{\mu \nu}^t [R_{\kappa \lambda}] &= - (\partial_t \mu \nu | \kappa \lambda) R_{\kappa \lambda} \\ K_{\mu \nu}^t [R_{\kappa \lambda}] &= - (\partial_t \mu \kappa | \nu \lambda) R_{\kappa \lambda} \end{align}\end{split}\]

函数效率警告： 存在更高效的函数 get_jk。这里不介绍其调用方式。

[15]:

%%timeit -n 10
mf_grad.get_j()

33.5 ms ± 1.35 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

[16]:

%%timeit -n 10
mf_grad.get_jk()

35.2 ms ± 1.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

2.2.5. `_vhf.direct_mapdm`

该函数专门地用于计算双电子积分与密度矩阵的缩并。

[17]:

_vhf.direct_mapdm

[17]:

<function pyscf.scf._vhf.direct_mapdm(intor, aosym, jkdescript, dms, ncomp, atm, bas, env, vhfopt=None, cintopt=None, shls_slice=None)>

intor 双电子积分的类型
- 必须要指定积分形式，譬如 int2e_ip1 是不允许的。如果所求的分子积分是球谐形式，那么就需要在后面增加词尾 _sph，从而是 int2e_ip1_sph。
- 词尾的增加可以是 mol._add_suffix("int2e_ip1")。一般球谐是 _sph，笛卡尔是 _cart，其他两种是 _spinor, _ssc。
aosym 双电子积分的对称性
- 尽管可以从 intor 关键词推断，但必须要手动指定。使用时需要非常谨慎。如果不清楚用哪个，可以使用 s1 而保证结果不会错，但这会大大降低积分效率。
- 可能的选项是 s8, s4, s2ij, s2kl, s1, aa4, a4ij, a4kl, a2ij, a2kl。这在 scf._vhf 文件中有所说明。
- 所有的双电子积分假设维度是 \((t, i, j, k, l)\)。s2ij 表示互换 \(i, j\) 角标结果不变，即 \(g^t_{ij, kl} = g^t_{kl, ij}\)。a4ij 表示 \(i, j\) 反对称而 \(k, l\) 对称。a2ij 表示 \(i, j\) 反对称。
- 举例而言，\((\partial_t \mu \nu | \kappa \lambda)\) 具有二重对称性：\((\partial_t \mu \nu | \kappa \lambda) = (\partial_t \mu \nu | \lambda \kappa)\)。由于是后两个角标对称，因此对称性记为 s2kl。
- 它不代表密度矩阵的对称性。同时，密度矩阵是否对称不会很影响计算效率。
jkdescript 张量缩并方式
- 对于类库伦积分，lk->s1ij 表示 \((\partial_t \mu \nu | \kappa \lambda) R_{\kappa \lambda}\)。由于最终的结果对于 \(\mu, \nu\) 不对称，因此这里是 s1。
- 密度角标只支持 ji, lk, li, jk；结果角标只支持 _s1, _s2 与 kl, ij, kj, il 的组合。
dms 密度矩阵，形状必须是 2 维度方阵或 3 维度张量
ncomp 双电子积分的大小指标
- 尽管可以从 intor 关键词推断，但必须要手动指定。
- 对于 \((\partial_t \mu \nu | \kappa \lambda)\)，由于 \(t\) 是 \(x, y, z\) 三个方向，因此该值为 3。
atm, bas, env 均是分子自身的参量，一般不需要更改

下面观察一下输出维度。如果张量缩并方式有 2 种，密度矩阵有 5 个，双电子积分大小有 3 个，那么输出的结果也会是 (2, 5, 3, nao, nao) 大小的张量。

[18]:

R_rand = np.random.randn(5, mol.nao, mol.nao)

[19]:

res = _vhf.direct_mapdm("int2e_ip1_sph", "s2kl", ('lk->s1ij', 'jk->s1il'), R_rand, 3, mol._atm, mol._bas, mol._env)
np.array(res).shape

[19]:

(2, 5, 3, 15, 15)

[20]:

np.allclose(np.einsum("tijkl, Blk -> Btij", mol.intor("int2e_ip1"), R_rand), res[0]), \
np.allclose(np.einsum("tijkl, Bjk -> Btil", mol.intor("int2e_ip1"), R_rand), res[1])

[20]:

(True, True)

指定对称性对效率的改变巨大： 若电子积分是二重对称的，那么如果程序中降低对称性，效率会恰好低一倍。往往双电子积分与密度矩阵缩并是次要性能关键步；它对性能的影响不小。因此要谨慎处理。

[21]:

%%timeit -n 5
_vhf.direct_mapdm("int2e_ip1_sph", "s2kl", 'lk->s1ij', R_rand, 3, mol._atm, mol._bas, mol._env)

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

[22]:

%%timeit -n 5
_vhf.direct_mapdm("int2e_ip1_sph", "s1", 'lk->s1ij', R_rand, 3, mol._atm, mol._bas, mol._env)

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

效率随体系不同而不同： 在小体系下，直接在内存中储存所有双电子积分，并且不考虑对称性地进行积分反而效率更高。但如果体系扩大，使用 _vhf.direct_mapdm 的必要性就出来了。同时，由于我们无法承受 \(O(N^4)\) 大小的内存量，因此有必要时就使用 _vhf.direct_mapdm 或者其他 _vhf 函数。

[23]:

%%timeit -n 5
np.einsum("tuvkl, Bkl -> Btuv", mol.intor("int2e_ip1"), R_rand)

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

[24]:

mol_large = gto.Mole()
mol_large.atom = """
N  0.  0.  0.
H  1.5 0.  0.2
H  0.1 1.2 0.
H  0.  0.  1.
"""
mol_large.basis = "cc-pVTZ"
mol_large.verbose = 0
mol_large.build()
R_rand_large = np.random.randn(5, mol_large.nao, mol_large.nao)

[25]:

%%timeit -n 2
np.einsum("tuvkl, Bkl -> Btuv", mol_large.intor("int2e_ip1"), R_rand_large)

659 ms ± 9.3 ms per loop (mean ± std. dev. of 7 runs, 2 loops each)

[26]:

%%timeit -n 2
_vhf.direct_mapdm("int2e_ip1_sph", "s2kl", ('lk->s1ij', 'jk->s1il'), R_rand_large, 3, mol_large._atm, mol_large._bas, mol_large._env)

487 ms ± 9.11 ms per loop (mean ± std. dev. of 7 runs, 2 loops each)

2.3. 能量的一阶梯度

在写可以实用化的一阶梯度程序时，需要考虑到的最重要因素之一，是内存的大小。

内存大小不能超过平方级别，甚至不允许是 \((n_\mathrm{atom}, n_\mathrm{AO}, n_\mathrm{AO})\) 大小。因此，处理时需要尽可能将原子分离开，更不能出现双电子积分。

2.3.1. pyxdh 实现方式

pyxdh 的实现方式是：

\[\frac{\partial E_\mathrm{tot}}{\partial A_t} = 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} + \frac{\partial E_\mathrm{nuc}}{\partial A_t}\]

[27]:

so = gradh.so
D = gradh.D

[28]:

(
    +        np.einsum("Auv, uv -> A", gradh.H_1_ao, D)
    + 0.5  * np.einsum("Auvkl, uv, kl -> A", gradh.eri1_ao, D, D)
    - 0.25 * np.einsum("Aukvl, uv, kl -> A", gradh.eri1_ao, D, D)
    - 2    * np.einsum("ij, Aij -> A", gradh.F_0_mo[so, so], gradh.S_1_mo[:, so, so])
    + mf_grad.grad_nuc().flatten()
).reshape(mol.natm, 3)

/home/a/miniconda3/lib/python3.8/site-packages/pyxdh/DerivOnce/deriv_once_scf.py:309: UserWarning: eri1_ao: 4-idx tensor ERI should be not used!
  warnings.warn("eri1_ao: 4-idx tensor ERI should be not used!")

[28]:

array([[-0.1408, -0.1166, -0.0278],
       [ 0.0947,  0.0102,  0.0289],
       [ 0.0195,  0.0815,  0.0225],
       [ 0.0266,  0.025 , -0.0236]])

2.3.2. Hamilton Core 贡献项

这个拆分是非常容易的：

[29]:

hcore_deriv = mf_grad.hcore_generator(mol)

[30]:

contrib_hcore = np.zeros((mol.natm, 3))
for A in range(mol.natm):
    contrib_hcore[A] += np.einsum("tuv, uv -> t", hcore_deriv(A), D)

[31]:

np.allclose(contrib_hcore.flatten(), np.einsum("Auv, uv -> A", gradh.H_1_ao, D))

[31]:

True

2.3.3. J 积分

依据对称性，我们注意到

\[\partial_{A_t} E_\mathrm{tot} \leftarrow \frac{1}{2} \partial_{A_t} (\mu \nu | \kappa \lambda) D_{\mu \nu} D_{\kappa \lambda} = 2 (\partial_{A_t} \mu \nu | \kappa \lambda) D_{\mu \nu} D_{\kappa \lambda} = - 2 (\partial_t \mu_A \nu | \kappa \lambda) D_{\mu_A \nu} D_{\kappa \lambda} = 2 J_{\mu_A \nu} [D_{\kappa \lambda}] D_{\mu_A \nu}\]

[32]:

mol_slice = gradh.mol_slice

[33]:

vj = - _vhf.direct_mapdm("int2e_ip1_sph", "s2kl", 'lk->s1ij', D, 3, mol._atm, mol._bas, mol._env)

[34]:

contrib_vj = np.zeros((mol.natm, 3))
for A in range(mol.natm):
    sA = mol_slice(A)
    contrib_vj[A] += 2 * np.einsum("tuv, uv -> t", vj[:, sA, :], D[sA, :])

[35]:

np.allclose(contrib_vj.flatten(), 0.5 * np.einsum("Auvkl, uv, kl -> A", gradh.eri1_ao, D, D))

[35]:

True

2.3.4. K 积分

非常类似地，我们可以得到

\[\partial_{A_t} E_\mathrm{tot} \leftarrow - \frac{1}{4} \partial_{A_t} (\mu \nu | \kappa \lambda) D_{\mu \kappa} D_{\nu \lambda} = - K_{\mu_A \kappa} [D_{\nu \lambda}] D_{\mu_A \kappa}\]

[36]:

vk = - _vhf.direct_mapdm("int2e_ip1_sph", "s2kl", 'jk->s1il', D, 3, mol._atm, mol._bas, mol._env)

[37]:

contrib_vk = np.zeros((mol.natm, 3))
for A in range(mol.natm):
    sA = mol_slice(A)
    contrib_vk[A] += - np.einsum("tuv, uv -> t", vk[:, sA, :], D[sA, :])

[38]:

np.allclose(contrib_vk.flatten(), - 0.25 * np.einsum("Aukvl, uv, kl -> A", gradh.eri1_ao, D, D))

[38]:

True

2.3.5. 能量加权部分

\[\partial_{A_t} E_\mathrm{tot} \leftarrow - 2 F_{ij} S_{ij}^{A_t} = - (2 F_{ij} C_{\mu i} C_{\nu j}) \partial_{A_t} (\mu | \nu) = 2 R_{\mu_A \nu} [F_{ij}] (\partial_t \mu_A | \nu)\]

[39]:

dme0 = mf_grad.make_rdm1e(mf_scf.mo_energy, mf_scf.mo_coeff, mf_scf.mo_occ)
s1 = mf_grad.get_ovlp(mol)

[40]:

contrib_dme0 = np.zeros((mol.natm, 3))
for A in range(mol.natm):
    sA = mol_slice(A)
    contrib_dme0[A] += - 2 * np.einsum("tuv, uv -> t", s1[:, sA, :], dme0[sA, :])

[41]:

np.allclose(contrib_dme0.flatten(), - 2 * np.einsum("ij, Aij -> A", gradh.F_0_mo[so, so], gradh.S_1_mo[:, so, so]))

[41]:

True

2.3.6. 最终加和

[42]:

contrib_hcore + contrib_vj + contrib_vk + contrib_dme0 + mf_grad.grad_nuc()

[42]:

array([[-0.1408, -0.1166, -0.0278],
       [ 0.0947,  0.0102,  0.0289],
       [ 0.0195,  0.0815,  0.0225],
       [ 0.0266,  0.025 , -0.0236]])