REST-DIIS 使用说明

1. 安装说明

该程序是库，没有主程序。若要在其他项目中使用该库，直接在其他库的 Cargo.toml 的 dependencies 中引入该库即可。

目前该程序没有发布到 crates.io 的计划。请将该库下载到本地。

该库不需要额外的编译；但在使用时，请作以下准备工作：

1. 准备 hdf5。这是 cargo crate hdf5-metno 的要求；在 Ubuntu 系统中请执行

   sudo apt install libhdf5-dev
   

   对于不适合 apt install 的操作系统或集群环境，需要自行尝试解决。

2. 准备 OpenBLAS。我们要求使用编译了 CBLAS、LAPACKE 的并行 OpenBLAS 后端 (pthread 或 OpenMP 均可)。并且您需要在项目中实现下述两者其一 (目的是将 libopenblas.so 和 OpenMP runtime 链入您的程序)：

   • build.rs 中引入如下语句：

     #![allow(unused)]
     fn main() {
     println!("cargo:rustc-link-lib=openblas");
     println!("cargo:rustc-link-lib=gomp"); // or other valid omp runtime library
     }

   • RUSTFLAGS = -l openblas -l gomp

   这是 RSTSR 本身还有待发展导致的情况

   必须要使用 OpenBLAS，是因为 RSTSR 目前还没有实现 Faer 后端的线性代数接口、或者 OpenBLAS 以外 BLAS 后端。

   DIIS 需要用到线性代数功能 (求本征值)，但计算性能需求并不大，所以理想的情况是不需要借助任何 FFI 后端实现 DIIS，即未来我们希望默认使用 Faer 后端。

2. 使用说明

2.1 HDF5 临时文件及其路径

本项目的 DIISSemiIncore 将会在 DIIS 外推过程中使用 HDF5 硬盘存储空间。该空间是临时存储，一般来说用户是不需要获得的 (以后也许我们会实现 DIIS 文件存读，但目前这不是亟需实现的功能)。

DIIS 的临时文件路径由环境变量 REST_TMPDIR 或 TMPDIR 指定。REST_TMPDIR 优先级更高。

对于临时文件的声明周期问题：

• 如果程序正常退出，那么 DIIS 临时文件会自然地删除。
• 如果程序被用户终止或崩溃，DIIS 临时文件很有可能会存留在硬盘中。用户需要手动删除这些文件。

2.2 程序使用大致说明

下述不完整的代码是程序演示。该 DIIS 使用 DIISIncore，即所有存储均在内存进行。

若想查看一个可以工作的完整代码，我们也在一个 简单的玩具问题示例 中展示了 DIIS 的应用 (参考 后一节)。

#![allow(unused)]
fn main() {
// 简化的张量类型记号
type Tsr<T> = Tensor<T, DeviceOpenBLAS, IxD>;

// 默认的 DIIS 选项
let diis_flags = rstsr_diis::DIISIncoreFlagsBuilder::default().build().unwrap();
// 生成 DIIS 实例
let mut diis_driver = rstsr_diis::DIISIncore::<f64>::new(diis_flags, &device);

// 初猜向量 (譬如 SCF 下 Fock 矩阵 F_uv；CCSD 下的振幅 t1、t2)
let vec_init: Tsr<f64> = ...;
// 误差向量 (譬如某种定义下的 F_ai)
// 如果误差不好定义可以直接传 None (譬如 CCSD 的情形)
let err_init: Option<Tsr<f64>> = ...;
// 向 DIIS 实例插入初猜向量
// 参数表：
// - 初猜 (或待迭代) 向量
// - 初猜 (或待迭代) 误差 (有最好，没有可以传 None)
// - 迭代步数标记 (这里没有，建议传 None)
diis_driver.update(vec_init, err_init, None)

// 实际迭代
for niter in 0..max_iter {
    ...
    // 更新的待外推向量
    let vec_new: Tsr<f64> = ...;
    // 当前向量的误差
    let err_new: Option<Tsr<f64>> = ...;
    // 向 DIIS 实例插入待外推向量，并得到外推后的向量
    let vec_extrapolated = diis_driver.update(vec_new, err_new, None)
    ...
}
}

2.3 程序功能

