PyG edge index 转换回 邻接矩阵

PyG的edge index形式是$[(nod{e}_1,nod{e}_2),(nod{e}_1,nod{e}_3)...]$这种edge pair。

naive

直接for循环，吧edge index里面的位置填充1：

import torch  
  
def edge_index_to_adjacency_matrix(edge_index, num_nodes):  
    # 创建大小为 (num_nodes, num_nodes) 的二维张量  
    adjacency_matrix = torch.zeros(num_nodes, num_nodes)  
      
    # 根据边索引填充邻接矩阵的元素  
    for i, j in zip(*edge_index):  
        adjacency_matrix[i, j] = 1  
        adjacency_matrix[j, i] = 1  
      
    return adjacency_matrix

效率很低

利用传播机制

用PyTorch的广播机制，通过将边索引直接作为索引，一次性将对应的邻接矩阵元素设置为1，避免了使用for循环进行逐个元素的填充。这种方法在大规模图形上具有更高的效率。

import torch  
  
def edge_index_to_adjacency_matrix(edge_index, num_nodes):  
    # 构建一个大小为 (num_nodes, num_nodes) 的零矩阵  
    adjacency_matrix = torch.zeros(num_nodes, num_nodes, dtype=torch.uint8)  
      
    # 使用索引广播机制，一次性将边索引映射到邻接矩阵的相应位置上  
    adjacency_matrix[edge_index[0], edge_index[1]] = 1  
    adjacency_matrix[edge_index[1], edge_index[0]] = 1  
      
    return adjacency_matrix