一个使用ALIGNN神经网络对材料性能预测的深度学习案例解读

案例：使用更先进的图神经网络（ALIGNN）结合Materials Project API进行材料带隙预测

在这个案例中，我们将使用一种更先进且性能更优的图神经网络模型------ALIGNN（Atomistic Line Graph Neural Network），结合Materials Project API（MP API）的数据，预测材料的带隙（band gap）。ALIGNN是一种专门为材料科学设计的GNN模型，它通过结合原子图和键角信息，能够更精确地捕捉原子间的相互作用和几何特性。相比SchNet，ALIGNN在多个材料性能预测任务中表现出更高的准确性。我们将详细讲解每个步骤，包括代码实现、参数设置及其含义，帮助你全面理解如何使用ALIGNN进行带隙预测。

环境准备

步骤解读

在开始之前，需要安装必要的Python库以支持数据获取、预处理和模型训练。ALIGNN依赖PyTorch和DGL（Deep Graph Library），并需要额外的工具来处理晶体结构。

代码实现

pip install mp_api pymatgen torch dgl alignn jarvis-tools

mp_api：用于访问Materials Project数据库，获取晶体结构和带隙数据。
pymatgen：处理晶体结构的强大工具，用于解析和转换材料数据。
torch：PyTorch深度学习框架，是ALIGNN的核心依赖。
dgl：Deep Graph Library，用于构建和处理图数据，ALIGNN基于此实现。
alignn：ALIGNN模型的官方实现库，提供预训练模型和工具。
jarvis-tools：JARVIS框架的工具集，与ALIGNN配合使用，支持数据转换和分析。

参数设置与含义

无需额外参数设置，只需确保安装的库版本兼容（建议Python 3.8+，PyTorch 1.12+，DGL 1.0+）。
你需要从Materials Project官网注册并获取一个API密钥，用于数据访问。

数据获取

步骤解读

使用MP API从Materials Project数据库中获取一组材料的晶体结构和带隙数据。这里我们以50种材料为例，增加数据量以提升模型训练效果。

代码实现

from mp_api.client import MPRester

# 替换为你的API密钥
API_KEY = "your_api_key"

# 使用MPRester获取数据
with MPRester(API_KEY) as mpr:
    materials = mpr.materials.summary.search(
        fields=["material_id", "structure", "band_gap"],
        num_elements=(1, 5),  # 限制元素数量在1到5之间
        limit=50  # 获取50种材料
    )

# 提取材料ID、结构和带隙
material_ids = [mat.material_id for mat in materials]
structures = [mat.structure for mat in materials]
band_gaps = [mat.band_gap for mat in materials]

# 输出验证
for mid, bg in zip(material_ids[:5], band_gaps[:5]):  # 只打印前5个以验证
    print(f"Material ID: {mid}, Band Gap: {bg} eV")

参数设置与含义

fields=["material_id", "structure", "band_gap"]：
- 指定返回的字段：材料ID、晶体结构和带隙。
num_elements=(1, 5)：
- 限制材料的元素数量在1到5之间，简化数据集，同时保持多样性。
limit=50：
- 获取50种材料，比SchNet案例中的20种更多，增加数据量有助于提高模型泛化能力。

数据预处理：将晶体结构转换为ALIGNN图数据

步骤解读

ALIGNN需要将晶体结构转换为图格式，包括原子图（描述原子和键）和线图（描述键和键角）。我们使用jarvis-tools提供的工具进行转换。

代码实现

from jarvis.core.atoms import Atoms
from jarvis.core.graphs import Graph
import torch

def structure_to_alignn_graph(structure):
    """
    将pymatgen的Structure对象转换为ALIGNN所需的图数据。
    参数:
        structure: pymatgen的Structure对象
    返回:
        graph: DGL图对象，包含原子和键角信息
    """
    # 转换为JARVIS的Atoms格式
    atoms = Atoms.from_pymatgen(structure)
    # 生成ALIGNN所需的图数据
    graph = Graph.atom_dgl_multigraph(
        atoms,
        cutoff=8.0,  # 截断距离
        max_neighbors=12  # 每个原子的最大邻居数
    )
    return graph

# 将所有结构转换为图数据
graph_data_list = [structure_to_alignn_graph(struct) for struct in structures]

# 为每个图数据添加带隙标签
for graph, bg in zip(graph_data_list, band_gaps):
    graph.ndata['target'] = torch.tensor([bg], dtype=torch.float32)

参数设置与含义

cutoff=8.0：
- 含义：截断距离（单位：Å），决定构建图时考虑的原子间最大距离。
- 设置建议：默认8.0 Å适用于大多数晶体结构，若材料具有较大的周期性单元，可增加到10.0 Å。
max_neighbors=12：
- 含义：每个原子的最大邻居数，限制图的连边数量。
- 设置建议：12是一个经验值，适合大多数材料；若材料结构稀疏，可减小到8，若稠密，可增加到16。
ndata['target']：
- 将带隙值作为节点属性存储，便于后续训练。

构建ALIGNN模型

步骤解读

ALIGNN结合了原子图和线图，通过多层GNN更新特征，最终输出材料属性。我们直接使用alignn库提供的实现。

代码实现

from alignn.models.alignn import ALIGNN
import torch.nn as nn

