CODEWITHSUNDEEP
pythonnetworkxpytorch-geometric
Ssong
Ssong

Reputation: 452

How can I map the node index in pytorch geometric graph?

I'd like to mapping node index to the original name in pytorch geometric graph for extracting node embedding.

import numpy as np
import torch
import pandas as pd

data = {'source': ['123', '2323', '545', '4928', '398'],
       'target': ['2323', '398', '958', '203', '545']}
df = pd.DataFrame(data)
df
#source target
#0  123 2323
#1  2323    398
#2  545 958
#3  4928    203
#4  398 545

Note that I inserted additional nodes without edges.

import networkx as nx
G = nx.from_pandas_edgelist(df, 'source', 'target')
G = nx.relabel_nodes(G, { n:str(n) for n in G.nodes()})

G.add_nodes_from(['1', '309', '6749'])
G.number_of_nodes()

Then I converted the networkx graph to pytorch geometric graph.

from torch_geometric.utils.convert import from_networkx

pyg_graph = from_networkx(G)
print(pyg_graph)
#Data(edge_index=[2, 10], num_nodes=10)

Finally, I get the below edge index that need to mapping index to name.
(It is normal situation that PyG graph has bidirectional graph automatically)

pyg_graph.edge_index
# tensor([[0, 1, 1, 2, 2, 3, 3, 4, 5, 6],
#         [1, 0, 2, 1, 3, 2, 4, 3, 6, 5]])

# Expected result
# 123   0
# 2323  1
# 398   2
# 545   3
# 958   4
# 4928  5
# 203   6
# 1     7 
# 309   8
# 6749  9

Upvotes: 0

Views: 942

Answers (1)

Enes Altınışık
Enes Altınışık

Reputation: 327

If you set the node label as an attribute, pyg_graph keeps the label as a feature.

nx.set_node_attributes(G,{n:{'label':n} for n in G.nodes()})
pyg_graph = from_networkx(G)
print(pyg_graph.label)
# or for detailed print
for node,lbl in enumerate(pyg_graph.label):
    print(lbl, node)

Upvotes: 0

Related Questions