REST Workshop：Rust 计算化学程序开发

该仓库以 RHF 与 RCCSD(T) 为例子，用于演示 Rust 语言下，计算化学程序的快速开发与高性能实现的一种策略与工程实践。

该仓库将使用 RSTSR (REST 子项目) 作为数学库工具，libcint (Rust binding and wrapper) 作为电子积分引擎。我们将展示 RSTSR 在单一语言框架 (Rust) 下，可以同时兼顾开发效率与运行效率。

编译与测试

cargo build
cargo test --lib -- playground_ri_ccsdt --exact --nocapture

该项目由于开发过程也需要关注效率，因此在 debug 下仍然开启了 opt-level=2 优化选项。上述编译后的程序效率与 release 没有明显差异。

上述程序测试的是水分子 def2-TZVP 基组下的 RI-RCCSD(T) 能量计算。

该编译流程会从 github 上下载电子积分库 libcint 并作静态链接 (static linking)。如果 github 下载较慢，可以先下载好 libcint 的源码，并声明到环境变量 CINT_SRC。

前言

计算化学程序开发，与其他科学计算领域类似，有快速验证与高效实现的需求。

基于 REST 平台，我们在 Rust 语言下开发了 RSTSR 数学库工具，提供便利的程序接口，应对计算化学常见高维度张量的高效运算。

作为例子，我们快速高效地实现了 RHF 到 RCCSD(T) 一系列方法，展示 Rust 语言下计算化学程序 (或类似的科学计算任务) 开发的可能性。

目录

• 1. Rust 语言与计算化学
  • 1.1 程序开发的两个需求：快速验证与高效实现
  • 1.2 Rust 语言：作为计算化学程序开发者的特色与不足
• 2. 数学库 RSTSR
  • 2.1 RSTSR 简介
  • 2.2 所有权问题
• 3. 约定俗成
  • 3.1 记号
  • 3.2 张量与后端类型
• 4. 必要的性能分析背景
  • 4.1 简单矩阵代数的 FLOPs 分析
  • 4.2 计算设备的理想性能
• 5. 前言结语

1. Rust 语言与计算化学

这里的核心问题是，为何 Rust 是适用于计算化学程序开发的语言？

1.1 程序开发的两个需求：快速验证与高效实现

计算化学程序用户比较关心程序功能是否全面、运行能力是否强大、程序是否易于使用。其中，运行能力是指是否可以高效地模拟较大的体系；这要求了程序性能必须足够高。

计算化学程序开发者还同时面临一些额外的问题：计算化学方法，特别是涉及到后自洽场、响应性质、激发态等问题的公式经常比较复杂。如何快速验证程序的正确性、并进一步作性能改善，是程序开发者面临的困难。

如何以客观的标准，评估程序语言或框架，是否能满足计算化学所需要的快速验证与高效实现？我们认为，

• 快速验证：因程序开发者经验与习惯差异而有不同看法；我们以代码行数作为客观标准，认为较短的代码比较契合快速开发的需求。理想情况下，一行公式能对应一行代码。
• 高效实现：在算法清晰、浮点数计算量 (FLOPs) 可以大致算出的前提下，我们可以对比计算设备的理论浮点计算能力 (FLOP/sec) 与程序的实际浮点计算效率，确认实现是否确实高效。

随着硬件与计算机语言的发展，早年时代表现好的程序，或许会在未来失去竞争力。但是，上面两个标准本身，不随计算机语言、硬件水平而不同；任何时代下的程序与框架应都可以此标准衡量。

1.2 Rust 语言：作为计算化学程序开发者的特色与不足

作为计算化学程序开发者，Rust 语言吸引我们的特色包括

• 与 C 同等级的性能；只要算法实现得当，几乎不存在因 Rust 语言本身的不足引起的性能问题；
• 现代函数式语言的表达能力、基于 trait 的泛型能力、强类型等适合于大型项目的语言特性；
• 内存安全性与无畏并行，能保证程序编写者调用失误的可能性大大下降；
• (除科学计算之外的) 程序生态丰富，在硬件层与底层应用上有诸多出色的项目。

Rust 语言的不足包括

• 作为编译语言，尽管支持增量编译，但其编译或链接时间仍然较长 (与 C++ 项目同等级，涉及到模板特化的耗时较多)；
• REPL 支持有限，难以像 Python, Matlab, Mathematica, Julia 脚本式地运行程序¹；
• 科学计算与 AI 生态与 GPU 支持还有待发展²。

Rust 语言具有争议的地方包括

• Rust 完全否定面向对象；必须通过类似于 C#、Java 的接口 (trait) 实现多态 (polymorphism)；
• 语法上对重载 (overload) 的支持有限³。

Rust 语言与其他语言的对比：

• C++ 的表达能力非常自由，但程序编写者更容易写出错误的代码；C++ 语言本身复杂，编写程序时需要考虑引用、右值引用、移动等语义。
• Python 语言尽管可以写出效率很高的程序，但需要与其他语言 (C, C++, Rust) 或特殊的 Python 方言 (如 Numba) 结合才能达到高效率。Python 语言的类型检查较弱，非常便于快速验证、但不便于大型程序开发。
• Julia 语言尽管是少有的同时实现 REPL 与编译功能的语言，但其在科学计算之外的程序生态薄弱、有影响力的项目较少。
• Mojo, Moonbit 等语言吸收了 Python 与 Rust 等语言的特性，但目前仍在发展，现阶段不适合作为主力语言。

在语言上容易产生误解的地方包括

• 作为编译语言，Rust 必须要用很多代码实现脚本语言同样的功能：这个项目会表明，在合适的数学库框架下，Rust 语言可以用不太长的代码实现计算化学关心的问题。Python 下能用 $n$ 行代码实现计算化学的算法，Rust 经常有办法在 $1.5n$ 行以内实现；Python 下不能高效实现的算法，Rust 可以实现。
• Rust 语言学习曲线陡峭：C++ 语言的学习难度未必低于 Rust，但这个问题因人而异。不考虑工程与架构，计算化学的算法实现所遇到的问题，大多不涉及 Rust 语法中困难的部分。尽管我们需要了解和使用 Rust 的生命周期，但我们的工作中不涉及循环链表和复杂的异步，因此并不需要对生命周期的所有细节都有所掌握。

2. 数学库 RSTSR

2.1 RSTSR 简介

该项目使用到数学库 RSTSR。

RSTSR 是 Rust 语言下，可用于处理张量存储与索引、以及对矩阵作线性代数运算的库。

• 功能与 API 对标 NumPy 与 SciPy，可以原生地在 Rust 语言下开发科学计算程序。对于习惯 NumPy 的用户，可以参考 NumPy-RSTSR 对照表 以快速熟悉 RSTSR 的使用。
• 支持多后端。目前 RSTSR 支持 OpenBLAS 与 Faer (纯 Rust 的数值代数库) 后端。未来会引入 MKL 与 CUDA 等其他后端。
• 涉及到矩阵线性代数的部分，性能由后端决定；其余的部分尽可能充分利用张量连续性与并行，性能较高。
• 同时支持行优先与列优先。

2.2 所有权问题

RSTSR 在使用上与 NumPy 有很多相似之处，但有一个非常大的区别在于数据所有权。

Rust 有非常严格的所有权规则；除非您的数据类型是 RC (reference counting) 智能指针，否则必须要区分数据是占有的还是引用的、对于引用类型还需要指定其生命周期。

RSTSR 的所有权规则参考自 Rust 库 ndarray；其简要说明可以参考 用户文档。调用的函数的过程中，

• 如果是不改变张量数据的操作 (比如索引、转置等)，函数签名中没有 into_ 前缀的将返回引用类型张量视窗 (如函数 i, transpose, swapaxes)；若有 into_ 前缀的则返回 self 自身的类型 (如函数 into_reverse_axes)。
• 张量维度更改 (reshape) 是可能产生新数据的操作。reshape 将返回占有或视窗的张量 (copy-on-write 类型)；into_shape 将返回占有数据的张量。调用 reshape 或 into_shape 时，如果原先的张量数据足够连续，可以避免数据的复制，从而以近乎为零的开销实现维度更改。

3. 约定俗成

3.1 记号

• 本项目统一使用行优先 (row-major)。
• 本项目仅考虑闭壳层基态，不引入赝势 (ECP)。占据轨道数严格等于实际体系电子数的一半。
• 本项目仅考虑实函数。

指标规则为

3.2 张量与后端类型

本工作中，为了方便起见，后端类型声明为 DeviceTsr。若用户可以通过 cargo feature 指定 Faer 后端或 OpenBLAS 后端。

#![allow(unused)]
fn main() {
// prelude.rs
#[cfg(not(feature = "use_openblas"))]
pub type DeviceTsr = DeviceFaer;
#[cfg(feature = "use_openblas")]
pub type DeviceTsr = DeviceOpenBLAS;
}

我们总是在 64 位浮点数、固定后端下进行开发；为了简化代码，张量类型则分别定义为

• 占有张量 Tsr
• 不可变张量视窗 TsrView
• 可变张量视窗 TsrMut

我们会在非常少数的情景下，使用固定维度张量 (相对于动态维度而言)，因此保留维度泛型参数。

#![allow(unused)]
fn main() {
// prelude.rs
pub type Tsr<D = IxD> = Tensor<f64, DeviceTsr, D>;
pub type TsrView<'a, D = IxD> = TensorView<'a, f64, DeviceTsr, D>;
pub type TsrMut<'a, D = IxD> = TensorMut<'a, f64, DeviceTsr, D>;
}

4. 必要的性能分析背景

4.1 简单矩阵代数的 FLOPs 分析

在本项目中，我们会对 FLOPs 作简单分析。

• 对于矩阵乘法 ($\mathbf{A}\in {\mathbb{F}}^{m\times k}$，$\mathbf{B}\in {\mathbb{F}}^{k\times n}$，$\mathbf{C}\in {\mathbb{F}}^{m\times n}$，其中 $\mathbb{F}$ 代表数据的类型)

  $$C_{ij}=\sum _a{A}_{ia}{B}_{aj}\text{or}\mathbf{C}=\mathbf{AB}$$

  则该计算的 FLOPs 总量是 $2mnk$；其中的 $2$ 系数来源与一次加法与一次乘法。

• 对于矩阵乘法

  $$C_{ij}=\sum _a{A}_{ia}{B}_{ja}\text{or}\mathbf{C}=\mathbf{A}{\mathbf{A}}^{†}$$

  则该计算的 FLOPs 近似为 $n^{2}k$ (由于 $C_{ij}=C_{ji}$，因此只需要作一半矩阵乘法，剩下一半通过转置得到)。

• 对于上三角或下三角矩阵 $\mathbf{A}\in {\mathbb{F}}^{n\times n}$ 以及长方形矩阵 $\mathbf{B}\in {\mathbb{F}}^{n\times m}$，作下述三角矩阵乘法或求解

  $$\begin{array}{rl}\mathbf{C}& =\mathbf{AB}\text{(triangular matrix multiply)}\\ \mathbf{C}& ={\mathbf{A}}^{-1}\mathbf{B}\text{(triangular matrix solve)}\end{array}$$

  则该计算的 FLOPs 近似为 $n^{2}m$。

在以矩阵乘法为性能瓶颈的问题中，其他 FLOPs 通常是可以作为小量忽略的。但也要留意，除了矩阵乘法外，

