Matplotlib: wrong colors on lineplot legend when using colormap

Question

I am trying to plot multiple lines, each with a unique color selected from a colormap. The lines are plotted with the correct color, but in the legend, each line is labeled with the same color. Here is some sample code:

from matplotlib import pyplot as plt
import numpy as np

fig = plt.figure()
ax = fig.add_subplot(111)
#values associated with each line
values=(-4,-3,-2,-1, 0,1,2,3,4)
#get an array of colors associated with the list of values
col=plt.cm.viridis(0.125*np.array(values)+0.5)
j=0
x=np.linspace(-5,5,101)
for val in values:
    y=get_associated_data(val)
    ax.plot(x,y,color=col[j],label='val='+str(val))
    j+=1
handles,labels = ax.get_legend_handles_labels()
ax.legend(np.unique(labels))
plt.show()

The above code produces a plot like this

How do I fix this? Also note that the labels are out of order, which I'd also like to fix.
EDIT: Here is the exact code I am using to make the plot, including @William Miller's fix for the out of order legend.

from matplotlib import pyplot as plt
import numpy as np

rc('font', size=16)
rc('text', usetex=True)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_xlabel(r'$N^\prime$')
ax.set_ylabel(r'$\Gamma$')
ax.set_title(r'Particle flux at $U=U^{\prime\prime}=0,\varepsilon=20$')
gradlist=(-4,-3,-2,-1, 0,1,2,3,4)
col=plt.cm.viridis(0.125*np.array(gradlist)+0.5)
x=[i/10. for i in range(-50,51)]
j=0
for grad in gradlist:
    average_flux=np.load('flux_'+str(grad)+'.npy')
    ax.plot(x,average_flux,color=col[j],label=r'$U^\prime=$'+str(grad))
    time.sleep(1)
    print(grad)
    j+=1
handles,labels = ax.get_legend_handles_labels()
labels = np.array(labels)
ax.legend(labels[np.sort(np.unique(labels, return_index=True)[1])])
plt.show()

You'll need the following dataset: https://www.dropbox.com/sh/2l2pot21f5sp6cw/AAD1xJcl-FLVf79ylpf7SZiTa?dl=0

Answer 1

You can resolve the sorting according to this answer by using a combination of np.unique with np.sort to preserve the ordering.

handles, labels = ax.get_legend_handles_labels()
labels = np.array(labels)
ax.legend(labels[np.sort(np.unique(labels, return_index=True)[1])])
plt.show()

which will give you

The coloring issue is caused by only the first 9 handles of ~900 being used, instead of the handles corresponding to the correct 9 labels, you can resolve this by selecting the correct indices from both handles and labels, something like this

handles, labels = ax.get_legend_handles_labels()
idx = np.sort(np.unique(np.array(labels), return_index=True)[1])
ax.legend(np.array(handles)[idx], np.array(labels)[idx])
plt.show()

Should give you the correct result,