I would like to make a heatmap for a matrix of data such that all positions that are 1 will be red, all positions that are 2 will be white, and etc. with an arbitrary specification. Ideally this should handle the case where all of the values are the same, plotting just a uniform color.
The best solution I have come up with is using:
from matplotlib import colors
import matplotlib.pyplot as plt
import numpy as np
cmap = colors.ListedColormap(['white', 'blue', 'red', 'purple'])
data = np.array([[0, 0,0,0], [1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]])
plt.imshow(data, interpolation='none', aspect='auto', origin='upper', cmap=cmap)
which successfully prints a stripe of each color. However, if I plot instead:
dat2 = np.array([[0, 0,0,0], [1, 1, 1, 1]])
plt.imshow(dat2, interpolation='none', aspect='auto', origin='upper', cmap=cmap)
instead, it plots white and purple rather than white and blue. If the data contains only one of the numbers, it will only plot white.