• 同时支持全内存 (incore) 与半内存 (semi-incore) 功能。其中，半内存是指输入的、以及 DIIS 外推输出的向量，都是完全存于内存的；但 DIIS 内部的存储空间使用了硬盘作为媒介。

  • 全内存的 DIIS 类型是 DIISIncore，设置参数的构造程序类型是 DIISIncoreFlagsBuilder。
  • 半内存的 DIIS 类型是 DIISSemiIncore，设置参数的构造程序类型是 DIISSemiIncoreFlagsBuilder。
  • 目前我们还没有实现全磁盘 (outcore) 的功能。
  • 由于目前数学库 RSTSR 还不支持 GPU 等设备，因此目前还不具备异构功能。

• 重要的设置参数 (flags)：

  • 不同的 DIIS 示例具有不同的设置参数类型，这些参数类型目前还互不相通。
  • space：DIIS 空间大小。该设置一方面并非越大越好，另一方面该数值越大、DIIS 内部存储的向量也会越多，对内存或磁盘的需求会越紧张。特别是对 CCSD 等参数量较大的问题，不建议该值设置地太大。
  • min_space：DIIS 初始外推空间大小。该数值需要不小于 1。
  • pop_strategy：DIIS 弹出策略。DIIS 可以看作是一种队列 (queue)，其队列大小是 space 参数所决定的。一般来说 DIIS 是作为双向队列 (dqueue) 使用的，即先入先出 (DIISPopStrategy::Iteration)；但既然 DIIS 同时还需要计算输入向量的误差，那么我们也可以每次弹出误差最大的向量 (DIISPopStrategy::ErrDiagonal)。但由于如果最后一次输入的向量误差较大 (这种情况一般来说是收敛不太成功的情况)，弹出该向量容易导致外推出来的向量总是有较大误差；因而这种情形下，我们会退回到先入先出的策略上。
  • chunk：对于较长的向量，其计算是分批次进行的。该值对走硬盘的 DIIS 更为关键。需要注意，该值的设置并非越大越好，特别是硬盘占用空间小于磁盘的缓冲区域的情景。我们认为一个 chunk 的大小小于 1 MB (128 k FP64) 是比较合适的。
  • scratch_path：DIIS 缓存文件的文件夹路径。该选项仅对走硬盘的 DIIS 提供。

• Cargo feature

  • sequential_io：DIIS 计算误差重叠矩阵与外推时，不使用异步 I/O。我们在实现 DIIS 时，默认是使用异步 I/O 的，因此该选项一般不建议开启。

2.4 Undefined Behavior 注意事项

我们在使用 DIIS 时，一般仅使用其 diis.update(vec, err, iter) -> extrapolated_vec 函数。在使用该函数时，有一些情况我们不会仔细检查，有可能会导致 UB。这需要用户自己保证，包括：

1. 每次传入的 vec, err 的向量最好是一维 (动态维度 IxD) 向量，且必须保证每次传入的向量长度一致。
2. 如果你的计算中，有任何一次 err 无法给出 (即需要传 None)，那么请自始至终都在 err 参数中传入 None。err 传入有效值或 None 在整个程序中是两种处理逻辑。
3. 一般来说 iter 传入 None 即可，除非有特殊需求。DIIS 是双向队列时，具体要弹出哪个向量是通过迭代步数给出的；如果用户传入不合理的 iter，那么弹出的过程可能会出现问题。

玩具模型上的示例

该文档是针对示例文件 toy_model.rs 的补充说明。

1. 问题描述

该问题是一个简单的 4x4 矩阵求解问题：

$$\mathbf{A}\mathbf{x}=\mathbf{b}$$

一般来说，该问题是 $O(N^3)$ 复杂度的，因为该问题需要先对 $\mathbf{A}$ 分解为上三角、对角、下三角矩阵的乘积 (Cholesky 若对称半正定，SVD 或 QR 若非对称，等等)，将问题化为若干 $O(N^2)$ 三角矩阵求解问题；但矩阵的分解本身是 $O(N^3)$ 的。

现在，我们注意到 $\mathbf{A}$ 矩阵本身具有特征：其矩阵数值上，对角矩阵元相比于非对角元大许多。譬如示例中用到的矩阵

