Search code examples
nodesnetworkxdata-conversionedgespytorch-geometric

Converting a PyG graph to a NetworkX graph


I am trying to convert my PyG graph to a NetworkX graph using to_networkx

According to the docs I can optionally pass node and edge attributes as str iterables, in addition to the Data object.

Below are by node and edge attribute lists, with values converted to strings:

Nodes:  ['3.3375725746154785', '2.0086510181427',..., '1.5960148572921753', '3.621992349624634']

Edges:  ['0.9940207804344958', '0.48573804411542043', ..., '0.7245483440145621', '0.24117984598949904']

to_networkx runs fine when I only pass it the Data object. However, when I also pass these attribute lists, I get the below error:

G[u][v][key] = values[key][i]
KeyError: '0.30194718370332896'

I've looked at the source code, but can't make out what it is doing. Could someone please help explain what is wrong with my attribute lists and what I need to change for them to be accepted.

What I can make out is that this error is specifically referring to my edge attributes. If I remove them, I get the following similar error related to the node attributes:

feat_dict.update({key: values[key][i]})
KeyError: '0.0'

How I construct my graph and pass it to to_networkx:

n1 = np.repeat(np.array([0,1,2,3,4,5,6]),5)
n2 = np.array([0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4])
cat = np.stack((n1,n2), axis=0)
e = torch.tensor(cat, dtype=torch.long)
edge_index = e.t().clone().detach()
edge_attr = torch.tensor(np.random.rand(35,1))

x = torch.tensor([[0], [0], [0], [0], [0], [1], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index.t().contiguous(), edge_attr = edge_attr)

Before I pass the node and edge attributes, I do the string conversion to conform to the str iterable requirment:

networkx_node_values = list(map(str, data.x.t()[0].tolist()))
networkx_edge_values = list(map(str, edge_attr.t()[0].tolist()))
    
networkX_graph = to_networkx(data, node_attrs = networkx_node_values, edge_attrs = networkx_edge_values)

Solution

  • You need to pass the names of the attributes as a list:

    to_networkx(<PyTorchGeometricDataObject>, node_attrs=[<Name of Node Attribute 1>, <Name of Node Attributes 2>, ... ], edge_attr=[<Edge Attribute 1>, ...])
    

    Or in context, based on your given minimal example:

    import numpy as np
    import torch
    from torch_geometric.data import Data
    from torch_geometric.utils import to_networkx
    
    n1 = np.repeat(np.array([0,1,2,3,4,5,6]),5)
    n2 = np.array([0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4])
    cat = np.stack((n1,n2), axis=0)
    e = torch.tensor(cat, dtype=torch.long)
    edge_index = e.t().clone().detach()
    edge_attr = torch.tensor(np.random.rand(35,1))
    
    x = torch.tensor([[0], [0], [0], [0], [0], [1], [1]], dtype=torch.float)
    
    data = Data(x=x, edge_index=edge_index.t().contiguous(), edge_attr = edge_attr)
    print(data)
    # Data(edge_attr=[35, 1], edge_index=[2, 35], x=[7, 1])
    
    networkX_graph = to_networkx(data, node_attrs=["x"], edge_attrs=["edge_attr"])
    
    print(networkX_graph.nodes(data=True))
    # [(0, {'x': 0.0}), (1, {'x': 0.0}),...
    print(networkX_graph.edges(data=True))
    # [(0, 0, {'edge_attr': 0.3412137594357493}), ...