• 在理想的程序实现、以及特定的算法下，内存带宽也会占用计算资源；
• 不理想的程序实现，会导致各种性能损耗，包括但不限于多余的内存资源分配与复制、低效率的指针索引、并行问题等等。

4.2 计算设备的理想性能

本文档着重于 CPU。CPU 厂商通常不会直接将 64 位浮点数矩阵乘法峰值能力标在产品说明上，我们通常需要自己作一些简单的换算。GPU 也可以作类似的分析，不过其峰值性能通常标注于产品说明，性能的比较会更加直观。

对于矩阵乘法问题，CPU 计算设备的理想性能是

$$\text{Ideal GFLOP/sec}=n_{\text{core}}{n}_{\text{FMA}}{n}_{\text{SIMD}}{n}_{\text{unit}}{f}_{\text{CPU}}$$

• $n_{\text{core}}$ 物理内核数 (关闭超线程下的最大有效线程数)；
• $n_{\text{FMA}}$ 乘加融合系数，固定为 2；
• $n_{\text{SIMD}}$ 单个指令允许最大的数据大小 (对于 64 位浮点，x86 AVX-512 为 8、AVX2 为 4，ARM Neon 为 2)；
• $n_{\text{unit}}$ 为 FMA 指令通道数 (AMD Zen AVX-512 为 1、AVX2 为 2，Intel Xeon 与 ARM Neon 通常为 2)；
• $f_{\text{CPU}}$ 为满负荷运行时，以 GHz 为单位的 CPU 频率。

本项目主要在个人电脑 AMD 7945HX 上测试。该设备的理想性能是 1.1 TFLOP/sec 或 1100 GFLOP/sec。

5. 前言结语

这份文档希望讨论程序的具体细节，以追求比较理想的程序性能。Rust 语言不仅能帮助我们较容易地达到理想性能，而且其具有较高的开发效率的潜力，适合计算化学的程序开发。

但也要留意，程序性能的提升经常只是改进；它当然也很重要，但真正的突破还需要依靠计算化学的算法、方法、或认知论的发展。

1. Rust 实际上有 REPL 库 evcxr，可以使用但体验与 Python 还是有差距的。作为参考，C++ xeus-cling 是类似的情况。 ↩

2. 一方面除了 C/C++ 外，其他所有计算机语言对 GPU 的支持事实上都不足 (Fortran 有支持但较有限)，大多都需要通过 FFI (foreign function interface) 实现。另一方面，只从 CUDA 生态而言，涉及 GPU 的代码应看作 CUDA 语言而非传统的 C/C++ (不是相同的语言)；只是大多数工具链 (toolchain) 允许 CUDA 与 C/C++ 混合编译，使用上与语法上非常接近。从工程量上来看，使用 C/C++ 不一定就比使用其他语言更少。 ↩

3. 一种常见的误解是 Rust 不能重载。实际上 Rust 具有重载能力 (基于 trait)，RSTSR 也大量使用了 Rust 重载；它或许可以看作支持泛型的 C11 _Generic，与 Python 的 named-arguments 不同。但重载在 Rust 语言中的支持确实不足，导致现在很多 RSTSR 函数的参数必须打两个括号才能使用。 ↩

RHF：极简演示

• 程序：rhf_slow.rs (gitee, github)
• 实现内容：RHF (restricted Hartree-Fock, conventional 4c-2e)
• 不关注程序效率
• 该程序总共约 50 行

目录

• 1. RHF 回顾
• 2. 程序
  • 2.1 输入与输出
  • 2.2 原子核互斥能计算
  • 2.3 自洽场计算
  • 2.4 RHF 能量计算

1. RHF 回顾

RHF 方法是分子体系计算化学中最基础的方法。在 Bohn-Oppenheimer 近似下，假定原子核不运动，则分子体系的能量 $E^{\mathsf{tot}}$ 可以视为原子核互斥能 $E^{\mathsf{nuc}}$ 与电子能量 $E^{\mathsf{elec}}$ 之和：

$$E^{\mathsf{tot}}=E^{\mathsf{nuc}}+E^{\mathsf{elec}}$$

其中，原子核互斥能量决定于原子核的相对位置 ${\mathbf{R}}_A$ 与其电荷数 $Z_A$：

$$E^{\mathsf{nuc}}=\sum _{A<B}^{n_{\mathrm{atom}}}\frac{Z_A{Z}_B}{|{\mathbf{R}}_A-{\mathbf{R}}_B|}$$

Hartree-Fock 近似要求电子的波函数是 Slater 行列式。在闭壳层下 ($\alpha$ 自旋与 $\beta$ 自旋的分子轨道与能级，在空间上对应相等)，其最终表达式可以写为¹

$$\begin{array}{rl}{E}^{\mathsf{elec}}[\mathbf{D}]& =\sum _{\mu \nu }{D}_{\mu \nu }\left(h_{\mu \nu }+\frac{1}{2}{J}_{\mu \nu }-\frac{1}{4}{K}_{\mu \nu }\right)\\ & =\sum _{\mu \nu }\left(h_{\mu \nu }+\sum _{\kappa \lambda }{D}_{\kappa \lambda }\left(\frac{1}{2}{g}_{\mu \nu ,\kappa \lambda }-\frac{1}{4}{g}_{\mu \kappa ,\nu \lambda }\right)\right)\end{array}$$

其中，

• $h_{\mu \nu }$ 是 core Hamiltonian，它包含了动能 $t_{\mu \nu }$ 与电子-原子核势 $v_{\mu \nu }$。
• $g_{\mu \nu ,\kappa \lambda }$ 是。
• $D_{\mu \nu }$ 是电子云密度矩阵，它是由占据轨道系数 $C_{\mu p}$ 构成；其表达式中的 $2$ 是指闭壳层下 $\alpha ,\beta$ 轨道有两份贡献。

以上的量中，$h_{\mu ,\nu },t_{\mu \nu },v_{\mu \nu },g_{\mu \nu ,\kappa \lambda }$ 是随体系确定的，$D_{\mu \nu }$ 是待求解的。记号 $E^{\mathsf{elec}}[\mathbf{D}]$ 代表我们将电子能量 $E^{\mathrm{elec}}$ 视作关于密度矩阵 $\mathbf{D}$ 的函数。

$$\begin{array}{rl}{h}_{\mu \nu }& =t_{\mu \nu }+v_{\mu \nu }\\ t_{\mu \nu }& =-\frac{1}{2}\int \mathrm{d}\mathbf{r}{\varphi }_{\mu }(\mathbf{r}){\nabla }^2{\varphi }_{\nu }(\mathbf{r})\\ v_{\mu \nu }& =\sum _A^{n_{\mathrm{atom}}}\int \mathrm{d}\mathbf{r}{\varphi }_{\mu }(\mathbf{r})\frac{-Z_A}{|\mathbf{r}-{\mathbf{R}}_A|}{\varphi }_{\nu }(\mathbf{r})\\ g_{\mu \nu ,\kappa \lambda }& =\int \mathrm{d}{\mathbf{r}}_1\int \mathrm{d}{\mathbf{r}}_2{\varphi }_{\mu }({\mathbf{r}}_1){\varphi }_{\nu }({\mathbf{r}}_1)\frac{1}{|{\mathbf{r}}_1-{\mathbf{r}}_2|}{\varphi }_{\kappa }({\mathbf{r}}_2){\varphi }_{\lambda }({\mathbf{r}}_2)\\ D_{\mu \nu }& =2\sum _i^{n_{\mathrm{occ}}}{C}_{\mu i}{C}_{\nu i}\text{or}\mathbf{D}=2\mathbf{C}{\mathbf{C}}^{†}\end{array}$$

RHF 要求密度矩阵 $D_{\mu \nu }$ 对应的占据轨道系数 $C_{\mu i}$ 满足正交归一条件

$$\sum _{\mu \nu }{\mathrm{C}}_{\mu i}{S}_{\mu \nu }{C}_{\nu j}={\delta }_{ij}\text{or}{\mathbf{C}}^{†}\mathbf{SC}=\mathbf{I}$$

以此正交归一条件为前提，电子能量 $E^{\mathsf{elec}}[\mathbf{D}]$ 取条件极小，即为体系的基态电子能量。上式中，重叠矩阵 $S_{\mu \nu }$ 是

$$S_{\mu \nu }=\int \mathrm{d}\mathbf{r}{\varphi }_{\mu }(\mathbf{r}){\varphi }_{\nu }(\mathbf{r})$$

定义能量对密度矩阵的导数矩阵 (该矩阵也称为 Fock 矩阵)

$$F_{\mu \nu }[\mathbf{D}]=\frac{\partial E^{\mathsf{elec}}}{\partial D_{\mu \nu }}=h_{\mu \nu }+J_{\mu \nu }-\frac{1}{2}{K}_{\mu \nu }$$

引入 Lagrangian 乘子 $2\mathbf{\epsilon }$ (由 Koopmans 的诠释，$\mathbf{\epsilon }$ 一定条件下具有轨道能的物理意义)，对能量作关于轨道系数 $\mathbf{C}$ 的条件极小，得到 Hartree-Fock-Roothaan 方程：

$$\begin{array}{rl}0& =\frac{1}{2}\frac{\partial \left(E^{\text{elec}}-2{\mathbf{C}}^{†}\mathbf{SC}\mathbf{\epsilon }\right)}{\partial {\mathbf{C}}^{†}}=\frac{1}{2}\frac{\partial E^{\mathsf{elec}}}{\partial \mathbf{D}}\frac{\partial \mathbf{D}}{\partial {\mathbf{C}}^{†}}-\mathbf{SC}\mathbf{\epsilon }\\ & =\mathbf{FC}-\mathbf{SC}\mathbf{\epsilon }\end{array}$$

该问题恰好与第一类广义本征值问题等价²。当作为广义本征值问题时，求解得到的系数矩阵是 $n_{\mathrm{basis}}\times n_{\mathrm{basis}}$ 即基组平方大小的；在计算密度矩阵 $D_{\mu \nu }$ 时，需要取本征值最低的 $n_{\mathrm{occ}}$ 个系数，重新构成 $C_{\mu i}$ 的 $n_{\mathrm{basis}}\times n_{\mathrm{occ}}$ 大小的矩阵。

2. 程序

2.1 输入与输出

整个程序中，涉及到的函数包括

#![allow(unused)]
fn main() {
pub fn get_energy_nuc(cint_data: &CInt) -> f64;
pub fn minimal_rhf(cint_data: &CInt) -> RHFResults
}

其中，输入参数 cint_data: &CInt 是定义了分子坐标与基组信息的类型 libcint::CInt。其初始化可以通过下述程序进行：

#![allow(unused)]
fn main() {
let cint_data = CInt::from_json(&mol_file);
}

这里的 mol_file 是定义了分子坐标与基组信息的 json 文件。该文件与 PySCF 输出的格式一致：

# python generation of `mol_file`
from pyscf import gto

mol = gto.Mole(atom="O; H 1 0.94; H 1 0.94 2 104.5", basis="def2-TZVP").build()
with open("h2o-tzvp.json", "w") as f:
    f.write(mol.dumps())

RHF 主程序是 minimal_rhf；其输出 RHFResults 定义为

#![allow(unused)]
fn main() {
pub struct RHFResults {
    pub mo_coeff: Tsr,
    pub mo_energy: Tsr,
    pub dm: Tsr,
    pub e_nuc: f64,
    pub e_elec: f64,
    pub e_tot: f64,
}
}

