How to visualize a list of strings on a colorbar in matplotlib

问题

I have a dataset like

x = 3,4,6,77,3
y = 8,5,2,5,5
labels = "null","exit","power","smile","null"

Then I use

from matplotlib import pyplot as plt
plt.scatter(x,y)
colorbar = plt.colorbar(labels)
plt.show()

to make a scatter plot, but cannot make colorbar showing labels as its colors.

How to get this?

回答1:

I'm not sure, if it's a good idea to do that for scatter plots in general (you have the same description for different data points, maybe just use some legend here?), but I guess a specific solution to what you have in mind, might be the following:

from matplotlib import pyplot as plt

# Data
x = [3, 4, 6, 77, 3]
y = [8, 5, 2, 5, 5]
labels = ('null', 'exit', 'power', 'smile', 'null')

# Customize colormap and scatter plot
cm = plt.cm.get_cmap('hsv')
sc = plt.scatter(x, y, c=range(5), cmap=cm)
cbar = plt.colorbar(sc, ticks=range(5))
cbar.ax.set_yticklabels(labels)
plt.show()

This will result in such an output:

The code combines this Matplotlib demo and this SO answer.

Hope that helps!

EDIT: Incorporating the comments, I can only think of some kind of label color dictionary, generating a custom colormap from the colors, and before plotting explicitly grabbing the proper color indices from the labels.

Here's the updated code (I added some additional colors and data points to check scalability):

from matplotlib import pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np

# Color information; create custom colormap
label_color_dict = {'null': '#FF0000',
                    'exit': '#00FF00',
                    'power': '#0000FF',
                    'smile': '#FF00FF',
                    'addon': '#AAAAAA',
                    'addon2': '#444444'}
all_labels = list(label_color_dict.keys())
all_colors = list(label_color_dict.values())
n_colors = len(all_colors)
cm = LinearSegmentedColormap.from_list('custom_colormap', all_colors, N=n_colors)

# Data
x = [3, 4, 6, 77, 3, 10, 40]
y = [8, 5, 2, 5, 5, 4, 7]
labels = ('null', 'exit', 'power', 'smile', 'null', 'addon', 'addon2')

# Get indices from color list for given labels
color_idx = [all_colors.index(label_color_dict[label]) for label in labels]

# Customize colorbar and plot
sc = plt.scatter(x, y, c=color_idx, cmap=cm)
c_ticks = np.arange(n_colors) * (n_colors / (n_colors + 1)) + (2 / n_colors)
cbar = plt.colorbar(sc, ticks=c_ticks)
cbar.ax.set_yticklabels(all_labels)
plt.show()

And, the new output:

Finding the correct middle point of each color segment is (still) not good, but I'll leave this optimization to you.

来源：https://stackoverflow.com/questions/58620640/how-to-visualize-a-list-of-strings-on-a-colorbar-in-matplotlib

标签

python-3.x

pandas

numpy

matplotlib

colorbar