#![allow(unused)]
fn main() {
[
    [1.0, 0.1, 0.3, 0.2],
    [0.1, 1.5, 0.2, 0.1],
    [0.3, 0.2, 1.8, 0.2],
    [0.2, 0.1, 0.2, 1.3],
]
}

在这种情况下，我们记矩阵 $\mathbf{D}=\mathrm{diag}(\mathbf{A})\simeq \mathbf{A}$ 是矩阵 $\mathbf{A}$ 的对角元构成的矩阵；其对角值向量是 $\mathbf{d}$。

这种具有对角值大特性的矩阵，是否有办法避免 $O(N^3)$ 复杂度的计算？DIIS 提供一种 (不保证一定求解成功，但在计算化学中屡试不爽) 的策略。

2. 简单迭代策略

在具体应用中，DIIS 迭代需要以简单迭代算法为基础。对于当前问题，简单迭代算法设计如下。

首先，初猜向量可以通过下述表达式给出： $$\mathbf{b}=\mathbf{A}\mathbf{x}\simeq \mathbf{D}\mathbf{x}\Rightarrow \mathbf{x}\simeq {\mathbf{D}}^{-1}\mathbf{b}=\mathbf{b}÷\mathbf{d}$$ 这里的除号是依元素除以的意思。因此，我们就以 ${\mathbf{x}}_0={\mathbf{D}}^{-1}\mathbf{b}$ 作为初猜。

其次，对于迭代过程中的向量 ${\mathbf{x}}_n$，为了得到下一次的向量 ${\mathbf{x}}_{n+1}$，我们注意并联立下述两式： $$\begin{array}{rl}{\mathbf{b}}_n& =\mathbf{A}{\mathbf{x}}_n\\ \mathbf{b}& =\mathbf{A}\mathbf{x}\end{array}$$ 得到 $$\mathbf{x}={\mathbf{A}}^{-1}(\mathbf{b}-{\mathbf{b}}_n)+{\mathbf{x}}_n\simeq {\mathbf{D}}^{-1}(\mathbf{b}-{\mathbf{b}}_n)+{\mathbf{x}}_n=(\mathbf{b}-{\mathbf{b}}_n)÷\mathbf{d}+{\mathbf{x}}_n$$ 我们就将该近似式定义为 ${\mathbf{x}}_{n+1}$。这样就得到了简单迭代策略。

3. 简单迭代程序实现

程序实现如下。

#![allow(unused)]
fn main() {
let f = |x: &Tsr<f64>| &a % x - &b;

// the initial guess
let mut x = &b / &d;
let mut x0;

// the number of iterations and the tolerance
let mut niter = 0;
let maxiter = 20;
let tol = 1e-7;

while niter < maxiter && f(&x).l2_norm_all() > tol {
    niter += 1;
    x0 = x;
    let b0 = &a % &x0;
    x = (&b - b0) / &d + &x0;
}
}

def f(x):
    return a @ x - b

# the initial guess
x = b / d
x0 = np.zeros_like(x)

# the number of iterations and the tolerance
niter = 0
maxiter = 20
tol = 1e-7

while np.linalg.norm(f(x)) > tol and niter < maxiter:
    niter += 1
    x0, x = x, x0
    b0 = a @ x0
    x = (b - b0) / d + x0

最终实际迭代次数是 19 次。

4. DIIS 实现

DIIS 相对于简单迭代情景，所需要增加的代码很少：

• 初始化 DIIS 实例需要 1--2 行；
• 将初猜提供给 DIIS 需要 1 行；
• 迭代过程中 DIIS 外推更新需要 1 行。

#![allow(unused)]
fn main() {
let f = |x: &Tsr<f64>| &a % x - &b;

// the initial guess
let mut x = &b / &d;
let mut x0;

// DIIS incore driver
let diis_flags = DIISIncoreFlagsBuilder::default().build().unwrap(); // <==
let mut diis = DIISIncore::<f64>::new(diis_flags, &device);          // <==
// initial guess to DIIS
diis.update(x.to_owned(), None, None);                               // <==

// the number of iterations and the tolerance
let mut niter = 0;
let maxiter = 20;
let tol = 1e-7;

while niter < maxiter && f(&x).l2_norm_all() > tol {
    niter += 1;
    x0 = x;
    let b0 = &a % &x0;
    x = (&b - b0) / &d + &x0;
    x = diis.update(x, None, None);                                  // <==
}
}

