Python NetworkX -- set node color automatically based on number of attribute options

Question

I am using NetworkX to analyze and visualize social networks. Often, the nodes within the network have attribute information associated with them, such as division. However, I do not always know how many options there might be for the attribute division. For example, sometimes there might be only 3 divisions represented as node attributes within the network, other times there might be 30 divisions.

I've figured out how to set node colors based on node attributes (see code below). However, in this example, I knew how different options there were for the node attribute group (5 options), and set each color automatically.

When there are only 3 or 5 options for node attributes, it's not difficult to pick the node attribute colors, but this becomes unrealistic when there are many more options.

What I would like to figure out is how can I automatically pick the optimal node attribute colors based on the number of node attribute options provided.

Sometime I have have 5 options for the attribute to color by, other times I might have 30 options for the node attribute to color by, and I don't want to set each node color individually.

I'm not sure if this is something that I should be able to do with the colormap function, or if that is only for color nodes by numeric measures such as degree centrality.

NETWORKX CODE

import networkx as nx

pylab inline

# create an empty graph
g = nx.Graph()

# open csv edgelist and read edges into graph
for line in open('phils_network_edgelist.csv', 'rb'):
    edge = line.rstrip().split(',')
    g.add_edge(edge[0], edge[1])

# draw network without node color
nx.draw(g, with_labels=False, node_size=25)

network without node attribute color

# read in node attributes as list of tuples
group_attr = []
for line in open('phils_network_attribute_group.csv', 'rb'):
    group_attr.append(tuple(line.rstrip().split(',')))

# convert list of tuples into a dict
group_attr_dict = dict(set(sorted(group_attr)))

# set nodes attributes
nx.set_node_attributes(g, "group", group_attr_dict)

# create empty list for node colors
node_color = []

# for each node in the graph
for node in g.nodes(data=True):

    # if the node has the attribute group1
    if 'group1' in node[1]['group']:
        node_color.append('blue')

    # if the node has the attribute group1
    elif 'group2' in node[1]['group']:
        node_color.append('red')

    # if the node has the attribute group1
    elif 'group3' in node[1]['group']:
        node_color.append('green')

    # if the node has the attribute group1
    elif 'group4' in node[1]['group']:
        node_color.append('yellow')

    # if the node has the attribute group1
    elif 'group5' in node[1]['group']:
        node_color.append('orange')  

# draw graph with node attribute color
nx.draw(g, with_labels=False, node_size=25, node_color=node_color)

network with node attribute color

NETWORK DATA

In[58]: 

g.nodes(data=True)

Out[58]:

[('BD', {'group': 'group5'}),
 ('WC', {'group': 'group3'}),
 ('BA', {'group': 'group4'}),
 ('WM', {'group': 'group3'}),
 ('JR', {'group': 'group1'}),
 ('JS', {'group': 'group3'}),
 ('JL', {'group': 'group4'}),
 ('JM', {'group': 'group2'}),
 ('JK', {'group': 'group2'}),
 ('JF', {'group': 'group2'}),
 ('JG', {'group': 'group2'}),
 ('JA', {'group': 'group2'}),
 ('JB', {'group': 'group4'}),
 ('JC', {'group': 'group4'}),
 ('RR', {'group': 'group3'}),
 ('RS', {'group': 'group3'}),
 ('TTI', {'group': 'group3'}),
 ('RB', {'group': 'group1'}),
 ('RL', {'group': 'group3'}),
 ('RO', {'group': 'group4'}),
 ('LHA', {'group': 'group2'}),
 ('LHI', {'group': 'group1'}),
 ('GF', {'group': 'group2'}),
 ('GB', {'group': 'group4'}),
 ('EM', {'group': 'group2'}),
 ('HR', {'group': 'group5'}),
 ('BS', {'group': 'group3'}),
 ('HH', {'group': 'group4'}),
 ('HA', {'group': 'group1'}),
 ('PS', {'group': 'group1'}),
 ('PW', {'group': 'group1'}),
 ('PB', {'group': 'group1'}),
 ('PC', {'group': 'group5'}),
 ('MFR', {'group': 'group4'}),
 ('JMA', {'group': 'group5'}),
 ('PN', {'group': 'group4'}),
 ('PL', {'group': 'group3'}),
 ('ZL', {'group': 'group4'}),
 ('EB', {'group': 'group2'}),
 ('ET', {'group': 'group3'}),
 ('EW', {'group': 'group1'}),
 ('ER', {'group': 'group3'}),
 ('MF', {'group': 'group3'}),
 ('MA', {'group': 'group4'}),
 ('MM', {'group': 'group2'}),
 ('MN', {'group': 'group4'}),
 ('MH', {'group': 'group3'}),
 ('MK', {'group': 'group2'}),
 ('JLA', {'group': 'group2'}),
 ('MP', {'group': 'group1'}),
 ('MS', {'group': 'group4'}),
 ('MR', {'group': 'group4'}),
 ('FI', {'group': 'group5'}),
 ('CJ', {'group': 'group4'}),
 ('CO', {'group': 'group5'}),
 ('CM', {'group': 'group4'}),
 ('CB', {'group': 'group2'}),
 ('CG', {'group': 'group2'}),
 ('CF', {'group': 'group5'}),
 ('CD', {'group': 'group3'}),
 ('CS', {'group': 'group2'}),
 ('CP', {'group': 'group2'}),
 ('CV', {'group': 'group2'}),
 ('KC', {'group': 'group1'}),
 ('KB', {'group': 'group3'}),
 ('SY', {'group': 'group2'}),
 ('KF', {'group': 'group2'}),
 ('KD', {'group': 'group3'}),
 ('KH', {'group': 'group1'}),
 ('SW', {'group': 'group1'}),
 ('KL', {'group': 'group2'}),
 ('KP', {'group': 'group3'}),
 ('KW', {'group': 'group1'}),
 ('SM', {'group': 'group2'}),
 ('SB', {'group': 'group4'}),
 ('DJ', {'group': 'group2'}),
 ('DD', {'group': 'group2'}),
 ('DV', {'group': 'group5'}),
 ('BJ', {'group': 'group3'}),
 ('DR', {'group': 'group2'}),
 ('KWI', {'group': 'group4'}),
 ('TW', {'group': 'group2'}),
 ('TT', {'group': 'group2'}),
 ('LH', {'group': 'group3'}),
 ('LW', {'group': 'group3'}),
 ('TM', {'group': 'group3'}),
 ('LS', {'group': 'group3'}),
 ('LP', {'group': 'group2'}),
 ('TG', {'group': 'group3'}),
 ('JCU', {'group': 'group2'}),
 ('AL', {'group': 'group1'}),
 ('AP', {'group': 'group3'}),
 ('AS', {'group': 'group3'}),
 ('IM', {'group': 'group4'}),
 ('AW', {'group': 'group3'}),
 ('HHI', {'group': 'group1'})]

In [59]:

g.edges(data=True)

Out[59]:

[('BD', 'ZL', {}),
 ('BD', 'JCU', {}),
 ('BD', 'DJ', {}),
 ('BD', 'BA', {}),
 ('BD', 'CB', {}),
 ('BD', 'CG', {}),
 ('BD', 'AS', {}),
 ('BD', 'MH', {}),
 ('BD', 'AP', {}),
 ('BD', 'HH', {}),
 ('BD', 'TM', {}),
 ('BD', 'CF', {}),
 ('BD', 'CP', {}),
 ('BD', 'DR', {}),
 ('BD', 'CV', {}),
 ('BD', 'EB', {}),
 ('WC', 'JCU', {}),
 ('WC', 'JS', {}),
 ('BA', 'JR', {}),
 ('BA', 'JB', {}),
 ('BA', 'RR', {}),
 ('BA', 'RS', {}),
 ('BA', 'LH', {}),
 ('BA', 'PC', {}),
 ('BA', 'TTI', {}),
 ('BA', 'PL', {}),
 ('BA', 'JCU', {}),
 ('BA', 'CF', {}),
 ('BA', 'EB', {}),
 ('BA', 'GF', {}),
 ('BA', 'AS', {}),
 ('BA', 'IM', {}),
 ('BA', 'BJ', {}),
 ('BA', 'CS', {}),
 ('BA', 'KH', {}),
 ('BA', 'SW', {}),
 ('BA', 'MH', {}),
 ('BA', 'MR', {}),
 ('BA', 'HHI', {}),
 ('WM', 'EM', {}),
 ('WM', 'JCU', {}),
 ('WM', 'CO', {}),
 ('WM', 'LP', {}),
 ('WM', 'AW', {}),
 ('WM', 'KD', {}),
 ('WM', 'TT', {}),
 ('WM', 'JS', {}),
 ('WM', 'PB', {}),
 ('WM', 'JM', {}),
 ('WM', 'MFR', {}),
 ('WM', 'RB', {}),
 ('WM', 'MR', {}),
 ('WM', 'DV', {}),
 ('WM', 'TG', {}),
 ('WM', 'JF', {}),
 ('WM', 'JMA', {}),
 ('WM', 'FI', {}),
 ('WM', 'JB', {}),
 ('JR', 'GF', {}),
 ('JR', 'MFR', {}),
 ('JR', 'KH', {}),
 ('JR', 'JB', {}),
 ('JS', 'EM', {}),
 ('JS', 'PS', {}),
 ('JS', 'MF', {}),
 ('JS', 'JCU', {}),
 ('JS', 'KD', {}),
 ('JS', 'MH', {}),
 ('JS', 'TTI', {}),
 ('JS', 'RB', {}),
 ('JS', 'TG', {}),
 ('JL', 'KB', {}),
 ('JL', 'MN', {}),
 ('JL', 'LW', {}),
 ('JL', 'CS', {}),
 ('JL', 'ET', {}),
 ('JL', 'ER', {}),
 ('JM', 'EM', {}),
 ('JM', 'PS', {}),
 ('JM', 'KD', {}),
 ('JM', 'CD', {}),
 ('JM', 'JK', {}),
 ('JM', 'TG', {}),
 ('JM', 'RO', {}),
 ('JM', 'CV', {}),
 ('JK', 'HR', {}),
 ('JK', 'PS', {}),
 ('JF', 'EM', {}),
 ('JF', 'PS', {}),
 ('JF', 'LP', {}),
 ('JF', 'LHA', {}),
 ('JF', 'CD', {}),
 ('JF', 'RB', {}),
 ('JF', 'JG', {}),
 ('JF', 'KF', {}),
 ('JG', 'CJ', {}),
 ('JG', 'SY', {}),
 ('JG', 'KF', {}),
 ('JG', 'LHA', {}),
 ('JG', 'CD', {}),
 ('JG', 'RB', {}),
 ('JG', 'BS', {}),
 ('JA', 'CS', {}),
 ('JB', 'KC', {}),
 ('JB', 'JCU', {}),
 ('JB', 'MA', {}),
 ('JB', 'AW', {}),
 ('JB', 'KWI', {}),
 ('JB', 'KH', {}),
 ('JB', 'CF', {}),
 ('JB', 'EB', {}),
 ('JB', 'PB', {}),
 ('JB', 'MFR', {}),
 ('JB', 'KW', {}),
 ('JB', 'RB', {}),
 ('JB', 'MR', {}),
 ('JB', 'RL', {}),
 ('JB', 'FI', {}),
 ('JB', 'JMA', {}),
 ('JC', 'SM', {}),
 ('RR', 'MS', {}),
 ('RR', 'SW', {}),
 ('RR', 'LH', {}),
 ('RS', 'LH', {}),
 ('TTI', 'JCU', {}),
 ('TTI', 'SW', {}),
 ('TTI', 'CF', {}),
 ('RB', 'EM', {}),
 ('RB', 'PS', {}),
 ('RB', 'SY', {}),
 ('RB', 'JCU', {}),
 ('RB', 'KD', {}),
 ('RB', 'CF', {}),
 ('RB', 'LHI', {}),
 ('RB', 'CD', {}),
 ('RB', 'MH', {}),
 ('RB', 'CJ', {}),
 ('RB', 'TG', {}),
 ('RB', 'EB', {}),
 ('RO', 'PS', {}),
 ('LHA', 'CJ', {}),
 ('LHA', 'SY', {}),
 ('LHA', 'KF', {}),
 ('LHA', 'CD', {}),
 ('LHI', 'PS', {}),
 ('LHI', 'CJ', {}),
 ('GF', 'KC', {}),
 ('GF', 'MA', {}),
 ('GB', 'HR', {}),
 ('GB', 'MM', {}),
 ('GB', 'LS', {}),
 ('EM', 'LP', {}),
 ('EM', 'DV', {}),
 ('EM', 'TG', {}),
 ('HR', 'MM', {}),
 ('HR', 'MH', {}),
 ('HR', 'EB', {}),
 ('HR', 'LS', {}),
 ('BS', 'CD', {}),
 ('HH', 'ZL', {}),
 ('HH', 'CB', {}),
 ('HH', 'CP', {}),
 ('HH', 'DR', {}),
 ('HH', 'CV', {}),
 ('HA', 'SM', {}),
 ('PS', 'KD', {}),
 ('PS', 'CF', {}),
 ('PS', 'TG', {}),
 ('PW', 'CM', {}),
 ('PW', 'TW', {}),
 ('PW', 'TT', {}),
 ('PW', 'MH', {}),
 ('PW', 'AL', {}),
 ('PW', 'MP', {}),
 ('PW', 'CS', {}),
 ('PW', 'HHI', {}),
 ('PW', 'EW', {}),
 ('PB', 'CO', {}),
 ('PB', 'KH', {}),
 ('PB', 'CF', {}),
 ('PB', 'MFR', {}),
 ('PB', 'AW', {}),
 ('PB', 'MA', {}),
 ('PC', 'CS', {}),
 ('PC', 'JCU', {}),
 ('PC', 'SW', {}),
 ('MFR', 'KC', {}),
 ('MFR', 'JCU', {}),
 ('MFR', 'KH', {}),
 ('MFR', 'MH', {}),
 ('MFR', 'MR', {}),
 ('JMA', 'KWI', {}),
 ('JMA', 'AW', {}),
 ('PN', 'SB', {}),
 ('PL', 'HHI', {}),
 ('PL', 'MK', {}),
 ('PL', 'LH', {}),
 ('ZL', 'CB', {}),
 ('ZL', 'AP', {}),
 ('ZL', 'CP', {}),
 ('ZL', 'DR', {}),
 ('ZL', 'CV', {}),
 ('EB', 'JCU', {}),
 ('EB', 'DJ', {}),
 ('EB', 'CM', {}),
 ('EB', 'SW', {}),
 ('EB', 'MM', {}),
 ('EB', 'LS', {}),
 ('EB', 'CS', {}),
 ('EB', 'CP', {}),
 ('EB', 'CV', {}),
 ('ET', 'LW', {}),
 ('ET', 'ER', {}),
 ('ET', 'KB', {}),
 ('EW', 'TW', {}),
 ('EW', 'TT', {}),
 ('EW', 'HHI', {}),
 ('EW', 'AL', {}),
 ('ER', 'LW', {}),
 ('ER', 'KB', {}),
 ('MA', 'KW', {}),
 ('MA', 'AW', {}),
 ('MA', 'MR', {}),
 ('MM', 'LS', {}),
 ('MH', 'JCU', {}),
 ('MH', 'SY', {}),
 ('MH', 'DJ', {}),
 ('MH', 'CM', {}),
 ('MH', 'AL', {}),
 ('MH', 'SW', {}),
 ('MH', 'CF', {}),
 ('MH', 'LS', {}),
 ('MH', 'CS', {}),
 ('MH', 'TG', {}),
 ('MH', 'CP', {}),
 ('MH', 'CV', {}),
 ('MK', 'LH', {}),
 ('MK', 'KL', {}),
 ('MK', 'JLA', {}),
 ('MK', 'MS', {}),
 ('MK', 'CS', {}),
 ('JLA', 'CM', {}),
 ('JLA', 'KL', {}),
 ('JLA', 'MS', {}),
 ('JLA', 'CS', {}),
 ('JLA', 'SB', {}),
 ('JLA', 'HHI', {}),
 ('MP', 'TW', {}),
 ('MP', 'TT', {}),
 ('MP', 'HHI', {}),
 ('MS', 'CS', {}),
 ('MS', 'HHI', {}),
 ('FI', 'KW', {}),
 ('FI', 'AW', {}),
 ('FI', 'CF', {}),
 ('CJ', 'SY', {}),
 ('CJ', 'DD', {}),
 ('CJ', 'CD', {}),
 ('CO', 'AW', {}),
 ('CM', 'TW', {}),
 ('CM', 'TT', {}),
 ('CM', 'AL', {}),
 ('CM', 'CS', {}),
 ('CB', 'DJ', {}),
 ('CB', 'CP', {}),
 ('CB', 'CV', {}),
 ('CG', 'CF', {}),
 ('CF', 'JCU', {}),
 ('CF', 'AW', {}),
 ('CF', 'KH', {}),
 ('CF', 'LH', {}),
 ('CF', 'AP', {}),
 ('CF', 'AS', {}),
 ('CF', 'KW', {}),
 ('CF', 'CS', {}),
 ('CF', 'CV', {}),
 ('CD', 'SY', {}),
 ('CD', 'LP', {}),
 ('CD', 'KF', {}),
 ('CS', 'JCU', {}),
 ('CS', 'TW', {}),
 ('CS', 'TT', {}),
 ('CS', 'AS', {}),
 ('CS', 'LH', {}),
 ('CS', 'SB', {}),
 ('CS', 'HHI', {}),
 ('CP', 'DJ', {}),
 ('CP', 'AP', {}),
 ('CP', 'DR', {}),
 ('CP', 'CV', {}),
 ('CV', 'DJ', {}),
 ('CV', 'AP', {}),
 ('CV', 'DR', {}),
 ('KB', 'LW', {}),
 ('SY', 'KF', {}),
 ('KF', 'AP', {}),
 ('KD', 'TG', {}),
 ('SW', 'BJ', {}),
 ('SW', 'IM', {}),
 ('SW', 'LH', {}),
 ('KL', 'TT', {}),
 ('KP', 'TM', {}),
 ('KW', 'JCU', {}),
 ('SB', 'AL', {}),
 ('DJ', 'TG', {}),
 ('BJ', 'IM', {}),
 ('KWI', 'AW', {}),
 ('TW', 'TT', {}),
 ('TW', 'AL', {}),
 ('TW', 'HHI', {}),
 ('TT', 'AL', {}),
 ('TT', 'HHI', {}),
 ('LH', 'JCU', {}),
 ('JCU', 'AP', {}),
 ('JCU', 'AS', {}),
 ('AL', 'HHI', {})]

