Build networkx/dgl graph with from numpy representations of feature and adjacency matrix

Question

Description

Generate a graph object (either DGL or NetworkX) from an Adjacency matrix and allow for node features to be established.

Result

I generate my solution below. However, other answers are encouraged.

Answer 1

Code

import numpy as np
import dgl
import networkx as nx

def numpy_to_graph(A,type_graph='dgl',node_features=None):
    '''Convert numpy arrays to graph

    Parameters
    ----------
    A : mxm array
        Adjacency matrix
    type_graph : str
        'dgl' or 'nx'
    node_features : dict
        Optional, dictionary with key=feature name, value=list of size m
        Allows user to specify node features

    Returns

    -------
    Graph of 'type_graph' specification
    '''
    
    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]
    
    if type_graph == 'nx':
        return G
    
    G = G.to_directed()
    
    if node_features != None:
        node_attrs = list(node_features.keys())
    else:
        node_attrs = []
        
    g = dgl.from_networkx(G, node_attrs=node_attrs, edge_attrs=['weight'])
    return g

Example

The adjacency matrix is passed to the function. Additionally, other features (i.e. Feature vectors, Labels, etc.) can be passed in node_features

# mxm adjacency matrix
A = np.array([[0,0,0],
              [2,0,0],
              [5,1,0]])

# Each m row is a feature vector for node m
F = np.array([[1,0,1,4,4],
             [2,4,0,12,4],
             [5,1,-4,2,9]])


G = numpy_to_graph(A,type_graph='nx',node_features={'feat':F})

import matplotlib.pyplot as plt
pos=nx.spring_layout(G) # pos = nx.nx_agraph.graphviz_layout(G)
nx.draw_networkx(G,pos)
labels = nx.get_edge_attributes(G,'weight')
nx.draw_networkx_edge_labels(G,pos,edge_labels=labels)
plt.show()