Search code examples
python-3.xmatplotlibplotcolorbarcolormap

discrete colorbar with discrete colormesh


I want to realize a colour plot of a matrix. The rules for this kind of plot are:

  • if -5 < my_matrix[i,j] < 8 use white.
  • if 8 < my_matrix[i,j] < 20 use black.
  • if 20 < my_matrix[i,j] < 30 use red.
  • if my_matrix[i,j] > 30 use green.

Here is the code, as you can see the matrix is random. I have also printed a text plot in order to check if the print is correct.

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

Z_1 = 8
Z_2 = 20
Z_3 = 30

period = 5
size = 10

np.random.seed(1)

my_matrix = np.random.randint(-4,50,[period,size])

my_matrix[0,0] = 33
my_matrix[1,0] = 31

print(" ")
print(my_matrix.T)
print(" ")


bounds =  [-5,Z_1,Z_2,Z_3]    
cmap = colors.ListedColormap(['white', 'black','red','green' ]).with_extremes(over='green')
pm=plt.pcolormesh(my_matrix.T, cmap = cmap)
plt.colorbar(pm, boundaries = bounds ,extend = 'max', extendfrac='auto' ,ticks=bounds)

for i in range(len(my_matrix[:,0])):
  for j in range(len(my_matrix[0,:])):
    plt.text(i,j,str(my_matrix[i,j]), va='center_baseline', ha='left',color = 'blue')
#plt.clim(0, 31)
plt.gca().invert_yaxis()

plt.show()

I think the command .with_extremes(over='...') is useful for this kind of plot. Unfortunately I get the following picture. The green colour is not present and some bound are not respected (see my_matrix(0)(1) or my_matrix(1)(0)).

enter image description here

I want to have a plot with the aforementioned rules which has all the related colours in the colorbar, also taking into account those values greater than 30.


Solution

  • You can use a BoundaryNorm to set where the boundaries are. Also, you don't need to include green as a colour in the list of colours. For example:

    bounds =  [-5, Z_1, Z_2, Z_3]    
    cmap = colors.ListedColormap(['white', 'black', 'red']).with_extremes(over='green')
    norm = colors.BoundaryNorm(bounds, cmap.N)
    
    pm = plt.pcolormesh(my_matrix.T, cmap=cmap, norm=norm)
    plt.colorbar(pm, extend='max', extendfrac='auto')
    
    

    enter image description here