AI4S图机器学习：DGL图构建接口的PyG替换

def make_full_graph(xyz, pair, idx, top_k=64, kmin=9):
        B, L = xyz.shape[:2]
        device = xyz.device

        # 确保xyz形状正确 
        if xyz.dim() > 3:
                xyz_flat = xyz[:,:,1] if xyz.shape[2] == 3 else xyz.reshape(B, L, 3)
        else:
                xyz_flat = xyz

        # 计算序列分离
        sep = idx[:,None,:] - idx[:,:,None] 
        b,i,j = torch.where(sep.abs() > 0)

        # 构建PyG图所需的边索引 
        src = b*L+i
        tgt = b*L+j

        # 创建图对象
        G = graph((src, tgt), num_nodes=B*L).to(device)

        # 计算相对位置
        rel_pos = xyz_flat[b,j,:] - xyz_flat[b,i,:]
        if rel_pos.dim() > 2 and rel_pos.shape[-1] == 3:
                rel_pos = rel_pos.reshape(-1, 3)
        G.edata['rel_pos'] = rel_pos.detach()

        # 处理边特征
        edge_feats = pair[b,i,j] 
        if edge_feats.dim() == 1:
                edge_feats = edge_feats.unsqueeze(-1)
        if edge_feats.dim() == 2:
                edge_feats = edge_feats.unsqueeze(-1) 

        # 归一化特征减少实现差异
        edge_feats = torch.tanh(edge_feats / 10.0) * 10.0

        return G, edge_feats