晶体图卷积神经网络

本软件包实现了晶体图卷积神经网络（CGCNN），该网络能够以任意晶体结构为输入，预测材料性质。

该软件包提供两大功能：

• 使用自定义数据集训练CGCNN模型。
• 使用预训练的CGCNN模型预测新晶体的材料性质。

以下论文详细介绍了CGCNN框架：

用于准确且可解释的材料性质预测的晶体图卷积神经网络

目录

• 如何引用
• 先决条件
• 使用方法
  • 定义自定义数据集
  • 训练CGCNN模型
  • 使用预训练的CGCNN模型预测材料性质
• 数据
• 作者
• 许可证

如何引用

如需使用CGCNN，请引用以下文献：

@article{PhysRevLett.120.145301,
  title = {用于准确且可解释的材料性质预测的晶体图卷积神经网络},
  author = {Xie, Tian 和 Grossman, Jeffrey C.},
  journal = {Phys. Rev. Lett.},
  volume = {120},
  issue = {14},
  pages = {145301},
  numpages = {6},
  year = {2018},
  month = {Apr},
  publisher = {美国物理学会},
  doi = {10.1103/PhysRevLett.120.145301},
  url = {https://link.aps.org/doi/10.1103/PhysRevLett.120.145301}
}

先决条件

本软件包需要以下依赖：

• PyTorch
• scikit-learn
• pymatgen

如果您是Python新手，安装这些先决条件最简单的方式是通过conda。安装conda后，运行以下命令创建一个名为cgcnn的新环境，并安装所有先决条件：

conda upgrade conda
conda create -n cgcnn python=3 scikit-learn pytorch torchvision pymatgen -c pytorch -c conda-forge

*注：此代码已在PyTorch v1.0.0及以上版本上测试通过，由于一些破坏性更改，不兼容v0.4.0以下的版本。

这将创建一个用于运行CGCNN的conda环境。在使用CGCNN之前，通过以下命令激活该环境：

source activate cgcnn

然后，在cgcnn目录下，您可以通过运行以下命令来测试所有先决条件是否已正确安装：

python main.py -h
python predict.py -h

这应显示main.py和predict.py的帮助信息。如果没有错误信息，则说明先决条件已正确安装。

使用完CGCNN后，通过以下命令退出环境：

source deactivate

使用方法

定义自定义数据集

要将晶体结构输入CGCNN，您需要定义一个自定义数据集。请注意，无论是训练还是预测，都必须定义自定义数据集。

在定义自定义数据集之前，您需要：

• CIF文件，记录您感兴趣的晶体结构。
• 每个晶体的目标属性（仅在预测时不需要，但您仍需在id_prop.csv中填入一些随机数）。

您可以通过创建一个名为root_dir的目录来构建自定义数据集，其中包含以下文件：

1. id_prop.csv：一个CSV文件，包含两列。第一列记录每个晶体的唯一ID，第二列记录目标属性的值。如果您想用predict.py预测材料性质，可以在第二列中填入任意数字。（第二列仍然是必需的。）

2. atom_init.json：一个JSON文件，存储每种元素的初始化向量。atom_init.json的一个示例是data/sample-regression/atom_init.json，对于大多数应用来说应该足够。

3. ID.cif：一个CIF文件，记录晶体结构，其中ID是该晶体的唯一ID。

root_dir的结构应如下：

root_dir
├── id_prop.csv
├── atom_init.json
├── id0.cif
├── id1.cif
├── ...

仓库中提供了两个自定义数据集的示例：data/sample-regression用于回归任务，data/sample-classification用于分类任务。

面向高级PyTorch用户

上述创建自定义数据集的方法使用了cgcnn.data中的CIFData类。如果您希望以更灵活的方式输入晶体结构，PyTorch有一个很好的教程，教您如何编写自己的数据集类。

训练CGCNN模型

在训练新的CGCNN模型之前，您需要：

• 在root_dir中定义自定义数据集，以存储您感兴趣的结构-属性关系。

然后，在cgcnn目录下，您可以使用以下命令为您的自定义数据集训练CGCNN模型：

python main.py root_dir

您可以通过--train-size、--val-size和--test-size设置训练、验证和测试数据的数量。或者，您也可以使用--train-ratio、--val-ratio和--test-ratio标志。请注意，比率标志不能与大小标志同时使用。例如，data/sample-regression共有10个数据点。您可以按以下方式训练模型：

python main.py --train-size 6 --val-size 2 --test-size 2 data/sample-regression

或者：

python main.py --train-ratio 0.6 --val-ratio 0.2 --test-ratio 0.2 data/sample-regression

您还可以通过--task classification标志训练分类模型。例如，您可以使用data/sample-classification进行训练：

python main.py --task classification --train-size 5 --val-size 2 --test-size 3 data/sample-classification

训练完成后，您将在cgcnn目录下获得三个文件：

• model_best.pth.tar：存储具有最佳验证准确率的CGCNN模型。
• checkpoint.pth.tar：存储最后一个epoch的CGCNN模型。
• test_results.csv：存储测试集中每个晶体的ID、目标值和预测值。

使用预训练的 CGCNN 模型预测材料性质

在开始预测材料性质之前，您需要：

• 在 root_dir 目录下定义一个自定义数据集，用于存放所有您希望进行预测的晶体结构。
• 获取一个名为 pre-trained.pth.tar 的预训练 CGCNN 模型。

随后，在 cgcnn 目录下，您可以对 root_dir 中的晶体进行性质预测：

python predict.py pre-trained.pth.tar root_dir

例如，您可以预测 data/sample-regression 目录中晶体的形成能：

python predict.py pre-trained/formation-energy-per-atom.pth.tar data/sample-regression

此外，您还可以预测 data/sample-classification 目录中的晶体是金属（1）还是半导体（0）：

python predict.py pre-trained/semi-metal-classification.pth.tar data/sample-classification

需要注意的是，对于分类任务，test_results.csv 文件中的预测值是一个介于 0 和 1 之间的概率，表示该晶体被分类为 1（即金属，如上述示例）的可能性。

预测完成后，您将在 cgcnn 目录下得到一个文件：

• test_results.csv：存储测试集中每种晶体的 ID、目标值和预测值。其中，目标值只是您在 id_prop.csv 文件中定义数据集时设定的任意数值，与实际意义无关。

数据

为了复现我们的论文，您可以按照说明下载相应的数据集。

作者

本软件主要由 Tian Xie 编写，并在 Jeffrey Grossman 教授 的指导下完成。

许可证

CGCNN 采用 MIT 许可证发布。

CGCNN 快速上手指南

环境准备

• 操作系统：Linux / macOS / Windows（推荐 Linux）
• Python ≥ 3.6
• PyTorch ≥ 1.0.0（不支持 0.4.0 以下版本）
• 依赖包：scikit-learn、pymatgen

安装步骤

1. 安装 conda（国内用户可用清华镜像：https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/）
2. 创建并激活环境

conda upgrade conda
conda create -n cgcnn python=3 scikit-learn pytorch torchvision pymatgen -c pytorch -c conda-forge
source activate cgcnn   # Windows 用 `conda activate cgcnn`

3. 验证安装

git clone https://github.com/txie-93/cgcnn.git
cd cgcnn
python main.py -h
python predict.py -h

无报错即安装成功。

基本使用

1. 准备数据集

在任意目录（如 mydata）创建以下文件：

mydata
├── id_prop.csv      # 两列：ID, 目标值（预测时可填 0）
├── atom_init.json   # 元素初始化向量，可直接复制 data/sample-regression/atom_init.json
├── 10001.cif        # 晶体结构文件，文件名与 id_prop.csv 第一列对应
├── 10002.cif
└── ...

2. 训练模型（回归示例）

python main.py --train-ratio 0.7 --val-ratio 0.1 --test-ratio 0.2 mydata

训练完成后得到：

• model_best.pth.tar：验证集最佳模型
• test_results.csv：测试集预测结果

3. 使用预训练模型预测

# 下载官方预训练模型到 pre-trained 目录
python predict.py pre-trained/formation-energy-per-atom.pth.tar mydata

输出 test_results.csv 包含预测值。