CODEWITHSUNDEEP
pythonmatplotlib
user
user

Reputation: 2093

Getting the correct colours for pcolormesh

I am using pcolormesh to show the results from some algorithm I am running. I create the usual mesh for some x_min, x_max, etc., and define my color map

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)

and then do plt.pcolormesh(xx, yy, Z, alpha=0.7, cmap=cmap), where Z is the result of my prediction (it could be any 0, 1 or 2 value, it does not matter).

Say Z = numpy.zeros(xx.shape) I should see everything blue, Z = numpy.ones(xx.shape) I should see green, and Z = 2*numpy.ones(xx.shape) I should see red, or so I thought. Instead, I always see blue. If I add these lines:

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

everything works as expected. It looks as if the result does not have all the possible results (0,1,2) then it defaults to use only the first color, blue, even if the result is all 2s, and I want red.

How can I force it to have the colours I desire, i.e. blue for 0, green for 1 and red for 2, in every case?

Upvotes: 2

Views: 2291

Answers (2)

Serenity
Serenity

Reputation: 36635

You have to normalize 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()

enter image description here

Z array:

[[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]]

Upvotes: 2

Georgia S
Georgia S

Reputation: 622

You can use clim to fix your colorbar. 

plt.clim(0, 3)

That will force 0 to be blue, 1 to be green, 2 to be red.

Upvotes: 3

Related Questions