matplotlib color N-slope-normalization

Question

Based from the source code of the TwoSlopeNorm color normalization of matplotlib, I tried to implement a more general normalization that can handle any number of "breakpoints" (instead of just one in the middle of the colormap - at 0.5). In other words, this a N-slope normalization. It produces the expected output, and the normalzation responds correctly by updating the image/colormap/colorbar when one of its attributes in modified programmaticaly as with im.norm.vcenters = [0.2, 0.6].

I really did not modify much from the source example. Basically just replaced vcenter from the TwoSlopeNorm by a list vcenters that store the v-values (those between vmin and vmax). Accordingly, a list of "boundaries" allows to control the position in the colormap dynamic (whereas the value 0.5 was hardcorded in the 2-slope norm).

One thing I am wondering is if this N-slope normalization could be implemented using FuncNorm. The end goal being to improve this N-slope norm into a N-slope-percentage so that it does not relies on value-space limits (vcenters), but rather on percentiles of the image (so vcenters will be replaced by a percentiles list).

# %matplotlib qt # uncomment if in notebook
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib.colors import Normalize

cmap = plt.cm.get_cmap('bwr').copy()
cmap.set_over('purple')
cmap.set_under('green')

class NSlopeNorm(Normalize):
    def __init__(self, vcenters, boundaries=None, vmin=None, vmax=None):

        super().__init__(vmin=vmin, vmax=vmax)

        if boundaries is None:
            boundaries = np.linspace(0, 1, len(vcenters)+2)
            boundaries = boundaries[1:-1]
        if not len(vcenters)==(len(boundaries)):
            raise ValueError('Incompatible length between vs and pcs')

        self._vcenters = vcenters
        self._boundaries = boundaries

        if vmin is not None and not np.all(np.diff(np.concatenate((np.atleast_1d(self.vmin), self.vcenters)))>0):
            raise ValueError('vs must be in ascending order')
        if vmax is not None and not np.all(np.diff(np.concatenate((self.vcenters, np.atleast_1d(self.vmax))))>0):
            raise ValueError('vs must be in ascending order')


    @property
    def vcenters(self):
        return np.array(self._vcenters)
    
    @property
    def boundaries(self):
        return [0, *self._boundaries, 1]
    
    @boundaries.setter
    def boundaries(self, value):
        if value != self._boundaries:
            self._boundaries = value
            self._changed()

    @vcenters.setter
    def vcenters(self, value):
        if value != self._vcenters:
            self._vcenters = value
            self._changed()

    def autoscale_None(self, A):
        """
        Get vmin and vmax.

        If all vcenters isn't in the range [vmin, vmax], either vmin or vmax
        is expanded.
        """
        super().autoscale_None(A)
        if self.vmin >= np.min(self.vcenters):
            self.vmin = np.min(self.vcenters)
        if self.vmax <= np.max(self.vcenters):
            self.vmax = np.max(self.vcenters)

    def __call__(self, value, clip=None):
        """
        Map value to the interval [0, 1]. The *clip* argument is unused.
        """
        result, is_scalar = self.process_value(value)
        self.autoscale_None(result)  # sets self.vmin, self.vmax if None

        if not (np.all(self.vmin <= self.vcenters) & np.all(self.vcenters <= self.vmax)):
            raise ValueError("vmin, vcenter, vmax must increase monotonically")
        # note that we must extrapolate for tick locators:
        result = np.ma.masked_array(
            np.interp(result, [self.vmin, *self.vcenters, self.vmax],
                      self.boundaries, left=-np.inf, right=np.inf),
            mask=np.ma.getmask(result))
        if is_scalar:
            result = np.atleast_1d(result)[0]
        return result

    def inverse(self, value):
        if not self.scaled():
            raise ValueError("Not invertible until both vmin and vmax are set")
        (vmin,), _ = self.process_value(self.vmin)
        (vmax,), _ = self.process_value(self.vmax)
        vcenters = []
        for v in self.vcenters:
            (val,), _ = self.process_value(v)
            vcenters.append(val)
        #(vcenters,), _ = self.process_value(self.vcenters)
        result = np.interp(value, self.boundaries, [vmin, *vcenters, vmax],
                           left=-np.inf, right=np.inf)
        return result


# say you data ranges from 20°C to 50°C
data = np.linspace(20, 50, 100).reshape(10,10) # toy data
fig, ax = plt.subplots(figsize=(8, 8))
PARAMS = {
    'vmin':25,                  # no interest below 25°C
    'vcenters':[34, 40],       # "zoom" between 34 and 40...
    'boundaries':[0.2, 0.8],    # ...by allowing most of the dynamic
    'vmax':42,                  # no intereset above 42
}
ax.set_title(f'{PARAMS}')
im = ax.imshow(
    data, 
    cmap=cmap,
    norm=NSlopeNorm(**PARAMS),
) 
plt.colorbar(im, extend='both')
plt.tight_layout()

# those attributes can be set/get interactively
print(im.norm.boundaries, im.norm.vcenters, im.norm.vmin, im.norm.vmax)

# for interpretation purpose
for i in range(data.shape[0]):
    for j in range(data.shape[1]):
        text = ax.text(j, i, f'{data[i, j]:.2f}', ha='center', va='center', color='yellow')

