FermiNet：费米子神经网络

FermiNet 是一种基于变分蒙特卡洛方法的神经网络，用于学习原子和分子的高精度基态波函数。

本仓库包含了首次发表于《使用深度神经网络求解多电子薛定谔方程的从头计算法》（David Pfau、James S. Spencer、Alex G de G Matthews 和 W.M.C. Foulkes，《物理评论·研究》2, 033429 (2020)）中的算法实现及实验，并涵盖了后续的研究与进展。

警告：这是一个 JAX 实现的研究级版本，目前仍在积极开发中。原始的 TensorFlow 实现可在 tf 分支中找到。

安装

运行 pip install -e . 将安装所有必需的依赖项。建议在 虚拟环境 中进行此操作。

virtualenv ~/venv/ferminet
source ~/venv/ferminet/bin/activate
pip install -e .

如果您有 GPU 可用（强烈推荐以加快训练速度），则可以安装支持 CUDA 的 JAX，例如：

pip install --upgrade jax jaxlib==0.1.57+cuda110 -f
https://storage.googleapis.com/jax-releases/jax_releases.html

请注意，jaxlib 的版本必须与您希望使用的现有 CUDA 安装相匹配。更多详细信息请参阅 JAX 文档。

测试最简单的方法是使用 pytest：

pip install -e '.[testing]'
python -m pytest

使用

ferminet 使用来自 ml_collections 的 ConfigDict 来配置系统。ferminet/configs/ 目录下包含了一些示例脚本，这些脚本主要用于测试，因此在生产级计算中可能需要额外的设置。

ferminet --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

或者

python3 ferminet/main.py --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

将训练 FermiNet 以寻找锂原子的基态波函数，使用 1024 个 MCMC 配置（在变分蒙特卡洛语言中称为“游走者”）作为批次大小，并进行 100 次预训练迭代（对于如此小的系统，默认的 1000 次预训练显然过于冗余）。可以通过修改配置文件或（对于一次性更改，更推荐的方式）使用命令行标志来控制系统和超参数。有关标志语法的更多信息，请参阅 ml_collections 的文档。所有可用配置选项的详细信息请参见 ferminet/base_config.py。

其他系统也可以轻松设置，只需创建一个新的配置文件或使用 ferminet 工具，或者编写自定义的训练脚本。例如，要对 H2 分子进行计算，您可以创建一个包含以下内容的配置文件：

from ferminet import base_config
from ferminet.utils import system

# 配置文件中的设置通过执行 get_config 函数加载。
def get_config():
  # 获取默认选项。
  cfg = base_config.default()
  # 设置分子
  cfg.system.electrons = (1,1)
  cfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]

  # 设置训练超参数
  cfg.batch_size = 256
  cfg.pretrain.iterations = 100

  return cfg

然后使用以下命令运行：

ferminet --config /path/to/h2_config.py

或者等效地编写如下脚本（或交互式执行）：

import sys

from absl import logging
from ferminet.utils import system
from ferminet import base_config
from ferminet import train

# 可选，用于将训练进度打印到标准输出。
# 如果运行脚本，也可以直接使用 --alsologtostderr 标志。
logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)

# 定义 H2 分子
cfg = base_config.default()
cfg.system.electrons = (1,1)  # (α电子，β电子)
cfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]

# 设置训练参数
cfg.batch_size = 256
cfg.pretrain.iterations = 100

train.train(cfg)

此外，您还可以直接传入一个 PySCF 的 'Molecule' 对象。可以使用以下代码创建 PySCF 分子：

from pyscf import gto
mol = gto.Mole()
mol.build(
    atom = 'H  0 0 1; H 0 0 -1',
    basis = 'sto-3g', unit='bohr')

一旦有了这个分子，就可以将其直接传递给配置对象，运行如下代码：

from ferminet import base_config
from ferminet import train

# 添加 H2 分子
cfg = base_config.default()
cfg.system.pyscf_mol = mol

# 设置训练参数
cfg.batch_size = 256
cfg.pretrain.iterations = 100

train.train(cfg)

注意：对于较大原子和分子的训练，尤其是当批次大小较大时，多 GPU 并行化至关重要。这可通过 JAX 的 pmap 功能实现。如果有多个 GPU 可用，它们将被自动检测并使用。

