CODEWITHSUNDEEP
pythonnetworkx
Nitin
Nitin

Reputation: 2874

Colorbar for edges in networkx

I am trying to get a colorbar for edges in a networkx graph. Here is a code snippet

import networkx as nx
import matplotlib.colors as colors
import matplotlib.cm as cmx

n = 12 # Number of clusters
w = 21 # Number of weeks

m = Basemap(
    projection='merc',
    ellps = 'WGS84',
    llcrnrlon=-98.5,
    llcrnrlat=25,
    urcrnrlon=-60,
    urcrnrlat=50,
    lat_ts=0,
    resolution='i',
    suppress_ticks=True)

mx, my = m(list(ccentroids['lon']), list(ccentroids['lat']))

# The NetworkX part
# put map projection coordinates in pos dictionary
G = nx.DiGraph()
G.add_nodes_from(range(n))
for i in range(n):
    for j in range(n):
       if P_opt[i,j] > 0.5 and i != j:
           G.add_edge(i,j, weight = P_opt[i,j])

pos = {i : (mx[i], my[i]) for i in range(n)}

# Add a color_map for the edges
jet = cm = plt.get_cmap('jet')
cNorm  = colors.Normalize(vmin=0, vmax=np.max(P_opt))
scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=jet)
colorList = []
weights_list = []

for i in G.edges():
   a, b = i
   colorVal = scalarMap.to_rgba(G.edge[a][b]['weight'])
   colorList.append(colorVal)
   weights_list.append(G.edge[a][b]['weight'])

plt.clf()
fig = plt.figure()
ax = fig.add_subplot(111, axisbg='w', frame_on=False)
fig.set_size_inches(18.5, 10.5)

# draw the network
nodes = nx.draw_networkx_nodes(G, pos, node_size=100, node_color=q[:,t], cmap =     plt.cm.jet, 
                               font_size=8, with_labels=False, label='Cluster centroids')
edges = nx.draw_networkx_edges(G, pos, edge_color=colorList)

m.drawcountries()    
m.bluemarble()

This gives me the following image: enter image description here

Now I want to add a colorbar for the edges. I tried doing something like,

plt.sci(edges)
edges.set_array(np.array(weights_list))
plt.colorbar(shrink = 0.8)

This gives me an image like: enter image description here

The colours of the arrows and edges seem to differ. How can I correct that? Thanks.

EDIT: I tried to use the following code by modifying the edge line:

edges = nx.draw_networkx_edges(G, pos, edge_color=colorList, edge_cmap = plt.cm.jet)
plt.colorbar(edges)

This gives me an error TypeError: You must first set_array for mappable

Changing the edge_color to be the weights_list I get the following picture:

enter image description here

Upvotes: 6

Views: 5854

Answers (2)

warped
warped

Reputation: 9481

I know this is an old one, but I spent some time figuring this out, and maybe it is still useful for someone.

nx.draw_networkx_edges

returns 

matplotlib.collection.LineCollection

if there are no arrows.

returns 

list of matplotlib.patches.FancyArrowPatch

if there are arrows.

docs

For my specific case I could get a colorbar like in Aric's answer only if I set arrows to False:

edges = nx.draw_networkx_edges(G, edge_color=colors, arrows=False)
plt.colorbar(edges)

Upvotes: 1

Aric
Aric

Reputation: 25319

It seems you are getting the colomap/colorbar for the nodes. Here is how to set a colormap and draw a colorbar for edge colors:

![import networkx as nx
import matplotlib.pyplot as plt
G = nx.star_graph(20)  
pos = nx.spring_layout(G)
colors = range(20) 
nodes = nx.draw_networkx_nodes(G,pos,node_color='k', with_labels=False)
edges = nx.draw_networkx_edges(G,pos,edge_color=colors,width=4,
                               edge_cmap=plt.cm.Blues)
plt.colorbar(edges)
plt.axis('off')

enter image description here

Upvotes: 8

Related Questions