How to get an array of strings plotted with plt.matshow and a colored mesh grid?-CodePudding

I have this kind of data (array of strings)

data = [['x', '1', '0'],['x', '1', 'x'],['0', '0', '1']]

and I'd like to have something similar to this code snippet

import matplotlib.pyplot as plt
import numpy as np

# a 2D array with linearly increasing values on the diagonal
a = np.diag(range(15))

plt.matshow(a)

plt.show()

with the following details:

'x's are colored darkred
'1's are colored white
'0's are colored orange
colored mesh between pixels
starting ticks from 1

Any help is appreciated!

CodePudding user response：

You could combine np.unique() and a ListedColormap as follows:

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

data = [['x', '1', '0'], ['x', '1', 'x'], ['0', '0', '1']]
data = np.array(data)
unique_chars, matrix = np.unique(data, return_inverse=True)
color_dict = {'x': 'darkred', '1': 'white', '0': 'orange'}
plt.matshow(matrix.reshape(data.shape), cmap=ListedColormap([color_dict[char] for char in unique_chars]))
plt.xticks(np.arange(data.shape[1]), np.arange(data.shape[1])   1)
plt.yticks(np.arange(data.shape[0]), np.arange(data.shape[0])   1)
plt.show()

You could also try seaborn's heatmap:

import seaborn as sns

sns.set(font_scale=2)
ax = sns.heatmap(matrix.reshape(data.shape), annot=data, annot_kws={'fontsize': 30}, fmt='',
                 linecolor='dodgerblue', lw=5, square=True, clip_on=False,
                 cmap=ListedColormap([color_dict[char] for char in unique_chars]),
                 xticklabels=np.arange(data.shape[1])   1, yticklabels=np.arange(data.shape[0])   1, cbar=False)
ax.tick_params(labelrotation=0)