Non-linear scaling of a colormap to enhance contrast

The following python code creates a heatmap of a开发者_开发问答 matrix that contains normally distributed values

import numpy as np
from matplotlib import pylab as plt


np.random.seed(123) #make sure we all have same data
m = np.random.randn(200).reshape(10, 20)
plt.imshow(m, cmap='RdYlGn', interpolation='nearest')
plt.colorbar()

This is the output of this code

I would like to enhance the contrast of this image by "fading out" the values close to zero. I can easily do this by using disigmoid scaling of the original data as follows:

def disigmoidScaling(values, steepnessFactor=1, ref=None):
    ''' Sigmoid scaling in which values around a reference point are flattened
    arround a reference point

    Scaled value y is calculated as 
        y = sign(v - d)(1 - exp(-((x - d)/s)**2)))
    where v is the original value,  d is the referenc point and s is the 
    steepness factor
    '''
    if ref is None:
        mn = np.min(values)
        mx = np.max(values)
        ref = mn + (mx - mn) / 2.0

    sgn = np.sign(values - ref)
    term1 = ((values - ref)/steepnessFactor) ** 2
    term2 = np.exp(- term1) 
    term3 = 1.0 - term2 
    return sgn * term3


plt.imshow(disigmoidScaling(m, 4), cmap='RdYlGn', interpolation='nearest')
plt.colorbar()

Here is the output.

I'm pleased with the result, except the fact that in this version the original values have been exchanged for scaled ones.

Is there a way to perform a non-linear mapping of values to colormap?

A colormap contains a dictionary of red, green and blue values mapped over the interval [0,1]. The Linear Segmented Colormap class docs give the example

cdict = {'red':   [(0.0,  0.0, 0.0),
               (0.5,  1.0, 1.0),
               (1.0,  1.0, 1.0)],

     'green': [(0.0,  0.0, 0.0),
               (0.25, 0.0, 0.0),
               (0.75, 1.0, 1.0),
               (1.0,  1.0, 1.0)],

     'blue':  [(0.0,  0.0, 0.0),
               (0.5,  0.0, 0.0),
               (1.0,  1.0, 1.0)]}

"Each row in the table for a given color is a sequence of x, y0, y1 tuples. In each sequence, x must increase monotonically from 0 to 1. For any input value z falling between x[i] and x[i+1], the output value of a given color will be linearly interpolated between y1[i] and y0[i+1]:"

The RdYlGn colormap has 11 x values for each color going from 0 to 1.0 in steps of 0.1. You can get the cdict values by calling

plt.cm.RdYlGn._segmentdata

You can then change the x values to whatever steps you want (as long as they are monotonically increasing and range from 0 to 1) and get a new colormap by calling matplotlib.colors.LinearSegmentedColormap on your new cdict. There are several great examples of this in the Matplotlib Cookbook.