class BandGapALIGNN(nn.Module):
    def __init__(self, hidden_features=256, edge_features=128, triplet_input_features=64):
        super().__init__()
        self.alignn = ALIGNN(
            hidden_features=hidden_features,  # 隐藏层特征维度
            edge_features=edge_features,      # 边特征维度
            triplet_input_features=triplet_input_features,  # 三体特征维度
            output_features=1                 # 输出带隙值（标量）
        )

    def forward(self, graph):
        # ALIGNN前向传播
        out = self.alignn(graph)
        return out

参数设置与含义

hidden_features=256：
- 含义：隐藏层特征的维度，决定模型的表示能力。
- 设置建议：默认256，若数据集较小可减至128，若需更高精度可增至512。
edge_features=128：
- 含义：边特征的维度，用于表示键的特性。
- 设置建议：默认128，适用于大多数任务；若键类型复杂，可增至256。
triplet_input_features=64：
- 含义：三体（键角）特征的初始维度，用于线图更新。
- 设置建议：默认64，若模型性能不足，可增至128。
output_features=1：
- 含义：输出层的维度，因为带隙是标量，故设为1。

数据集划分和训练准备

步骤解读

将数据分为训练集（80%）和测试集（20%），并使用DGL的DataLoader加载数据。

代码实现

from dgl.dataloader import GraphDataLoader

# 划分训练集和测试集
train_data = graph_data_list[:40]
test_data = graph_data_list[40:]

# 创建数据加载器
train_loader = GraphDataLoader(train_data, batch_size=8, shuffle=True)
test_loader = GraphDataLoader(test_data, batch_size=8, shuffle=False)

参数设置与含义

batch_size=8：
- 含义：每个批次包含8个图，平衡计算效率和内存使用。
- 设置建议：根据GPU内存选择，4~16常见，默认8适合中等规模数据。
shuffle=True：
- 含义：训练时打乱数据顺序，提升模型泛化能力。
- 设置建议：训练时始终启用，测试时关闭。

模型训练

步骤解读

使用均方误差（MSE）作为损失函数，训练ALIGNN模型以预测带隙。

代码实现

# 初始化模型、优化器和损失函数
model = BandGapALIGNN(hidden_features=256, edge_features=128, triplet_input_features=64)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)
criterion = nn.MSELoss()

# 训练函数
def train():
    model.train()
    total_loss = 0
    for batch in train_loader:
        optimizer.zero_grad()
        out = model(batch)
        loss = criterion(out, batch.ndata['target'].view(-1, 1))
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(train_loader)

# 训练模型
for epoch in range(200):
    loss = train()
    if (epoch + 1) % 20 == 0:
        print(f"Epoch {epoch+1}, Loss: {loss:.4f}")

参数设置与含义

lr=0.0005：
- 含义：学习率，控制参数更新的步长。
- 设置建议：ALIGNN推荐较小的学习率（0.0001~0.001），此处设为0.0005以平衡收敛速度和稳定性。
criterion=nn.MSELoss()：
- 含义：均方误差损失，用于回归任务。
- 设置建议：带隙预测是回归问题，MSE是标准选择。
epochs=200：
- 含义：训练轮数，决定模型训练的充分程度。
- 设置建议：200~500常见，需根据损失收敛调整。

模型评估

步骤解读

在测试集上评估模型性能，使用平均绝对误差（MAE）衡量预测精度。

代码实现

def test(loader):
    model.eval()
    total_error = 0
    with torch.no_grad():
        for batch in loader:
            out = model(batch)
            error = torch.abs(out - batch.ndata['target'].view(-1, 1)).sum().item()
            total_error += error
    return total_error / len(loader.dataset)

# 测试模型
test_error = test(test_loader)
print(f"Test MAE: {test_error:.4f} eV")

参数设置与含义

无需额外参数，MAE直接反映预测误差的平均值，单位为eV。

预测新材料的带隙

步骤解读

使用训练好的模型预测一个新材料的带隙，例如mp-150。

代码实现

# 获取新材料
with MPRester(API_KEY) as mpr:
    new_material = mpr.materials.summary.search(material_ids=["mp-150"], fields=["structure"])
new_structure = new_material[0].structure

# 转换为图数据
new_graph = structure_to_alignn_graph(new_structure)

# 预测
model.eval()
with torch.no_grad():
    prediction = model(new_graph)
print(f"Predicted band gap for mp-150: {prediction.item():.4f} eV")

参数设置与含义

无需额外参数，确保new_graph与训练数据格式一致。

案例总结与参数设置指南

总结

通过这个案例，我们展示了如何使用更先进的ALIGNN模型结合MP API预测材料带隙。ALIGNN通过原子图和线图的联合建模，捕捉了更丰富的几何信息，相比SchNet具有更高的预测精度。以下是关键参数的总结和设置建议：

ALIGNN模型参数

hidden_features=256：
- 含义：隐藏层维度，控制模型容量。
- 建议：128~512，默认256适合中等数据集。
edge_features=128：
- 含义：边特征维度，影响键信息建模。
- 建议：64~256，默认128。
triplet_input_features=64：
- 含义：三体特征维度，影响线图更新。
- 建议：32~128，默认64。

数据预处理参数

cutoff=8.0：
- 含义：截断距离，决定图的连边范围。
- 建议：6.0~10.0 Å，默认8.0。
max_neighbors=12：
- 含义：最大邻居数，限制图的稠密程度。
- 建议：8~16，默认12。

训练参数

batch_size=8：
- 含义：批次大小，影响训练效率。
- 建议：4~16，默认8。
lr=0.0005：
- 含义：学习率，控制优化步长。
- 建议：0.0001~0.001，默认0.0005。
epochs=200：
- 含义：训练轮数。
- 建议：100~500，视收敛情况调整。