第4章：创建和操作DGL图对象

# 手动指定节点数
g = dgl.graph((u, v), num_nodes=5)
print(g)

4.2 给图添加特征

图建好了，下一步就是塞特征。DGL用ndata和edata两个字典来挂载节点特征和边特征，非常直观。

# 节点特征：假设每个节点有一个32维向量
g.ndata['feat'] = torch.randn(3, 32)
g.ndata['label'] = torch.tensor([0, 1, 0])

# 边特征：每条边有一个16维向量
g.edata['weight'] = torch.randn(3, 16)

print(g.ndata['feat'].shape)
print(g.edata['weight'].shape)

有一点我刚开始总搞混：ndata里每个特征的第一维长度必须等于num_nodes()，edata同理。如果你图有3个节点却传了4个特征向量，DGL会直接抛异常。这个设计其实挺好，帮你提前发现数据对齐的问题。

4.3 dgl.heterograph：处理异构图

药物图神经网络里，图很少是同构的。你想，一个药物分子、一个蛋白质靶点、再加上两者之间的“结合”关系——节点类型就两种，边类型也有“化学结合”和“相互作用”之分。这就得上dgl.heterograph了。

它的输入是一个字典，键是三元组('源节点类型', '边类型', '目标节点类型')，值是(u, v)元组。我来给你写个药物-靶点网络的例子：

# 构建异构图
graph_data = {
    ('drug', 'binds', 'protein'): (torch.tensor([0, 1, 2]), torch.tensor([0, 1, 0])),
    ('drug', 'similar', 'drug'): (torch.tensor([0, 1]), torch.tensor([1, 2]))
}

hg = dgl.heterograph(graph_data)
print(hg)
print('节点类型:', hg.ntypes)
print('边类型:', hg.canonical_etypes)

输出里ntypes会列出所有节点类型，canonical_etypes则给出规范化的边类型三元组。DGL内部会给每种节点类型分配独立的ID空间，所以drug的节点0和protein的节点0是两码事。

为了帮你快速对比，我列了个表格：

4.4 异构图特征操作

给异构图加特征，需要在ndata和edata前面加上节点/边类型前缀。写法是g.nodes['drug'].data['feat']，而不是g.ndata['feat']。

# 给药物节点加特征
hg.nodes['drug'].data['feat'] = torch.randn(3, 64)
hg.nodes['protein'].data['feat'] = torch.randn(2, 128)

# 给特定边类型加特征
hg.edges['binds'].data['affinity'] = torch.randn(3, 1)

print(hg.nodes['drug'].data['feat'].shape)
print(hg.edges['binds'].data['affinity'].shape)

我记得第一次用异构图时，写了个hg.ndata['feat']，结果报错说ndata找不到。后来才意识到，异构图上每个节点类型的特征维度可能不同，DGL不允许你混着来。这个设计很合理——药物节点和蛋白质节点的特征维度不一样，本来就是常事。

4.5 实用操作与避坑指南

我整理了三个最常用的操作，以及我曾经踩过的坑，希望对你有帮助。

4.5.1 获取近邻节点

不管是同构图还是异构图，g.successors()和g.predecessors()都是一把好手。但对于异构图，要指定边类型：

# 同构图：获取节点0的所有后继
neighbors = g.successors(0)

# 异构图：获取drug节点0通过binds边连接的protein节点
neighbors = hg.successors(0, etype='binds')
print(neighbors)

4.5.2 构建子图

我经常需要根据某些条件筛选节点，然后只在这个子集上做计算。DGL提供了node_subgraph和edge_subgraph：

# 提取节点0和1构成的子图
sub_g = dgl.node_subgraph(g, [0, 1])
print(sub_g)

# 异构图也支持：提取drug类型的[0,1]和protein类型的[0]
sub_hg = dgl.node_subgraph(hg, {'drug': [0, 1], 'protein': [0]})
print(sub_hg)

4.5.3 图的序列化

训练好的图数据或模型中间结果，一般用dgl.save_graphs和dgl.load_graphs保存和读取：

# 保存
dgl.save_graphs('./demo_graph.bin', [g, hg])

# 加载
loaded_graphs, _ = dgl.load_graphs('./demo_graph.bin')
loaded_g = loaded_graphs[0]
loaded_hg = loaded_graphs[1]

我曾经在保存图时忘了检查特征的数据类型，结果加载出来发现float64变成了float32。DGL默认会做类型转换，如果你对精度有要求，最好在保存前用.type(torch.float32)统一一下。