python - 如何控制 matplotlib.pyplot.imshow 中像素的颜色？

Question

我想使用 matplotlib.pyplot.imshow() 表示一个 2X2 矩阵（数组）。它工作正常，但我想自己控制每个像素的颜色，而不是由函数来控制。就像，我有一个数组说：

for i in range(N):
    for j in range(N):
        x = np.random.random()
        if x <= 0.4:
            lat[i, j] = 0
        elif 0.4 < x <= 0.5:
            lat[i, j] = 1            
        elif 0.5 < x <= 0.6:
            lat[i, j] = 2           
        else:
            lat[i, j] = 3

这会生成我想要的矩阵。现在在使用该plt.imshow()函数时，如果矩阵元素具有特定值（在本例中为 0、1、2 或 3），我想使用特定颜色。我怎样才能做到这一点？

标签： pythonmatplotlib

Answer 1

您可以LinearSegmentedColormap使用所需的颜色创建一个：

from matplotlib import pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np

N = 5
lat = np.empty((N, N), dtype=np.int)
for i in range(N):
    for j in range(N):
        x = np.random.random()
        if x <= 0.4:
            lat[i, j] = 0
        elif 0.4 < x <= 0.5:
            lat[i, j] = 1
        elif 0.5 < x <= 0.6:
            lat[i, j] = 2
        else:
            lat[i, j] = 3

my_colors = ['crimson', 'lime', 'dodgerblue', 'gold'] # colors for 0, 1, 2 and 3
cmap = LinearSegmentedColormap.from_list('', my_colors, len(my_colors))
plt.imshow(lat, cmap=cmap, vmin=0, vmax=len(my_colors) - 1, alpha=0.4)
for i in range(lat.shape[0]):
    for j in range(lat.shape[1]):
        plt.text(j, i, lat[i, j])
plt.show()

PS：注意numpy也有digitize自动化第一步代码的功能：

x = np.random.random((N, N))
lat = np.digitize(x, [0.4, 0.5, 0.6], right=True)

my_colors = ['fuchsia', 'lime', 'turquoise', 'gold']
cmap = LinearSegmentedColormap.from_list('', my_colors, len(my_colors))
plt.imshow(lat, cmap=cmap, vmin=0, vmax=len(my_colors) - 1)
for i in range(lat.shape[0]):
    for j in range(lat.shape[1]):
        plt.text(j, i, f'{x[i, j]:.2f}\n{lat[i, j]}', ha='center', va='center')
plt.show()