score 42 · Accepted Answer · edited Jan 24 '23 at 11:51

42

Here is an example of how to use a colormap. It's a little tricky. If you want a customized discrete colormap you can try this SO answer Matplotlib discrete colorbar

import matplotlib.pyplot as plt
# create number for each group to allow use of colormap
from itertools import count
# get unique groups
groups = set(nx.get_node_attributes(g,'group').values())
mapping = dict(zip(sorted(groups),count()))
nodes = g.nodes()
colors = [mapping[g.nodes[n]['group']] for n in nodes]

# drawing nodes and edges separately so we can capture collection for colobar
pos = nx.spring_layout(g)
ec = nx.draw_networkx_edges(g, pos, alpha=0.2)
nc = nx.draw_networkx_nodes(g, pos, nodelist=nodes, node_color=colors, node_size=100, cmap=plt.cm.jet)
plt.colorbar(nc)
plt.axis('off')
plt.show()

enter image description here

edited Jan 24 '23 at 11:51

Dos

2,250
1
29
39

answered Mar 07 '15 at 14:50

Aric

24,511
5
78
77

1

do you know if there is anyway to create a categorical legend for the attribute colors compared to the colorbar option that you showed? I ask because not all of my attributes are numeric, and associating a non-numeric attribute to a numeric value is just one more step in interpreting the graph. – CurtLH Mar 12 '15 at 00:20
How would you do the same thing but with multiple graphs to be plotted at once? How do you adjust the colorbar? – FaCoffee Jun 27 '17 at 12:41
`colors = [mapping[g.node[n]['group']] for n in nodes]]` has one too many `]` at the end. Can't propose as edit ;p – JJJ Jun 18 '19 at 12:40
@Aric This coming up with == KeyError: 'group'. Any idea why? – Strayhorn Jul 05 '21 at 17:24
2

