Matplotlib cmap - custom color definition-CodePudding

I am plotting a graph using matplotlib.

Here is the code:

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_title("Grid search results - "   model_name)
ax.set_xlabel("Log10(Wight decay)")
ax.set_ylabel("Log10(Learning rate)")
ax.set_zlabel("Batch size")
ax.set_xticks(weightdecay)
ax.set_yticks(learningrate)
ax.set_zticks(trainbatchsize)

scat_plot = ax.scatter(xs=weightdecay, ys=learningrate, zs=trainbatchsize, c=f1, cmap="bwr")
ax.text(top_score[0], top_score[1], top_score[2], top_score[3], color="black")

cb = plt.colorbar(scat_plot, pad=0.2)
cb.ax.set_xlabel('F1 score')

plt.plot(top_score[0], top_score[1], top_score[2], marker="o", markersize=15, markerfacecolor="yellow")

path = Path(output_dir)
plt.savefig(str(path.absolute())   '/grid_search_plot_'   model_name   ".pdf")
plt.show()

The graph I am getting looks like:

What I would like to do is to use a more granular color-bar. For example for my F1-score (colour-bar), show in:

color1 scores < 0.5
color2 scores 0.5 - 0.75
color3 scores 0.75 - 0.80
color4 scores 0.8 - 0.85
color5 scores 0.85-1

I was trying to re-use some code to create a custom cmap but nothing was working as expected.

CodePudding user response：

One cheap/quick solution might be to create a "categorical color value", like this:

import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import ListedColormap
import numpy as np

N = 40
x = np.random.uniform(0, 1, N)
y = np.random.uniform(0, 1, N)
z = np.random.uniform(0, 1, N)
# color values
c = np.random.uniform(0, 1, N)
# new color values
new_col = c.copy()
new_col[c < 0.5] = 0
new_col[(c >= 0.5) & (c < 0.75)] = 1
new_col[(c >= 0.75) & (c < 0.8)] = 2
new_col[(c >= 0.8) & (c < 0.85)] = 3
new_col[c >= 0.85] = 4
new_col = new_col / new_col.max()

fig = plt.figure()
ax = fig.add_subplot(projection="3d")
scatter = ax.scatter(x, y, z, c=new_col, cmap=cm.get_cmap("tab10", 5))
cb = fig.colorbar(scatter)
cb.ax.set_yticklabels([0, 0.5, 0.75, 0.80, 0.85, 1])
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_zlim(0, 1)