Color heatmap according to value

I would like to plot a heatmap and specify a color for each value:

import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
data = pd.DataFrame()
data['percentile'] = pd.DataFrame({'percentile':[95] * 3 + [98] * 3 + [99] * 3})
data['time_scale'] = pd.DataFrame({'time_scale' : [1 , 7, 30 ] * 3 })
data['drivers'] = pd.DataFrame({'drivers' : [ 1 , 3 , 5 , 3 ,1 , 4 , 4  ,5 , 1] })

data = data.pivot_table(values='drivers', index='percentile', columns='time_scale').to_numpy()
im = plt.imshow(data, interpolation='none', origin='lower')

How to specify the colors depending on the value of drivers?
1 : grey
2 : black
3 : blue, etc…

I tried this :

cmap = matplotlib.cm.get_cmap('Set2', 5)

But I don’t know how to specify the exact colors I want.
Thank you.

>Solution :

You can use LinearSegmentedColormap from matplotlib.colors :

cmap = matplotlib.colors.LinearSegmentedColormap.from_list('', ['grey', 'black', 'blue', 'purple', 'red'])
im = plt.imshow(data, interpolation='none', origin='lower', cmap=cmap)

You can also add annotations to check if the colors are arranged correctly:

for i in range(len(data)):
    for j in range(len(data)):
        text = plt.text(j, i, data[i, j],
                       ha="center", va="center", color="w")