2.2 原子核互斥能计算

原子核互斥能的计算代价较小，也有多种实现方案。下面的实现方案并非是性能上最快的，但代码长度少，也比较直观。

首先，通过指标对称性，我们更改原子核互斥能的表达式为

$$E^{\mathsf{nuc}}=\frac{1}{2}\sum _{A\ne B}^{n_{\mathrm{atom}}}\frac{Z_A{Z}_B}{|{\mathbf{R}}_A-{\mathbf{R}}_B|}$$

进一步地，我们对下述矩阵作定义：

$$\begin{array}{rl}{P}_{AB}& =|{\mathbf{R}}_A-{\mathbf{R}}_B|\\ P_{AB}^{\prime }& =\{\begin{array}{ll}{P}_{AB}& A\ne B\\ \infty & A=B\end{array}\end{array}$$

那么原子核互斥能的表达式会更加简洁：

$$E^{\mathsf{nuc}}=\frac{1}{2}\sum _{AB}^{n_{\mathrm{atom}}}\frac{Z_A{Z}_B}{P_{AB}^{\prime }}$$

我们需要得到原子核电荷 $Z_A$。这是一个 $n_{\mathrm{atom}}$ 的向量。

我们可以用 rt::asarray，生成一个占有数据的张量：

#![allow(unused)]
fn main() {
let atom_charges = rt::asarray((cint_data.atom_charges(), &device)); // Tsr (natm, )
}

我们需要得到原子核坐标 ${\mathbf{R}}_A$。这是一个 $n_{\mathrm{atom}}\times 3$ 的矩阵。

CInt::atom_coords 函数输出的是 Vec<[f64; 3]> 类型；它不能由 rt::asarray 直接读取为 f64 类型张量。

#![allow(unused)]
fn main() {
let coords = cint_data.atom_coords(); // Vec<[f64; 3]> (natm)
}

其中一种办法是，先将长度为 $n_{\mathrm{atom}}$ 的 Vec<[f64; 3]> 类型转换为 $3n_{\mathrm{atom}}$ Vec<f64> 类型，

#![allow(unused)]
fn main() {
let coords = coords.into_iter().flatten().collect::<Vec<f64>>(); // Vec<f64> (3 * natm)
}

随后代入到 rt::asarray 中构成张量类型 Tsr，并通过 into_shape 重塑其形状为 $n_{\mathrm{atom}}\times 3$：

#![allow(unused)]
fn main() {
rt::asarray((coords, &device)).into_shape((-1, 3)) // Tsr (natm, 3)
}

另一种办法是，RSTSR 提供了函数 into_unpack_array；它可以将 [f64; 3] 类型的数据转为 [f64] 类型。其输入的参数是新增加出来的维度的位置 (在当前的例子中，我们希望维度增加在第 1 个位置)：

#![allow(unused)]
fn main() {
let atom_coords = rt::asarray((cint_data.atom_coords(), &device)).into_unpack_array(1);
}

对于 $P_{AB}=|{\mathbf{R}}_A-{\mathbf{R}}_B|$，它实际上就是两个原子之间的 Euclidean 距离。在 RSTSR 的 rt::sci 模块，我们有 cdist 函数以实现该功能 (类比 SciPy cdist 或 Matlab pdist2)。该函数目前还是要求传入张量视窗；视窗可以直接从张量变量后加 .view() 得到。

#![allow(unused)]
fn main() {
let mut dist = rt::sci::distance::cdist((atom_coords.view(), atom_coords.view()));
}

对于 $P_{AB}^{\prime }=\{\begin{array}{ll}{P}_{AB}& A\ne B\\ \infty & A=B\end{array}$，它相对于 $P_{AB}$ 是在对角线上的值设为无穷大。下述代码中，diagonal_mut 是获得矩阵 $P_{AB}$ 的对角线的可变引用，并随后调取 fill 以设置具体的数值。

#![allow(unused)]
fn main() {
dist.diagonal_mut(None).fill(f64::INFINITY);
}

最后，计算总原子核排斥能 $E^{\mathsf{nuc}}=\frac{1}{2}{\sum }_{AB}^{n_{\mathrm{atom}}}{Z}_A{Z}_B/P_{AB}^{\prime }$ 的代码为

#![allow(unused)]
fn main() {
0.5 * ((atom_charges.i((.., None)) * &atom_charges) / dist).sum()
}

该代码用到了张量广播 (broadcast) 的特性。上述表达式在作求和之前，生成的是该张量： $$T_{AB}=Z_A{Z}_B/P_{AB}^{\prime }$$ 但该数乘计算需要先进行张量的维度对齐才能进行；这是指在 $Z_A$ 后增加一个维度、并在 $Z_B$ 前增加一个维度： $$T_{AB}=Z_{A\square }{Z}_{\square B}/P_{AB}^{\prime }$$ 不过对于 row-major，$Z_{\square B}$ 前面增加出来的维度在程序中是可以被识别出来的。因此，我们只需要对 $Z_A$ 作维度对齐到 $Z_{A\square }$ 即可。

在 RSTSR 中，与 NumPy 一样地，为对齐而扩充维度，可以通过基础索引中插入 None 实现。

2.3 自洽场计算

首先，我们需要获得电子积分。这些量与密度无关，可以事先存储。我们在程序 util.rs (gitee, github) 中封装了电子积分调取函数 intor_row_major，可以在一个函数调用中直接给出 Tsr 张量类型的电子积分。

下述代码中，hcore 为 $h_{\mu \nu }$、ovlp 为 $S_{\mu \nu }$、int2e 为 $g_{\mu \nu ,\kappa \lambda }$。

#![allow(unused)]
fn main() {
let hcore = util::intor_row_major(cint_data, "int1e_kin") + util::intor_row_major(cint_data, "int1e_nuc");
let ovlp = util::intor_row_major(cint_data, "int1e_ovlp");
let int2e = util::intor_row_major(cint_data, "int2e");
}

随后是自洽场循环本身。我们回顾到自洽场问题是为了求解占据轨道系数矩阵 $C_{\mu i}$ 与电子云密度矩阵 $D_{\mu \nu }$；但由于求解过程中 Fock 矩阵也是随密度有依赖关系的，因此不能一次性计算得到，而需要通过迭代求解得到。该迭代关系总结下来是：

1. 通过密度矩阵得到 Fock 矩阵 $F_{\mu \nu }$： $$F_{\mu \nu }[\mathbf{D}]=h_{\mu \nu }+J_{\mu \nu }-\frac{1}{2}{K}_{\mu \nu }=h_{\mu \nu }+\sum _{\kappa \lambda }{D}_{\kappa \lambda }\left(g_{\mu \nu ,\kappa \lambda }-\frac{1}{2}{g}_{\mu \kappa ,\nu \lambda }\right)$$
2. 对 Fock 矩阵求解本征问题得到轨道系数 $C_{\mu p}$，该轨道系数需要依能级自小到大地排列： $$\sum _{\nu }{F}_{\mu \nu }{C}_{\nu p}=\sum _{\nu }{S}_{\mu \nu }{C}_{\nu p}{\epsilon }_p\text{or}\mathbf{FC}=\mathbf{SC}\mathbf{\epsilon }$$
3. 取本征值最低的占据轨道数量 $n_{\mathrm{occ}}$ 个本征向量构成占据轨道系数 $C_{\mu i}$，进而获得密度矩阵 $D_{\mu \nu }$： $$D_{\mu \nu }=2\sum _i^{n_{\mathrm{occ}}}{C}_{\mu i}{C}_{\nu i}\text{or}\mathbf{D}=2\mathbf{C}{\mathbf{C}}^{†}$$

自洽场既然是三行公式可以说明的问题，那么程序上也可以用三行写出来：

#![allow(unused)]
fn main() {
let mut dm = rt::zeros(([nao, nao], &device));
let mut mo_coeff = rt::zeros(([nao, nao], &device));
let mut mo_energy = rt::zeros(([nao], &device));
for _ in 0..NITER {
    let fock = &hcore + ((1.0_f64 * &int2e - 0.5_f64 * int2e.swapaxes(1, 2)) * &dm).sum_axes([-1, -2]);
    (mo_energy, mo_coeff) = rt::linalg::eigh((&fock, &ovlp)).into();
    dm = 2.0_f64 * mo_coeff.i((.., ..nocc)) % mo_coeff.i((.., ..nocc)).t();
}
}

尽管实际上该程序有 8 行，但其中 3 行是对可变变量 dm, mo_coeff, mo_energy 作定义，2 行是构建循环体。因此，真正用于自洽场算法实现的，确实只有 3 行。

#![allow(unused)]
fn main() {
let fock = &hcore + ((1.0_f64 * &int2e - 0.5_f64 * int2e.swapaxes(1, 2)) * &dm).sum_axes([-1, -2]);
}

Fock 矩阵计算中，

• int2e.swapaxes(1, 2) 是将 4-D 张量 $g_{\mu \nu ,\kappa \lambda }$ 中的 $\nu ,\kappa$ 维度作交换，得到 $g_{\mu \kappa ,\nu \lambda }$： $${\text{T1}}_{\mu \nu ,\kappa \lambda }=g_{\mu \kappa ,\nu \lambda }$$ 在 RSTSR 中，交换两个维度 (swapaxes) 或转置 (transpose) 都是没有计算与内存代价的，但可能对未来的计算效率产生影响；这与 NumPy 或 Julia 是一样的。
• 0.5_f64 * int2e.swapaxes(1, 2) 中，我们需要明确定义 0.5_f64 是 f64 位浮点，而不能仅写为 0.5。这是由于 Rust 在处理二目运算类型推断时，左右表达式的地位不等价、以及 RSTSR 内部实现本身有关。该表达式与 int2e.swapaxes(1, 2) * 0.5 是等价的实现，但后者的写法可以成功地进行类型推断。
• sum_axes([-1, -2]) 是对最后两个维度作求和。具体来说，下述代码碎片

  #![allow(unused)]
  fn main() {
  (1.0_f64 * &int2e - 0.5_f64 * int2e.swapaxes(1, 2)) * &dm
  }

  执行的是 broadcast 数乘运算，得到了 4-D 张量： $${\text{T2}}_{\mu \nu ,\kappa \lambda }=\left(g_{\mu \nu ,\kappa \lambda }-\frac{1}{2}{\text{T1}}_{\mu \nu ,\kappa \lambda }\right)D_{\square \square \kappa \lambda }$$ 随后对 ${\text{T2}}_{\mu \nu ,\kappa \lambda }$ 的最后两个维度 $\lambda ,\kappa$ 作求和： $${\text{T3}}_{\mu \nu }=\sum _{\lambda \kappa }{\text{T2}}_{\mu \nu ,\kappa \lambda }$$
• 最后的 Fock 矩阵则通过 $F_{\mu \nu }=h_{\mu \nu }+{\text{T3}}_{\mu \nu }$ 得到。

#![allow(unused)]
fn main() {
(mo_energy, mo_coeff) = rt::linalg::eigh((&fock, &ovlp)).into();
}

广义本征值求解可以通过 RSTSR 的函数 rt::linalg::eigh 实现。

• 该函数返回的值是 EighResult 类型，需要通过 .into() 转换为 (Tsr, Tsr) 的 tuple 类型。
• 这里由于是改写可变变量 mo_energy, mo_coeff，因此赋值语句左边不需要用 Rust 的 let。

#![allow(unused)]
fn main() {
dm = 2.0_f64 * mo_coeff.i((.., ..nocc)) % mo_coeff.i((.., ..nocc)).t();
}

密度矩阵计算中，

• .i((.., ..nocc)) 是基础索引语句，它是指对第一个维度作全索引、第二个维度取前 $n_{\mathrm{occ}}$ 个指标，构成新的矩阵。具体来说，本征值求解得到的矩阵是 $n_{\mathrm{basis}}\times n_{\mathrm{basis}}$ 大小的 $C_{\mu p}$；但实际的密度矩阵需要前 $n_{\mathrm{occ}}$ 个本征值最低的本征向量构成的 $n_{\mathrm{basis}}\times n_{\mathrm{occ}}$ 大小的 $C_{\mu i}$。
• .t() 与函数 .to_reverse_axes() 等价；它是生成一个转置的张量视窗。
• % 是 RSTSR 特有的语法糖；它的意义是矩阵乘法 (或矩阵-向量乘法、向量-矩阵乘法)。RSTSR 中，不论输入的矩阵是 row-major 或 col-major，都可以正常地调用 BLAS 的 GEMM 函数，不会因为输入矩阵是否转置了而产生效率损失。输出矩阵的连续性由编译时设置的 cargo feature (或后端的默认设定) 决定。

2.4 RHF 能量计算

对于 RHF 方法，能量计算的代码与 Fock 矩阵的实现是非常相似的，只是在 4c-2e 双电子积分 $g_{\mu \nu ,\kappa \lambda }$ 上的系数有差异。

#![allow(unused)]
fn main() {
let eng_scratch = &hcore + ((0.5_f64 * &int2e - 0.25_f64 * int2e.swapaxes(1, 2)) * &dm).sum_axes([-1, -2]);
let e_elec = (&dm * &eng_scratch).sum();
let e_tot = e_nuc + e_elec;
}

1. 这里的表达式可以参考 Szabo & Ostlund, eqs (3.154, 3.184)；但需要注意记号上的差异。 ↩

2. 这里第第一类是程序意义上的，参考 Lapack dsygv (参数 ITYPE) 或 SciPy eigh (参数 type) 程序文档。 ↩

RI-RHF：极简演示

• 程序：ri_rhf_slow.rs (gitee, github)
• 实现内容：RI-RHF (restricted resolution-of-identity Hartree-Fock)
• 不关注程序效率
• 该程序总共约 50 行

在本节中，我们经常将密度矩阵 $D_{\mu \nu }$ 看作 $n_{\mathrm{basis}}^2$ 长度的 1-D 向量，而非 $n_{\mathrm{basis}}\times n_{\mathrm{basis}}$ 的 2-D 矩阵。类似地，4c-2e 双电子积分经常看作是 $n_{\mathrm{basis}}^2\times n_{\mathrm{basis}}^2$ 的 2-D 矩阵，而非 $n_{\mathrm{basis}}\times n_{\mathrm{basis}}\times n_{\mathrm{basis}}\times n_{\mathrm{basis}}$ 的 4-D 张量。

目录

• 1. RI-RHF 回顾
  • 1.1 4c-2e 双电子积分分解
  • 1.2 J 积分
  • 1.3 K 积分
• 2. 程序
  • 2.1 4c-2e 双电子积分分解
  • 2.2 J 积分
  • 2.3 K 积分

1. RI-RHF 回顾

相比于普通的 RHF，RI-RHF 方法不需要计算庞大的 4c-2e 双电子积分，而只需要计算较小的 3c-2e 双电子积分。同时，通过张量分解，其 Fock 矩阵的计算量明显比 RHF 方法小很多。

以 J 积分 $J_{\mu \nu }$ 为例，RI-V 近似下的 RI-RHF 的理解思路可以大致用下图表示：

1.1 4c-2e 双电子积分分解

RI-V 近似除了普通的分子轨道基组外，还需要设置辅助基组。辅助基数量越大、精度越高但计算消耗也越多。一般来说，辅助基组的数量 $n_{\mathrm{aux}}$ 大约是 $3n_{\mathrm{basis}}$ 大小。从矩阵分解的角度来看，如果将 4c-2e 双电子积分看作 $n_{\mathrm{basis}}^2\times n_{\mathrm{basis}}^2$ 的矩阵，那么辅助基的大小 $n_{\mathrm{aux}}$ 数量级上应该要接近 4c-2e 双电子积分的秩 (rank)。

3c-2e 双电子积分与 2c-2e 双电子积分定义如下：

$$\begin{array}{rl}{g}_{\mu \nu ,P}& =\int \mathrm{d}{\mathbf{r}}_1\int \mathrm{d}{\mathbf{r}}_2{\varphi }_{\mu }({\mathbf{r}}_1){\varphi }_{\nu }({\mathbf{r}}_1)\frac{1}{|{\mathbf{r}}_1-{\mathbf{r}}_2|}{\varphi }_P({\mathbf{r}}_2)\\ g_{PQ}& =\int \mathrm{d}{\mathbf{r}}_1\int \mathrm{d}{\mathbf{r}}_2{\varphi }_P({\mathbf{r}}_1)\frac{1}{|{\mathbf{r}}_1-{\mathbf{r}}_2|}{\varphi }_Q({\mathbf{r}}_2)\end{array}$$

随后对 2c-2e 双电子积分作 Cholesky 分解，分解得到的 $L_{PQ}$ 是下三角矩阵：

$$g_{PQ}=\sum _L^{n_{\mathrm{aux}}}{L}_{PR}{L}_{QR}\text{or}{\mathbf{g}}^{\text{2c}}=\mathbf{L}{\mathbf{L}}^{†}$$

考虑到 4c-2e 双电子积分在 RI-V 下的近似表达式为

$${\mathbf{g}}^{\text{4c}}\simeq {\mathbf{g}}^{\text{3c}}({\mathbf{g}}^{\text{2c}}{)}^{-1}({\mathbf{g}}^{\text{3c}}{)}^{†}={\mathbf{g}}^{\text{3c}}({\mathbf{L}}^{†}{)}^{-1}{\mathbf{L}}^{-1}({\mathbf{g}}^{\text{3c}}{)}^{†}$$

定义 Cholesky 分解量 $B_{P,\mu \nu }$ 为 (在我们当前的项目中，需要注意其指标顺序与 3c-2e 双电子积分 $g_{\mu \nu ,P}$ 不同)

$$\mathbf{B}={\mathbf{L}}^{-1}({\mathbf{g}}^{\text{3c}}{)}^{†}\text{or}{B}_{P,\mu \nu }=\sum _Q^{n_{\mathrm{aux}}}(L^{-1}{)}_{PQ}{g}_{\mu \nu ,Q}$$

那么 4c-2e 双电子积分的近似表达式也可以写为

$${\mathbf{g}}^{\text{4c}}\simeq {\mathbf{B}}^{†}\mathbf{B}\text{or}{g}_{\mu \nu ,\kappa \lambda }\simeq \sum _P^{n_{\mathrm{aux}}}{B}_{P,\mu \nu }{B}_{P,\kappa \lambda }$$

1.2 J 积分

J 积分的计算近似为

$$J_{\mu \nu }=\sum _{\kappa \lambda }^{n_{\mathrm{basis}}}{g}_{\mu \nu ,\kappa \lambda }{D}_{\kappa \lambda }\simeq \sum _P^{n_{\mathrm{aux}}}{B}_{P,\mu \nu }\sum _{\kappa \lambda }^{n_{\mathrm{basis}}}{B}_{P,\kappa \lambda }{D}_{\kappa \lambda }$$

如果用矩阵-向量乘法来表述，则是

$$\mathbf{J}\simeq {\mathbf{B}}^{†}\mathbf{BD}$$

1.3 K 积分

在 RI 近似下，K 积分的计算应代入分子轨道系数以提升计算效率。在闭壳层下，占据轨道总是双占据的。

$$K_{\mu \nu }=\sum _{\kappa \lambda }^{n_{\mathrm{basis}}}{g}_{\mu \kappa ,\nu \lambda }{D}_{\kappa \lambda }\simeq 2\sum _P^{n_{\mathrm{aux}}}\sum _{\kappa \lambda }^{n_{\mathrm{basis}}}\sum _i^{n_{\mathrm{occ}}}{B}_{P,\mu \kappa }{B}_{P,\nu \lambda }{C}_{\kappa i}{C}_{\lambda i}$$

当我们定义半转换 (half-transformed) 的 $B_{P,i\mu }$ 为

$$B_{P,i\mu }=\sum _{\kappa }^{n_{\mathrm{basis}}}{B}_{P,\mu \kappa }{C}_{\kappa i}$$

我们可以将 $K_{\mu \nu }$ 进一步写为

$$K_{\mu \nu }\simeq 2\sum _P^{n_{\mathrm{aux}}}\sum _i^{n_{\mathrm{occ}}}{B}_{P,i\mu }{B}_{P,i\nu }$$

我们再定义记号 ${\bar{B}}_{Pi,\mu }=B_{P,i\mu }$，这两个张量完全等价。但当作矩阵看待时，$B_{P,i\mu }$ 是 $n_{\mathrm{aux}}\times n_{\mathrm{occ}}{n}_{\mathrm{basis}}$ 大小，而 ${\bar{B}}_{Pi,\mu }$ 是 $n_{\mathrm{aux}}{n}_{\mathrm{occ}}\times n_{\mathrm{basis}}$ 大小。定义 ${\bar{B}}_{Pi,\mu }$ 的方便之处在于，K 积分可以用简单的矩阵乘法形式表达出来 (上标 half 是指该张量是经过半转换的)：

$$\mathbf{K}\simeq 2{{\bar{\mathbf{B}}}^{\text{half}}}^{†}{\bar{\mathbf{B}}}^{\text{half}}$$

2. 程序

2.1 4c-2e 双电子积分分解

在获得电子积分后，我们的主要目标是得到 $B_{P,\mu \nu }$。总共需要作两步计算：Cholesky 分解与矩阵求解： $$\begin{array}{rl}{\mathbf{g}}^{\text{2c}}& =\mathbf{L}{\mathbf{L}}^{†}\\ \mathbf{B}& ={\mathbf{L}}^{-1}({\mathbf{g}}^{\text{3c}}{)}^{†}\end{array}$$

在程序中的实现如下：

#![allow(unused)]
fn main() {
let int2c2e_l = rt::linalg::cholesky((int2c2e.view(), Lower));
let cderi = rt::linalg::solve_triangular((int2c2e_l.view(), int3c2e.reshape([nao * nao, naux]).t(), Lower));
let cderi = cderi.into_shape([naux, nao, nao]).into_contig(RowMajor);
}

• 关于 Cholesky 分解，代码相信已经很直观了；
• 关于矩阵求解，需要注意到这里的 $\mathbf{L}$ 是下三角矩阵；通过三角矩阵求解函数 solve_triangular 会比通常的矩阵求解函数 solve_general 或者求逆 inv 函数，性能上要快许多。
• 最后，我们对 $B_{P,\mu \nu }$ 作维度修改 (从 $n_{\mathrm{aux}}\times n_{\mathrm{basis}}^2$ 的 2-D 矩阵修改到 $n_{\mathrm{aux}}\times n_{\mathrm{basis}}\times n_{\mathrm{basis}}$ 的 3-D 张量)，并从效率角度出发要求该张量需要是 row-major 的 (需要让辅助基指标 $P$ 是最不连续的维度，这在 K 积分计算中是比较关键的)。

2.2 J 积分

J 积分的过程非常简单，它就是两次矩阵-向量乘法。但需要注意乘法的顺序，否则效率会非常糟糕：

$$\mathbf{J}\simeq {\mathbf{B}}^{†}(\mathbf{BD})$$

#![allow(unused)]
fn main() {
let cderi_flat = cderi.reshape([naux, nao * nao]);
let dm_flat = dm.reshape([nao * nao]);
(cderi_flat.t() % (&cderi_flat % dm_flat)).into_shape([nao, nao])
}

上面的代码花了一些功夫在维度更改上。

2.3 K 积分

K 积分的计算是两次矩阵乘法：

$$\begin{array}{rl}{B}_{P,i\mu }& =\sum _{\kappa }^{n_{\mathrm{basis}}}{B}_{P,\mu \kappa }{C}_{\kappa i}\\ \mathbf{K}& \simeq 2{{\bar{\mathbf{B}}}^{\text{half}}}^{†}{\bar{\mathbf{B}}}^{\text{half}}\end{array}$$

#![allow(unused)]
fn main() {
let occ_coeff = mo_coeff.i((.., ..nocc));
let scr = (occ_coeff.t() % cderi.reshape([naux * nao, nao]).t()).into_shape([nocc, naux, nao]);
let scr_flat = scr.reshape([naux * nocc, nao]);
2.0_f64 * (scr_flat.t() % &scr_flat)
}

上述程序的

• 第 1 行是从分子轨道系数中取出占据轨道的部分 $C_{\mu i}$；
• 第 2 行是半转换得到 $B_{i,P\mu }$；
• 第 3 行是重塑维度得到 ${\bar{B}}_{iP,\mu }$；
• 第 4 行是进行 $\mathbf{K}\simeq 2{{\bar{\mathbf{B}}}^{\text{half}}}^{†}{\bar{\mathbf{B}}}^{\text{half}}$ 的实际计算。

这里的程序实现与先前 RI-RHF 回顾中用到的公式还有些许差别。我们需要作一些额外的展开。

第 2 行中，我们实际执行的矩阵乘法是

$$B_{i,P\mu }=\sum _{\kappa }^{n_{\mathrm{basis}}}{C}_{\kappa i}{B}_{P\mu ,\kappa }$$

该矩阵乘法对 $B_{P,\mu \kappa }$ 重塑维度，并对参与乘法的两个矩阵都作转置，才能得到目标结果。这么做非常费劲，其目的是将半转换后指标 $i$ 与 $P$ 联系在一起，方便下面得到矩阵 $K_{\mu \nu }$ 时可以矩阵乘法一步到位。

由于 $\mathbf{K}\simeq 2{{\bar{\mathbf{B}}}^{\text{half}}}^{†}{\bar{\mathbf{B}}}^{\text{half}}$ 中，指标 $i$ 与 $P$ 是并在一起的；因此半转换得到的是 $B_{i,P\mu }$ 还是 $B_{P,i\mu }$ 并不重要，只要 $i$ 与 $P$ 联系在一起即可。

RI-RHF：高效率实现

• 程序：ri_rhf.rs (gitee, github)
• 实现内容：RI-RHF (restricted resolution-of-identity Hartree-Fock)
• 性能尚可
  • 体系 (H2O)₁₀ cc-pVDZ ($n_{\mathrm{basis}}=240$, $n_{\mathrm{aux}}=1160$)
  • 计算设备 Ryzen HX7945，16 cores，算力约 1.1 TFLOP/sec
  • OpenBLAS 1.4 sec / 16 iter (RSTSR v0.3 的矩阵-向量乘法用 GEMM 而非 GEMV 实现，J 积分较慢)
  • Faer 1.3 sec / 16 iter
• 该程序总共约 110 行
• 该程序包含 DIIS 迭代

目录

• 1. RI-RHF 性能需要考虑的因素
• 2. 3c-2e 双电子积分的获得
• 3. Cholesky 分解张量的获得
• 4. J 积分计算
• 5. K 积分计算

1. RI-RHF 性能需要考虑的因素

提升 RI-RHF 主要的因素是 3c-2e 双电子积分 $g_{\mu \nu ,P}$ 与 Cholesky 分解张量 $B_{P,\mu \nu }$ 具有对称性： $$g_{\mu \nu ,P}=g_{\nu \mu ,P}$$ 利用对称性可以节省一半内存，以及减小 J 积分的浮点计算量、以及 K 积分计算时的 DRAM 带宽。

利用对称性当然也可以减少电子积分本身的计算时间；但对于本项目的 RI-RHF 而言，我们已经假定电子积分可以全部存入内存，那么自洽过程就不需要额外计算电子积分了。同时，3c-2e 电子积分计算时间一般不大于一次 SCF 迭代，因此电子积分耗时本来就不是大头。是否节省这点计算量，对于当前项目的算法而言，意义并不大。

Conventional 4c-2e 双电子积分的 RHF 方法体系稍大一些时，内存是无法存入所有电子积分的；因此每次 SCF 迭代一定会重新算电子积分，从而电子积分的耗时变得关键。这与 RI-RHF 的情形并不相同。

2. 3c-2e 双电子积分的获得

我们对 3c-2e 双电子积分有两个考量因素：

• 我们希望 3c-2e 双电子积分只存一半；对于 row-major 情形，我们希望存下三角部分 $\mu ⩾\nu$。
• 考虑到后面 solve_triangular 函数的调用，我们希望生成 row-major 下连续的 $g_{P,\mu \nu }$，即辅助基的指标是内存上最不连续的。

libcint 给出的 3c-2e 电子积分始终是 $(\mu \nu |P)$ 形式的，必须要进行一次转置才能得到 $(P|\mu \nu )$。既然我们希望得到的是 row-major 下的 $g_{P,\mu \nu }$，那么反过来想，我们可以从 col-major 的 $g_{\nu \mu ,P}=g_{\mu \nu ,P}$ 作一次转置，就可以得到 row-major 下的 $g_{P,\mu \nu }$ 了。

综合这些考量因素，我们给出如下的 3c-2e 双电子积分计算代码：

#![allow(unused)]
fn main() {
use libcint::prelude::*;

let int3c2e = {
    let (out, shape) = CInt::integrate_cross("int3c2e", [cint_data, cint_data, aux_cint_data], "s2ij", None).into();
    rt::asarray((out, shape.f(), &device)).into_reverse_axes()
};
}

• CInt::integrate_cross 是可以对多个分子基组信息作交叉电子积分计算的函数，其结果是 col-major 的。参数 "s2ij" 表示其对称性。
• 在函数 rt::asarray 调用时，使用到的 shape.f() 是指，在构成张量时，我们希望使用 col-major 布局 (fortran-contiguous)。随后作 into_reverse_axes 将给出 col-major 到 row-major 的张量转置，这一步不会产生额外的内存复制。
• 最后返回得到的 int3c2e $g_{P,\mu \nu }$，其布局是 row-major，维度是 $n_{\mathrm{aux}}\times \frac{1}{2}{n}_{\mathrm{basis}}(n_{\mathrm{basis}}+1)$，关于指标 $(\mu ,\nu )$ 下三角存储。

整个过程中，除了电子积分本身的耗时与内存占用之外，其余部分的资源消耗是近乎为零的。

3. Cholesky 分解张量的获得

$$B_{P,\mu \nu }=\sum _Q^{n_{\mathrm{aux}}}(L^{-1}{)}_{PQ}{g}_{Q,\mu \nu }\text{or}\mathbf{B}={\mathbf{L}}^{-1}{\mathbf{g}}^{\text{3c}}$$

#![allow(unused)]
fn main() {
let cderi = rt::linalg::solve_triangular((int2c2e_l.view(), int3c2e, Lower));
}

这里需要留意的是，我们传入的 int3c2e 是占有数据的 Tsr 张量类型，而不是视窗引用。这意味着变量 int3c2e 在调用 solve_triangular 函数后就无法使用了。这与上一节 ri_rhf_slow.rs 程序的情况不同；上一节的程序的第二个参数 int3c2e.reshape([nao * nao, naux]).t() 是张量视窗。

这里里用到了 RSTSR 的参数重载。当第二个参数 (被求解的长方形矩阵) 是 Tsr 或 TsrMut、而非 TsrView 时，求解得到的结果会被写入到原来的张量中。这种情况下，只要保证传入 solve_triangular 存在一个维度是连续的，不论它是行连续还是列连续，solve_triangular 求解时都不会额外分配内存¹。

4. J 积分计算

考虑到对称性，J 积分可以写为

$$\begin{array}{rl}{J}_{\mu \nu }& =\sum _P^{n_{\mathrm{aux}}}{B}_{P,\mu \nu }\sum _{\kappa \lambda }^{n_{\mathrm{basis}}}{B}_{P,\kappa \lambda }{D}_{\kappa \lambda }\\ & =\sum _P^{n_{\mathrm{aux}}}{B}_{P,\mu \nu }\sum _{\kappa ⩾\lambda }^{n_{\mathrm{basis}}}(2-{\delta }_{\kappa \lambda })B_{P,\kappa \lambda }{D}_{\kappa \lambda }\end{array}$$

由此，我们可以考虑构建一个缩放后的密度矩阵 ${\bar{D}}_{\kappa \lambda }$，并在程序实现时只取其下三角部分 $\kappa ⩾\lambda$：

$${\bar{D}}_{\kappa \lambda }=(2-{\delta }_{\kappa \lambda })D_{\kappa \lambda }$$

#![allow(unused)]
fn main() {
let mut dm_scaled = 2.0_f64 * dm.to_owned();
let dm_diag = dm.diagonal(None);
dm_scaled.diagonal_mut(None).assign(dm_diag);
let dm_scaled_tp = dm_scaled.pack_tril();
}

随后进行普通的矩阵-向量乘法，得到下三角的 $J_{\mu \nu }$ $(\mu ⩾\nu )$：

#![allow(unused)]
fn main() {
let j_tp = (&cderi % dm_scaled_tp) % &cderi;
}

最后将下三角的 J 积分展开到方形矩阵：

#![allow(unused)]
fn main() {
j_tp.unpack_tril(FlagSymm::Sy)
}

5. K 积分计算

在这里，K 积分的计算与上一节有些许差别。

K 积分计算的核心目的仍然不变，是求取半转换的 Cholesky 分解张量时，尽量将指标辅助基指标 $P$ 与占据轨道指标 $i$ 放到一起。但是在放置顺序上，上一篇文档的程序是 $B_{i,P\mu }$；这里程序的策略则是 $B_{P,i\mu }$。

当前程序的伪代码可以写为

• 对指标 $P$ 作并行；
  • 对特定的指标 $P$，取得子下三角矩阵 $b_{\mu \nu }=B_{P,\mu \nu }$ $(\mu ⩾\nu )$；
  • 对子下三角矩阵 $b_{\mu \nu }=b_{\nu \mu }$ 展开到方形矩阵；
  • 求取 $b_{i\mu }={\sum }_{\nu }{C}_{\nu i}{b}_{\nu \mu }$
  • 对特定的指标 $P$，写入 $b_{i\mu }$ 到半转换张量 $B_{P,i\mu }$
• 将半转换张量的维度重置为 $B_{Pi,\mu }$
• 通过矩阵乘法得到 $K_{\mu \nu }={\sum }_{Pi}{B}_{Pi,\mu }{B}_{Pi,\nu }$

#![allow(unused)]
fn main() {
let occ_coeff = mo_coeff.i((.., ..nocc));
let scr_xob = unsafe { rt::empty(([naux, nocc, nao], &device)) };
(0..naux).into_par_iter().for_each(|p| {
    let mut scr_xob = unsafe { scr_xob.force_mut() };
    let cderi_ob = cderi.i(p).unpack_tril(FlagSymm::Sy);
    scr_xob.i_mut(p).matmul_from(&occ_coeff.t(), &cderi_ob, 1.0, 0.0);
});
let scr_xob_flat = scr_xob.reshape([naux * nocc, nao]);
2.0_f64 * (scr_xob_flat.t() % &scr_xob_flat)
}

程序上，比较特殊的地方有

• 对指标 $P$ 并行，在 C++/OpenMP 中可以写为

  #pragma omp parallel for // if don't want parallel, comment this line
  for (size_t p = 0; p < naux; ++p) { ... }
  

  相对而言，在 Rust 下，使用 Rayon 库，语法上有些许不同，但思路是一样的：

  #![allow(unused)]
  fn main() {
  // if don't want parallel, change `into_par_iter` to `into_iter`
  (0..naux).into_par_iter().for_each(|p| { ... });
  }

• 第一个 unsafe 是 rt::empty 的要求。在 Rust 中，声明一块没有初始化的内存是不安全的行为。对于 f64 类型而言，这个不安全性至多影响到数据安全问题，这在科学计算上不是很重要。但如果是对一个具有析构函数 (trait Drop implemented) 的类型不作初始化，那么很容易会出现非有效内存错误。

• 第二个 unsafe，从程序的角度来说，是强制在每个线程中，对声明为不可变的 scr_xob ($B_{Pi,\mu }$) 取到其可变的引用。

  这从一般的程序开发角度来看，是比较危险的 unsafe。但在科学计算问题里，这是非常常见的情况。这个 unsafe 其实在 C++/OpenMP 中是声明某个变量为 shared：

  #pragma omp parallel for shared(scr_xob)
  

  如果您作为 C++ 编程者，认为这里的 omp shared 语句是可以接受的，那么这里的 unsafe 您也应该容易接受。

• 第二个 unsafe 确实是有更安全的解决方案，但它可能会导致代码变得不那么直观：

  #![allow(unused)]
  fn main() {
  let mut scr_xob_iter = scr_xob.axes_iter_mut(0);
  (0..naux).into_par_iter().zip(scr_xob_iter).for_each(|(p, mut scr_ob)| {
      let cderi_ob = cderi.i(p).unpack_tril(FlagSymm::Sy);
      scr_ob.matmul_from(&occ_coeff.t(), &cderi_ob, 1.0, 0.0);
  });
  }

  其中，axes_iter_mut(0) 函数是指对第 0 个维度 (指标 $P$ 所在的维度) 返回 mutable 子张量的迭代器。这个迭代器可以用于并行的环境。

• i_mut(p) 与 i 函数一样是基础索引，但它将返回张量的可变视窗。

• matmul_from 与 % 或者 rt::matmul 一样，都是张量的矩阵乘法。但不同的是，% 与 rt::matmul 会为矩阵乘法的结果新开辟一片堆空间内存；但 matmul_from 则是在已有内存上进行矩阵乘法，且允许指定 $\mathbf{C}=\alpha \mathbf{AB}+\beta \mathbf{C}$ 的系数 $\alpha$ 与 $\beta$。使用 matmul_from 或许从阅读上看来并不直观；但从效率与内存占用上，matmul_from 确实是更好的选择。

• 我们在 Rayon 并行区域内调用了矩阵乘法。在 RSTSR 中，对于 BLAS 后端，

  • 如果在 Rayon 并行区域外调用 BLAS，则会使用并行版本的 BLAS，并控实际使用的制线程数量。
  • 如果在 Rayon 并行区域内调用 BLAS，则会通过 openblas_set_num_threads 或 omp_set_num_threads 控制线程数为 1，从而线程内使用串行的 BLAS、但线程外由 Rayon 控制任务调度，达到并行的效果。
  • 对于 pthreads 并行的 OpenBLAS，如果您希望在其他地方直接使用 OpenBLAS 函数前，最好重新通过 openblas_set_num_threads 重新设置一下线程数量。Rayon 并行区域内调用 BLAS 可能会破坏 OpenBLAS 的全局并行参数。对于 OpenMP 并行的 OpenBLAS 应没有上述问题。

• 注意到最后实现矩阵乘法 $\mathbf{K}=2{\bar{\mathbf{B}}}^{†}\bar{\mathbf{B}}$ 的代码

  #![allow(unused)]
  fn main() {
  2.0_f64 * (scr_xob_flat.t() % &scr_xob_flat)
  }

  括号影响了计算顺序，即我们需要先进行 ${\bar{\mathbf{B}}}^{†}\bar{\mathbf{B}}$，随后作 2 倍的乘法。如果不加括号，考虑到 * 与 % 是同优先级算符，计算顺序是会发生变化的。同时，对两个相同矩阵的转置乘法 (类似于 ${\mathbf{A}}^{†}\mathbf{A}$ 或 $\mathbf{A}{\mathbf{A}}^{†}$)，RSTSR 会作一步额外优化：

  • BLAS 后端会调用 SYRK 而非 GEMM 函数计算；
  • Faer 后端会使用仅输出三角矩阵的矩阵乘法模块。 理想情况下，这类型矩阵乘法的浮点计算量相比于普通矩阵乘法，可以节省一半。

1. 是否要分配额外的内存，这实际上依赖于后端的实现。不过当前实现了 solve_triangular 的后端 Faer 与 OpenBLAS 一般都不会额外分配内存。 ↩

RI-RCCSD：高效率实现

• 程序：ri_rccsd.rs (gitee, github)
• 实现内容：RI-RCCSD (restricted resolution-of-identity coupled-cluster (singles and doubles))
• 性能较高
  • 体系 (H2O)₁₀ cc-pVDZ ($n_{\mathrm{basis}}=240$, $n_{\mathrm{aux}}=840$，无冻结轨道)
  • 计算设备 Ryzen HX7945，16 cores，算力约 1.1 TFLOP/sec
  • 一次 CCSD 迭代约 24.0 sec，CPU 性能利用率约 49%
  • 本文档的所有性能分析均基于 OpenBLAS 后端讨论
• 该程序总共约 700 行
• 该程序包含 DIIS 迭代
• 程序测评脚本

  cargo build --release --features="use_openblas"
  target/release/showcase-workshop-rstsr-ricc ri-rccsd \
      -mol assets/h2o_10-pvdz.json -aux assets/h2o_10-pvdz_jk.json --aux-ri assets/h2o_10-pvdz_ri.json
  

目录

• 1. RI-RCCSD 实现前言
• 2. Cholesky 分解积分
  • 2.1 公式
  • 2.2 程序
  • 2.3 性能上的考量
  • 3.4 性能评价
• 3. RCCSD 初猜与 RI-RMP2 能量
  • 3.1 输出类型 RCCSDResults
  • 3.2 公式
  • 3.3 程序
  • 3.4 性能评价
• 4. RI-RCCSD 中间变量：第一部分
  • 4.1 RI-RCCSD 中间变量
  • 4.2 公式
  • 4.3 程序
  • 4.4 性能上的考量
  • 4.4 性能评价
• 5. RI-RCCSD 能量计算
• 6. RI-RCCSD 中间变量：第二部分
  • 6.1 公式
  • 6.2 程序
• 7. RI-RCCSD 一次激发 rhs1
  • 7.1 公式
  • 7.2 程序
• 8. RI-RCCSD 二次激发 rhs2：$t_{ij}^{ab}$ 与类微扰 Fock 矩阵缩并部分
  • 8.1 公式
  • 8.2 程序
  • 8.3 性能评价
• 9. RI-RCCSD 二次激发 rhs2：直积贡献
  • 9.1 公式
  • 9.2 程序
  • 9.3 性能评价
• 10. RI-RCCSD 二次激发 rhs2：HHL
  • 10.1 公式
  • 10.2 程序
  • 10.3 性能评价
• 11. RI-RCCSD 二次激发 rhs2：PHL
  • 11.1 公式
  • 11.2 程序
  • 11.3 性能与程序结构上的考量
  • 11.4 性能评价
• 12. RI-RCCSD 二次激发 rhs2：PPL
  • 12.1 公式
  • 12.2 最简程序实现
  • 12.3 程序及其性能考量
    • 12.3.1 指标 $(a,b)$ 的分批
    • 12.3.2 计算 $W_{abcd}$
    • 12.3.3 ${\mathcal{T}}_{ij}^{ab}$ 的计算
    • 12.3.4 对 ${\text{RHS2}}_{ij}^{ab}$ 作更新
  • 12.4 性能评价
• 13. RI-RCCSD 其他计算过程与总结

1. RI-RCCSD 实现前言

相比于 RHF、RI-RHF、RI-CCSD(T) 的 triplet 微扰部分都只有 100 行左右，RI-RCCSD 的 600 行程序长度有明显的提升。这其中直接的原因是

• RI-RCCSD 本身的公式数量相当多；
• 需要用工程化的思路组织程序，这引入了许多不执行具体计算的业务逻辑代码 (大量的 getter/setter)；
• 为了提升效率，一些情况下我们确实需要引入更多的代码；
• RSTSR 目前暂不支持 einsum 式的张量缩并¹，因此需要用矩阵乘法实现张量缩并。

如果从提升实现效率的角度而言，还有额外的原因是

• 为了将张量缩并问题化为矩阵乘法，我们需要经常思考如何排布张量的指标顺序²。因此，公式推导和程序的实现复杂程度还会进一步提升。

但相对于 RI-RCCSD(T) 的 triplet 微扰，RI-RCCSD 相对来说更简单的部分是

• 绝大多数计算都可以用数学库的标准函数高效实现，不需要作额外的 hack。

我们知道耦合簇方法，从公式上，可以认为是联立下述关于波函数态 $\xi$ 方程组，求解 $\hat{T}={\sum }_{\xi }{t}_{\xi }\hat{\xi }$ 中的激发系数 $t_{\xi }$： $$⟨{\Phi }_{\xi }|e^{-\hat{T}}\hat{H}{e}^{\hat{T}}|{\Phi }_0⟩={\delta }_{\xi 0}{E}^{\text{CC}}$$ 其中，

• $\hat{H}$ 是体系 Hamiltonian，
• $|{\Phi }_0⟩$ 是参考基态 (一般是 Hartree-Fock 基态)，
• $|{\Phi }_{\xi }⟩=\hat{\xi }|{\Phi }_0⟩$ 是参考方法能级为 $\xi$ 的态波函数。$\hat{\xi }$ 这里指代激发算符；在二次量子化下，一次激发 $|{\Phi }_i^a⟩=i^{†}a|{\Phi }_0⟩$，二次激发 $|{\Phi }_{ij}^{ab}⟩=i^{†}{j}^{†}ba|{\Phi }_0⟩$。

上式的 $\hat{T}$ 原则上应该包含任意阶次的激发算符 (最多到参考态占据电子数)，但大多数时候 $\hat{T}$ 截断到 2-3 次激发就达到一般化学需求 (或者达到可计算能力上限)。当 $\hat{T}$ 截断到二次激发时，称为 CCSD 方法；进一步通过引入 $\hat{T}$ 的特定三次微扰时，则称为 CCSD(T) 方法。

耦合簇理论很优美，但程序的实现没办法做得这么漂亮。即使是概述 RI-RCCSD 的实现，也不可避免地涉及大量公式细节³。因此，这一节中，我们会拆解 RI-RCCSD 到数个子任务，对每个子任务进行公式、程序实现、以及性能提升策略的叙述。

2. Cholesky 分解积分

• 程序函数：get_riccsd_intermediates_cderi

2.1 公式

与 RI-RHF 的情形不同，在 RI-RCCSD 计算时，我们希望辅助基指标在最连续的维度上，且我们希望得到分子轨道基下的 Cholesky 分解积分：

$$B_{pq,P}=\sum _{\mu \nu }{B}_{\mu \nu ,P}{C}_{\mu p}{C}_{\nu q}$$

2.2 程序

下述程序从分子基组信息开始，

• 首先获得原子轨道基下的 2c-2e 双电子积分 $g_{PQ}$ 与 row-major 的 3c-2e 双电子积分 $g_{\mu \nu ,P}$；

  #![allow(unused)]
  fn main() {
  let int3c2e = util::intor_3c2e_row_major(cint_data, aux_cint_data, "int3c2e");
  let int2c2e = util::intor_row_major(aux_cint_data, "int2c2e");
  }

• 对原子轨道基下的积分作转换，得到 Cholesky 分解积分 $B_{\mu \nu ,P}$，

  #![allow(unused)]
  fn main() {
  let int3c2e_trans = int3c2e.into_shape([nao * nao, naux]).into_reverse_axes();
  let int2c2e_l = rt::linalg::cholesky((int2c2e.view(), Lower));
  let cderi = rt::linalg::solve_triangular((int2c2e_l.view(), int3c2e_trans, Lower));
  let cderi_uvp = cderi.into_reverse_axes().into_shape([nao, nao, naux]);
  }

• 作第一次半转换得到 $B_{p\nu ,P}$： $$B_{p\nu ,P}\equiv B_{p,\nu P}=\sum _{\mu }{C}_{\mu p}{B}_{\mu ,\nu P}$$

  #![allow(unused)]
  fn main() {
  let cderi_svp = (mo_coeff.t() % cderi_uvp.reshape((nao, nao * naux))).into_shape((nmo, nao, naux));
  }

• 对三种情形分别作第二次半转换，得到最终的分子轨道基下的 Cholesky 分解积分；下述计算利用了 RSTSR 所支持的 broadcast 矩阵乘法： $$\begin{array}{rl}{B}_{ij,P}\equiv B_{i,jP}& =\sum _{\nu }{C}_{\square ,\nu j}{B}_{i,\nu P}\\ B_{ia,P}\equiv B_{i,aP}& =\sum _{\nu }{C}_{\square ,\nu a}{B}_{i,\nu P}\\ B_{ab,P}\equiv B_{a,bP}& =\sum _{\nu }{C}_{\square ,\nu b}{B}_{a,\nu P}\end{array}$$

  #![allow(unused)]
  fn main() {
  let so = slice!(0, nocc);
  let sv = slice!(nocc, nocc + nvir);
  let b_oo = mo_coeff.i((.., so)).t() % cderi_svp.i(so);
  let b_ov = mo_coeff.i((.., sv)).t() % cderi_svp.i(so);
  let b_vv = mo_coeff.i((.., sv)).t() % cderi_svp.i(sv);
  }

2.3 性能上的考量

1. 与 RI-RHF 实现不同，3c-2e 双电子积分我们要求使用 row-major 的版本得到 $g_{\mu \nu ,P}$，而不是使用 col-major 的版本再作一次转置得到 $g_{P,\mu \nu }$。这是因为 RI-RCCSD 的许多计算中，辅助基 $P$ 是最连续的维度会比较方便。那么不论是 3c-2e 电子积分还是 Cholesky 分解积分，我们都要求辅助基指标 $P$ 是最连续的。

2. rt::linalg::solve_triangular 支持 inplace solve (上述代码的 cderi 复用了 int3c2e_trans 的内存数据，效率更高)，因此对 int3c2e $g_{\mu \nu ,P}$ 的维度更改操作，始终使用 into_ 函数 (输入输出都是占据数据的张量)，而不是使用视窗。

3. 我们需要存储三个分子轨道基 Cholesky 分解积分 $B_{ij,P}$、$B_{ia,P}$、$B_{ab,P}$。一个自然的疑问是，为什么我们不存储完整的分子轨道空间积分 $B_{pq,P}$？这会牺牲不是很多的内存，但可以减少程序用到的变量数量，简化程序编写。

   但考虑到后续的计算过程中，经常要作 reshape 维度更改 (譬如 $B_{ij,P}\equiv B_{i,jP}$)。而维度更改要求了数据的连续性 (该例子中，如果要求维度更改不会产生新的内存，$i$ 与 $j$ 指标未必是连续的，但 $j$ 和 $P$ 之间必须是连续的)。

   这意味着，如果存储为 row-major 连续的 $B_{ij,P}$，那么取 $B_{i,jP}$ 时是不会发生内存复制的：

   #![allow(unused)]
   fn main() {
   b_oo.reshape((nocc, nocc * naux)) // DataCow::Ref, data not cloned
   }

   如果存储为 $B_{pq,P}$、而在用到子张量 $B_{i,jP}$ 时使用索引加 reshape，是会发生内存复制的：

   #![allow(unused)]
   fn main() {
   b_all.i((so, so, ..)).reshape((nocc, nocc * naux)) // DataCow::Owned, data cloned
   }

   考虑到 $B_{ij,P}$ 到 $B_{i,jP}$ 这类 reshape 经常会用到，我们不能只存一个变量 $B_{pq,P}$ 而需要分开存三个，为运行效率对代码复杂度作妥协。

3.4 性能评价

• Cholesky 分解积分的浮点效率是 0.45 TFLOP/sec，性能利用率 40%，其性能是比较高的；
• 原子轨道到分子轨道变换的浮点效率是 0.16 TFLOP/sec，性能利用率 14%，该步骤性能有待改进；但考虑到 CCSD 中分子轨道变换占总计算消耗很少，因此仅针对当前任务，不是非常有必要对该过程进一步优化：
  • 上述计算过程没有利用电子积分的对称性，内存利用上有少许提升空间；
  • 上述代码的实现较为简洁，但诸如 cderi_svp 等中间张量其实是可以避免生成、从而避免较大的内存带宽消耗。原子轨道到分子轨道变换的性能有进一步提升的空间。

3. RCCSD 初猜与 RI-RMP2 能量

• 程序函数：get_riccsd_initial_guess
• 输出类型：RCCSDResults

CCSD 的初猜一般是 MP2 的激发振幅；在真正开始 CCSD 计算前，由于 MP2 能量计算并不额外消耗太多计算资源，因此 RI-RMP2 能量可以顺便给出。

3.1 输出类型 RCCSDResults

#![allow(unused)]
fn main() {
pub struct RCCSDResults {
    pub e_corr: f64,
    pub t1: Tsr,
    pub t2: Tsr,
}
}

上表中，$\Delta E^{\mathsf{SD}}$ 是指 CCSD 的相关能；只有与 Hartree-Fock 参考态能量求和，才得到完整的 CCSD 体系能量。

3.2 公式

首先，为计算 MP2 的二次激发张量 $t_{ij}^{ab}{}^{(2)}$，需要获得分子轨道下的 4c-2e 双电子积分 $g_{ij}^{ab}$，求得多体微扰分母 ${\Delta }_{ij}^{ab}$： $$\begin{array}{rl}{g}_{ij}^{ab}& =\sum _P{B}_{ia,P}{B}_{jb,P}\\ {\Delta }_{ij}^{ab}& ={\epsilon }_i+{\epsilon }_j-{\epsilon }_a-{\epsilon }_b\\ t_{ij}^{ab}{}^{\text{(2)}}& =g_{ij}^{ab}/{\Delta }_{ij}^{ab}\end{array}$$

随后可以计算 MP2 能量 (下式的 OS 表示 opposite-spin、SS 表示 same-spin，bi1 与 bi2 表示闭壳层 bi-orthogonal 下的两个分量)： $$\begin{array}{rl}\Delta E^{\text{bi1}}& =\sum _{ijab}{t}_{ij}^{ab}{}^{\text{(2)}}{g}_{ij}^{ab}\\ \Delta E^{\text{bi2}}& =\sum _{ijab}{t}_{ij}^{ab}{}^{\text{(2)}}{g}_{ij}^{ba}\\ \Delta E^{\text{(2)-OS}}& =\Delta E^{\text{bi1}}\\ \Delta E^{\text{(2)-SS}}& =\Delta E^{\text{bi1}}-\Delta E^{\text{bi2}}\\ \Delta E^{\text{(2)}}& =\Delta E^{\text{(2)-OS}}+\Delta E^{\text{(2)-SS}}\end{array}$$

3.3 程序

MP2 的激发张量与能量可以在 3-5 行公式内写出，意味着程序也可以在 3-5 行代码 (或 3-5 个 statement) 实现 (下述代码片段中张量的指标顺序是 $(i,a,j,b)$，与正文其他部分不同)：

#![allow(unused)]
fn main() {
let g2 = b_ov.reshape((nocc * nvir, -1)) % b_ov.reshape((nocc * nvir, -1)).t();
let g2 = g2.into_shape((nocc, nvir, nocc, nvir));
let d2 = mo_energy.i((so, None, None, None))
    - mo_energy.i((None, sv, None, None))
    + mo_energy.i((None, None, so, None))
    - mo_energy.i((None, None, None, sv));
let t2 = &g2 / &d2;
let e_corr_mp2 = ((2.0_f64 * &t2 - t2.swapaxes(1, 3)) * &g2).sum();
}

这段程序尽管非常直观短小，但计算量与内存消耗都较大。实际运行的程序是

#![allow(unused)]
fn main() {
let d_ov = mo_energy.i((so, None)) - mo_energy.i((None, sv));
let t1 = rt::zeros(([nocc, nvir], &device));
let t2 = rt::zeros(([nocc, nocc, nvir, nvir], &device));

let e_corr_oo = rt::zeros(([nocc, nocc], &device));
(0..nocc).into_par_iter().for_each(|i| {
    (0..i + 1).into_par_iter().for_each(|j| {
        let g2_ab = b_ov.i(i) % b_ov.i(j).t();
        let d2_ab = d_ov.i((i, .., None)) + d_ov.i((j, None, ..));
        let t2_ab = &g2_ab / &d2_ab;
        let mut t2 = unsafe { t2.force_mut() };
        t2.i_mut((i, j)).assign(&t2_ab);
        if i != j {
            t2.i_mut((j, i)).assign(&t2_ab.t());
        }
        let e_bi1 = (&t2_ab * &g2_ab).sum_all();
        let e_bi2 = (&t2_ab * &g2_ab.swapaxes(-1, -2)).sum_all();
        let e_corr_ij = 2.0 * e_bi1 - e_bi2;
        let mut e_corr_oo = unsafe { e_corr_oo.force_mut() };
        *e_corr_oo.index_mut([i, j]) = e_corr_ij;
        *e_corr_oo.index_mut([j, i]) = e_corr_ij;
    });
});

let e_corr = e_corr_oo.sum_all();
}

上述代码可以利用输出张量 $t_{ij}^{ab}=t_{ji}^{ba}$ 的对称性，从而只对 $i⩾j$ 的下三角情形作计算；而上三角的部分复制下三角的部分即可，减少一半计算量。

3.4 性能评价

• 该计算总浮点计算量大于 70.6 GFLOPs；
• 计算耗时约 180 msec，实际运行效率不低于 0.39 TFLOP/sec；
• CPU 性能利用率不少于 36%。

这个性能利用率可以接受但不算很高，说明如下：

• 在确定计算量时，我们没有统计入其他 $O(N^4)$ 的部分；
• 如果仅仅是计算 MP2 能量，那么 MP2 激发张量 $t_{ij}^{ab}{}^{\text{(2)}}$ 不需要写入到内存；但作为 CCSD 计算任务的初猜过程，这个较大张量的写入是必须的，从而有额外的内存读写的开销。这应当是主要的效率影响因素。若不需要在内存写入 $t_{ij}^{ab}{}^{\text{(2)}}$ 张量，则 CPU 性能利用率可能达到 50%。
• 在计算浮点计算量 FLOPs 数值时，我们只考虑了 4c-2e 双电子积分 $g_{ij}^{ab}$ 的部分，而没有考虑 MP2 能量计算；因此上面的 CPU 性能利用率是被低估的。同时，像 &t2_ab * &g2_ab 与 &t2_ab * &g2_ab.swapaxes(-1, -2) 这类运算会开辟新的内存空间；但我们的目标其实是求和得到能量，这里开辟出新的空间是浪费内存带宽的。因此这里的 MP2 能量计算性能有进一步提升的空间，但提升空间并不大 (可能到 60%)。

4. RI-RCCSD 中间变量：第一部分

• 程序函数：get_riccsd_intermediates_1

4.1 RI-RCCSD 中间变量

RI-RCCSD 有许多中间变量；在一次 CCSD 迭代中，这些中间变量会被多次使用到，因此需要找内存空间存储起来。

从程序实现角度，RI-RCCSD 的中间变量分为两部分：一部分是计算 CCSD 相关能时需要用到的中间量，另一部分是只在更新 $t_i^a$、$t_{ij}^{ab}$ 时才会用的中间量。

中间变量的命名规则是

• 中间变量字母记为 $M$；
• 数字 1、2 分别代表中间变量中包含的激发次数；
• 如果有多种意义不同、但包含激发次数相同的中间变量，则以 a, b 等进行区分；
• RI-RCCSD 用到的中间变量大多是 3-D 张量，但也有其他维度张量。变量最后的字母 o 表示占据轨道维度，v 表示非占轨道维度，j 表示辅助基维度。

4.2 公式

$$\begin{array}{rl}{M}_P^{\text{1}}& =\sum _{ia}{t}_i^a{B}_{ia,P}\\ M_{ij,P}^{\text{1}}& =\sum _a{t}_i^a{B}_{ja,P}\\ M_{ia,P}^{\text{2a}}& =\sum _{jb}(2t_{ij}^{ab}-t_{ij}^{ba})B_{jb,P}=\sum _{jb}{\tilde{t}}_{ij}^{ab}{B}_{jb,P}\end{array}$$

需要留意，$M_{ij,P}^{\text{1}}$ 不具有占据轨道上的对称性，即一般来说 $M_{ij,P}^{\text{1}}\ne M_{ji,P}^{\text{1}}$。

4.3 程序

#![allow(unused)]
fn main() {
let m1_j = t1.reshape(-1) % b_ov.reshape((-1, naux));

let m1_oo: Tsr = rt::zeros(([nocc, nocc, naux], &device));
(0..nocc).into_par_iter().for_each(|j| {
    let mut m1_oo = unsafe { m1_oo.force_mut() };
    *&mut m1_oo.i_mut((.., j)) += &t1 % &b_ov.i(j);
});

let m2a_ov: Tsr = rt::zeros(([nocc, nvir, naux], &device));
(0..nocc).into_par_iter().for_each(|i| {
    let mut m2a_ov = unsafe { m2a_ov.force_mut() };
    let scr_jba: Tsr = 2 * t2.i((.., i)) - t2.i(i);
    *&mut m2a_ov.i_mut(i) += scr_jba.reshape((-1, nvir)).t() % b_ov.reshape((-1, naux));
});
}

4.4 性能上的考量

• 矩阵乘法的实现：上述计算都使用 % 即普通的矩阵乘法得到新的张量后，通过 *&mut var += new_mat 进行赋值。这种赋值语句可能比较直观，但由于产生额外的内存开销，性能上会比 var.matmul_from(a, b, alpha, beta) 要低一些。

• 特定情景下不使用 broadcast 矩阵乘法：计算 m1_oo $M_{ij,P}^{\text{1}}$ 时，对指标 $j$ 作手动并行。这里使用 broadcast 矩阵乘法会导致内存连续性问题，不适合用一行代码直接实现该矩阵乘法。

• 避免大量中间内存的考量：计算 m2a_ov $M_{ia,P}^{\text{2a}}$ 时，对指标 $i$ 作手动并行。这里并行的目的是避免直接生成 4-D 张量 $t_{ij}^{ab}=2t_{ij}^{ab}-t_{ij}^{ba}$，而是在每个线程中生成一个 3-D 的、固定了指标 $i$ 数值的张量 $t_{ij}^{ab}$，较大的内存开销。

  • 尽管在 RI-RCCSD 中间张量计算中，我们可以通过这种方式避免大量中间内存；但在后续计算过程中，若要使得效率最大化，我们似乎不太能避免产生一些新的 $n_{\mathrm{occ}}^2{n}_{\mathrm{vir}}^2$ 浮点数大小的中间张量。

• 为性能考量而更改公式表达：计算 m2a_ov $M_{ia,P}^{\text{2a}}$ 时，程序与公式的结果等价，但具体的做法上有差异。

  • 更改中间张量指标顺序：我们注意到 ${\tilde{t}}_{ij}^{ab}$ 在固定特定的指标 $i$ 时，正常情况下应该是 $(j,a,b)$ 的 3-D 张量：

    #![allow(unused)]
    fn main() {
    let scr_jab: Tsr = 2 * t2.i(i) - t2.i(i).swapaxes(2, 3);
    }

    但考虑到后续要将其与 $B_{jb,P}$ 作张量缩并，其被缩并的指标是 $(j,b)$，那么将张量缩并归约到矩阵乘法时，应该要是 ${\sum }_{jb}{\text{T}}_{jb,a}{B}_{jb,P}$ 或 ${\sum }_{jb}{\text{T}}_{jb,a}{B}_{jb,P}$；但这要求的是一个 ${\text{T}}_{jba}$ 即 $(j,b,a)$ 或 ${\text{T}}_{ajb}$ 即 $(a,j,b)$ 的 3-D 张量。无论是 ${\text{T}}_{jba}$ 还是 ${\text{T}}_{ajb}$，这都与 ${\tilde{t}}_{ij}^{ab}$ 通常存储的顺序 $(j,a,b)$ 不一致。

    既然 ${\tilde{t}}_{ij}^{ab}$ 只是一个中间张量，我们不见得就要严格按照既定的公式与程序约定俗成来实现该变量。在当前程序实现里，我们直接实现转置了指标的 $(a,b)$ 张量 ${\text{T}}_{jba}$：

    $${\text{T}}_{jba}=2t_{ij}^{ab}-t_{ij}^{ba}\text{(index i selected)}$$

    #![allow(unused)]
    fn main() {
    let scr_jba: Tsr = 2 * t2.i(i).swapaxes(2, 3) - t2.i(i);
    }

  • 考虑 Elementwise 操作的连续性：我们还注意到上述实现与实际的程序实现并不相同：

    #![allow(unused)]
    fn main() {
    // let scr_jba: Tsr = 2 * t2.swapaxes(0, 1).i(i) - t2.i(i); // equilvant code to below
    let scr_jba: Tsr = 2 * t2.i((.., i)) - t2.i(i);
    }

    实际程序中使用到的公式是

    $${\text{T}}_{jba}=2t_{ji}^{ba}-t_{ij}^{ba}\text{(index i selected)}$$

    这是里用到了 $t_{ij}^{ab}=t_{ji}^{ba}$ 的二重对称性。注意到在 row-major 下，$t_{ji}^{ba}$ 与 $t_{ij}^{ba}$ 在内存中，共同最大连续维度是指标 $(b,a)$ 对应的维度，其长度是 $n_{\mathrm{vir}}^2$；处理这两个张量之间的 elementwise 运算时，由于内存的连续性，其性能较高。而 $t_{ij}^{ab}$ 与 $t_{ij}^{ba}$ 在内存中，没有共同的连续维度；处理这两个张量之间的 elementwise 运算，性能会比较低。

4.4 性能评价

• 该计算总浮点计算量约 142 GFLOPs；
• 计算耗时约 320 msec，实际运行效率为 0.44 TFLOP/sec；
• CPU 性能利用率为 40%。

5. RI-RCCSD 能量计算

• 程序函数：get_riccsd_energy

在有了一些中间变量后，我们可以计算 RI-RCCSD 能量：

$$\Delta E^{\text{SD}}=2\sum _P{M}_P^{\text{1}}{M}_P^{\text{1}}-\sum _{ijP}{M}_{ij,P}^{\text{1}}{M}_{ji,P}^{\text{1}}+\sum _{iaP}{B}_{ia,P}{M}_{ia,P}^{\text{2a}}$$

在合理的 reshape 下，程序实现也比较直观：

#![allow(unused)]
fn main() {
let e_t1_j = 2.0_f64 * (m1_j % m1_j);
let e_t1_k = -(m1_oo.reshape(-1) % m1_oo.swapaxes(0, 1).reshape(-1));
let e_t2 = b_ov.reshape(-1) % m2a_ov.reshape(-1);
let e_corr = e_t1_j + e_t1_k + e_t2;
e_corr.to_scalar()
}

需要留意，在 RSTSR 下，向量-向量点乘得到的结果仍然是张量，需要通过 to_scalar 函数返回一个 f64 浮点数。

该计算耗时约为 5 msec，相比于其他过程完全是可以忽略的，没有必要讨论其计算性能。

6. RI-RCCSD 中间变量：第二部分

• 程序函数：get_riccsd_intermediates_2

6.1 公式

第二部分中间变量不会在能量计算中用到，但会在更新激发张量 $t_i^a$ 与 $t_{ij}^{ab}$ 时会用到。

$$\begin{array}{rl}{M}_{ia,P}^{\text{1a}}& =\sum _j{t}_j^a{B}_{ij,P}\\ M_{ia,P}^{\text{1b}}& =\sum _b{t}_i^b{B}_{ab,P}\\ M_{ab,P}^{\text{1}}& =\sum _i{t}_i^a{B}_{ib,P}\\ M_{ia,P}^{\text{2b}}& =\sum _{jb}{t}_j^a{t}_i^b{B}_{jb,P}=\sum _j{t}_j^a{M}_{ij,P}^{\text{1}}\end{array}$$