基于材料热电数据集构建含属性的图以实现GNN预测

核心逻辑 #

将Pymatgen的Structure对象转换成GNN框架可识别的图结构，关键是提取节点属性（原子特征）和边属性（原子间连接特征），再适配目标GNN库的数据格式。

1. 提取节点属性 #

节点对应结构中的单个原子，常用特征可从元素基础属性中提取：

• 原子序数
• 原子量
• Pauling电负性
• 价电子数
• 共价半径

示例代码：

from pymatgen.core.periodic_table import Element

def extract_node_features(structure):
    node_features = []
    for site in structure:
        elem = Element(site.species.elements[0].symbol)
        feats = [
            elem.Z,
            elem.atomic_mass,
            elem.X,
            elem.number_of_valence_electrons,
            elem.covalent_radius or 0.0  # 部分元素无共价半径时补0
        ]
        node_features.append(feats)
    return node_features

2. 提取边属性（连接特征） #

边代表原子间的连接关系，先通过截断距离判断成键，再提取边的特征：

• 原子间距离
• 电负性差值
• 共价半径差值

示例代码：

import numpy as np

def extract_edge_features(structure, cutoff=3.0):
    # 计算原子间距离矩阵
    dist_matrix = structure.distance_matrix
    # 筛选符合截断距离的原子对（排除自身连接）
    edge_indices = np.where((dist_matrix < cutoff) & (dist_matrix > 1e-8))
    edge_indices = np.array(edge_indices).T  # 转为[num_edges, 2]格式
    
    edge_features = []
    for i, j in edge_indices:
        elem_i = Element(structure[i].species.elements[0].symbol)
        elem_j = Element(structure[j].species.elements[0].symbol)
        # 构建边特征
        dist = dist_matrix[i][j]
        en_diff = abs(elem_i.X - elem_j.X)
        radius_diff = abs((elem_i.covalent_radius or 0.0) - (elem_j.covalent_radius or 0.0))
        edge_features.append([dist, en_diff, radius_diff])
    
    return edge_indices, edge_features

3. 适配PyTorch Geometric数据格式 #

以主流的PyTorch Geometric为例，将特征封装为Data对象：

import torch
from torch_geometric.data import Data

def structure_to_graph(structure, target_value):
    # 转换节点特征为张量
    x = torch.tensor(extract_node_features(structure), dtype=torch.float)
    # 转换边索引和边特征为张量
    edge_index, edge_attr = extract_edge_features(structure)
    edge_index = torch.tensor(edge_index, dtype=torch.long).T  # PyG要求边索引为[2, num_edges]格式
    edge_attr = torch.tensor(edge_attr, dtype=torch.float)
    # 目标属性张量（比如热电ZT值）
    y = torch.tensor([target_value], dtype=torch.float)
    
    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)

4. 批量处理数据集 #

遍历你的DataFrame，批量生成图数据列表：

graph_dataset = []
# 假设df是原始数据集，包含目标属性列（比如'zt_value'）
for idx, row in df1.iterrows():
    target = df.loc[df['mpid'] == row['material_id'], 'zt_value'].values[0]
    graph = structure_to_graph(row['structure'], target)
    graph_dataset.append(graph)

注意事项 #

• 截断距离cutoff需根据材料体系调整：共价键材料可选1.8-2.5Å，金属体系可选3.0-4.0Å
• 节点特征可按需扩充，比如添加电离能、电子亲和能等元素属性
• 若使用DGL框架，仅需将边索引和特征转换为DGLGraph格式，核心特征提取逻辑一致

内容的提问来源于stack exchange，提问作者Delchere Dontsa