matplotlib · alvarosg · Dec 16, 2016 · Dec 17, 2016 · Jan 15, 2017 · Jan 15, 2017
diff --git a/examples/color/colormap_normalizations_funcnorm.py b/examples/color/colormap_normalizations_funcnorm.py
@@ -0,0 +1,85 @@
+"""
+=====================================================================
+Examples of normalization using  :class:`~matplotlib.colors.FuncNorm`
+=====================================================================
+
+This is an example on how to perform a normalization using an arbitrary
+function with :class:`~matplotlib.colors.FuncNorm`. A logarithm normalization
+and a square root normalization will be use as examples.
+
+"""
+
+import matplotlib.cm as cm
+import matplotlib.colors as colors
+import matplotlib.pyplot as plt
+
+import numpy as np
+
+
+def main():
+    fig, axes = plt.subplots(3, 2, gridspec_kw={
+              'width_ratios': [1, 3.5]}, figsize=plt.figaspect(0.6))
+
+    # Example of logarithm normalization using FuncNorm
+    norm_log = colors.FuncNorm(f='log10', vmin=0.01)
+    # The same can be achieved with
+    # norm_log = colors.FuncNorm(f=np.log10,
+    #                            finv=lambda x: 10.**(x), vmin=0.01)
+
+    # Example of root normalization using FuncNorm
+    norm_sqrt = colors.FuncNorm(f='sqrt', vmin=0.0)
+    # The same can be achieved with
+    # norm_sqrt = colors.FuncNorm(f='root{2}', vmin=0.)
+    # or with
+    # norm_sqrt = colors.FuncNorm(f=lambda x: x**0.5,
+    #                             finv=lambda x: x**2, vmin=0.0)
+
+    normalizations = [(None, 'Regular linear scale'),
+                      (norm_log, 'Log normalization'),
+                      (norm_sqrt, 'Root normalization')]
+
+    for i, (norm, title) in enumerate(normalizations):
+        X, Y, data = get_data()
+
+        # Showing the normalization effect on an image
+        ax2 = axes[i][1]
+        cax = ax2.imshow(data, cmap=cm.afmhot, norm=norm)
+        ticks = cax.norm.ticks(5) if norm else np.linspace(0, 1, 6)
+        fig.colorbar(cax, format='%.3g', ticks=ticks, ax=ax2)
+        ax2.set_title(title)
+        ax2.axes.get_xaxis().set_ticks([])
+        ax2.axes.get_yaxis().set_ticks([])
+
+        # Plotting the behaviour of the normalization
+        ax1 = axes[i][0]
+        d_values = np.linspace(cax.norm.vmin, cax.norm.vmax, 100)
+        cm_values = cax.norm(d_values)
+        ax1.plot(d_values, cm_values)
+        ax1.set_xlabel('Data values')
+        ax1.set_ylabel('Colormap values')
+
+    plt.show()
+
+
+def get_data(_cache=[]):
+    if len(_cache) > 0:
+        return _cache[0]
+    x = np.linspace(0, 1, 300)
+    y = np.linspace(-1, 1, 90)
+    X, Y = np.meshgrid(x, y)
+
+    data = np.zeros(X.shape)
+
+    def gauss2d(x, y, a0, x0, y0, wx, wy):
+        return a0 * np.exp(-(x - x0)**2 / wx**2 - (y - y0)**2 / wy**2)
+    N = 15
+    for x in np.linspace(0., 1, N):
+        data += gauss2d(X, Y, x, x, 0, 0.25 / N, 0.25)
+
+    data = data - data.min()
+    data = data / data.max()
+    _cache.append((X, Y, data))
+
+    return _cache[0]
+
+main()
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -960,6 +960,229 @@ def scaled(self):
         return (self.vmin is not None and self.vmax is not None)
 
 
