How to create a discrete colorbar with matplotlib ?

There are several ways to create a discrete colorbar for visualizations using matplotlib. Here are some examples:

Create data

Before we begin, let's make a two-dimensional function that can be evaluated on a grid:

import numpy as np
import random

def f(x1,x2):
    return np.exp(-(x1**2+x2**2))

x1_min = -2.0
x1_max = 2.0
x2_min = -2.0
x2_max = 2.0

x1, x2 = np.meshgrid(np.arange(x1_min,x1_max, 0.1), np.arange(x2_min,x2_max, 0.1))

y = f(x1,x2)

Create discrete colorbars from continuous ones

We can use the "imshow" function along with the "hot" color map to visualize the dataset created above:

from pylab import figure, cm

import matplotlib.pyplot as plt
import matplotlib

plt.imshow(y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cm.hot, origin='lower')

plt.colorbar()

plt.title("How to evaluate a 2D function using a python grid?" , fontsize=8)

plt.savefig("create_discrete_colorbar_01.png", bbox_inches='tight')

plt.show()

To convert a continuous colorbar into a discrete one, we can utilize the get_cmap() function. This function requires the name of the desired colormap (in this case, 'hot') and the number of discrete values we want (let's use 10 as an example):

cmap = cm.get_cmap('hot', 10)

color_list = []
for i in range(cmap.N):
    rgba = cmap(i)
    print(matplotlib.colors.rgb2hex(rgba))
    color_list.append(matplotlib.colors.rgb2hex(rgba))

Output

#0b0000
#550000
#9f0000
#ea0000
#ff3500
#ff8000
#ffca00
#ffff20
#ffff8f
#ffffff

We can create a discrete colormap using ListedColormap() by choosing the colors from a list of hexadecimal values and then plot the data:

cmap = color_list

cmap = matplotlib.colors.ListedColormap(cmap)

plt.imshow(y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cmap, origin='lower')

plt.colorbar()

plt.title("How to create a discrete colorbar with matplotlib ?" , fontsize=8)

plt.savefig("create_discrete_colorbar_02.png", bbox_inches='tight')

plt.show()

Add a color to the discrete colorbar

The benefit of using the approach mentioned is that it allows for simple addition of color. For instance, if we desire to replace the small value with the color gray, we can accomplish this by using the following code:

Create new color list:

cmap = cm.get_cmap('hot', 10)    # PiYG

color_list = ['#808080'] # Gray

for i in range(cmap.N):
    rgba = cmap(i)
    # rgb2hex accepts rgb or rgba
    #print(matplotlib.colors.rgb2hex(rgba))
    color_list.append(matplotlib.colors.rgb2hex(rgba))

Plot the data:

cmap = color_list

cmap = matplotlib.colors.ListedColormap(cmap)

plt.imshow(y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cmap, origin='lower')

plt.colorbar()

plt.title("How to create a discrete colorbar with matplotlib ?" , fontsize=8)

plt.savefig("create_discrete_colorbar_03.png", bbox_inches='tight')

plt.show()

Edit colorbar labels

Finally, we can define the labels on the discrete colorbar as follows:

cmap = color_list

cmap = matplotlib.colors.ListedColormap(cmap)

bounds = [i/10 for i in range(11)]

norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)

img = plt.imshow(y,extent=[x1_min,x1_max,x2_min,x2_max], cmap=cmap, origin='lower')

cbar_bounds = bounds
cbar_ticks =  [(cbar_bounds[i+1]-cbar_bounds[i])/2.0+cbar_bounds[i] for i in range( len(cbar_bounds) - 1 )] 
cbar_labels = ['No Fire', 'Smoldering','Flaming','Saturated', 'Bad Data']

cbar = plt.colorbar(img, cmap=cmap, norm=norm, boundaries=cbar_bounds, ticks=cbar_ticks)

plt.title("How to create a discrete colorbar with matplotlib ?" , fontsize=8)

plt.savefig("create_discrete_colorbar_04.png", bbox_inches='tight')

plt.show()

References