ENH: Add func norm

jklymak · jklymak · commit 989c966cddbd · 2020-12-19T19:36:20.000-08:00
diff --git a/doc/api/colors_api.rst b/doc/api/colors_api.rst
@@ -30,6 +30,7 @@ Classes
    PowerNorm
    SymLogNorm
    TwoSlopeNorm
+   FuncNorm
 
 Functions
 ---------
diff --git a/examples/userdemo/colormap_normalizations.py b/examples/userdemo/colormap_normalizations.py
@@ -77,7 +77,6 @@
                        shading='nearest')
 fig.colorbar(pcm, ax=ax[1], extend='both')
 
-
 ###############################################################################
 # Custom Norm: An example with a customized normalization.  This one
 # uses the example above, and normalizes the negative data differently
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1464,6 +1464,38 @@ def inverse(self, value):
         *bound_init_signature.parameters.values()])
     return Norm
 
+@_make_norm_from_scale(
+    scale.FuncScale,
+    init=lambda functions, vmin=None, vmax=None, clip=False: None)
+class FuncNorm(Normalize):
+    """
+    Arbitrary normalization using functions for the forward and inverse.
+
+    Parameters
+    ----------
+    functions : (callable, callable)
+        two-tuple of the forward and inverse functions for the normalization.
+        The forward function must be monotonic.
+
+        Both functions must have the signature ::
+
+           def forward(values: array-like) -> array-like
+
+    vmin, vmax : float or None
+        If *vmin* and/or *vmax* is not given, they are initialized from the
+        minimum and maximum value, respectively, of the first input
+        processed; i.e., ``__call__(A)`` calls ``autoscale_None(A)``.
+
+    clip : bool, default: False
+        If ``True`` values falling outside the range ``[vmin, vmax]``,
+        are mapped to 0 or 1, whichever is closer, and masked values are
+        set to 1.  If ``False`` masked values remain masked.
+
+        Clipping silently defeats the purpose of setting the over, under,
+        and masked colors in a colormap, so it is likely to lead to
+        surprises; therefore the default is ``clip=False``.
+    """
+
 
 @_make_norm_from_scale(functools.partial(scale.LogScale, nonpositive="mask"))
 class LogNorm(Normalize):
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -526,6 +526,25 @@ def test_Normalize():
     assert 0 < norm(1 + 50 * eps) < 1
 
 
+def test_FuncNorm():
+    def _forward(x):
+        return (x**2)
+    def _inverse(x):
+        return np.sqrt(x)
+
+    norm = mcolors.FuncNorm((_forward, _inverse), vmin=0, vmax=10)
+    expected = np.array([0, 0.25, 1])
+    assert_array_almost_equal(norm([0, 5, 10]), expected)
+
+    def _forward(x):
+        return np.log10(x)
+    def _inverse(x):
+        return 10**x
+    norm = mcolors.FuncNorm((_forward, _inverse), vmin=0.1, vmax=10)
+    lognorm = mcolors.LogNorm(vmin=0.1, vmax=10)
+    assert_array_almost_equal(norm([0.2, 5, 10]), lognorm([0.2, 5, 10]))
+
+
 def test_TwoSlopeNorm_autoscale():
     norm = mcolors.TwoSlopeNorm(vcenter=20)
     norm.autoscale([10, 20, 30, 40])
diff --git a/tutorials/colors/colormapnorms.py b/tutorials/colors/colormapnorms.py
@@ -169,14 +169,16 @@
 X, Y = np.mgrid[0:3:complex(0, N), 0:2:complex(0, N)]
 Z1 = (1 + np.sin(Y * 10.)) * X**2
 
-fig, ax = plt.subplots(2, 1)
+fig, ax = plt.subplots(2, 1, constrained_layout=True)
 
 pcm = ax[0].pcolormesh(X, Y, Z1, norm=colors.PowerNorm(gamma=0.5),
                        cmap='PuBu_r', shading='auto')
 fig.colorbar(pcm, ax=ax[0], extend='max')
+ax[0].set_title('PowerNorm()')
 
 pcm = ax[1].pcolormesh(X, Y, Z1, cmap='PuBu_r', shading='auto')
 fig.colorbar(pcm, ax=ax[1], extend='max')
+ax[1].set_title('Normalize()')
 plt.show()
 
 ###############################################################################
@@ -274,17 +276,46 @@
 # Simple geographic plot, set aspect ratio beecause distance between lines of
 # longitude depends on latitude.
 ax.set_aspect(1 / np.cos(np.deg2rad(49)))
+ax.set_title('TwoSlopeNorm(x)')
 fig.colorbar(pcm, shrink=0.6)
 plt.show()
 
 
+###############################################################################
+# FuncNorm: Arbitrary function normalization
+# ------------------------------------------
+#
+# If the above norms do not provide the normalization you want, and you
+# can express your normalization in terms of function, you can use
+# `~.colors.FuncNorm`.  Note that this example is the same as
+# `PowerNorm` with a power of 0.5:
+
+def _forward(x):
+    return np.sqrt(x)
+
+
+def _inverse(x):
+    return x**2
+
+N = 100
+X, Y = np.mgrid[0:3:complex(0, N), 0:2:complex(0, N)]
+Z1 = (1 + np.sin(Y * 10.)) * X**2
+fig, ax = plt.subplots()
+
+norm = colors.FuncNorm((_forward, _inverse), vmin=0, vmax=20)
+pcm = ax.pcolormesh(X, Y, Z1, norm=norm, cmap='PuBu_r', shading='auto')
+ax.set_title('FuncNorm(x)')
+fig.colorbar(pcm, shrink=0.6)
+plt.show()
+
 ###############################################################################
 # Custom normalization: Manually implement two linear ranges
 # ----------------------------------------------------------
 #
 # The `.TwoSlopeNorm` described above makes a useful example for
 # defining your own norm.
 
+
 class MidpointNormalize(colors.Normalize):
     def __init__(self, vmin=None, vmax=None, vcenter=None, clip=False):
         self.vcenter = vcenter
@@ -303,5 +334,6 @@ def __call__(self, value, clip=None):
 pcm = ax.pcolormesh(longitude, latitude, topo, rasterized=True, norm=midnorm,
                     cmap=terrain_map, shading='auto')
 ax.set_aspect(1 / np.cos(np.deg2rad(49)))
+ax.set_title('Custom norm')
 fig.colorbar(pcm, shrink=0.6, extend='both')
 plt.show()
