Matplotlib - custom cmap not in accordance with graph colors if the some data is not available

Question

I have the following code:

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xticks(weightdecay)
ax.set_yticks(learningrate)
ax.set_zticks(trainbatchsize)

arr = numpy.array(f1)
new_col = arr.copy()

new_col[arr < 0.5] = 0
new_col[(arr >= 0.5) & (arr < 0.75)] = 1
new_col[(arr >= 0.75) & (arr < 0.8)] = 2
new_col[(arr >= 0.8) & (arr < 0.85)] = 3
new_col[arr >= 0.85] = 4
new_col = new_col / new_col.max()

cmap = ListedColormap(["magenta", "green", "blue", "orange", "red"])
scat_plot = ax.scatter(xs=weightdecay, ys=learningrate, zs=trainbatchsize, c=new_col, cmap=cmap)

cb = fig.colorbar(scat_plot, pad=0.2)
cb.ax.set_yticklabels([0, 0.5, 0.75, 0.80, 0.85, 1])

However, the colors within the color-bar are not the same as the colors within the graph. This is due to the fact that within some of my data-sets, some data corresponding to a specific color are not defined. Any idea on how to solve this. So I would like to always have the same color-bar even if data represented by some colors is missing.

Answer 1

I'd assume the problem lies in the usage of ListedColormap, did you try LinearSegmentedColormap (an example tutorial can be found in the docs)?

Answer 2

I would recommend using the color sequence directly. According to the documentation, the c argument can be: "A sequence of colors of length n".

So you can use something like:

colors = ["magenta", "green", "blue", "orange", "red"]
bins = [0.5, 0.75, 0.8, 0.85]
color_indices = numpy.digitize(f1, bins)
c = [colors[i] for i in color_indices]

In this case, it will be more tricky to create the colorbar without explicit cmap yet you can add a simple legend.