Misaraty
全部 标签 分类 关于
Misaraty

目录

mMEGNet

 收录于 学术
   约 1335 字   预计阅读 3 分钟 
misaraty 更新 | 2026-10-06
前言
下载：mMEGNet。

mMEGNet

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

代码保留了 MEGNet 的主要组成部分，包括周期性晶体图、原子嵌入、高斯距离展开、边/节点/全局状态更新、残差 MEGNet 模块、原子与键的双 Set2Set 池化以及回归头。代码使用 pymatgen 解析 CIF 并构建周期性晶体图，不依赖 TensorFlow、Keras 或原始 MEGNet 仓库。

环境依赖

1

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

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

1

pip install optuna

数据准备

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

1
2
3
4
5
6

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

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

cif bandgap
1 1.23
2 0.87
3 2.15

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

使用方法

1

python mMEGNet_v2.py

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

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

1
2
3
4
5
6
7
8

MEGNet_v2/
|-- MEGNet_best.pt
|-- figure/
|-- dat/
|-- table/
|-- log/
|-- split/
|-- cache/

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

v1

该脚本实现了一个基于MEGNet图神经网络的预测流程，用于预测晶体材料的超导临界温度Tc。程序从data.xlsx读取结构ID和对应的Tc标签，从cif目录加载相应结构文件，并可选择加入元素电负性均值和标准差作为全局特征。数据可采用train/test划分或KFold交叉验证方式进行训练，并支持EarlyStopping和学习率调度。脚本输出RMSE、MAE和R2指标，保存预测散点图；在megnet_tuned模式下，使用Optuna优化网络宽度和学习率，并自动保存每次trial和fold的日志及模型文件。

运行命令：python mMEGNet_v14.py。

  • TARGET_COL：预测目标列名称（例如'tc'表示超导临界温度）。

  • method：训练模式选择（'megnet_default' 为固定结构，'megnet_tuned' 为使用Optuna进行超参数优化）。

  • batch_size：每次梯度更新的样本数，影响收敛稳定性和显存占用。

  • lr：初始学习率，决定优化步长。

  • train_ratio / test_ratio：训练集和测试集划分比例。

  • USE_EN_GLOBAL：是否加入基于电负性的全局状态特征。

  • n1, n2, n3：控制MEGNet网络宽度（模型容量）；数值越大模型越复杂，但可能增加过拟合风险。

  • epochs：最大训练轮数。

  • n_folds：交叉验证折数，用于控制模型评估的稳健性（推荐5，或 'none' 表示单次划分）。

请将 ~\anaconda3\Lib\site-packages\megnet\data\graph.py 中的 get_atom_features 函数替换为以下实现：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

    @staticmethod
    def get_atom_features(structure) -> List[Any]:
        z_list = []
        for site in structure:
            try:
                if getattr(site, "is_ordered", True):
                    z = int(site.specie.Z)
                else:
                    items = list(site.species.items())  # [(Element, occ), ...]
                    items.sort(key=lambda kv: (float(kv[1]), getattr(kv[0], "Z", 0)), reverse=True)
                    z = int(getattr(items[0][0], "Z", 0))
            except Exception:
                tok = str(site.species_string).split()[0]
                try:
                    from pymatgen.core.periodic_table import Element
                    z = int(Element(tok).Z)
                except Exception:
                    z = 0
            z_list.append(z)
        return np.array(z_list, dtype="int32").tolist()

该修改通过在分数占位（无序）结构中选择占据比例最高的元素，避免在读取包含部分占位原子的CIF文件时出现报错。

引用

原始 MEGNet 文献：

本工作：

v2:

论文正式发表后补充。

v1:

 机器学习
返回 | 主页