Implemented feedback from @QuLogic and @story645

matplotlib · alvarosg · Dec 16, 2016 · Dec 17, 2016 · Jan 15, 2017 · Jan 15, 2017
commit 44658d4ce621291c74226dd7e082687dc5c4a634
diff --git a/examples/color/colormap_normalizations_funcnorm.py b/examples/color/colormap_normalizations_funcnorm.py
@@ -15,71 +15,60 @@
 
 import numpy as np
 
+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)
 
-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()
+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')]
+
+# Fabricating some data
+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)
+
+for x in np.linspace(0., 1, 15):
+    data += gauss2d(X, Y, x, x, 0, 0.25 / 15, 0.25)
+
+data -= data.min()
+data /= data.max()
+
+# Using the custom normalizations to plot the data
+fig, axes = plt.subplots(3, 2, sharex='col',
+                         gridspec_kw={'width_ratios': [1, 3.5]},
+                         figsize=plt.figaspect(0.6))
+
+for (ax_left, ax_right), (norm, title) in zip(axes, normalizations):
+
+    # Showing the normalization effect on an image
+    cax = ax_right.imshow(data, cmap=cm.afmhot, norm=norm, aspect='auto')
+    ticks = cax.norm.ticks(5) if norm else np.linspace(0, 1, 6)
+    fig.colorbar(cax, format='%.3g', ticks=ticks, ax=ax_right)
+    ax_right.set_title(title)
+    ax_right.xaxis.set_ticks([])
+    ax_right.yaxis.set_ticks([])
+
+    # Plotting the behaviour of the normalization
+    d_values = np.linspace(cax.norm.vmin, cax.norm.vmax, 100)
+    cm_values = cax.norm(d_values)
+    ax_left.plot(d_values, cm_values)
+    ax_left.set_ylabel('Colormap values')
+
+ax_left.set_xlabel('Data values')
+
+plt.show()
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -965,30 +965,43 @@ 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.
+    values in the ``[0,1]`` range.
     """
 
-    def __init__(self, f, finv=None, **normalize_kw):
+    def __init__(self, f, finv=None, vmin=None, vmax=None, clip=False):
         """
         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)).
+        calculated as (f(x)-f(vmin))/(f(vmax)-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.
+            'cubic', 'x**{p}', 'sqrt', 'cbrt', 'root{p}(x)', 'log', 'log10',
+            'log2', 'log{p}(x)', 'log(x+{p}) 'log10(x+{p})', 'log{p}(x+{p})]
+            can be used, replacing 'p' by the corresponding value of the
+            parameter, when present.
         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`.
+            Inverse function of `f` that satisfies finv(f(x))==x.
+            Optional/ignored when `f` is a string.
+        vmin : float or None, optional
+            Value assigned to the lower limit of the colormap. If None, it
+            will be assigned to the minimum value of the data provided.
+            Default None.
+        vmax : float or None, optional
+            Value assigned to the upper limit of the colormap. If None, it
+            will be assigned to the maximum value of the data provided.
+            Default None.
+        clip : bool, optional
+            If True, any value below `vmin` will be clipped to `vmin`, and
+            any value above `vmax` will be clip to `vmin`. This effectively
+            defeats the purpose of setting the over and under values of the
+            color map. If False, values below `vmin` and above `vmax` will
+            be set to -0.1 and 1.1 respectively, after the normalization.
+            Default False.
 
         Examples
         --------
@@ -997,14 +1010,15 @@ def __init__(self, f, finv=None, **normalize_kw):
         >>> import matplotlib.colors as colors
         >>> norm = colors.FuncNorm(f='log10', vmin=0.01, vmax=2)
 
-        Or doing it manually:
+        Or 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)
 
         """
+        super(FuncNorm, self).__init__(vmin=vmin, vmax=vmax, clip=clip)
 
         if isinstance(f, six.string_types):
             func_parser = cbook._StringFuncParser(f)
@@ -1021,32 +1035,30 @@ def __init__(self, f, finv=None, **normalize_kw):
         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.
+        # to accommodate 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.
+        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.
+            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,
+            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.
 
@@ -1057,6 +1069,7 @@ def __call__(self, value, clip=None):
         result, is_scalar = self.process_value(value)
         self.autoscale_None(result)
 
+        self._check_vmin_vmax()
         vmin = float(self.vmin)
         vmax = float(self.vmax)
 
@@ -1070,7 +1083,7 @@ def __call__(self, value, clip=None):
             resultnorm = result.copy()
             mask_over = result > vmax
             mask_under = result < vmin
-            mask = (result >= vmin) * (result <= vmax)
+            mask = ~(mask_over | mask_under)
             # Since the non linear function is arbitrary and may not be
             # defined outside the boundaries, we just set obvious under
             # and over values
@@ -1079,72 +1092,38 @@ def __call__(self, value, clip=None):
             resultnorm[mask] = (self._f(result[mask]) - self._f(vmin)) / \
                                (self._f(vmax) - self._f(vmin))
 
-        return np.ma.array(resultnorm)
+        resultnorm = np.ma.array(resultnorm)
+        if is_scalar:
+            return resultnorm[0]
+        else:
+            return resultnorm
 
     def inverse(self, value):
         """
-        Performs the inverse normalization from the `[0.0, 1.0]` into the
-        `[vmin, vmax]` interval and returns it.
+        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.
+            Data in the ``[0.0, 1.0]`` interval.
 
         Returns
         -------
         result : float or ndarray of floats
             Data before normalization.
 
         """
+        self._check_vmin_vmax()
         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:
+    def _check_vmin_vmax(self):
+        if self.vmin >= self.vmax:
             raise ValueError("vmin must be smaller than vmax")
 
     def ticks(self, nticks=13):

diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -169,6 +169,20 @@ 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_clip_true(self):
+        norm = mcolors.FuncNorm(f='log10', vmin=0.01, vmax=2.,
+                                clip=True)
+        assert_array_equal(norm([0.0, 2.5]), [0.0, 1.0])
+
+    def test_clip_false(self):
+        norm = mcolors.FuncNorm(f='log10', vmin=0.01, vmax=2.,
+                                clip=False)
+        assert_array_equal(norm([0.0, 2.5]), [-0.1, 1.1])
+
+    def test_clip_default_false(self):
+        norm = mcolors.FuncNorm(f='log10', vmin=0.01, vmax=2.)
+        assert_array_equal(norm([0.0, 2.5]), [-0.1, 1.1])
+
     def test_intermediate_values(self):
         norm = mcolors.FuncNorm(f='log10')
         assert_array_almost_equal(norm([0.01, 0.5, 2]),