1. 简单理解 PySCF 临时文件 chkfile 使用

这份文档会介绍 PySCF 的 chkfile 功能。

chkfile 类似于 Gaussian 的 checkpoint 文件。对于自洽场方法，它用于储存分子信息、自洽场轨道等等。

PySCF 是通过 Python 的 h5py 实现；换言之，PySCF 的 chkfile 相当于 h5py 的高级接口。

1.1. 初始化

from pyscf import gto, scf, hessian, lib
import numpy as np
import h5py

np.set_printoptions(5, suppress=True, linewidth=150)

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

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

scf_eng = scf.RHF(mol).run()
scf_hess = hessian.RHF(scf_eng).run()

1.2. 介绍

1.2.1. h5py 文件与内容结构

临时文件的位置如下：

scf_eng.chkfile

'/home/a/Documents/2021-03-01-MP2_pyscf_prop/tmp94oic47h'

该文件可以使用 h5py 直接进行读写。我们先了解文件目录结构。

该文件的顶层目录是

with h5py.File(scf_eng.chkfile, "r") as f:
    print(f.keys())

<KeysViewHDF5 ['mol', 'scf', 'scf_f1ao', 'scf_mo1']>

对于 scf 键值，其子目录是

with h5py.File(scf_eng.chkfile, "r") as f:
    print(f["scf"].keys())

<KeysViewHDF5 ['e_tot', 'mo_coeff', 'mo_energy', 'mo_occ']>

事实上，我们可以用 scf/e_tot 直接获得能量的结果；但需要注意，在最后需要加 [()] 以获得结果，否则得到的是 h5py.Dataset 类实例：

with h5py.File(scf_eng.chkfile, "r") as f:
    print(f["scf/e_tot"][()])

-56.02979155465495

最后，我们可以使用下面的小程序，打印出完整的 chkfile 文件结构：

def print_h5py_group_dir(f):
    spl = f.name.split("/")
    level = len(spl) - 1 if len(f.name) > 1 else 0
    name = spl[-1]
    if isinstance(f, h5py.Group):
        if not isinstance(f, (h5py.File)):
            print("  "*(level-1) + "|- " + name)
        for k in f.keys():
            print_h5py_group_dir(f[k])
    else:
        val = f[()]
        if not isinstance(val, np.ndarray):
            print("  "*(level-1) + "|- " + name + ": " + str(type(val)))
        else:
            print("  "*(level-1) + "|- " + name + ": " + str(type(val)) + ", dtype: " + str(val.dtype))

with h5py.File(scf_eng.chkfile, "r") as f:
    print_h5py_group_dir(f)

|- mol: <class 'bytes'>
|- scf
  |- e_tot: <class 'numpy.float64'>
  |- mo_coeff: <class 'numpy.ndarray'>, dtype: float64
  |- mo_energy: <class 'numpy.ndarray'>, dtype: float64
  |- mo_occ: <class 'numpy.ndarray'>, dtype: float64
|- scf_f1ao
  |- 0: <class 'numpy.ndarray'>, dtype: float64
  |- 1: <class 'numpy.ndarray'>, dtype: float64
  |- 2: <class 'numpy.ndarray'>, dtype: float64
  |- 3: <class 'numpy.ndarray'>, dtype: float64
|- scf_mo1
  |- 0: <class 'numpy.ndarray'>, dtype: float64
  |- 1: <class 'numpy.ndarray'>, dtype: float64
  |- 2: <class 'numpy.ndarray'>, dtype: float64
  |- 3: <class 'numpy.ndarray'>, dtype: float64

1.2.2. 读取数据

PySCF 具有自己的读取 chkfile 的方式。作为高级 API，它确实更方便一些。

读取特定的数组或结果，可以直接用下述代码实现：

lib.chkfile.load(scf_eng.chkfile, "scf/mo_coeff").shape

(15, 15)

如果读取的是某个目录结构，那么它会将所有的子目录或子数据结果递归地转换为字典，储存到内存：

lib.chkfile.load(scf_eng.chkfile, "scf").keys()

dict_keys(['e_tot', 'mo_coeff', 'mo_energy', 'mo_occ'])

1.2.3. 储存数据

PySCF 支持三种储存的方法，单独的结果、列表与字典。

val_float = 0.4
val_list = [1.4, 1.8]
val_dict = {"foo": 2.4, "bar": 2.8}

lib.chkfile.dump(scf_eng.chkfile, "val_float", val_float)
lib.chkfile.dump(scf_eng.chkfile, "val_list", val_list)
lib.chkfile.dump(scf_eng.chkfile, "val_dict", val_dict)

其数据结构如下：

with h5py.File(scf_eng.chkfile, "r") as f:
    print_h5py_group_dir(f)

|- mol: <class 'bytes'>
|- scf
  |- e_tot: <class 'numpy.float64'>
  |- mo_coeff: <class 'numpy.ndarray'>, dtype: float64
  |- mo_energy: <class 'numpy.ndarray'>, dtype: float64
  |- mo_occ: <class 'numpy.ndarray'>, dtype: float64
|- scf_f1ao
  |- 0: <class 'numpy.ndarray'>, dtype: float64
  |- 1: <class 'numpy.ndarray'>, dtype: float64
  |- 2: <class 'numpy.ndarray'>, dtype: float64
  |- 3: <class 'numpy.ndarray'>, dtype: float64
|- scf_mo1
  |- 0: <class 'numpy.ndarray'>, dtype: float64
  |- 1: <class 'numpy.ndarray'>, dtype: float64
  |- 2: <class 'numpy.ndarray'>, dtype: float64
  |- 3: <class 'numpy.ndarray'>, dtype: float64
|- val_dict
  |- bar: <class 'numpy.float64'>
  |- foo: <class 'numpy.float64'>
|- val_float: <class 'numpy.float64'>
|- val_list__from_list__
  |- 000000: <class 'numpy.float64'>
  |- 000001: <class 'numpy.float64'>

其中，列表的情况是比较特殊的。在 PySCF 中，列表名称以 __from_list__ 结尾。其读取也会生成列表而非字典。

lib.chkfile.load(scf_eng.chkfile, "val_list")

[1.4, 1.8]

储存数据的值是可以覆盖的。譬如更改 val_float 的数值：

lib.chkfile.dump(scf_eng.chkfile, "val_float", 10.5)
lib.chkfile.load(scf_eng.chkfile, "val_float")

10.5

1.2.4. 直接对 File 对象写入数据

作为一种特殊情况，我们现在拥有的不是 chkfile 的名称 (譬如这里的 scf_eng.chkfile)，而是一个可以读写的实际文件：

f = h5py.File(scf_eng.chkfile, "a")

此时可以用下述方式直接进行写入操作：

f.create_dataset("grp1/grp2", data="val3")

<HDF5 dataset "grp2": shape (), type "|O">

f["grp1/grp2"][()]

b'val3'

这种情况在使用 lib.H5TmpFile() 生成的 hdf5 文件时比较实用。