Получение правильных цветов для pcolormesh

Question

Я использую pcolormesh, чтобы показать результаты по некоторому алгоритму, который я запускаю. создать обычную сетку для некоторых x_min, x_max и т.д., и определить мой цвет карты

h = (x_max - x_min)/1000. 
xx, yy = numpy.meshgrid(numpy.arange(x_min, x_max, h), numpy.arange(y_min, y_max, h)) 
colours = ("blue", "green", "red") 
cmap = colors.ListedColormap(colours) 

, а затем сделать plt.pcolormesh(xx, yy, Z, alpha=0.7, cmap=cmap), где Z является результатом моего предсказания (это может быть любой 0, 1 или 2 значение, Это не имеет значения).

Скажите Z = numpy.zeros(xx.shape) Я должен видеть все синее, Z = numpy.ones(xx.shape) Я должен видеть зеленый, и Z = 2*numpy.ones(xx.shape) Я должен видеть красный цвет, или так я думал. Вместо этого я всегда вижу синий. Если добавить эти строки:

Z[0] = 0 
    Z[1] = 1 
    Z[2] = 2 

все работает, как ожидалось. Похоже, что результат не имеет всех возможных результатов (0,1,2), тогда он по умолчанию использует только первый цвет, синий, даже если результат - все 2s, и я хочу красный.

Как я могу заставить его иметь желаемые цвета, т. Е. Синий для 0, зеленый для 1 и красный для 2, в каждом случае?

Answer 1

Вы можете использовать clim для исправления цвета.

plt.clim(0, 3) 

Это заставило бы 0 быть синим, 1 - зеленым, 2 - красным.

Answer 2

Вы должны нормализовать Colormap:

import matplotlib.pylab as plt 
import numpy as np 
from matplotlib.colors import ListedColormap, BoundaryNorm 

x_max = 100. 
x_min = 0. 
y_max = 100. 
y_min = 0. 
h = (x_max - x_min)/5. 
xx, yy = np.meshgrid(np.arange(x_min, x_max+h, h), np.arange(y_min, y_max+h, h)) 
Z = np.random.randint(3, size=(5,5)) 
# define color map & norm it 
colours = (["blue", "green", "red"]) 
cmap = ListedColormap(colours) 
bounds=[0,1,2,np.max(Z)+1] # discrete values of Z 
norm = BoundaryNorm(bounds, cmap.N) 
# for colorbar 
ticks = [.5,1.5,2.5] 
labels = ['0','1','2'] 
# plot 
pcm = plt.pcolormesh(xx, yy, Z, alpha=0.7, cmap=cmap, norm=norm) 
cb = plt.colorbar(pcm, cmap=cmap, norm=norm, ticks=ticks) 
cb.set_ticklabels(labels) 
plt.show() 

Z массив:

[[0 0 1 0 0] 
[0 0 0 1 1] 
[1 0 0 0 1] 
[1 1 2 2 0] 
[1 1 1 2 2]]