def f(x):
    return a @ x - b

# the initial guess
x = b / d
x0 = np.zeros_like(x)

# DIIS driver
diis = pyscf.lib.diis.DIIS()                                         ## <==
diis.update(x)                                                       ## <==

# the number of iterations and the tolerance
niter = 0
maxiter = 20
tol = 1e-7

while np.linalg.norm(f(x)) > tol and niter < maxiter:
    niter += 1
    x0, x = x, x0
    b0 = a @ x0
    x = (b - b0) / d + x0
    x = diis.update(x)                                               ## <==

经过 DIIS 加速，该问题可以在 5 次迭代中完成。

5. 同时允许 semi-incore (disk I/O) 或 fully incore 实现的 DIIS

在实际的化学问题中，我们可能需要同时处理 semi-incore 与 fully incore 两种 DIIS：在大量计算的较小体系的情景下，我们希望用 fully incore DIIS 以增加效率；在计算较大体系时，则使用 semi-incore DIIS 以保证内存空间充足。

在这种情形下，DIISAPI 作为 dyn-compatible trait 提供了一种方式，以使得 semi-incore 与 fully incore 能同时处理：

#![allow(unused)]
fn main() {
// DIIS incore driver
let semi_incore = x.size() < 3; // assuming x > 3 requires semi-incore (disk I/O)
let mut diis: Box<dyn DIISAPI<Tsr<f64>>> = match semi_incore {
    false => Box::new({
        let diis_flags = DIISIncoreFlagsBuilder::default().build().unwrap();
        DIISIncore::<f64>::new(diis_flags, &device)
    }),
    true => Box::new({
        let diis_flags = DIISSemiIncoreFlagsBuilder::default().build().unwrap();
        DIISSemiIncore::<f64>::new(diis_flags, &device)
    }),
};
// initial guess to DIIS
diis.update(x.to_owned(), None, None);
}

所涉及到的编程实践

1. 异步 I/O

RSTSR-DIIS 在处理 semi-incore 情景时，使用标准库实现异步 I/O (以 semi_incore_inner_dot 函数为例)。

首先，由于我们用到的变量通常有生命周期，因此不能直接使用 std::thread::spawn 函数开辟新的线程。其解决方案是，需要使用 std::thread::scope 先圈住声明周期，随后进行双线程。

以下述伪代码为例。假设我们要执行 func_a 与 func_b (两者是 FnOnce，不返回任何结果)

#![allow(unused)]
fn main() {
for i in (0..niter) {
    func_a();
    func_b();
}
}

实际情景中，func_b 会依赖于 func_a，即不能乱序执行；但 func_b 不依赖于 func_a。若现在要进行异步问题，我们可以这么操作：

#![allow(unused)]
fn main() {
std::thread::scope(|scope| {  // 圈住生命周期
    let mut task = scope.spawn(|| {});
    for i in (0..niter) {
        func_a();             // 此时 i 步的 func_a 与 i-1 步的 func_b 同时进行
        task.join().unwrap(); // 相当于 barrier，保证现在执行 func_b 时已经完成 func_a
        task = scope.spawn(move || {
            func_b();         // func_b 并不在 master thread 上生成，而是平行进行的
        })
    }
})
}

很可能我们在 func_b 时需要用到一些可变引用；碰到这种情况会比较麻烦，需要在 spawn 前传入不可变引用，然后在 spawn 内部用一些 unsafe 技巧弄成可变的。这其实有点像 OpenMP 并行时，一些变量是 shared 即所有线程可见的；这在 Rust 中其实是 unsafe 的，但这样的程序非常好写，而且大多数时候不会真的造成 race，那我的看法是不如就大方地用 unsafe (= trust me)。

同时，我们指出上面的做法是 post-process 的异步。另一种异步策略是 prefetch：

# 伪代码
func_a(0)
for i in range(1, niter)
    barrier()
    if i < niter - 1:
        task = spawn(func_a(i + 1))  # prefetch
    func_b(i)

但 prefetch 的策略我感觉不是非常好写，因为会对第一次循环读取需要在循环外进行、会增加一些判断语句去确定是否有下一轮循环。

被异步到分支线程的任务不一定必须是 I/O，也可以是计算过程本身。