define custom Normalisation function in matplotlib when using plt.colorbar()

Question

to use matplotlib colorbar, one has to specify a matplotlib.cm.ScalarMappable with an object from a subclass of matplotlib.colors.Normalize, from which the colorbar can know how to normalise the data into [0,1] float value.

there are only a few normalisation process provided by matplotlib, linear normalisation, log, power law etc. but in practice, we might want to use other normalisation function written by ourselves.

we can normalise the data array into [0,1] using whatever function, but without Scalarmappable built with Nomalization subclass, the colorbar won't have ticks and labels right.

I am wondering is it my understanding to matplotlib colorbar right or there is other way to do it quite easily? or perhaps we have to manually write a subclass to wrap the custom normalisation function?

Answer 1

You can easily subclass matplotlib.colors.Normalize for this purpose. Here's an example of a piece-wise normalization class I wrote for a previous SO question:

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

class PiecewiseNorm(Normalize):
    def __init__(self, levels, clip=False):
        # input levels
        self._levels = np.sort(levels)
        # corresponding normalized values between 0 and 1
        self._normed = np.linspace(0, 1, len(levels))
        Normalize.__init__(self, None, None, clip)

    def __call__(self, value, clip=None):
        # linearly interpolate to get the normalized value
        return np.ma.masked_array(np.interp(value, self._levels, self._normed))

    def inverse(self, value):
        return 1.0 - self.__call__(value)

For example:

y, x = np.mgrid[0.0:3.0:100j, 0.0:5.0:100j]
H = 50.0 * np.exp( -(x**2 + y**2) / 4.0 )
levels = [0, 1, 2, 3, 6, 9, 20, 50]

H1 = -50.0 * np.exp( -(x**2 + y**2) / 4.0 )
levels1 = [-50, -20, -9, -6, -3, -2, -1, 0]

fig, ax = plt.subplots(2, 2, gridspec_kw={'width_ratios':(20, 1), 'wspace':0.05})

im0 = ax[0, 0].contourf(x, y, H, levels, cmap='jet', norm=PiecewiseNorm(levels))
cb0 = fig.colorbar(im0, cax=ax[0, 1])
im1 = ax[1, 0].contourf(x, y, H1, levels1, cmap='jet', norm=PiecewiseNorm(levels1))
cb1 = fig.colorbar(im1, cax=ax[1, 1])

plt.show()