2

I have a sample scatterplot via matplotlib via the code below.

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0, 100, 501)
y = np.sin(x)

label = 'xy data sample'

plt.scatter(x, y, cmap='plasma', c=x, label=label)
legend_dict = dict(ncol=1, loc='best', scatterpoints=4, fancybox=True, shadow=True)
plt.legend(**legend_dict)
plt.show()

Running the code above produces the plot below.

enter image description here

The colormap was successfully plotted, but the legend shows points that are all blue rather than points in a color that correspond to the chosen colormap. Why does this happen?

I tried putting cmap='plasma' in legend_dict, but it results in the error below.

File "/Users/.../
site-packages/matplotlib/axes/_axes.py", line 550, in legend
    self.legend_ = mlegend.Legend(self, handles, labels, **kwargs)
TypeError: __init__() got an unexpected keyword argument 'cmap'

EDIT:

My desired output is to have the four dots represented in the legend to be a different color via the chosen colormap. Ideally, cmap='plasma' in this example could produce a legend using something similar to a blue dot, then a purple dot, then an orange-red dot, then a yellow dot. Although a colorbar could make for a possible alternative, I have yet to look through any documentation about colorbars.

  • The reason that happens is that matplotlib is designed in that way. Which other color would you like the points to have? Please write clearly what you want to achieve. – ImportanceOfBeingErnest Mar 15 '18 at 10:57
  • 1
    Would a colorbar be more suited to this? Otherwise, you'd have a lot of legend entries. Something like [this](https://stackoverflow.com/questions/6063876/matplotlib-colorbar-for-scatter) perhaps? – DavidG Mar 15 '18 at 10:58
  • I updated the post to address this. –  Mar 15 '18 at 11:02

1 Answers1

2

A colorbar can be achieved via plt.colorbar(). This would allow to directly see the values corresponding to the colors.

Having the points in the legend show different colors is of course also nice, although it would not allow to give any quantitative information.

Unfortunately matplotlib does not provide any inbuilt way to achieve this. So one way would be to subclass the legend handler used to create the legend handle and implement this feature.

Here we create a ScatterHandler with a custom create_collection method, in which we create the desired PathCollection and use this by specifying it in the legend_map dictionary of the legend.

handler_map={ type(sc) : ScatterHandler()}

The following code seems a bit complicated at first sight, however you may simply copy the class without understanding it completely and use it in your code.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerRegularPolyCollection

class ScatterHandler(HandlerRegularPolyCollection):
    def update_prop(self, legend_handle, orig_handle, legend):
        legend._set_artist_props(legend_handle)
        legend_handle.set_clip_box(None)
        legend_handle.set_clip_path(None)

    def create_collection(self, orig_handle, sizes, offsets, transOffset):
        p = type(orig_handle)([orig_handle.get_paths()[0]],
                              sizes=sizes, offsets=offsets,
                              transOffset=transOffset,
                              cmap=orig_handle.get_cmap(),
                              norm=orig_handle.norm )

        a = orig_handle.get_array()
        if type(a) != type(None):
            p.set_array(np.linspace(a.min(),a.max(),len(offsets)))
        else:
            self._update_prop(p, orig_handle)
        return p


x = np.linspace(0, 100, 501)
y = np.sin(x)*np.cos(x/50.)

sc = plt.scatter(x, y, cmap='plasma', c=x, label='xy data sample')

legend_dict = dict(ncol=1, loc='best', scatterpoints=4, fancybox=True, shadow=True)
plt.legend(handler_map={type(sc) : ScatterHandler()}, **legend_dict)

plt.show()

enter image description here

ImportanceOfBeingErnest
  • 321,279
  • 53
  • 665
  • 712
  • I'm slowly getting more comfortable with classes; I posted this question to incorporate the answer into a class that outputs a plot of the desired type (`scatter`, `bar`, etc) based on inputs (rather than manually specifying `plt.xlim(...)`, `plt.xlabel(...)`, etc and various kwargs used in matplotlib functions like `label` and `color`. Since my intention is to have my plot class inherit your class (or somehow redefine it within my plot class), would you mind explaining the `ScatterHandler` class - specifically about what the various types of handles are and `orig_handle.get_paths()[0]`? –  Mar 15 '18 at 22:13
  • I don't think you want to change `ScatterHandler` at all. And inheriting the handler class does not make sense. `orig_handle.get_paths()[0]` is the shape of the scatter (i.e. a circle in this case). – ImportanceOfBeingErnest Mar 15 '18 at 22:25
  • That makes sense, that approach works much simpler. I tried using `orig_handle.get_paths()[0]` in a `plt.bar` routine unsuccessfully. I'm having trouble finding relevant documentation for adjusting legends for various plots. Do you know where I can find relevant docs? –  Mar 16 '18 at 11:20
  • 2
    `plt.bar()` returns a [`BarContainer`](https://matplotlib.org/api/container_api.html#module-matplotlib.container), which has an attribute `patches`. The strategy is always to look at the return types and locate them in the documentation. – ImportanceOfBeingErnest Mar 16 '18 at 11:45