CODEWITHSUNDEEP
pythonmatplotlib
dylkot
dylkot

Reputation: 2475

Make a heatmap with a specified discrete color mapping with matplotlib in python

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.

Upvotes: 1

Views: 2692

Answers (1)

James
James

Reputation: 36781

I believe the colormap is renormalizing to your data. Passing values that range from 0 to 1 gives a colormapping of white:0, blue:0.333, red:0.666, purple:1.0.

You can prevent this behavior by passing vmin and vmax to the plot.

plt.imshow(dat2, interpolation='none', aspect='auto', origin='upper',
           cmap=cmap, vmin=0, vmax=4)

(I don't have python running on my computer right now, so I am unable to test this.)

Upvotes: 5

Related Questions