Arbirtrary non-linear colorbar using Matplotlib

Question

I want to color edges of networks using Networkx and Matplotlib, where each edge (i,j) is given a value G[i][j]['label'] between 0 and 1.

However, often, these values are either very close to zero, or very close to 1. It's then difficult to visualize the variations of color, since everything is either very red or very blue (using a coolwarm colormap).

Then, my idea is to apply a filter filtR like one of these ones :

It's simply a polynomial function which provides a bijection from [0,1] to [0,1], and stretches more values around 0 or 1. If needed, the inverse is tractable.

For now, I just apply it to the value of the edge, in order to define its color :

cm        = plt.get_cmap('coolwarm') 
cNorm     = colors.Normalize(vmin=0., vmax=1.)
scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cm)
colorList = []

# The color is defined by filtR(G[i][j]['label'])
val_map   = {(i,j): filtR(G[i][j]['label']) for (i,j) in G.edges()}
values    = [scalarMap.to_rgba(val_map[e]) for e in G.edges()]
edges     = nx.draw_networkx_edges(G,edge_color=values,edge_cmap=plt.get_cmap('coolwarm'))


# Definition of the colorbar :-(
sm = cmx.ScalarMappable(cmap=cmx.coolwarm)
sm.set_array(values)
plt.colorbar(sm)

The question is now : I would like to define the corresponding colorbar.

For now, it shows the evaluation of my edges by the filtR function, which is meaningless : the only purpose of the filter is to modify the repartition of colors on the [0,1] interval in order to improve the readibility of the graph.

For instance, I get :

I'm happy with the left part, but not the right one, where the colorbar should be something like:

Here the filter function is clearly not the best one, but it should provide to you a better illustration.

I tried to redefine values just before the definition of the colorbar :

# Definition of the colorbar :-(

new_val_map = {(i,j): filtR(G[i][j]['label']) for (i,j) in G.edges()}
new_values  = [scalarMap.to_rgba(val_map[e]) for e in G.edges()]

sm = cmx.ScalarMappable(cmap=cmx.coolwarm)
sm.set_array(new_values)
plt.colorbar(sm)

But nothing changes.

My understanding of Matplotlib is kind of limited, and the presented code is already a patchwork of stack overflow answers.

Answer 1

Essentially you do not want to change the colormap at all. Instaed you want to create your custom normalization. To this end, you can subclass matplotlib.colors.Normalize and let it return the values of your custom function. The function would need to take values between vmin and vmax as input and return values in the range [0,1].

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as mcolors


class MyNormalize(mcolors.Normalize):
    def __call__(self, value, clip=None):
        # function to normalize any input between vmin and vmax linearly to [0,1]
        n = lambda x: (x-self.vmin)/(self.vmax-self.vmin)
        # nonlinear function between [0,1] and [0,1]
        f = lambda x,a: (2*x)**a*(2*x<1)/2. +(2-(2*(1-1*x))**a)*(2*x>=1)/2.
        return np.ma.masked_array(f(n(value),0.5))


fig, (ax,ax2) = plt.subplots(ncols=2)

x = np.linspace(-0.3,1.2, num=101)
X = (np.sort(np.random.rand(100))*1.5-0.3)

norm=  MyNormalize(vmin=-0.3, vmax=1.2)

ax.plot(x,norm(x))
im = ax2.imshow(X[::-1,np.newaxis], norm=norm, cmap="coolwarm", aspect="auto")
fig.colorbar(im)

plt.show()

The image of the desired colorbar rather suggests a partially linear function like the following beeing used.

class MyNormalize2(mcolors.Normalize):
    def __call__(self, value, clip=None):
        n = lambda x: self.vmin+(self.vmax-self.vmin)*x
        x, y = [self.vmin, n(0.2), n(0.8), self.vmax], [0, 0.48,0.52, 1]
        return np.ma.masked_array(np.interp(value, x, y))

Answer 2

You have your favorite colormap (let say coolwarm), and you want to distort it according a filtR function :

Nb : this function is the inverse of the one suggested in the initial question.

Thanks to Serenity's enlightments : the work has to be done on the colormap definition :

def distortColorMap(cm,inv = lambda x:x):
    """Inspired from 'make_colormap' in Serenity's answer.

    Inputs : a pre-existing colormap cm, 
             the distorsion function inv

    Output : the distorted colormap"""

    def f(color,inv):
        """In the sequence definition, modifies the position of stops tup[0] according the transformation function.

           Returns the distorted sequence."""
        return map(lambda tup:(inv(tup[0]),tup[1],tup[2]),color)

    # Extract sequences from cm, apply inv
    C = cm.__dict__['_segmentdata']
    cdict = {'red': f(C['red']  ,inv), 'green': f(C['green'],inv), 'blue': f(C['blue'] ,inv)}

    name = 'new_'+cm.__dict__['name']
    return colors.LinearSegmentedColormap(name, cdict)

Then, that's very easy to use :

cm        = plt.get_cmap('coolwarm')
cm        = distortColorMap(cm,inv = filtR) # all the job is done here

cNorm     = colors.Normalize(vmin=0., vmax=1.)
scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cm)

# The color is the natural value G[i][j]['label']
val_map   = {(i,j): G[i][j]['label'] for (i,j) in G.edges()}
values    = [scalarMap.to_rgba(val_map[e]) for e in G.edges()]
edges     = nx.draw_networkx_edges(G,edge_color=values,edge_cmap=plt.get_cmap('coolwarm'))


# Definition of the colorbar : just use the new colormap
sm = cmx.ScalarMappable(cmap=cm)
sm.set_array(values)
plt.colorbar(sm)

And we get then the corresponding colorbar :

Which is cool, because you don't need anymore to define the whole color sequence (everything is now done from the definition of the distorsion function), and because you can still use the fancy colormaps provided by Matplotlib !

EDIT

More info about the filtR function, and my motivations.

In this example, the filtR is defined as :

exponent = 7. 
filtR    = lambda y: ((2*y-1)**(1./exponent)+1.)/2.

With different values for exponent, we have a class of functions (with more or less smooth behaviour). Being able to jump from one definition to an other can be helpful to determine the best visualization.

Actually, for any e (even odd), Python does not like to deal with x**1/e when x is negative. But that's not a big deal, we just define properly the 7-root (or any other odd exponent).

It's, however, not the hot point : we just need a mathematical bijection from [0,1] to [0,1]. We can then take the one which fit the most our needs.

For instance, we could also want to define the filtR function as filtR = lambda y: y**4, because we want to have a better readibility on the lowest values. We would get then :

It should also work for log, piecewise, or staircase function...

I wanted a general and flexible tool, which could allow me to focus quickly on some specific areas. I don't want to create by hand sequences with stops and color values for each test of visualization.

I also want to be able to reuse this work for other projects if needed.

Answer 3

You have to define your own custom colormap and use it in custom cbar:

import matplotlib.pylab as plt
from matplotlib import colorbar, colors

def make_colormap(seq, name='mycmap'):
    """Return a LinearSegmentedColormap
    seq: a sequence of floats and RGB-tuples. The floats should be increasing
    and in the interval (0,1).
    """
    seq = [(None,) * 3, 0.0] + list(seq) + [1.0, (None,) * 3]
    cdict = {'red': [], 'green': [], 'blue': []}
    for i, item in enumerate(seq):
        if isinstance(item, float):
            r1, g1, b1 = seq[i - 1]
            r2, g2, b2 = seq[i + 1]
            cdict['red'].append([item, r1, r2])
            cdict['green'].append([item, g1, g2])
            cdict['blue'].append([item, b1, b2])
    return colors.LinearSegmentedColormap(name, cdict)

def generate_cmap(lowColor, highColor, lowBorder, highBorder):
    """Apply edge colors till borders and middle is in grey color"""
    c = colors.ColorConverter().to_rgb
    return make_colormap([c(lowColor), c('grey'),l owBorder, c('grey'), .5, \
     c('grey'), highBorder ,c('grey'), c(highColor)])

fig = plt.figure()
ax = fig.add_axes([.05, .05, .02, .7]) # position of colorbar

cbar = colorbar.ColorbarBase(ax, cmap=generate_cmap('b','r',.15,.85),
  norm=colors.Normalize(vmin=.0, vmax=1)) # set min, max of colorbar
ticks = [0.,.1,.2,.3,.4,.5,.6,.7,.8,.9,1.]
cbar.set_ticks(ticks) # add ticks
plt.show()