Search code examples
pythonmatplotlibcolorslegendlegend-properties

Matplotlib Scatter Plot Legend Creation Mystery


I have the following snipped of code (values for c, s, x, y are mockups, but the real lists follow the same format, just much bigger. Only two colors are used - red and green though. All lists are of the same size)

The issue is that the color legend fails to materialize. I am completely at loss as to why. Code snippets for legend generation is basically a cut-n-paste from docs, i.e. (https://matplotlib.org/3.1.1/gallery/lines_bars_and_markers/scatter_with_legend.html#sphx-glr-gallery-lines-bars-and-markers-scatter-with-legend-py)

Anyone has any idea??

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
%matplotlib inline

c = [ 'g', 'r', 'r', 'g', 'g', 'r', 'r', 'r', 'g', 'r']
s = [ 10, 20, 10, 40, 60, 90, 90, 50, 60, 40]
x = [ 2.4, 3.0, 3.5, 3.5, 3.5, 3.5, 3.5, 2.4, 3.5, 3.5]
y = [24.0, 26.0, 20.0, 19.0, 19.0, 21.0, 20.0, 23.0, 20.0, 20.0]

fig, ax = plt.subplots()

scatter = plt.scatter(x, y, s=s, c=c, alpha=0.5)

# produce a legend with the unique colors from the scatter
handles, lables = scatter.legend_elements()
legend1 = ax.legend(handles, labels, loc="lower left", title="Colors")
ax.add_artist(legend1)

# produce a legend with a cross section of sizes from the scatter
handles, labels = scatter.legend_elements(prop="sizes", alpha=0.5)
legend2 = ax.legend(handles, labels, loc="upper right", ncol=2, title="Sizes")

plt.show()

Plot output:

enter image description here


Solution

  • It seems that legend_elements() is only meant to be used when c= is passed a numeric array to be mapped against a colormap. You can test by replacing c=c by c=s in your code, and you will get the desired output.

    Personally, I would have expected your code to work, and maybe it is worth bringing it up either as a bug or a feature request at matplotlib's github. EDIT: actually, there is already a discussion about this very issue on the issue tracker

    One way to circumvent this limitation is to replace your array of colors names with a numeric array and creating a custom colormap that maps each value in your array to the desired color:

    #c = [ 'g', 'r', 'r', 'g', 'g', 'r', 'r', 'r', 'g', 'r']
    c = [0, 1, 1, 0, 0, 1, 1, 1, 0, 1]
    cmap = matplotlib.colors.ListedColormap(['g','r'])
    s = [ 10, 20, 10, 40, 60, 90, 90, 50, 60, 40]
    x = [ 2.4, 3.0, 3.5, 3.5, 3.5, 3.5, 3.5, 2.4, 3.5, 3.5]
    y = [24.0, 26.0, 20.0, 19.0, 19.0, 21.0, 20.0, 23.0, 20.0, 20.0]
    
    fig, ax = plt.subplots()
    
    scatter = plt.scatter(x, y, s=s, c=c, alpha=0.5, cmap=cmap)
    
    # produce a legend with the unique colors from the scatter
    handles, labels = scatter.legend_elements()
    legend1 = ax.legend(handles, labels, loc="lower left", title="Colors")
    ax.add_artist(legend1)
    
    # produce a legend with a cross section of sizes from the scatter
    handles, labels = scatter.legend_elements(prop="sizes", alpha=0.5)
    legend2 = ax.legend(handles, labels, loc="upper right", ncol=2, title="Sizes")
    
    plt.show()
    

    enter image description here