Note that if you are using NetworkX 2.4 or higher then you should use `colors = [mapping[g.nodes[n]['group']] for n in nodes]` instead since g.nide has been deprecated [see here](https://networkx.org/documentation/stable/release/release_2.4.html#deprecations). Also, if you are using NetworkX 2.5 or higher you should remove `with_labels=False` from the call to `nx.draw_networkx_nodes()`. Apparently it was never used and in v2.5 a check was added for validity of keywords in `nx.draw` and friends [see here](https://networkx.org/documentation/stable/release/release_2.5.html#api-changes) – René Oct 28 '21 at 08:20
Given that the `with_labels` is no longer present in the function. How could one add the node labels? – Ananth_Rao May 01 '22 at 09:41

Ari · Answer 2 · 2022-08-28T20:35:56.153

1

This flexibly assigns node color by the the value for the attribute of each node:

import networkx as nx # my version of networkx is 2.8.5

ATTRIBUTE_NAME = 'group'
COLOR_SCHEME = "Set1" # try plt.cm.Blues if your attribute is a continuous variable
colors = [g.nodes[node][ATTRIBUTE_NAME] for node in list(g.nodes())]
nx.draw_networkx(g, node_color=colors, cmap=COLOR_SCHEME)

This creates the color bar:

ax=plt.gca()
PCM=ax.get_children()[0] # this index may vary
plt.colorbar(PCM, ax=ax)

edited Aug 28 '22 at 20:35

answered Aug 28 '22 at 20:27

Ari

51
1
5

1

can also use: `node_color=list(nx.get_node_attributes(g, ATTRIBUTE_NAME).values())` – Minstein Nov 07 '22 at 20:12

Python NetworkX -- set node color automatically based on number of attribute options

2 Answers2

Linked