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

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.

CodePudding user response：

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

CodePudding user response：

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.