Pcolormesh plot with colorbar such that colors correspond to strings

Question

Suppose I have four objects ('strings', can assign numbers to them (say, [0,2,4,6]) which I want to plot using pcolormesh with 4 colors representing the 4 string,, and even one string is missing from the data, it should appear on the colorbar with same label.

The attached code is what I was trying.

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

Cmap = ListedColormap(['white', 'green', 'blue','red']) ##Corresponding to 4 strings labelled as [0, 2, 4, 6]

x = np.arange(0, 10, 1); y = np.arange(0,5,1)
X, Y = np.meshgrid(x,y)
Z = X+Y
for ii in range (len(x)):
  for jj in range (len(y)):
    ztemp = Z[jj,ii]
    if (ztemp<3):
      Z[jj,ii] = 0 ##'Apple'
    elif (6>ztemp>=3):
      Z[jj,ii] = 2 ##'Orange'
    elif (10>ztemp>=6):
      Z[jj,ii] = 4 ##'Banana'
    elif (ztemp>=10):
      Z[jj,ii] = 6 ##'Nothing'
print(Z)

#Z[Z==0]=2. ## This should only fill up the white portion by green

plt.figure()
plt.pcolormesh(X, Y, Z, cmap=Cmap)
plt.colorbar()
plt.show()

Answer 1

You can use vmin and vmax to set a precise range for the colorbar. If the colormap contains 4 values, with vmin=-1, vmax=7 the first color will cover the range -1,1, the second 1,3, the third 3,5 and the fourth 5,7. The position for 0 will be nicely at the center of the first color range (it's similar for the other colors).

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

names = ['Apple', 'Orange', 'Banana', 'Nothing']
Cmap = ListedColormap(['white', 'green', 'blue', 'red'])  ##Corresponding to 4 strings labelled as [0, 2, 4, 6]

x = np.arange(0, 10, 1);
y = np.arange(0, 5, 1)
X, Y = np.meshgrid(x, y)
Z = X + Y
for ii in range(len(x)):
     for jj in range(len(y)):
          ztemp = Z[jj, ii]
          if (ztemp < 3):
               Z[jj, ii] = 0  ##'Apple'
          elif (6 > ztemp >= 3):
               Z[jj, ii] = 2  ##'Orange'
          elif (10 > ztemp >= 6):
               Z[jj, ii] = 4  ##'Banana'
          elif (ztemp >= 10):
               Z[jj, ii] = 6  ##'Nothing'
plt.figure()
plt.pcolormesh(X, Y, Z, cmap=Cmap, vmin=-1, vmax=7)
cbar = plt.colorbar()
cbar.set_ticks([0, 2, 4, 6])
cbar.set_ticklabels(names)
plt.show()

PS: Note that you get a warning MatplotlibDeprecationWarning: shading='flat' when X and Y have the same dimensions as C is deprecated since 3.3. Either specify the corners of the quadrilaterals with X and Y, or pass shading='auto', 'nearest' or 'gouraud'. By default, the X and Y set the borders of the cells, and the Z indicates the cell colors. If cells are colored uniform (default shading='flat'), there will be one column of cells less than there are X positions (and one row less than there are Y positions).

One way to tackle this, is to provide 1 more x and 1 more y position. Subtracting 0.5 would put the integer ticks positions in the centers of the cells:

plt.pcolormesh(np.arange(-0.5, 10), np.arange(-0.5, 5), Z, cmap=Cmap, vmin=-1, vmax=7)
plt.xticks(x)
plt.yticks(y)