Set norms using scale names.

anntzer · anntzer · commit d6f86c1df699 · 2022-07-03T19:58:13.000+02:00
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -16,12 +16,13 @@
 """
 
 from collections.abc import Mapping
+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
 
@@ -308,6 +309,34 @@ def unregister_cmap(name):
     return cmap
 
 
+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").
+
+    Note that ``make_norm_from_scale`` caches the generated norm classes
+    (not the instances) and reuses them for later calls.  For example,
+    ``type(_auto_norm_from_scale("log")) == LogNorm``.
+    """
+    # 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.
@@ -318,12 +347,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`
@@ -353,11 +383,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.
@@ -531,9 +561,18 @@ 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):
+            try:
+                scale_cls = scale._scale_mapping[norm]
+            except KeyError:
+                raise ValueError(
+                    "Invalid norm str name; the following values are "
+                    "supported: {}".format(", ".join(scale._scale_mapping))
+                ) from None
+            norm = _auto_norm_from_scale(scale_cls)()
 
         if norm is self.norm:
             # We aren't updating anything
@@ -555,7 +594,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
@@ -1598,21 +1598,37 @@ class norm_cls(Normalize):
     if base_norm_cls is None:
         return functools.partial(make_norm_from_scale, scale_cls, init=init)
 
+    if isinstance(scale_cls, functools.partial):
+        scale_args = scale_cls.args
+        scale_kwargs_items = tuple(scale_cls.keywords.items())
+        scale_cls = scale_cls.func
+    else:
+        scale_args = scale_kwargs_items = ()
+
     if init is None:
         def init(vmin=None, vmax=None, clip=False): pass
 
     return _make_norm_from_scale(
-        scale_cls, base_norm_cls, inspect.signature(init))
+        scale_cls, scale_args, scale_kwargs_items,
+        base_norm_cls, inspect.signature(init))
 
 
 @functools.lru_cache(None)
-def _make_norm_from_scale(scale_cls, base_norm_cls, bound_init_signature):
+def _make_norm_from_scale(
+    scale_cls, scale_args, scale_kwargs_items,
+    base_norm_cls, bound_init_signature,
+):
     """
     Helper for `make_norm_from_scale`.
 
-    This function is split out so that it takes a signature object as third
-    argument (as signatures are picklable, contrary to arbitrary lambdas);
-    caching is also used so that different unpickles reuse the same class.
+    This function is split out to enable caching (in particular so that
+    different unpickles reuse the same class).  In order to do so,
+
+    - ``functools.partial`` *scale_cls* is expanded into ``func, args, kwargs``
+      to allow memoizing returned norms (partial instances always compare
+      unequal, but we can check identity based on ``func, args, kwargs``;
+    - *init* is replaced by *init_signature*, as signatures are picklable,
+      unlike to arbitrary lambdas.
     """
 
     class Norm(base_norm_cls):
@@ -1631,15 +1647,18 @@ def __reduce__(self):
             except (ImportError, AttributeError):
                 pass
             return (_picklable_norm_constructor,
-                    (scale_cls, base_norm_cls, bound_init_signature),
+                    (scale_cls, scale_args, scale_kwargs_items,
+                     base_norm_cls, bound_init_signature),
                     vars(self))
 
         def __init__(self, *args, **kwargs):
             ba = bound_init_signature.bind(*args, **kwargs)
             ba.apply_defaults()
             super().__init__(
                 **{k: ba.arguments.pop(k) for k in ["vmin", "vmax", "clip"]})
-            self._scale = scale_cls(axis=None, **ba.arguments)
+            self._scale = functools.partial(
+                scale_cls, *scale_args, **dict(scale_kwargs_items))(
+                    axis=None, **ba.arguments)
             self._trf = self._scale.get_transform()
 
         __init__.__signature__ = bound_init_signature.replace(parameters=[
@@ -1693,12 +1712,12 @@ 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:
+        Norm.__name__ = f"{scale_cls.__name__}Norm"
+        Norm.__qualname__ = f"{scale_cls.__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__
 
@@ -1746,9 +1765,10 @@ def forward(values: array-like) -> array-like
     """
 
 
-@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."""
+LogNorm = make_norm_from_scale(
+    functools.partial(scale.LogScale, nonpositive="mask"))(Normalize)
+LogNorm.__name__ = LogNorm.__qualname__ = "LogNorm"
+LogNorm.__doc__ = "Normalize a given value to the 0-1 range on a log scale."
 
 
 @make_norm_from_scale(
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -982,8 +982,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)
 
 
@@ -2431,8 +2431,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
@@ -1410,3 +1410,25 @@ def test_large_image(fig_test, fig_ref, dim, size, msg, origin):
                        extent=(0, 1, 0, 1),
                        interpolation='none',
                        origin=origin)
+
+
+@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
+    axts = fig_test.subplots(1, 5)
+    axts[0].imshow(t, norm="log")
+    axts[1].imshow(t, norm="log", vmin=.2)
+    axts[2].imshow(t, norm="symlog")
+    axts[3].imshow(t, norm="symlog", vmin=.3, vmax=.7)
+    axts[4].imshow(t, norm="logit", vmin=.3, vmax=.7)
+    axrs = fig_ref.subplots(1, 5)
+    axrs[0].imshow(t, norm=colors.LogNorm())
+    axrs[1].imshow(t, norm=colors.LogNorm(vmin=.2))
+    # same linthresh as SymmetricalLogScale's default.
+    axrs[2].imshow(t, norm=colors.SymLogNorm(linthresh=2))
+    axrs[3].imshow(t, norm=colors.SymLogNorm(linthresh=2, vmin=.3, vmax=.7))
+    axrs[4].imshow(t, norm="logit", clim=(.3, .7))
+
+    assert type(axts[0].images[0].norm) == colors.LogNorm  # Exactly that class
+    with pytest.raises(ValueError):
+        axts[0].imshow(t, norm="foobar")