+class FuncNorm(Normalize):
+    """
+    Creates a normalizer using a custom function
+
+    The normalizer will be a function mapping the data values into colormap
+    values in the [0,1] range.
+    """
+
+    def __init__(self, f, finv=None, **normalize_kw):
+        """
+        Specify the function to be used, and its inverse, as well as other
+        parameters to be passed to `Normalize`. The normalization will be
+        calculated as (f(x)-f(vmin))/(f(max)-f(vmin)).
+
+        Parameters
+        ----------
+        f : callable or string
+            Function to be used for the normalization receiving a single
+            parameter, compatible with scalar values and ndarrays.
+            Alternatively a string from the list ['linear', 'quadratic',
+            'cubic', 'sqrt', 'cbrt','log', 'log10', 'power{a}', 'root{a}',
+            'log(x+{a})', 'log10(x+{a})'] can be used, replacing 'a' by a
+            number different than 0 when necessary.
+        finv : callable, optional
+            Inverse function of `f` that satisfies finv(f(x))==x. It is
+            optional in cases where `f` is provided as a string.
+        normalize_kw : dict, optional
+            Dict with keywords (`vmin`,`vmax`,`clip`) passed
+            to `matplotlib.colors.Normalize`.
+
+        Examples
+        --------
+        Creating a logarithmic normalization using the predefined strings:
+
+        >>> import matplotlib.colors as colors
+        >>> norm = colors.FuncNorm(f='log10', vmin=0.01, vmax=2)
+
+        Or doing it manually:
+
+        >>> import matplotlib.colors as colors
+        >>> norm = colors.FuncNorm(f=lambda x: np.log10(x),
+        ...                        finv=lambda x: 10.**(x),
+        ...                        vmin=0.01, vmax=2)
+
+        """
+
+        if isinstance(f, six.string_types):
+            func_parser = cbook._StringFuncParser(f)
+            f = func_parser.function
+            finv = func_parser.inverse
+        if not callable(f):
+            raise ValueError("`f` must be a callable or a string.")
+
+        if finv is None:
+            raise ValueError("Inverse function `finv` not provided.")
+        elif not callable(finv):
+            raise ValueError("`finv` must be a callable.")
+
+        self._f = f
+        self._finv = finv
+
+        super(FuncNorm, self).__init__(**normalize_kw)
+
+    def _update_f(self, vmin, vmax):
+        # This method is to be used by derived classes in cases where
+        # the limits vmin and vmax may require changing/updating the
+        # function depending on vmin/vmax, for example rescaling it
+        # to accomodate to the new interval.
+        return
+
+    def __call__(self, value, clip=None):
+        """
+        Normalizes `value` data in the `[vmin, vmax]` interval into
+        the `[0.0, 1.0]` interval and returns it.
+
+        Parameters
+        ----------
+        value : float or ndarray of floats
+            Data to be normalized.
+        clip : boolean, optional
+            Whether to clip the data outside the `[vmin, vmax]` limits.
+            Default `self.clip` from `Normalize` (which defaults to `False`).
+
+        Returns
+        -------
+        result : masked array of floats
+            Normalized data to the `[0.0, 1.0]` interval. If clip == False,
+            values smaller than `vmin` or greater than `vmax` will be clipped
+            to -0.1 and 1.1 respectively.
+
+        """
+        if clip is None:
+            clip = self.clip
+
+        result, is_scalar = self.process_value(value)
+        self.autoscale_None(result)
+
+        vmin = float(self.vmin)
+        vmax = float(self.vmax)
+
+        self._update_f(vmin, vmax)
+
+        if clip:
+            result = np.clip(result, vmin, vmax)
+            resultnorm = (self._f(result) - self._f(vmin)) / \
+                         (self._f(vmax) - self._f(vmin))
+        else:
+            resultnorm = result.copy()
+            mask_over = result > vmax
+            mask_under = result < vmin
+            mask = (result >= vmin) * (result <= vmax)
+            # Since the non linear function is arbitrary and may not be
+            # defined outside the boundaries, we just set obvious under
+            # and over values
+            resultnorm[mask_over] = 1.1
+            resultnorm[mask_under] = -0.1
+            resultnorm[mask] = (self._f(result[mask]) - self._f(vmin)) / \
+                               (self._f(vmax) - self._f(vmin))
+
+        return np.ma.array(resultnorm)
+
+    def inverse(self, value):
+        """
+        Performs the inverse normalization from the `[0.0, 1.0]` into the
+        `[vmin, vmax]` interval and returns it.
+
+        Parameters
+        ----------
+        value : float or ndarray of floats
+            Data in the `[0.0, 1.0]` interval.
+
+        Returns
+        -------
+        result : float or ndarray of floats
+            Data before normalization.
+
+        """
+        vmin = self.vmin
+        vmax = self.vmax
+        self._update_f(vmin, vmax)
+        value = self._finv(
+            value * (self._f(vmax) - self._f(vmin)) + self._f(vmin))
+        return value
+
+    @staticmethod
+    def _fun_normalizer(fun):
+        if fun(0.) == 0. and fun(1.) == 1.:
+            return fun
+        elif fun(0.) == 0.:
+            return (lambda x: fun(x) / fun(1.))
+        else:
+            return (lambda x: (fun(x) - fun(0.)) / (fun(1.) - fun(0.)))
+
+    def autoscale(self, A):
+        """
+        Autoscales the normalization based on the maximum and minimum values
+        of `A`.
+
+        Parameters
+        ----------
+        A : ndarray or maskedarray
+            Array used to calculate the maximum and minimum values.
+
+        """
+        self.vmin = float(np.ma.min(A))
+        self.vmax = float(np.ma.max(A))
+
+    def autoscale_None(self, A):
+        """
+        Autoscales the normalization based on the maximum and minimum values
+        of `A`, only if the limits were not already set.
+
+        Parameters
+        ----------
+        A : ndarray or maskedarray
+            Array used to calculate the maximum and minimum values.
+
+        """
+        if self.vmin is None:
+            self.vmin = float(np.ma.min(A))
+        if self.vmax is None:
+            self.vmax = float(np.ma.max(A))
+        self.vmin = float(self.vmin)
+        self.vmax = float(self.vmax)
+        if self.vmin > self.vmax:
+            raise ValueError("vmin must be smaller than vmax")
+
+    def ticks(self, nticks=13):
+        """
+        Returns an automatic list of `nticks` points in the data space
+        to be used as ticks in the colorbar.
+
+        Parameters
+        ----------
+        nticks : integer, optional
+            Number of ticks to be returned. Default 13.
+
+        Returns
+        -------
+        ticks : ndarray
+            1d array of length `nticks` with the proposed tick locations.
+
+        """
+        ticks = self.inverse(np.linspace(0, 1, nticks))
+        finalticks = np.zeros(ticks.shape, dtype=np.bool)
+        finalticks[0] = True
+        ticks = FuncNorm._round_ticks(ticks, finalticks)
+        return ticks
+
+    @staticmethod
+    def _round_ticks(ticks, permanenttick):
+        ticks = ticks.copy()
+        for i in range(len(ticks)):
+            if i == 0 or i == len(ticks) - 1 or permanenttick[i]:
+                continue
+            d1 = ticks[i] - ticks[i - 1]
+            d2 = ticks[i + 1] - ticks[i]
+            d = min([d1, d2])
+            order = -np.floor(np.log10(d))
+            ticks[i] = float(np.round(ticks[i] * 10**order)) / 10**order
+        return ticks
+
+
 class LogNorm(Normalize):
     """
     Normalize a given value to the 0-1 range on a log scale

diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -146,6 +146,47 @@ def test_BoundaryNorm():
     assert_true(np.all(bn(vals).mask))
 
 
+class TestFuncNorm(object):
+    def test_limits_with_string(self):
+        norm = mcolors.FuncNorm(f='log10', vmin=0.01, vmax=2.)
+        assert_array_equal(norm([0.01, 2]), [0, 1.0])
+
+    def test_limits_with_lambda(self):
+        norm = mcolors.FuncNorm(f=lambda x: np.log10(x),
+                                finv=lambda x: 10.**(x),
+                                vmin=0.01, vmax=2.)
+        assert_array_equal(norm([0.01, 2]), [0, 1.0])
+
+    def test_limits_without_vmin_vmax(self):
+        norm = mcolors.FuncNorm(f='log10')
+        assert_array_equal(norm([0.01, 2]), [0, 1.0])
+
+    def test_limits_without_vmin(self):
+        norm = mcolors.FuncNorm(f='log10', vmax=2.)
+        assert_array_equal(norm([0.01, 2]), [0, 1.0])
+
+    def test_limits_without_vmax(self):
+        norm = mcolors.FuncNorm(f='log10', vmin=0.01)
+        assert_array_equal(norm([0.01, 2]), [0, 1.0])
+
+    def test_intermediate_values(self):
+        norm = mcolors.FuncNorm(f='log10')
+        assert_array_almost_equal(norm([0.01, 0.5, 2]),
+                                  [0, 0.73835195870437, 1.0])
+
+    def test_inverse(self):
+        norm = mcolors.FuncNorm(f='log10', vmin=0.01, vmax=2.)
+        x = np.linspace(0.01, 2, 10)
+        assert_array_almost_equal(x, norm.inverse(norm(x)))
+
+    def test_ticks(self):
+        norm = mcolors.FuncNorm(f='log10', vmin=0.01, vmax=2.)
+        expected = [0.01, 0.016, 0.024, 0.04, 0.06,
+                    0.09, 0.14, 0.22, 0.3, 0.5,
+                    0.8, 1.3, 2.]
+        assert_array_almost_equal(norm.ticks(), expected)
+
+
 def test_LogNorm():
     """
     LogNorm ignored clip, now it has the same