Legend for networkx draw function

Question

So I have the following function to draw a problem im working on. Its basically a critical node detection problem or interdiction. I have some values x, and decision to attack to the node z. Basically I wanna color my graph with active and inactive nodes and nodes that are being treated/attack. Here is what I have so far.

def draw_solution(g, zsolution, xsolution, T, xmin=0, filename='test.pdf'):

# draw solution
pos = {n: ndata['coord'] for n,ndata in g.nodes_iter(data=True)}
ncolour=[]
for n,ndata in g.nodes_iter(data=True):
    if ndata['fuel_load'] < xmin:
        ncolour.append('gold')
    else:
        ncolour.append('yellowgreen')
pp = PdfPages(filename)
fig = plt.figure(figsize=(11.7,8.3))
fig.suptitle('full graph and initial fuel load')
plt.axis('off')
nx.draw_networkx(g, pos, font_size=9, node_color=ncolour)
fig.savefig(pp, format='pdf')
fig.clf()

for t in range(T):
    g_copy = g.copy()
    #node colour
    ncolour=[]
    for i in g.nodes_iter():
        if zsolution[i,t] > 0.99:
            ncolour.append('lightcoral')
            #g_copy.remove_edges_from(g.edges(i))
        elif xsolution[i,t] < xmin:
            ncolour.append('gold')
            g_copy.remove_edges_from(g.edges(i))
        else:
            ncolour.append('yellowgreen')

    plt.axis('off')
    fig.suptitle('t={}, before treatment'.format(t))
    nx.draw_networkx(g_copy, pos, font_size=9, node_color=ncolour, label=ncolour)
    fig.savefig(pp, format='pdf')
    fig.clf()

    g_copy = g.copy()
    #node colour
    ncolour=[]
    for i in g.nodes_iter():
        if zsolution[i,t] > 0.99:
            ncolour.append('lightcoral')
            g_copy.remove_edges_from(g.edges(i))
        elif xsolution[i,t] < xmin:
            ncolour.append('gold')
            g_copy.remove_edges_from(g.edges(i))
        else:
            ncolour.append('yellowgreen')

    plt.axis('off')
    fig.suptitle('t={}, after treatment'.format(t))
    nx.draw_networkx(g_copy, pos, font_size=9, node_color=ncolour, label=ncolour)
    fig.savefig(pp, format='pdf')
    fig.clf()

pp.close()
fig.clf()  

For some reason I cant get the legend right?

Please help! Thanks!

Answer 1

Is this what you're after?

import networkx as nx
import matplotlib.pyplot as plt

G = nx.fast_gnp_random_graph(20,0.2)

rednodes = [1,2,4,5]
bluenodes = [10,12]
greennodes = [3,6,9]
yellowgreennodes = [node for node in G.nodes() if
                    node not in rednodes + greennodes + bluenodes]
pos = nx.spring_layout(G)
nx.draw_networkx_nodes(G, pos=pos, nodelist=rednodes,
                       node_color='red', label='red nodes')
nx.draw_networkx_nodes(G, pos=pos, nodelist=bluenodes,
                       node_color='blue', label='blue nodes')
nx.draw_networkx_nodes(G, pos=pos, nodelist=greennodes,
                       node_color='green', label='green nodes')
nx.draw_networkx_nodes(G, pos=pos, nodelist=yellowgreennodes,
                       node_color='yellowgreen', label='yellowgreen nodes')
nx.draw_networkx_edges(G, pos=pos)
plt.legend(scatterpoints = 1)
plt.show()