Assigning arbitrary colors to values in python matplotlib

Question

I have a 2-dimensional array of float numbers in python with only 8 possible values which I input into a square grid colored plot which represents the values the array. I only know of assigning different colors by interpolation=nearest, but since my values can be similar, getting a clear image requires changing the values in the array for them to be different enough.

Is there a method such that given varaibles

val1,...,val8

Can I choose that this values will appear with 8 arbitrary colors?

Answer 1

You can use np.unique(..., return_inverse=True) to get all the unique values and the inverse. That inverse gives the corresponding index into the list of unique values for each input value. The inverse is a 1D array and needs to be made 2D again. Note that the array of unique values will be sorted from low to high.

Then, you can for example use seaborns heatmap to draw the heatmap. Optionally together with annotated values and/or a colorbar. (You can do something similar directly with matplotlib).

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

input_values = np.random.rand(8)  # 8 different values
m = np.random.choice(input_values, (10, 10))  # 100 cells, each with one o the 8 values
values, m_cat = np.unique(m, return_inverse=True)
m_cat = m_cat.reshape(m.shape)
colors = ['#e41a1c', '#377eb8', '#4daf4a', '#984ea3', '#ff7f00', '#ffff33', '#a65628', '#f781bf']  # 8 colors

ax = sns.heatmap(m_cat, cmap=ListedColormap(colors), annot=m, fmt='.2f', cbar=True,
                 cbar_kws={'ticks': np.linspace(0, len(values) - 1, 2 * len(values) + 1)[1::2]})
cbar_ax = ax.figure.axes[-1]  # last created ax
cbar_ax.set_yticklabels([f'{v:.6f}' for v in values])
cbar_ax.set_title('Values')
plt.show()