Setting arbitrary color in pseudocolor plot with matplotlib [duplicate]

Question

Possible Duplicate:
Add separate colors for two (or more) specific values in color plot and color bar

I've got an array of values and I'd like to plot these using pcolor in matplotlib. I'm using the color map 'YlOrRd' and it works fine - but I'd like to use the color map for all values other than zero. That is, all values not 0, should use the color map - I'd like 0 to be black.

currently I'm using, for my numpy array of values 'x'.

pcolor(x,cmap=cm.YlOrRd)

is there a way to arbitrarily fix all values in array x that are 0 to black?

Thanks, Dave.

Answer 1

Here are two ways to do it. One by creating your own colormap, and one by using a masked array. Say we have:

import matplotlib
from pylab import *

data = np.arange(-50, 50).reshape(10, 10)
data = np.abs(data)
pcolor(data, cmap=cm.YlOrRd)
show()

This produces: Now we do the same thing, but create a list called colors that has the same values as cm.YlOrRd, apart from the 0 entry, which we set to black (0,0,0 in rgb). Then we use LinearSegmentedColormap.from_list to actually make the colormap:

import matplotlib
from pylab import *

data = np.arange(-50, 50).reshape(10, 10)
data = np.abs(data)
colors = [(0,0,0)] + [(cm.YlOrRd(i)) for i in xrange(1,256)]
new_map = matplotlib.colors.LinearSegmentedColormap.from_list('new_map', colors, N=256)
pcolor(data, cmap=new_map)
savefig('map.png')
show()

This produces the same plot, but zero values are black:

Here is the other way using masked arrays, its a bit more involved, comments in the code explain the steps:

from pylab import *
import numpy.ma as ma

data=np.arange(-50,50).reshape(10,10)
data=np.abs(data)
#create a mask where only values=0 are true:
mask = data == 0
#create a masked array by combining our mask and data:
mx = ma.masked_array(data, mask)
#set masked values in cm.YlOrRd to 'black'
cm.YlOrRd.set_bad(color='black', alpha=None)
# pcolor(data,cmap=cm.YlOrRd)
#we must use pcolormesh instead of pcolor, as pcolor does not draw masked values at all
pcolormesh(mx,cmap=cm.YlOrRd)
show()

This produces the same plot as immediately above.

There is a potential difference between these methods, the first method will round down data values and apply the appropriate color, while the second method will only set values equal to 0 to black (ie. 0.001 will not be masked, so will be the appropriate cm.YlOrRd color). The main advantage of the second is that you can mask out entries totally arbitrarily.