Operating on color vectors
Ever wanted to reverse a colormap, or to desaturate one ? Here is a routine to apply a function to the look up table of a colormap:
1 def cmap_map(function,cmap):
2 """ Applies function (which should operate on vectors of shape 3:
3 [r, g, b], on colormap cmap. This routine will break any discontinuous points in a colormap.
4 """
5 cdict = cmap._segmentdata
6 step_dict = {}
7 # Firt get the list of points where the segments start or end
8 for key in ('red','green','blue'): step_dict[key] = map(lambda x: x[0], cdict[key])
9 step_list = sum(step_dict.values(), [])
10 step_list = array(list(set(step_list)))
11 # Then compute the LUT, and apply the function to the LUT
12 reduced_cmap = lambda step : array(cmap(step)[0:3])
13 old_LUT = array(map( reduced_cmap, step_list))
14 new_LUT = array(map( function, old_LUT))
15 # Now try to make a minimal segment definition of the new LUT
16 cdict = {}
17 for i,key in enumerate(('red','green','blue')):
18 this_cdict = {}
19 for j,step in enumerate(step_list):
20 if step in step_dict[key]:
21 this_cdict[step] = new_LUT[j,i]
22 elif new_LUT[j,i]!=old_LUT[j,i]:
23 this_cdict[step] = new_LUT[j,i]
24 colorvector= map(lambda x: x + (x[1], ), this_cdict.items())
25 colorvector.sort()
26 cdict[key] = colorvector
27
28 return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)
Lets try it out: I want a jet colormap, but lighter, so that I can plot things on top of it:
light_jet = cmap_map(lambda x: x/2+0.5, cm.jet)
x,y=mgrid[1:2,1:10:0.1]
imshow(y, cmap=light_jet)
As a comparison, this is what the original jet looks like: 
Operating on indices
OK, but what if you want to change the indices of a colormap, but not its colors.
1 def cmap_xmap(function,cmap):
2 """ Applies function, on the indices of colormap cmap. Beware, function
3 should map the [0, 1] segment to itself, or you are in for surprises.
4
5 See also cmap_xmap.
6 """
7 cdict = cmap._segmentdata
8 function_to_map = lambda x : (function(x[0]), x[1], x[2])
9 for key in ('red','green','blue'): cdict[key] = map(function_to_map, cdict[key])
10 cdict[key].sort()
11 assert (cdict[key][0]<0 or cdict[key][-1]>1), "Resulting indices extend out of the [0, 1] segment."
12
13
14 return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)
Discrete colormap
Here is how you can discretize a continuous colormap.
1 def cmap_discretize(cmap, N):
2 """Return a discrete colormap from the continuous colormap cmap.
3
4 cmap: colormap instance, eg. cm.jet.
5 N: Number of colors.
6
7 Example
8 x = resize(arange(100), (5,100))
9 djet = cmap_discretize(cm.jet, 5)
10 imshow(x, cmap=djet)
11 """
12
13 cdict = cmap._segmentdata.copy()
14 # N colors
15 colors_i = linspace(0,1.,N)
16 # N+1 indices
17 indices = linspace(0,1.,N+1)
18 for key in ('red','green','blue'):
19 # Find the N colors
20 D = array(cdict[key])
21 I = interpolate.interp1d(D[:,0], D[:,1])
22 colors = I(colors_i)
23 # Place these colors at the correct indices.
24 A = zeros((N+1,3), float)
25 A[:,0] = indices
26 A[1:,1] = colors
27 A[:-1,2] = colors
28 # Create a tuple for the dictionary.
29 L = []
30 for l in A:
31 L.append(tuple(l))
32 cdict[key] = tuple(L)
33 # Return colormap object.
34 return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)
So for instance, this is what you would get by doing cmap_discretize(cm.jet, 6). 