Aaditya Ura
Aaditya Ura

Reputation: 12679

Convert numpy adjacency matrix and features to torch geometric data object?

I have adj matrix and features per node in NumPy format:

num_of_nodes     = 14
num_of_feats     = 14
output_dim       = 32
    
    
adj              = np.random.randint(0,2,[num_of_nodes,num_of_nodes])
node_feats       = np.random.uniform(-1,1,[num_of_nodes, num_of_feats])

I want to convert it into Torch geometric dataset format. What I have tried:

I am converting the data first into networkx format and then pyg format

from torch_geometric.utils.convert import from_networkx
import numpy as np
import networkx as nx



num_of_nodes = 14
num_of_feats = 14
output_dim   = 32


def numpy_to_graph(A, node_features=None):
    
    G = nx.from_numpy_array(A)
    if node_features != None:
        for n in G.nodes():
            for k,v in node_features.items():
                G.nodes[n][k] = v[n]
    return G


def get_data():
    
    adj         = np.random.randint(0,2,[num_of_nodes,num_of_nodes])
    feats       = np.random.uniform(-1,1,[num_of_nodes, num_of_feats])
    G           = numpy_to_graph(adj, node_features={'x': feats})
    pyg_graph   = from_networkx(G)

    return pyg_graph

is it correct? and is there any better way to do this?

Upvotes: 1

Views: 421

Answers (1)

mindless
mindless

Reputation: 11

You can do this compactly and at higher speed by first converting the adj to a torch tensor and then using torch.nonzero method to it. This will return all row and column indices in 1D for all positions where the entry is nonzero. You can then stack the indexes returned from it using torch.cat method to build adjacency matrix in COO format as required by torch geometric.

Upvotes: 1

Related Questions