Search code examples
python-3.xcsvnetworkxtopology

Color nodes by networkx


I am generating a network topology diagram through the data in a csv file where s0..s2 and c1..c3 are nodes of the diagram.

network.csv:

source,port,destination

s1,1,c3

s2,1,c1

s0,1,c2

s1,2,s2

s2,2,s0

I need to make all the source to be blue and destinations to be green. How can I do it without overriding the source nodes?


Solution

  • Following is a working solution:

    import csv
    import networkx as nx
    from matplotlib import pyplot as plt
    
    with open('../resources/network.csv') as csvfile:
        reader = csv.DictReader(csvfile)
        edges = {(row['source'], row['destination']) for row in reader }
    print(edges) # {('s1', 'c3'), ('s1', 's2'), ('s0', 'c2'), ('s2', 's0'), ('s2', 'c1')}
    
    G = nx.DiGraph()
    source_nodes = set([edge[0] for edge in edges])
    G.add_edges_from(edges)
    for n in G.nodes():
        G.nodes[n]['color'] = 'b' if n in source_nodes else 'g'
    
    pos = nx.spring_layout(G)
    colors = [node[1]['color'] for node in G.nodes(data=True)]
    nx.draw_networkx(G, pos, with_labels=True, node_color=colors)
    plt.show()
    

    We first read the csv to an edge list, which is later used for the construction of G. For well defining the colors we set each source node with blue and the rest of the nodes as green (i.e., all destination nodes that also not source nodes).

    We also use nx.draw_networkx to get a more compact implementation for drawing the graph.

    The result should be something like:

    enter image description here