Set norms using scale names.

anntzer · anntzer · commit a44db64bc24d · 2021-12-21T21:16:00.000+01:00
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -16,12 +16,13 @@
 """
 
 from collections.abc import Mapping, 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
 
@@ -331,6 +332,35 @@ 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.
+      (`.colors.make_norm_from_scale` also memoizes, but we call it with a
+      `functools.partial` instances which always compare unequal, so the cache
+      doesn't get hit.)
+    """
+    # Actually try to construct an instance, to verify whether
+    # ``nonpositive="mask"`` is supported.
+    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)()
+    return type(norm)
+
+
 class ScalarMappable:
     """
     A mixin class to map scalar data to RGBA.
@@ -341,12 +371,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`
@@ -376,11 +407,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.
@@ -554,9 +585,13 @@ def norm(self):
 
     @norm.setter
     def norm(self, norm):
-        _api.check_isinstance((colors.Normalize, None), norm=norm)
+        _api.check_isinstance((colors.Normalize, str, None), norm=norm)
         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)()
 
         if norm is self.norm:
             # We aren't updating anything
@@ -578,7 +613,7 @@ def set_norm(self, norm):
 
         Parameters
         ----------
-        norm : `.Normalize` or None
+        norm : `.Normalize` or str or None
 
         Notes
         -----
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1587,12 +1587,15 @@ def autoscale_None(self, A):
             in_trf_domain = np.extract(np.isfinite(self._trf.transform(A)), A)
             return super().autoscale_None(in_trf_domain)
 
-    Norm.__name__ = (
-            f"{scale_cls.__name__}Norm" if base_norm_cls is Normalize
-            else base_norm_cls.__name__)
-    Norm.__qualname__ = (
-            f"{scale_cls.__qualname__}Norm" if base_norm_cls is Normalize
-            else base_norm_cls.__qualname__)
+    if base_norm_cls is Normalize:
+        name_source = (scale_cls.func
+                       if isinstance(scale_cls, functools.partial)
+                       else scale_cls)
+        Norm.__name__ = f"{name_source.__name__}Norm"
+        Norm.__qualname__ = f"{name_source.__qualname__}Norm"
+    else:
+        Norm.__name__ = base_norm_cls.__name__
+        Norm.__qualname__ = base_norm_cls.__qualname__
     Norm.__module__ = base_norm_cls.__module__
     Norm.__doc__ = base_norm_cls.__doc__
 
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -923,8 +923,8 @@ def test_imshow_norm_vminvmax():
     a = [[1, 2], [3, 4]]
     ax = plt.axes()
     with pytest.raises(ValueError,
-                       match="Passing parameters norm and vmin/vmax "
-                             "simultaneously is not supported."):
+                       match="Passing a Normalize instance simultaneously "
+                             "with vmin/vmax is not supported."):
         ax.imshow(a, norm=mcolors.Normalize(-10, 10), vmin=0, vmax=5)
 
 
@@ -2285,8 +2285,8 @@ def test_scatter_norm_vminvmax(self):
         x = [1, 2, 3]
         ax = plt.axes()
         with pytest.raises(ValueError,
-                           match="Passing parameters norm and vmin/vmax "
-                                 "simultaneously is not supported."):
+                           match="Passing a Normalize instance simultaneously "
+                                 "with vmin/vmax is not supported."):
             ax.scatter(x, x, c=x, norm=mcolors.Normalize(-10, 10),
                        vmin=0, vmax=5)
 
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -1376,3 +1376,21 @@ def test_rgba_antialias():
     # alternating red and blue stripes become purple
     axs[3].imshow(aa, interpolation='antialiased', interpolation_stage='rgba',
                   cmap=cmap, vmin=-1.2, vmax=1.2)
+
+
+@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, 5)
+    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[4].imshow(t, norm="logit", vmin=.3, vmax=.7)
+    axs = fig_ref.subplots(1, 5)
+    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))
+    axs[4].imshow(t, norm="logit", clim=(.3, .7))
