Change colorbar gradient in matplotlib

Question

I have a grid of weights (Y) evolving with time (X):

I cannot distinguish correctly the variations of weights as the distribution is asymmetric between positive and negative weights; the null weights should be recognized as it means the given variables are not used.

For these reasons, I would like to change the color gradient to get something like those (either a or b):

Any idea on how to approach this?

Answer 1

A colorbar in matplotlib maps number between 0 and 1 to a color. In order to map other numbers to colors you need a normalization to the range [0,1] first. This is usually done automatically from the minimum and maximum data, or by using vmin and vmax arguments to the respective plotting function. Internally a normalization instance matplotlib.colors.Normalize is used to perform the normalization and by default a linear scale between vmin and vmax is assumed.

Here you want a nonlinear scale, which (a) shifts the middle point to some specified value, and (b) squeezes the colors around that value.

The idea can now be to subclass matplotlib.colors.Normalize and let it return a a mapping which fulfills the criteria (a) and (b).

An option might be the combination of two root functions as shown below.

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

class SqueezedNorm(matplotlib.colors.Normalize):
    def __init__(self, vmin=None, vmax=None, mid=0, s1=2, s2=2, clip=False):
        self.vmin = vmin # minimum value
        self.mid  = mid  # middle value
        self.vmax = vmax # maximum value
        self.s1=s1; self.s2=s2
        f = lambda x, zero,vmax,s: np.abs((x-zero)/(vmax-zero))**(1./s)*0.5
        self.g = lambda x, zero,vmin,vmax, s1,s2: f(x,zero,vmax,s1)*(x>=zero) - \
                                             f(x,zero,vmin,s2)*(x<zero)+0.5
        matplotlib.colors.Normalize.__init__(self, vmin, vmax, clip)

    def __call__(self, value, clip=None):
        r = self.g(value, self.mid,self.vmin,self.vmax, self.s1,self.s2)
        return np.ma.masked_array(r)


fig, (ax, ax2, ax3) = plt.subplots(nrows=3, 
                                   gridspec_kw={"height_ratios":[3,2,1], "hspace":0.25})

x = np.linspace(-13,4, 110)
norm=SqueezedNorm(vmin=-13, vmax=4, mid=0, s1=1.7, s2=4)

line, = ax.plot(x, norm(x))
ax.margins(0)
ax.set_ylim(0,1)

im = ax2.imshow(np.atleast_2d(x).T, cmap="Spectral_r", norm=norm, aspect="auto")
cbar = fig.colorbar(im ,cax=ax3,ax=ax2, orientation="horizontal")

The function is chosen such that independent of its parameters it will map any range onto the range [0,1], such that a colormap can be used. The parameter mid determines which value should be mapped to the middle of the colormap. This would be 0 in this case. The parameters s1 and s2 determine how squeezed the colormap is in both directions.

Setting mid = np.mean(vmin, vmax), s1=1, s2=1 would recover the original scaling.

In order to choose good parameters, one may use some Sliders to see the live updated plot.

from matplotlib.widgets import Slider

midax = plt.axes([0.1, 0.04, 0.2, 0.03], facecolor="lightblue")
s1ax = plt.axes([0.4, 0.04, 0.2, 0.03], facecolor="lightblue")
s2ax = plt.axes([0.7, 0.04, 0.2, 0.03], facecolor="lightblue")

mid = Slider(midax, 'Midpoint', x[0], x[-1], valinit=0)
s1 = Slider(s1ax, 'S1', 0.5, 6, valinit=1.7)
s2 = Slider(s2ax, 'S2', 0.5, 6, valinit=4)


def update(val):
    norm=SqueezedNorm(vmin=-13, vmax=4, mid=mid.val, s1=s1.val, s2=s2.val)
    im.set_norm(norm)
    cbar.update_bruteforce(im) 
    line.set_ydata(norm(x))
    fig.canvas.draw_idle()

mid.on_changed(update)
s1.on_changed(update)
s2.on_changed(update)

fig.subplots_adjust(bottom=0.15)

Answer 2

You can use a custom normalizer. Conveniently, the example for this in the docs is already an 'alternative midpoint' normalizer. The example is made by Joe Kington, so all credits to him.

See the bottom of this page: https://matplotlib.org/users/colormapnorms.html

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

Custom normalize class:

class MidpointNormalize(mpl.colors.Normalize):
    ## class from the mpl docs:
    # https://matplotlib.org/users/colormapnorms.html

    def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
        self.midpoint = midpoint
        super().__init__(vmin, vmax, clip)

    def __call__(self, value, clip=None):
        # I'm ignoring masked values and all kinds of edge cases to make a
        # simple example...
        x, y = [self.vmin, self.midpoint, self.vmax], [0, 0.5, 1]
        return np.ma.masked_array(np.interp(value, x, y))    

The result:

data = np.linspace(-5,1,100)[None,:]

fig, axs = plt.subplots(2,1, figsize=(5,2), facecolor='w', subplot_kw=dict(xticks=[], yticks=[]))

props = dict(aspect=15, cmap=plt.cm.coolwarm)

axs[0].imshow(data, **props)
axs[1].imshow(data, norm=MidpointNormalize(midpoint=0), **props)

This is a relative simple example, but more complex scalings can be achieved in a similar matter.