推断

训练完成后，在参数固定的情况下，对能量和其他可观测量进行多次时间步长的计算，有助于积累低方差的物理量估计值。为此，只需使用与训练相同的命令，但添加 --config.optim.optimizer 'none' 标志即可。请确保 cfg.log.save_path 的值保持不变，或者将 cfg.log.restore_path 设置为原始训练运行时的 cfg.log.save_path 值。

此外，在推断时收集原始训练中未包含的可观测量统计数据也很有用。如果配置文件中未启用，可以通过在命令行中添加 --config.observables.s2、--config.observables.dipole 和 --config.observables.density 来跟踪自旋模量、偶极矩和密度矩阵。

激发态

系统的激发态性质可以使用以下两种方法之一来计算：VMC 的自然激发态算法 (NES-VMC) 或 系综惩罚方法。要计算系统的 k 个状态，只需在配置文件中设置 cfg.system.states=k 即可。默认情况下使用 NES-VMC，但若要启用系综惩罚方法，则需在配置中添加 cfg.optim.objective='vmc_overlap'。NES-VMC 没有需要设置的参数，而系综惩罚方法则允许自由选择能量和重叠惩罚项的权重，这些权重可以在 cfg.optim.overlap 中设定。如果配置中未为各状态的能量设置权重，系统会自动将其设为 1/k。我们发现，NES-VMC 通常比系综惩罚方法更为准确，但为了完整性，我们同时提供了这两种方法。介绍 NES-VMC 的论文中所有实验的配置文件都位于 configs/excited 文件夹中，所有实验都可以通过运行 tests/excited_test.py 在较小的网络上进行测试。

输出

结果目录包含 train_stats.csv 文件，其中记录了每次迭代的局域能量和 MCMC 接受概率；此外还有 checkpoints 目录，用于存储训练过程中生成的检查点。在计算激发态的可观测量或基态密度矩阵时，还会将 .npy 文件保存到同一文件夹中。每次优化迭代都会将一个 NumPy 数组保存到同一个文件中。分析这些输出的 Colab 笔记本示例位于 notebooks/excited_states_analysis.ipynb 中。

致谢

如果您在工作中使用了此代码，请引用相关论文。最初的论文详细介绍了架构及其在多种体系上的结果：

