mPotNet

mPotNet

mPotNet 是一个基于 PyTorch 和 PyTorch Geometric 的独立 PotNet 实现，用于根据 CIF 晶体结构预测材料带隙。

代码保留了 PotNet 的核心设计：将局域周期邻居图与包含全部原子对的长程图结合，并为长程边计算周期性 Coulomb 势、London 色散势和 Pauli 排斥势。局域距离与周期势经过径向基函数编码后，由门控消息传递层、全局平均池化和标量回归头完成带隙预测。

代码使用 pymatgen 解析 CIF 并构建周期图。周期求和函数直接使用 NumPy 和 SciPy 实现，不依赖 JARVIS、原始 PotNet 仓库、仓库内部模块或需要单独编译的 Cython/GSL 扩展。

环境依赖

1

pip install torch torch-geometric pymatgen scipy numpy pandas openpyxl scikit-learn matplotlib tqdm

仅当 USE_OPTUNA = True 时需要安装 Optuna：

1

pip install optuna

数据准备

在脚本所在目录中准备以下文件：

1
2
3
4
5
6

mPotNet_v1.py
data.xlsx
cif/
|-- 1.cif
|-- 2.cif
|-- 3.cif

脚本读取 data.xlsx 的前两列：

第一列中的 1 对应 ./cif/1.cif。当文件名不包含 .cif 后缀时，脚本会自动补充。缺失、重复、标签非有限值、无法解析或含无序占位的样本会被跳过，并在日志中记录原因。

使用方法

1

python mPotNet_v1.py

脚本会自动完成 CIF 检查、周期势预处理与缓存、固定的 80/10/10 训练集/验证集/测试集划分、仅基于训练集拟合的目标归一化、模型训练、早停、评估和绘图。当 CUDA 可用时自动使用 GPU，否则使用 CPU；CUDA 环境默认启用 BF16 自动混合精度。

所有结果保存在由 MODEL_NAME 和 RUN_VERSION 确定的目录中。按照默认设置，输出目录为：

1
2
3
4
5
6
7
8

PotNet_v1/
|-- PotNet_best.pt
|-- figure/
|-- dat/
|-- table/
|-- log/
|-- split/
|-- cache/

输出内容包括训练集、验证集和测试集的 MAE、RMSE 与 R2，各数据集及合并奇偶图和对应数据，RMSE 训练曲线及其数据，可复用的固定数据划分，PotNet 图缓存，按照验证集 RMSE 选择的最佳模型权重，以及完整训练日志。

脚本中的 load_trained_model() 和 predict_cifs() 可以被其他 Python 程序导入，用于加载最佳模型并预测新的 CIF 结构。

显存说明

PotNet 为每个结构构建约含 N^2 条原子对边的完整长程图，因此显存会随原子数快速增长。对于较大的晶胞，建议首先降低 BATCH_SIZE，随后依次降低 HIDDEN_DIM 和 POTENTIAL_RBF_BINS。在支持的 NVIDIA GPU 上建议保持 USE_AMP = True 和 AMP_DTYPE = "bfloat16"。

引用

原始 PotNet 文献：

• Lin Y, Yan K, Luo Y, et al. Efficient approximations of complete interatomic potentials for crystal property prediction. International conference on machine learning. PMLR, 2023: 21260-21287

本工作：

论文正式发表后补充。