Recreating Graph Visualization in Randomly Wired Neural Network

visualization.py

import networkx as nx


def to_directed(graph):
    G = nx.DiGraph(source=[], sink=[])
    for node in range(graph.number_of_nodes()):
        neighbors = list(graph.neighbors(node))
        neighbors.sort()
        if node < neighbors[0]:  # input nodes
            G.graph['source'].append(node)
            G.add_node(node, node_type=0, fillcolor='deepskyblue1')
        elif node > neighbors[-1]:  # output nodes
            G.graph['sink'].append(node)
            G.add_node(node, node_type=1, fillcolor='firebrick3')
        else:  # intermediate nodes
            G.add_node(node, node_type=-1, fillcolor='white')
        for neighbor in neighbors[::-1]:
            if node >= neighbor:
                break
            G.add_edge(node, neighbor)
    return G


def draw_graph(graph, filename):
    A = nx.nx_agraph.to_agraph(graph)        # convert to a graphviz graph
    A.node_attr.update(label=" ", shape='circle', style='filled')  # label = " " ensures label to be white
    A.add_subgraph(graph.graph['source'], rank='source')
    A.add_subgraph(graph.graph['sink'], rank='sink')
    A.draw(filename, format='png', prog='dot')


if __name__ == "__main__":
    n = 32
    K = 4
    P = 0.75

    G = nx.random_graphs.connected_watts_strogatz_graph(n, K, P, tries=200)
    D = to_directed(G)
    draw_graph(D, "directed.png")