Answer 1

if not (np.all(self.vmin <= self.vcenters) & np.all(self.vcenters <= self.vmax)):

Use boolean operators when possible. Bitwise operators don't short-circuit. Here you're just comparing two bools, so use the short-circuiting and/or.

Also, avoid negation when possible. Here if not (all and all) is logically equivalent to the more intuitive if any or any.

if vmin is not None and not np.all(np.diff(np.concatenate((np.atleast_1d(self.vmin), self.vcenters)))>0):

For readability, limit lines to <80 chars. You could either reformat or restructure. Here I'd probably extract the concatenation as v.

And for speed, v[1:] > v[:-1] would be ~2x faster than np.diff(v) > 0.

raise ValueError('vs must be in ascending order')

Be explicit about what "vs" means ("vmin, vcenters, and vmax"). In fact that's what matplotlib originally did, but you abbreviated it. Keep in mind that this error text is not just for you but also for end users, so please keep matplotlib's explicit wording for clarity.

raise ValueError('Incompatible length between vs and pcs')

Case in point, here "vs" means something different (only "vcenters"). Be explicit and write the actual keywords ("vcenters" and "boundaries").

And since this is part of an existing codebase, best practice is to mirror their style. Rephrase the message to match matplotlib's "must be" wording (e.g., "vcenters and boundaries must be the same length"), and change the single quotes to double quotes.

---

Revised NSlopeNorm:

(Note that I only reviewed the parts you had added/changed.)

import numpy as np
from matplotlib.colors import Normalize


class NSlopeNorm(Normalize):
    def __init__(self, vcenters, boundaries=None, vmin=None, vmax=None):
        """
        Normalize data with set centers.

        Useful when mapping data with an unequal rates of change around
        conceptual centers, e.g., data that range from -2 to 4,
        with 0 and 1 as the midpoints.

        Parameters
        ----------
        vcenters : array-like
            The data values that define the inner boundaries in the normalization.
        boundaries : array-like, optional
            The values that define the colormap breakpoints.
        vmin : float, optional
            The data value that defines ``0.0`` in the normalization.
            Defaults to the min value of the dataset.
        vmax : float, optional
            The data value that defines ``1.0`` in the normalization.
            Defaults to the max value of the dataset.
        """

        super().__init__(vmin=vmin, vmax=vmax)

        if boundaries is None:
            boundaries = np.linspace(0, 1, len(vcenters)+2)
            boundaries = boundaries[1:-1]
        if not len(vcenters)==(len(boundaries)):
            raise ValueError("vcenters and boundaries must be the same length")

        self._vcenters = vcenters
        self._boundaries = boundaries

        if vmin is not None:
            v = np.concatenate((np.atleast_1d(self.vmin), self.vcenters))
            if not np.all(v[1:] > v[:-1]):
                raise ValueError("vmin and vcenters must be in ascending order")
            
        if vmax is not None:
            v = np.concatenate(self.vcenters), (np.atleast_1d(self.vmax))
            if not np.all(v[1:] > v[:-1]):
                raise ValueError("vcenters and vmax must be in ascending order")

    @property
    def vcenters(self):
        return np.array(self._vcenters)
    
    @property
    def boundaries(self):
        return [0, *self._boundaries, 1]
    
    @boundaries.setter
    def boundaries(self, value):
        if value != self._boundaries:
            self._boundaries = value
            self._changed()

    @vcenters.setter
    def vcenters(self, value):
        if value != self._vcenters:
            self._vcenters = value
            self._changed()

    def autoscale_None(self, A):
        """
        Get vmin and vmax.

        If all vcenters aren't in the range [vmin, vmax], either vmin or vmax
        is expanded.
        """
        super().autoscale_None(A)
        if self.vmin >= np.min(self.vcenters):
            self.vmin = np.min(self.vcenters)
        if self.vmax <= np.max(self.vcenters):
            self.vmax = np.max(self.vcenters)

    def __call__(self, value, clip=None):
        """
        Map value to the interval [0, 1]. The *clip* argument is unused.
        """
        result, is_scalar = self.process_value(value)
        self.autoscale_None(result)  # sets self.vmin, self.vmax if None

        if np.any(self.vmin > self.vcenters) or np.any(self.vcenters > self.vmax):
            raise ValueError("vmin, vcenters, vmax must increase monotonically")
        # note that we must extrapolate for tick locators:
        result = np.ma.masked_array(
            np.interp(result, [self.vmin, *self.vcenters, self.vmax],
                      self.boundaries, left=-np.inf, right=np.inf),
            mask=np.ma.getmask(result))
        if is_scalar:
            result = np.atleast_1d(result)[0]
        return result

    def inverse(self, value):
        if not self.scaled():
            raise ValueError("Not invertible until both vmin and vmax are set")
        (vmin,), _ = self.process_value(self.vmin)
        (vmax,), _ = self.process_value(self.vmax)
        vcenters = []
        for v in self.vcenters:
            (val,), _ = self.process_value(v)
            vcenters.append(val)
        result = np.interp(value, self.boundaries, [vmin, *vcenters, vmax],
                           left=-np.inf, right=np.inf)
        return result