Legend based on edge color in networkx

Question

Is there a way to create a legend in networkx based on edge color (as opposed to by node color)?

This is my graph:

plt.figure(figsize = (15, 10))
G = nx.from_pandas_dataframe(df, 'From', 'To', ['Order', 'Colors'])
edge_labels = nx.get_edge_attributes(G, 'Order')
nx.draw_networkx(G, with_labels = False, node_color = 'black', alpha = 0.5, node_size = 3, linewidths = 1, edge_color = df['Colors'], edge_cmap = 
plt.cm.Set2)
plt.show()

In this, ['Order'] is a descriptor of the edge and ['Color'] is a unique integer mapped to each value in ['Order'], which is working to create the edge colors based on the Set2 colormap.

I can get the edge labels with something like: edge_labels = nx.get_edge_attributes(G, 'Order') but how do I put this into a legend?

I'm happy to share the data and complete code if helpful!

Answer 1

One way you can do it is in the spirit of this SO answer which uses proxies for each member of a LineCollection in the legend.

You can get the LineCollection by using the step-by-step graph drawing functions, drawing the nodes and edges separately (e.g. draw_networkx_nodes doc.)

import matplotlib.pyplot as plt
import networkx as nx
from matplotlib.lines import Line2D

# make a test graph
n = 7 # nodes
m = 5 # edges
G = nx.gnm_random_graph(n, m)
# and define some color strings (you'll get this array from the dataframe)
_c = 'rgbcmky' * m # way too many colors, trim after
clrs = [c for c in _c[:m]]

plt.ion()

plt.figure(figsize = (9, 7), num=1); plt.clf()
# draw the graph in several steps
pos = nx.spring_layout(G)
h1 = nx.draw_networkx_nodes(G, pos=pos, node_color = 'black',
                            alpha = 0.9, node_size = 300, linewidths=6)
# we need the LineCollection of the edges to produce the legend (h2)
h2 = nx.draw_networkx_edges(G, pos=pos, width=6, edge_color=clrs)

# and just show the node labels to check the labels are right!
h3 = nx.draw_networkx_labels(G, pos=pos, font_size=20, font_color='c')


#https://stackoverflow.com/questions/19877666/add-legends-to-linecollection-plot - uses plotted data to define the color but here we already have colors defined, so just need a Line2D object.
def make_proxy(clr, mappable, **kwargs):
    return Line2D([0, 1], [0, 1], color=clr, **kwargs)

# generate proxies with the above function
proxies = [make_proxy(clr, h2, lw=5) for clr in clrs]
# and some text for the legend -- you should use something from df.
labels = ["{}->{}".format(fr, to) for (fr, to) in G.edges()]
plt.legend(proxies, labels)

plt.show()

This produces something like: