Matplotlib cmap - 自定义颜色定义

Matplotlib cmap - custom color definition

我正在使用 matplotlib 绘制图表。

代码如下:

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()

我得到的图表如下所示:

我想做的是使用更细粒度的颜色条。例如,对于我的 F1 分数(颜色条),显示在:

我试图重新使用一些代码来创建自定义 cmap,但没有按预期工作。

一个 cheap/quick 解决方案可能是创建一个“分类颜色值”,如下所示:

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)

编辑 以容纳评论。以下应该能够处理类别没有任何元素的情况:

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)
# number of categories
NC = 5
# 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 / NC

fig = plt.figure()
ax = fig.add_subplot(projection="3d")
cmap = ListedColormap(["red", "green", "blue", "magenta", "cyan"])
scatter = ax.scatter(x, y, z, c=cmap(new_col))
cb = fig.colorbar(cm.ScalarMappable(cmap=cmap))
cb.ax.set_yticks(np.linspace(0, 1, NC+1), [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)