ENH: Add func norm

Author: jklymak

doc/api/colors_api.rst

@@ -30,6 +30,7 @@ Classes
    PowerNorm
    SymLogNorm
    TwoSlopeNorm
+   FuncNorm
 
 Functions
 ---------

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

lib/matplotlib/colors.py

@@ -1467,6 +1467,39 @@ def inverse(self, value):
     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):
     """Normalize a given value to the 0-1 range on a log scale."""

lib/matplotlib/tests/test_colors.py

@@ -537,6 +537,29 @@ 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])
+    input = np.array([0, 5, 10])
+    assert_array_almost_equal(norm(input), expected)
+    assert_array_almost_equal(norm.inverse(expected), input)
+
+    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]))
+    assert_array_almost_equal(norm.inverse([0.2, 5, 10]),
+                              lognorm.inverse([0.2, 5, 10]))
+
+
 def test_TwoSlopeNorm_autoscale():
     norm = mcolors.TwoSlopeNorm(vcenter=20)
     norm.autoscale([10, 20, 30, 40])

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,45 @@
 # 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, you can use
+# `~.colors.FuncNorm` to define your own.  Note that this example is the same
+# as `~.colors.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 +333,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()