Node attributes in for loops, NetworkX - networkx

Question

I'm trying to model voting dynamics on networks, and would like to be able to create a graph in NetworkX where I can iterate the voter process on nodes, having their colour change corresponding to their vote 'labels'.
I've managed to get this code to let me see the attributes for each node, but how do I go about using those in a for loop to designate colour?
H = nx.Graph()
H.add_node(1,vote='labour')
H.add_node(2,vote='labour')
H.add_node(3,vote='conservative')
h=nx.get_node_attributes(H,'vote')
h.items()
Gives me the result:
[(1, 'labour'), (2, 'labour'), (3, 'conservative')]
I've got a for loop to do this type of colour coding based on the node number as follows, but haven't managed to make it work for my 'vote' status.
S=nx.star_graph(10)
colour_map=[]
for node in S:
if node % 2 ==0:
colour_map.append('blue')
else: colour_map.append('yellow')
nx.draw(S, node_color = colour_map,with_labels = True)
plt.show()

Answer 1

You can iterate the node attributes with H.nodes(data=True) which returns the node name and the node attributes in a dictionary. Here's a full example using your graph.
import networkx as nx
import matplotlib.pyplot as plt
H = nx.Graph()
H.add_node(1, vote='labour')
H.add_node(2, vote='labour')
H.add_node(3, vote='conservative')
color_map = []
for node, data in H.nodes(data=True):
if data['vote'] == 'labour':
color_map.append(0.25) # blue color
elif data['vote'] == 'conservative':
color_map.append(0.7) # yellow color
nx.draw(H, vmin=0, vmax=1, cmap=plt.cm.jet, node_color=color_map, with_labels=True)
plt.show()
This code will draw a different layout of nodes each time you run it (some layouts, as e.g. draw_spring, are available here).
Regarding colors, I use 0.25 for blue and 0.7 for yellow. Note that I use the jet matplotlib colormap and that I set vmin=0 and vmax=1 so that the color values are absolute (and not relative to eachother).
Output of the code above:
UPDATE:
I wasn't aware that you could simply use color names in matplotlib. Here's the updated for loop:
for node, data in H.nodes(data=True):
if data['vote'] == 'labour':
color_map.append("blue")
elif data['vote'] == 'conservative':
color_map.append("yellow")
And the updated draw command:
nx.draw(H, node_color=color_map, with_labels=True)
Note that this way you get different shades of blue and yellow than in the image above.