@article{pfau2020ferminet,
  title={基于深度神经网络的多电子薛定谔方程从头计算解},
  author={D. Pfau 和 J.S. Spencer、A.G. de G. Matthews、W.M.C. Foulkes},
  journal={Phys. Rev. Research},
  year={2020},
  volume={2},
  issue = {3},
  pages={033429},
  doi = {10.1103/PhysRevResearch.2.033429},
  url = {https://link.aps.org/doi/10.1103/PhysRevResearch.2.033429}
}

一篇 NeurIPS 工作坊“机器学习与物理”会议论文描述了 JAX 实现：

@misc{spencer2020better,
  title={更好、更快的费米子神经网络},
  author={James S. Spencer、David Pfau、Aleksandar Botev 和 W. M.C. Foulkes},
  year={2020},
  eprint={2011.07125},
  archivePrefix={arXiv},
  primaryClass={physics.comp-ph},
  url={https://arxiv.org/abs/2011.07125}
}

PsiFormer 架构则在 2023 年 ICLR 论文中进行了详细说明：

@misc{vonglehn2023psiformer,
  title={一种用于从头量子化学的自注意力 Ansatz},
  author={Ingrid von Glehn、James S Spencer 和 David Pfau},
  journal={ICLR},
  year={2023},
}

周期性边界条件最初是在《物理评论快报》的一篇文章中提出的：

@article{cassella2023discovering,
  title={利用费米子神经网络发现量子相变},
  author={Cassella, Gino、Sutterud, Halvard、Azadi, Sam、Drummond, ND、Pfau, David、Spencer, James S、Foulkes, W Matthew C},
  journal={物理评论快报},
  volume={130},
  number={3},
  pages={036401},
  year={2023},
  publisher={APS}
}

Wasserstein QMC（感谢 Kirill Neklyudov）则在 2023 年 NeurIPS 论文中有所介绍：

@article{neklyudov2023wasserstein,
  title={Wasserstein 量子蒙特卡洛：求解量子多体薛定谔方程的一种新方法},
  author={Neklyudov, Kirill、Nys, Jannes、Thiede, Luca、Carrasquilla, Juan、Liu, Qiang、Welling, Max、Makhzani, Alireza},
  journal={NeurIPS},
  year={2023}
}

自然激发态则在这篇文章中被提出，这也是我们团队首次使用赝势的论文：

@article{pfau2024excited,
  title={利用神经网络精确计算量子激发态},
  author={Pfau, David、Axelrod, Simon、Sutterud, Halvard、von Glehn, Ingrid、Spencer, James S},
  journal={Science},
  volume={385},
  number={6711},
  pages={eadn0137},
  year={2024},
  url={https://doi.org/10.1126/science.adn0137},
}

本仓库可按如下方式引用：

@software{ferminet_github,
  author = {James S. Spencer、David Pfau 和 FermiNet 贡献者},
  title = {{FermiNet}},
  url = {http://github.com/deepmind/ferminet},
  year = {2020},
}

免责声明

本项目并非 Google 官方产品。

FermiNet 快速上手指南

FermiNet 是一个基于变分蒙特卡洛（VMC）方法的神经网络，用于学习原子和分子的高精度基态波函数。本项目提供了该算法的 JAX 实现。

环境准备

• 操作系统：推荐 Linux 或 macOS。
• Python：建议使用 Python 3.8+。
• 硬件加速：强烈推荐使用支持 CUDA 的 NVIDIA GPU 以加快训练速度。
• 依赖管理：建议在虚拟环境中安装，避免污染系统环境。

安装步骤

1. 创建并激活虚拟环境

virtualenv ~/venv/ferminet
source ~/venv/ferminet/bin/activate

2. 安装基础依赖

进入项目目录后，执行以下命令安装核心包：

pip install -e .

提示：国内用户若遇到下载缓慢，可添加清华或阿里镜像源，例如：pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 配置 GPU 支持（可选但推荐）

若需使用 GPU 加速，需安装带有 CUDA 支持的 jax 和 jaxlib。请根据您本地的 CUDA 版本选择对应的 jaxlib 版本（以下示例为 CUDA 11.0）：

pip install --upgrade jax jaxlib==0.1.57+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

注意：jaxlib 版本必须与您系统中安装的 CUDA 版本严格对应。更多细节请参考 JAX 官方文档。

4. 安装测试依赖（可选）

如需运行单元测试：

pip install -e '.[testing]'
python -m pytest

基本使用

FermiNet 使用 ConfigDict 进行配置。您可以通过命令行参数直接覆盖配置文件中的设置。

示例：计算锂原子 (Li) 的基态波函数

以下命令将训练 FermiNet 以寻找锂原子的基态波函数。设置了批量大小为 256，预训练迭代次数为 100（对于小系统，默认值 1000 次通常过多）。

ferminet --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

或者直接使用 Python 脚本运行：

python3 ferminet/main.py --config ferminet/configs/atom.py --config.system.atom Li --config.batch_size 256 --config.pretrain.iterations 100

自定义分子系统（如 H2）

您也可以通过编写简单的 Python 脚本来定义自定义分子（例如 H2）并启动训练：

import sys
from absl import logging
from ferminet.utils import system
from ferminet import base_config
from ferminet import train

# 配置日志输出到标准输出
logging.get_absl_handler().python_handler.stream = sys.stdout
logging.set_verbosity(logging.INFO)

# 初始化默认配置
cfg = base_config.default()

# 定义 H2 分子：(alpha 电子数，beta 电子数) 和 原子坐标
cfg.system.electrons = (1, 1)
cfg.system.molecule = [system.Atom('H', (0, 0, -1)), system.Atom('H', (0, 0, 1))]

# 设置训练超参数
cfg.batch_size = 256
cfg.pretrain.iterations = 100

# 开始训练
train.train(cfg)

多 GPU 并行

对于较大的原子或分子以及大批量训练，代码会自动检测并利用多个 GPU 进行并行计算（基于 JAX 的 pmap），无需额外配置。