graph_from_adjacency_mat.py

import unittest
import networkx as nx

class TestExtract(unittest.TestCase):

    def test_graph_from_adj_mat(self):
        X = np.array([
            [(0,), (1,), (2,)],
            [(1,), (0,), (0,)],
            [(2,), (0,), (0,)]], dtype=[('weight', '<i8')])
        tested = graph_from_adj_mat(X, labels=['A','B','C'])
        expected = {
            'A': {'C': {'weight': 2}, 'B': {'weight': 1}},
            'B': {'A': {'weight': 1}},
            'C': {'A': {'weight': 2}}}
        for k in tested:
            msg = '\nExpected:\n{}\nGot:\n{}'.format(expected, tested)
            self.assertEqual(expected[k], tested[k], msg=msg)

def graph_from_adj_mat(adj_mat, labels=None):
    G = nx.from_numpy_matrix(adj_mat)
    if labels:
        nx.relabel_nodes(G, dict(enumerate(labels)), copy=False)
    return G