mCrystalformer

mCrystalformer

mCrystalformer 是一个纯 PyTorch 的独立 Crystalformer 实现，用于根据 CIF 晶体结构预测材料带隙。

代码保留了 Crystalformer 的核心设计，包括全连接周期注意力、实空间与倒空间周期编码、依赖查询向量的高斯注意力、径向值编码、交替堆叠的 Latticeformer 编码器、T-Fixup 初始化、晶体级池化和标量回归头。代码使用 pymatgen 解析 CIF 并将晶体标准化为原胞。

该版本不依赖 PyTorch Geometric、DGL、JARVIS、CuPy、pytorch-pfn-extras、自定义 CUDA 源文件、JSON 参数文件或原始 Crystalformer 仓库。当 CUDA 可用时，标准 PyTorch 张量运算仍会在 GPU 上执行；由于没有使用原仓库的融合 CUDA kernel，纯 PyTorch 版本的运行速度可能较慢。

环境依赖

1

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

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

1

pip install optuna

数据准备

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

1
2
3
4
5
6

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

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

第一列必须是可以转换为整数的编号。第一列中的 1 会被转换为 1.cif，并对应 ./cif/1.cif；第二列为带隙，单位为 eV。

无效标签、缺失的 CIF、无法解析的结构、无效占位信息和奇异晶格会被跳过，并记录在训练日志中。有效结构数量不得少于 20 个。

使用方法

1

python mCrystalformer_v2.py

脚本会自动完成 CIF 检查与缓存、原胞标准化、基于 SEED = 42 的固定 80/10/10 训练集/验证集/测试集划分、仅使用训练集拟合的标签标准化、模型训练、基于验证集 RMSE 的早停、最佳模型重新加载、最终评估和绘图。当 CUDA 可用时自动使用 GPU，否则使用 CPU。

默认模型包含 4 个 128 维 Latticeformer 编码块和 8 个注意力头。DOMAIN = "real-reci" 表示在不同编码块之间交替使用实空间和倒空间周期注意力。训练采用 MSE 损失、AdamW、逆平方根学习率衰减和最多 300 个 epoch。是否启用 Optuna 由 USE_OPTUNA 控制。

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

Crystalformer_v2/
|-- Crystalformer_best.pt
|-- figure/
|   |-- Crystalformer_parity_train.jpg
|   |-- Crystalformer_parity_val.jpg
|   |-- Crystalformer_parity_test.jpg
|   |-- Crystalformer_parity_all.jpg
|   `-- Crystalformer_rmse_curve.jpg
|-- dat/
|-- table/
|-- log/
|-- split/
`-- cache/

输出内容包括训练集、验证集和测试集的 MAE、RMSE 与 R2，单独及合并的奇偶图和制表符分隔的预测数据，RMSE 训练曲线及其数据，可复用的固定数据划分，解析后的结构缓存，由验证集 RMSE 选择的最佳模型权重，以及完整训练日志。修改 RUN_VERSION 可以建立新的结果目录，避免覆盖以前的实验。

引用

原始 Crystalformer 文献：

• Taniai T, Igarashi R, Suzuki Y, et al. Crystalformer: Infinitely connected attention for periodic structure encoding. International Conference on Learning Representations. 2024, 2024: 45083-45105.

本工作：

论文正式发表后补充。