Python NetworkX — set node color automatically based on a list of values

Question

I generated a graph with networkx

import networkx as nx  
s = 5
G = nx.grid_graph(dim=[s,s])
nodes = list(G.nodes)
edges = list(G.edges)
p = []
for i in range(0, s):
    for j in range(0, s):
        p.append([i,j])
for i in range(0, len(nodes)):
    G.nodes[nodes[i]]['pos'] = p[i]

pos = {}
for i in range(0, len(nodes)):
        pos[nodes[i]] = p[i]

nx.draw(G, pos)

Now I would like to assign a value to each node between 0 and 4

from random import randint
val = []
for i in range(0, len(G.nodes())):
    val.append(randint(0,4)) 

And I would like to assign the color to each node base on the val list and plot something like shown here

Answer 1

networkx.draw has the node_color, vmin, vmax and cmap parameters:

cmap (Matplotlib colormap, optional (default=None)) – Colormap for mapping intensities of nodes

vmin,vmax (float, optional (default=None)) – Minimum and maximum for node colormap scaling

node_color (color string, or array of floats, (default=’#1f78b4’)) – Node color. Can be a single color format string, or a sequence of colors with the same length as nodelist. If numeric values are specified they will be mapped to colors using the cmap and vmin,vmax parameters. See matplotlib.scatter for more details.

You can write a list in it so your nodes will be colored (for example):

colors = [i/len(G.nodes) for i in range(len(G.nodes))]
...
...
nx.draw(G, pos, node_color=colors)

Answer 2

To set a node property, you can use:

nx.set_node_attributes(G, val, 'val')

Networkx draw calls down to draw_networkx_nodes, this takes a cmap and color list, so all you would have to do would be something like:

nx.draw(G, pos, node_color = nx.get_node_attributes(G,'val'), vmin=0, vmax=4, cmap = plt.cm.get_cmap('rainbow'))