简体   繁体   中英

Add colorbar as legend to matplotlib scatterplot (multiple subplots, multiple scatters)

I have several subplots to which I want to add a single colorbar. Each subplot consists of 7 scatters. I found advise on how to add colorbars, but they are mostly related to the value of each scatter-point and not to the row itself.

Representative sample code:

import numpy as np
from matplotlib import pyplot as plt

x = range(50)
scales = np.linspace(0, 2, 7)
locs = range(4)
cmap = plt.get_cmap("Spectral")
for s_plot in range(4):
    plt.subplot(2, 2, s_plot+1)
    color = iter(cmap(np.linspace(0, 1, len(scales))))
    for scale in scales:
        c = next(color)
        y = np.random.normal(loc=locs[s_plot], scale=scale, size=50)
        plt.scatter(x, y, c=c, s=5)
        plt.title("Mean = {:d}".format(locs[s_plot]))
plt.subplots_adjust(hspace=0.4)
plt.show()

The above example gives: 在此输入图像描述

My desired colorbar looks like this (fake, to be placed next to the plot):

在此输入图像描述

So the colorbar does not depict the value of my scatterpoints, but rather the different "rows" (in this case: different scales) that are iterated through. In the example that would help match the points to the scales.

What I tried is a simple

plt.colorbar()

which is called once after finishing each subplot. But I get TypeError: You must first set_array for mappable Also, since it is the different scales I want to create the colormap for, I also tried

plt.colorbar(scales) 

which returns: AttributeError: 'numpy.ndarray' object has no attribute 'autoscale_None' .

I am currently lacking orientation on how to proceed on this. Edit: I was marked as possible duplicate of matplotlib colorbar for scatter . I found that question already, but it didn't help with my problem. In my case, I need a colormap that is independent of a z-value, but will only indicate the "row number" or "scatter-row" or however you want to call it (equivalent to "lines" in a plt.plot ).

A colorbar needs a ScalarMappable as input. So if none of the things you create in your plot is suitable for that, you may create it yourself.

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.cm import ScalarMappable

x = range(50)
scales = np.linspace(0, 2, 7)
locs = range(4)
cmap = plt.get_cmap("Spectral")
norm = plt.Normalize(scales.min(), scales.max())

fig, axes = plt.subplots(2,2, constrained_layout=True, sharey=True)

for s_plot, ax in enumerate(axes.flat):
    for scale in scales:
        y = np.random.normal(loc=locs[s_plot], scale=scale, size=50)
        sc = ax.scatter(x, y, c=[cmap(norm(scale))], s=5)
        ax.set_title("Mean = {:d}".format(locs[s_plot]))

sm =  ScalarMappable(norm=norm, cmap=cmap)
sm.set_array([])
cbar = fig.colorbar(sm, ax=axes[:,1])
cbar.ax.set_title("scale")

plt.show()

在此输入图像描述

If I understand correctly then you have some range and want to plot a colormap for that (without some plot actually using the colormap). Basically you can plot a colormap in any axes using

import matplotlib
norm = matplotlib.colors.Normalize(vmin=0, vmax=50)

ax = plt.gca()
matplotlib.colorbar.ColorbarBase(ax, cmap='viridis', norm=norm)

where of course you can use any axes (or use inset_axes to place axes somewhere specific).

More tricky is getting colors for your scatter plots that match the colormap in the first place. I'm not sure if there is an easier way, but I convert the colors to RGB for plotting. Here's a full example:

import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import numpy as np

N = 10

# dummy data
x_ = [k/10*np.arange(10) for k in range(N)]

cmap = matplotlib.cm.get_cmap('viridis')
cmap_values = np.linspace(0., 1., N)
colors = cmap(cmap_values)

colors_rgb = ['#{0:02x}{1:02x}{2:02x}'.format(int(255*a), int(255*b), int(255*c)) for a, b, c, _ in colors]

plt.figure()

for x, c in zip(x_, colors_rgb):
    plt.plot(x, c=c)

norm = matplotlib.colors.Normalize(vmin=0, vmax=50)
ticks = np.arange(0, 60, 10)

# vertical colorbar
cbaxes = inset_axes(plt.gca(), width="3%", height="80%", loc=2)
cbar = matplotlib.colorbar.ColorbarBase(cbaxes, cmap=cmap, norm=norm, ticks=ticks)
cbar.set_label('scale')
cbar.ax.set_yticklabels(ticks, fontsize=12)

在此输入图像描述

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM