Search code examples
pythoncolorsnodesnetworkx

Node color Networkx Python 3.8


I'm creating a graph with Networkx on python 3.8 and I want to assign a color to every single node by the time i'm appending the node to the graph, as follow:

     if raw_output in occurrencies_dict_gold: 
       G.add_node(raw_output, color = 'g')  

     else: 
       G.add_node(raw_output, color = 'r') 
       print('RED STATE : ', raw_output) 

in which i check that a node doesn't exist in a particular dictionary previously created, and i consequently add the node with its color.

I do the same for edges

               if transition_key in transitions_dict_gold:
                 G.add_edge(previous_raw_output, raw_output, color = 'g')
               else:
                 G.add_edge(previous_raw_output, raw_output, color = 'r')

At printing stage i do the following:

edges = G.edges()
 nodes = G.nodes()
 e_colors = [G[u][v]['color'] for u,v in edges]
 n_colors = [G[u]['color'] for u in nodes]
 
 nx.draw(G, node_color=n_colors, edge_color=e_colors,  with_labels=True)
 plt.show()
 
 
 plt.savefig("filename.png") 

which works perfectly if i only do it for the edges, while if i try as written above i recieve the following message:

Traceback (most recent call last):
  File "result_iterative_parser.py", line 157, in <module>
    n_colors = [G[u]['color'] for u in nodes]
  File "result_iterative_parser.py", line 157, in <listcomp>
    n_colors = [G[u]['color'] for u in nodes]
  File "/anaconda3/lib/python3.8/site-packages/networkx/classes/coreviews.py", line 51, in __getitem__
    return self._atlas[key]
KeyError: 'color'

any idea?


Solution

  • Instead of:

    n_colors = [G[u]['color'] for u in nodes]
    

    You could do:

    n_colors = [G.nodes(data='color')[u] for u in nodes]
    

    But the following looks way cleaner:

    nodes = G.nodes(data="color")
    n_colors = [u[1] for u in nodes]
    

    You may be confused by u[1] but it is because G.nodes(data="color")returns a list of nodes as tuples (node, color).


    Your solution doesn't work because G[u] returns the edges(and edge info) of node u. Instead what you need is the information on the node so use .nodes(data=True) or .nodes(data='color')
    Here is a comparison between what each returns:

    In:
    G=nx.Graph()
    
    G.add_node(1, n_color="red")
    G.add_node(2, n_color="blue")
    G.add_node(3, n_color="green")
    G.add_node(4, n_color="red")
    G.add_edge(1, 3, e_color="yellow")
    G.add_edge(2, 4, e_color="yellow")
    G.add_edge(2, 3, e_color="black")
    
    
    print("G[3] returns\t\t\t  ",                  G[3])
    print("G.nodes(data=True)[3] returns\t  ",     G.nodes(data=True)[3])
    print("G.nodes(data='n_color')[3] returns ",   G.nodes(data='n_color')[3])
    print("G.nodes(data='e_color')[3] returns ",   G.nodes(data='e_color')[3])
    
    Out:
    G[3] returns                        {1: {'e_color': 'yellow'}, 2: {'e_color': 'black'}}
    G.nodes(data=True)[3]returns        {'n_color': 'green'}
    G.nodes(data='n_color')[3] returns  green
    G.nodes(data='e_color')[3] returns  None