使用 matplotlib.pyplot.imshow 的基于列的颜色图

Question

我想将我的输入和输出向量转换为一个数组，并使用可以正常工作的颜色图进行绘图。请参见下面的示例：

x = np.round(np.array([-1, 0,     1,  2, 3, 2, 0, 0, 1, 2, 3]))
y = np.round(np.array([-1, 1.2, -3.1, 2, 3, 2, 1, 1, 1, 3, 3]))

xrange = np.arange(np.min(x),np.max(x) + 1)
yrange = np.arange(np.min(y),np.max(y) + 1)
a = np.zeros((len(yrange), len(xrange)))
for i in range(len(x)):
    a[yrange == y[i],xrange == x[i]] = a[yrange == y[i],xrange == x[i]] + 1

fig, ax = plt.subplots()
im = ax.imshow(a,cmap='Wistia',origin="lower")
for i in range(len(yrange)):
    for j in range(len(xrange)):
        text = ax.text(j, i, str(int(a[i, j])),
                       ha="center", va="center", color="k")
ax.set(xlabel = 'Input', ylabel = 'Output')
fig.tight_layout()
plt.show()

这给了我这个：

唯一的问题是我喜欢基于列的颜色图。 这意味着例如在第一列中，我希望 1 被视为最高数字，并且与第二列中的 3 具有相同的颜色，这适用于所有列.

Answer 1

一种简单的方法是将 a 转换为 DataFrame 并使用以下函数创建按列的 cmap。

import pandas as pd
df = pd.DataFrame(a)

df.style.background_gradient(cmap="Wistia")

这样，您还可以使用“子集”参数对特定列进行更多控制。