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


class PNorm(Normalize):
    def __init__(self, vmin=None, vmax=None, clip=False, *, gamma):
        super().__init__(vmin=vmin, vmax=vmax, clip=clip)
        self.gamma = gamma

    def __call__(self, value, clip=None):
        if clip is None:
            clip = self.clip
        result, is_scalar = self.process_value(value)
        self.autoscale_None(result)
        vmin, vmax = self.vmin, self.vmax
        gamma = self.gamma
        if vmin > vmax:
            raise ValueError("minvalue must be less than or equal to maxvalue")
        elif vmin == vmax:
            result.fill(0)
        else:
            if clip:
                mask = np.ma.getmask(result)
                result = np.ma.array(np.clip(result.filled(vmax), vmin, vmax),
                                     mask=mask)
            resdat = result.data
            resdat = resdat ** gamma
            resdat -= min(vmin ** gamma, vmax ** gamma)
            resdat /= abs(vmax ** gamma - vmin ** gamma)
            result = np.ma.array(resdat, mask=result.mask, copy=False)
        if is_scalar:
            result = result[0]
        return result

    def inverse(self, value):
        if not self.scaled():
            raise ValueError("Not invertible until scaled")
        if np.iterable(value):
            return np.vectorize(self.inverse)(value)
        else:
            vmin, vmax = self.vmin, self.vmax
            gamma = self.gamma
            return (value * abs(vmax ** gamma - vmin ** gamma)
                    + min(vmin ** gamma, vmax ** gamma)) ** (1 / gamma)


fig, axs = plt.subplots(1, 2)
im = axs[0].imshow(np.arange(1., 101.).reshape((10, 10)), norm=PNorm(gamma=2))
fig.colorbar(im, ax=axs[0])
im = axs[1].imshow(np.arange(1., 101.).reshape((10, 10)), norm=PNorm(gamma=-2))
fig.colorbar(im, ax=axs[1])

plt.show()