Set norms using scale names.

anntzer · anntzer · commit 17aeeea64cae · 2021-07-28T01:25:23.000+02:00
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -16,12 +16,13 @@
 """
 
 from collections.abc import MutableMapping
+import functools
 
 import numpy as np
 from numpy import ma
 
 import matplotlib as mpl
-from matplotlib import _api, colors, cbook
+from matplotlib import _api, colors, cbook, scale
 from matplotlib._cm import datad
 from matplotlib._cm_listed import cmaps as cmaps_listed
 
@@ -235,6 +236,36 @@ def unregister_cmap(name):
     return _cmap_registry.pop(name)
 
 
+@functools.lru_cache(None)
+def _auto_norm_from_scale(scale_cls):
+    """
+    Automatically generate a norm class from *scale_cls*.
+
+    This differs from `.colors.make_norm_from_scale` in the following points:
+
+    - This function is not a class decorator, but directly returns a norm class
+      (as if decorating `.Normalize`).
+    - The scale is automatically constructed with ``nonpositive="mask"``, if it
+      supports such a parameter, to work around the difference in defaults
+      between standard scales (which use "clip") and norms (which use "mask").
+    - The returned norm class is memoized and reused for later calls.  (This
+      would be desirable in `.colors.make_norm_from_scale` as well, but
+      because that function sometimes gets called with `functools.partial`
+      instances which always compare unequal, a simple ``@lru_cache`` would be
+      insufficient.)
+    """
+    try:
+        norm = colors.make_norm_from_scale(
+            functools.partial(scale_cls, nonpositive="mask"))(
+            colors.Normalize)()
+    except TypeError:
+        norm = colors.make_norm_from_scale(scale_cls)(
+            colors.Normalize)()
+    norm_cls = type(norm)
+    norm_cls.__name__ = f"{scale_cls.__name__}Norm"
+    return norm_cls
+
+
 class ScalarMappable:
     """
     A mixin class to map scalar data to RGBA.
@@ -245,12 +276,13 @@ class ScalarMappable:
 
     def __init__(self, norm=None, cmap=None):
         """
-
         Parameters
         ----------
-        norm : `matplotlib.colors.Normalize` (or subclass thereof)
+        norm : `.Normalize` (or subclass thereof) or str or None
             The normalizing object which scales data, typically into the
             interval ``[0, 1]``.
+            If a `str`, a `.Normalize` subclass is dynamically generated based
+            on the scale with the corresponding name.
             If *None*, *norm* defaults to a *colors.Normalize* object which
             initializes its scaling based on the first data processed.
         cmap : str or `~matplotlib.colors.Colormap`
@@ -277,11 +309,11 @@ def _scale_norm(self, norm, vmin, vmax):
         """
         if vmin is not None or vmax is not None:
             self.set_clim(vmin, vmax)
-            if norm is not None:
+            if isinstance(norm, colors.Normalize):
                 raise ValueError(
-                    "Passing parameters norm and vmin/vmax simultaneously is "
-                    "not supported. Please pass vmin/vmax directly to the "
-                    "norm when creating it.")
+                    "Passing a Normalize instance simultaneously with "
+                    "vmin/vmax is not supported. Please pass vmin/vmax "
+                    "directly to the norm when creating it.")
 
         # always resolve the autoscaling so we have concrete limits
         # rather than deferring to draw time.
@@ -454,18 +486,22 @@ def set_norm(self, norm):
 
         Parameters
         ----------
-        norm : `.Normalize` or None
+        norm : `.Normalize` or str or None
 
         Notes
         -----
         If there are any colorbars using the mappable for this norm, setting
         the norm of the mappable will reset the norm, locator, and formatters
         on the colorbar to default.
         """
-        _api.check_isinstance((colors.Normalize, None), norm=norm)
+        _api.check_isinstance((colors.Normalize, str, None), norm=norm)
         in_init = self.norm is None
         if norm is None:
             norm = colors.Normalize()
+        elif isinstance(norm, str):
+            # case-insensitive, consistently with scale_factory.
+            scale_cls = scale._scale_mapping[norm.lower()]
+            norm = _auto_norm_from_scale(scale_cls)()
         self.norm = norm
         if not in_init:
             self.changed()  # Things are not set up properly yet.
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -1328,3 +1328,19 @@ def test_nonuniform_and_pcolor():
         ax.set_axis_off()
         # NonUniformImage "leaks" out of extents, not PColorImage.
         ax.set(xlim=(0, 10))
+
+
+@check_figures_equal(extensions=["png"])
+def test_str_norms(fig_test, fig_ref):
+    t = np.random.rand(10, 10) * .8 + .1  # between 0 and 1
+    axs = fig_test.subplots(1, 4)
+    axs[0].imshow(t, norm="log")
+    axs[1].imshow(t, norm="log", vmin=.2)
+    axs[2].imshow(t, norm="symlog")
+    axs[3].imshow(t, norm="symlog", vmin=.3, vmax=.7)
+    axs = fig_ref.subplots(1, 4)
+    axs[0].imshow(t, norm=colors.LogNorm())
+    axs[1].imshow(t, norm=colors.LogNorm(vmin=.2))
+    # same linthresh as SymmetricalLogScale's default.
+    axs[2].imshow(t, norm=colors.SymLogNorm(linthresh=2))
+    axs[3].imshow(t, norm=colors.SymLogNorm(linthresh=2, vmin=.3, vmax=.7))
