代码导读：从命令到一次参数更新¶

这篇按一次真实训练的调用顺序阅读代码。目标是看到每个模块在 VMC 流程中的位置，而不先陷进实现细节。

命令行入口¶

典型运行命令：

python -m nnqs_tutorial.cli.train \
  --ham qubit_op.data \
  --config nnqs_tutorial/examples/lih_12q.yaml

cli/train.py 做三件事：

解析命令行参数。
调用 load_config(...) 读取 YAML。
创建 VMCTrainer，再调用 train()。

这一层只负责把外部输入变成训练对象。

配置对象¶

config.py 把 YAML 转成 TutorialConfig。后续代码使用：

cfg.model.d_model

而少直接访问：

cfg["model"]["d_model"]

这样做有两个好处：字段集中在 dataclass 定义里，编辑器也能补全。对教学代码来说，这比灵活但松散的字典更容易读。

模型输出¶

models.py 的主接口是：

out = model(states)

返回：

out["log_amp"]
out["phase"]
out["psi"]

对应数学式：

\[ \psi_\theta(x) = \exp(\log A_\theta(x))\exp(i\phi_\theta(x)). \]

其中：

\[ \log A_\theta(x) = {1 \over 2}\sum_i\log P_\theta(t_i\mid t_{<i}). \]

1/2 来自 \(P_\theta(x)=|\psi_\theta(x)|^2\)。

采样器¶

如果直接保存十万条 sample，里面会有大量重复 state。VMC 估计只需要去重后的 state 以及出现次数，所以 sampling.py 返回：

unique_states, counts, psi_values

训练中使用：

\[ w(x) = {\mathrm{count}(x) \over \sum_{x'}\mathrm{count}(x')}. \]

采样器逐 token 生成 bitstring，并用 electron mask 保证最终满足 \(\alpha\)、\(\beta\) 电子数约束。

Hamiltonian 与局域能¶

Hamiltonian 是 Pauli term 的和：

\[ H = \sum_k c_kP_k . \]

一个 Pauli term 作用到 \(x\) 后，得到 connected state \(x'\) 和矩阵元因子。局域能是：

\[ E_{\rm loc}(x) = \sum_{x'}H_{xx'}{\psi_\theta(x')\over\psi_\theta(x)}. \]

ExactHamiltonian 逐项实现这个公式。它速度不一定最快，但最适合用来对照公式和调试。

Trainer 的一轮¶

trainer.py 中一轮 epoch 可以按下面的结构理解：

_sample_states()
  从当前模型采样 unique states 和 counts

_evaluate_energy()
  计算 psi、local energies、energy mean

_vmc_loss_proxy()
  构造 VMC 梯度代理

_optimizer_step()
  backward 并更新参数

能量估计：

\[ E_{\rm mean} = \sum_x w(x)E_{\rm loc}(x). \]

梯度代理：

\[ L_{\rm proxy} = 2\,\mathrm{Re}\left[ \sum_x w(x)\overline{\log\psi_\theta(x)} (E_{\rm loc}(x)-E_{\rm mean}) \right]. \]

这条 loss_proxy.backward() 的意义是产生 VMC 梯度；训练日志里真正要看的是能量、唯一样本数和数值是否有限。

最短调用链¶

cli/train.py
  -> config.py
  -> trainer.py
     -> sampling.py
     -> models.py
     -> hamiltonian.py
     -> checkpoint.py

读代码时可以先沿这条链走一遍，再回头看每个模块里的边界条件。