mSchNet

mSchNet

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

代码保留了 SchNet 的主要组成部分，包括原子序数嵌入、高斯径向基展开、连续滤波卷积、shifted-softplus 激活函数、残差交互模块、原子级输出和晶体级池化。代码使用 pymatgen 直接根据 CIF 构建周期性邻居。当前配置使用平滑余弦截断，也可以通过 USE_COSINE_CUTOFF 将其关闭，以便与 hard cutoff 对比。脚本不依赖 TensorFlow、PyTorch Geometric、ASE 数据库或原始 SchNet 仓库。

环境依赖

1

pip install torch pymatgen numpy pandas openpyxl scikit-learn matplotlib

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

1

pip install optuna

数据准备

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

1
2
3
4
5
6

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

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

第一列中的 1 对应 ./cif/1.cif。

使用方法

1

python mSchNet_v1.py

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

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

1
2
3
4
5
6
7
8

SchNet_v1_USE_COSINE_CUTOFF/
|-- SchNet_best.pt
|-- figure/
|-- dat/
|-- table/
|-- log/
|-- split/
|-- cache/

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

引用

原始 SchNet 文献：

• Schütt K, Kindermans P J, Sauceda Felix H E, et al. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 2017, 30.

本工作：

论文正式发表后补充。