title: Make a heatmap with a specified discrete color mapping with matplotlib in python
body: <p>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. </p>
  
  <p>The best solution I have come up with is using:</p>
  
  <pre><code>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)
  </code></pre>
  
  <p>which successfully prints a stripe of each color. However, if I plot instead:</p>
  
  <pre><code>dat2 = np.array([[0, 0,0,0], [1, 1, 1, 1]])
  plt.imshow(dat2, interpolation='none', aspect='auto', origin='upper', cmap=cmap)
  </code></pre>
  
  <p>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.</p>