Notes‎ > ‎

posted Feb 18, 2010, 10:04 PM by Johann Goetz   [ updated Aug 4, 2011, 6:00 AM ]
When plotting a two dimensional histogram, the color maps provided by matplotlib need some adjustment to bring out features in the data. Here is a bit of code that eases the creation of new color maps with constrast in the region you want:

```import math, copy
import numpy
from matplotlib import pyplot, colors, cm

'''
returns a new colormap based on the one given

newcmap = oldcmap**a
'''
if a < 0.:
return cmap
cdict = copy.copy(cmap._segmentdata)
fn = lambda x : (x[0]**a, x[1], x[2])
for key in ('red','green','blue'):
cdict[key] = map(fn, cdict[key])
cdict[key].sort()
assert (cdict[key][0]<0 or cdict[key][-1]>1), \
"Resulting indices extend out of the [0, 1] segment."
return colors.LinearSegmentedColormap('colormap',cdict,1024)

'''
returns a new colormap based on the one given
but adjusted so that the old center point higher
(>0.5) or lower (<0.5)
'''
if not (0. < center_ratio) & (center_ratio < 1.):
return cmap
a = math.log(center_ratio) / math.log(0.5)

'''
converts center to a ratio between 0 and 1 of the
range given and calls cmap_center_adjust(). returns
'''
if not ((range[0] < center) and (center < range[1])):
return cmap
abs(center - range[0]) / abs(range[1] - range[0]))

if __name__ == '__main__':
### create some 2D histogram-type data
def func3(x,y):
return (1- x/2 + x**5 + y**3)*numpy.exp(-x**2-y**2)
x = numpy.linspace(-3.0, 3.0, 60)
y = numpy.linspace(-3.0, 3.0, 60)
X,Y = numpy.meshgrid(x, y)
Z = func3(X, Y)
extent = [x[0],x[-1],y[0],y[-1]]

plotkwargs = {
'extent' : extent,
'origin' : 'lower',
'interpolation' : 'nearest',
'aspect' : 'auto'}

fig = pyplot.figure(figsize=(8,3))
ax = [fig.add_subplot(1,3,i) for i in range(1,4,1)]

cmap = cm.seismic

plt = ax[0].imshow(Z, cmap=cmap, **plotkwargs)
cb = ax[0].figure.colorbar(plt, ax=ax[0])
ax[0].set_title('cmap: seismic')

plt = ax[1].imshow(Z, cmap=cmap_center_adjust(cmap, 0.75), **plotkwargs)
cb = ax[1].figure.colorbar(plt, ax=ax[1])
ax[1].set_title('center raised by 25%')

plt = ax[2].imshow(Z,
**plotkwargs)
cb = ax[2].figure.colorbar(plt, ax=ax[2])
ax[2].set_title('center set to 0')

pyplot.show()

```

ċ