mComFormer

2026-10-06 约 871 字预计阅读 2 分钟

前言

mComFormer 是一个基于 PyTorch 和 PyTorch Geometric 的独立 iComFormer 与 eComFormer 实现，用于根据 CIF 晶体结构预测材料带隙。两个模型包含在同一个 Python 脚本中，通过 MODEL_NAME 选择。

代码使用 pymatgen 解析 CIF 并构建周期性晶体图，不依赖 DGL、JARVIS 或原始 ComFormer 仓库。

模型

模型	几何表示	主要组成部分
iComFormer	SE(3) 不变	原子间距离、三个规范晶格参考向量、夹角编码、四个节点 ComFormer 层和一个边更新层
eComFormer	SO(3) 等变	原子间位移向量、球谐函数、三个节点 ComFormer 层和一个等变张量积更新层

两个模型均使用可学习的 92 维原子嵌入、径向基距离编码、图级平均池化和标量回归头。

安装共同依赖：

1

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

eComFormer 还需要安装 e3nn：

1

pip install e3nn

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

1

pip install optuna

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

1
2
3
4
5
6


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

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

第一列中的 1 会转换为 1.cif，对应 ./cif/1.cif。包含无序或部分占位原子位点的 CIF 会被跳过，并在训练日志中记录。

在 mComFormer_v1.py 开头选择模型：

1

MODEL_NAME = "iComFormer"

或者：

1

MODEL_NAME = "eComFormer"

运行版本可以单独修改：

1

RUN_VERSION = "v1"

1

python mComFormer_v1.py

脚本会自动完成 CIF 检查、周期性晶体图构建与缓存、固定的 80/10/10 训练集/验证集/测试集划分、仅基于训练集的目标归一化、模型训练、早停、评估和绘图。当 CUDA 可用时自动使用 GPU，否则使用 CPU。

训练采用 MSE 损失、AdamW 和 ReduceLROnPlateau，并根据验证集 RMSE 选择最佳模型。当有效 CIF 集合保持不变时，脚本会复用已保存的数据划分。

所有结果保存在由 MODEL_NAME 和 RUN_VERSION 确定的目录中。例如：

1
2
3
4
5
6
7
8


iComFormer_v1/
|-- iComFormer_best.pt
|-- figure/
|-- dat/
|-- table/
|-- log/
|-- split/
|-- cache/

或者：

1
2
3
4
5
6
7
8


eComFormer_v1/
|-- eComFormer_best.pt
|-- figure/
|-- dat/
|-- table/
|-- log/
|-- split/
|-- cache/

输出内容包括训练集、验证集和测试集的 MAE、RMSE 与 R2，单独及合并奇偶图和数据文件，RMSE 训练曲线及其数据，固定数据划分，晶体图缓存，最佳模型权重，可选的 Optuna 结果和完整训练日志。

脚本还提供 load_trained_model() 和 predict_cifs()，用于载入已保存的模型并预测新 CIF 结构的带隙。

原始 ComFormer 文献：

本工作：

论文正式发表后补充。