SALTED: a symmetry-adapted machine-learning program for predicting electron-densities in molecules and materials

技术摘要：SALTED：用于预测电子密度的对称性适配机器学习程序

问题陈述
密度泛函理论（DFT）是模拟材料性质的标准方法，在准确性和计算成本之间取得了平衡。然而，随着数据驱动材料科学向更大规模系统和更长时间尺度迈进，即使是 DFT 相对较低的成本也变得难以承受。目前迫切需要能够以极低的计算成本、同时保持高精度的机器学习（ML）方法来预测 DFT 中最基本的物理量——实空间电子密度 $n(\mathbf{r})$。现有的机器学习方法往往在不同化学环境下的迁移能力较弱，或者需要海量的训练数据，而生成这些高质量电子结构方法的训练数据是非常困难的。

方法论
SALTED（对称性适配电子密度学习）通过特定的表示形式和学习算法解决了这些挑战：

• 原子基组表示法：与在 3D 网格上采样 $n(\mathbf{r})$ 的网格型机器学习模型不同，SALTED 采用了线性原子中心分解。电子密度被展开在一组以每个原子为中心的径向函数 $R^\lambda_n$ 和球谐函数 $Y^\lambda_\mu$ 之上：
  $$n(\mathbf{r}) = \sum_{i n \lambda \mu} c_{i n \lambda \mu} \sum_{\mathbf{u}} R^\lambda_n(|\mathbf{r} - \mathbf{r}_i - \mathbf{u}|) Y^\lambda_\mu\left(\widehat{\mathbf{r} - \mathbf{r}_i - \mathbf{u}}\right)$$
  其中 $i$ 表示原子索引，$\mathbf{u}$ 代表周期性系统的晶胞平移向量。系数 $c_{i n \lambda \mu}$ 是目标变量。这种表示法与基于原子轨道的电子结构代码天然兼容，允许通过解析度恒等式（Resolution of the Identity, RI）方法直接提取参考系数。

• 对称性适配高斯过程回归（GPR）：学习算法利用了对称适配扩展的高斯过程回归。SALTED 并非使用等变神经网络，而是计算球张量核（$K_\lambda$）以满足精确的 $O(3)$ 对称性。这些核是由 featomic 库生成的原子密度特征推导而来的。

  • 数据效率：由这些核定义的受限函数空间，结合特征的物理可解释性，有助于防止过拟合。这使得 GPR 在小样本数据机制（例如 $\sim 10^2-10^3$ 次计算）中表现尤为出色，这对于使用高水平电子结构数据进行训练至关重要。
  • 矢量场核：通过实现矢量场核函数，SALTED 可以学习电子密度对外加电场的响应（$\partial n(\mathbf{r})/\partial \mathbf{E}$），从而确保 3D 矢量场的正确变换特性。
  • 长程相互作用：使用 featomic 能够引入长程等变（LODE）特征，这对于捕捉高度极化系统（如金属框架）中的非局部密度重新分布至关重要。

• 软件架构：SALTED 是一个开源 Python 包，分为三个阶段的工作流：数据集准备、模型训练和预测。它与主要的电子结构代码（CP2K、FHI-aims 和 PySCF）对接，以统一训练数据。该工作流专为高性能计算（HPC）设计，具有检查点功能、用于描述符收缩和 Hessian 构造的 MPI 并行化，以及用于提升性能的 Numba 编译核。

核心贡献

1. 开源框架：SALTED 提供了一个实用的、可扩展的框架，用于构建电子密度及其线性响应的机器学习模型。
2. 互操作性：它与成熟的电子结构软件包（CP2K、FHI-aims、PySCF）无缝集成，允许研究人员利用其现有的工作流和设置，而无需更换代码。
3. 响应性质：该软件独特地支持预测对电场的密度响应，从而能够计算极化率张量和极化矢量。
4. 小数据机制下的高效性：通过利用具有对称适配核的 GPR，该方法与数据饥渴型的神经网络架构相比，能以显著更少的训练点实现高精度。

结果与应用
论文强调了 SALTED 集成到计算工作流中的几个成功应用案例：

• QM/MM 模拟：通过预测金属电极中的电荷密度响应，在计算模型离子电容器的 DFT 级静电作用力时实现了约 $10^3$ 倍的加速。
• 大规模电子结构：将预测的密度作为输入，用于单次 Kohn-Sham 哈密顿量对角化，从而实现了对超大超胞中扭曲双层 Moiré 材料电子性质的研究。
• 性质计算：该方法有助于精确评估库仑力、极化矢量和极化率张量。
• 晶体学：预测的 RI 密度系数可直接应用于来自 X 射线和电子衍射数据的晶体结构精修。

意义
SALTED 代表了在保留物理意义明确的可观测量的同时，加速电子结构计算的重要一步。通过直接预测电子密度，它提供了一条以极低于 DFT 的成本来计算众多下游材料性质的途径。该软件旨在支持可重复性研究，并作为机器学习与多尺度模拟工具之间的桥梁，既满足了寻求高效预测的用户，也服务于对扩展方法论感兴趣的方法开发者。其处理分子及凝聚态系统的能力，以及与标准电子结构代码的兼容性，使其成为计算化学和材料科学领域的通用工具。