Merge pull request #17709 from jpmattern/symnorm

jklymak · web-flow · commit 03c293b056d7 · 2020-10-22T07:01:04.000-07:00
Enh: SymNorm for normalizing symmetrical data around a center
diff --git a/doc/api/colors_api.rst b/doc/api/colors_api.rst
@@ -20,6 +20,7 @@ Classes
 
    BoundaryNorm
    Colormap
+   CenteredNorm
    LightSource
    LinearSegmentedColormap
    ListedColormap
diff --git a/doc/users/next_whats_new/centerednorm.rst b/doc/users/next_whats_new/centerednorm.rst
@@ -0,0 +1,35 @@
+New CenteredNorm for symmetrical data around a center
+-----------------------------------------------------
+In cases where data is symmetrical around a center, for example, positive and
+negative anomalies around a center zero, `~.matplotlib.colors.CenteredNorm`
+is a new norm that automatically creates a symmetrical mapping around the
+center. This norm is well suited to be combined with a divergent colormap which
+uses an unsaturated color in its center.
+
+  .. plot::
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from matplotlib.colors import CenteredNorm
+
+    np.random.seed(20201004)
+    data = np.random.normal(size=(3, 4), loc=1)
+
+    fig, ax = plt.subplots()
+    pc = ax.pcolormesh(data, cmap=plt.get_cmap('RdGy'), norm=CenteredNorm())
+    fig.colorbar(pc)
+    ax.set_title('data centered around zero')
+
+    # add text annotation
+    for irow, data_row in enumerate(data):
+        for icol, val in enumerate(data_row):
+            ax.text(icol + 0.5, irow + 0.5, f'{val:.2f}', color='C0',
+                    size=16, va='center', ha='center')
+    plt.show()
+
+If the center of symmetry is different from 0, it can be set with the *vcenter*
+argument. To manually set the range of `~.matplotlib.colors.CenteredNorm`, use
+the *halfrange* argument.
+
+See :doc:`/tutorials/colors/colormapnorms` for an example and more details
+about data normalization.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1255,6 +1255,101 @@ def __call__(self, value, clip=None):
         return result
 
 
+class CenteredNorm(Normalize):
+    def __init__(self, vcenter=0, halfrange=None, clip=False):
+        """
+        Normalize symmetrical data around a center (0 by default).
+
+        Unlike `TwoSlopeNorm`, `CenteredNorm` applies an equal rate of change
+        around the center.
+
+        Useful when mapping symmetrical data around a conceptual center
+        e.g., data that range from -2 to 4, with 0 as the midpoint, and
+        with equal rates of change around that midpoint.
+
+        Parameters
+        ----------
+        vcenter : float, default: 0
+            The data value that defines ``0.5`` in the normalization.
+        halfrange : float, optional
+            The range of data values that defines a range of ``0.5`` in the
+            normalization, so that *vcenter* - *halfrange* is ``0.0`` and
+            *vcenter* + *halfrange* is ``1.0`` in the normalization.
+            Defaults to the largest absolute difference to *vcenter* for
+            the values in the dataset.
+
+        Examples
+        --------
+        This maps data values -2 to 0.25, 0 to 0.5, and 4 to 1.0
+        (assuming equal rates of change above and below 0.0):
+
+            >>> import matplotlib.colors as mcolors
+            >>> norm = mcolors.CenteredNorm(halfrange=4.0)
+            >>> data = [-2., 0., 4.]
+            >>> norm(data)
+            array([0.25, 0.5 , 1.  ])
+        """
+        self._vcenter = vcenter
+        # calling the halfrange setter to set vmin and vmax
+        self.halfrange = halfrange
+        self.clip = clip
+
+    def _set_vmin_vmax(self):
+        """
+        Set *vmin* and *vmax* based on *vcenter* and *halfrange*.
+        """
+        self.vmax = self._vcenter + self._halfrange
+        self.vmin = self._vcenter - self._halfrange
+
+    def autoscale(self, A):
+        """
+        Set *halfrange* to ``max(abs(A-vcenter))``, then set *vmin* and *vmax*.
+        """
+        A = np.asanyarray(A)
+        self._halfrange = max(self._vcenter-A.min(),
+                              A.max()-self._vcenter)
+        self._set_vmin_vmax()
+
+    def autoscale_None(self, A):
+        """Set *vmin* and *vmax*."""
+        A = np.asanyarray(A)
+        if self.vmax is None and A.size:
+            self.autoscale(A)
+
+    @property
+    def vcenter(self):
+        return self._vcenter
+
+    @vcenter.setter
+    def vcenter(self, vcenter):
+        self._vcenter = vcenter
+        if self.vmax is not None:
+            # recompute halfrange assuming vmin and vmax represent
+            # min and max of data
+            self._halfrange = max(self._vcenter-self.vmin,
+                                  self.vmax-self._vcenter)
+            self._set_vmin_vmax()
+
+    @property
+    def halfrange(self):
+        return self._halfrange
+
+    @halfrange.setter
+    def halfrange(self, halfrange):
+        if halfrange is None:
+            self._halfrange = None
+            self.vmin = None
+            self.vmax = None
+        else:
+            self._halfrange = abs(halfrange)
+
+    def __call__(self, value, clip=None):
+        if self._halfrange is not None:
+            # enforce symmetry, reset vmin and vmax
+            self._set_vmin_vmax()
+        return super().__call__(value, clip=clip)
+
+
 def _make_norm_from_scale(scale_cls, base_norm_cls=None, *, init=None):
     """
     Decorator for building a `.Normalize` subclass from a `.Scale` subclass.
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -390,6 +390,56 @@ def test_BoundaryNorm():
     assert_array_equal(cmshould(mynorm(x)), cmref(refnorm(x)))
 
 
+def test_CenteredNorm():
+    np.random.seed(0)
+
+    # Assert equivalence to symmetrical Normalize.
+    x = np.random.normal(size=100)
+    x_maxabs = np.max(np.abs(x))
+    norm_ref = mcolors.Normalize(vmin=-x_maxabs, vmax=x_maxabs)
+    norm = mcolors.CenteredNorm()
+    assert_array_almost_equal(norm_ref(x), norm(x))
+
+    # Check that vcenter is in the center of vmin and vmax
+    # when vcenter is set.
+    vcenter = int(np.random.normal(scale=50))
+    norm = mcolors.CenteredNorm(vcenter=vcenter)
+    norm.autoscale_None([1, 2])
+    assert norm.vmax + norm.vmin == 2 * vcenter
+
+    # Check that halfrange input works correctly.
+    x = np.random.normal(size=10)
+    norm = mcolors.CenteredNorm(vcenter=0.5, halfrange=0.5)
+    assert_array_almost_equal(x, norm(x))
+    norm = mcolors.CenteredNorm(vcenter=1, halfrange=1)
+    assert_array_almost_equal(x, 2 * norm(x))
+
+    # Check that halfrange input works correctly and use setters.
+    norm = mcolors.CenteredNorm()
+    norm.vcenter = 2
+    norm.halfrange = 2
+    assert_array_almost_equal(x, 4 * norm(x))
+
+    # Check that prior to adding data, setting halfrange first has same effect.
+    norm = mcolors.CenteredNorm()
+    norm.halfrange = 2
+    norm.vcenter = 2
+    assert_array_almost_equal(x, 4 * norm(x))
+
+    # Check that manual change of vcenter adjusts halfrange accordingly.
+    norm = mcolors.CenteredNorm()
+    assert norm.vcenter == 0
+    # add data
+    norm(np.linspace(-1.0, 0.0, 10))
+    assert norm.vmax == 1.0
+    assert norm.halfrange == 1.0
+    # set vcenter to 1, which should double halfrange
+    norm.vcenter = 1
+    assert norm.vmin == -1.0
+    assert norm.vmax == 3.0
+    assert norm.halfrange == 2.0
+
+
 @pytest.mark.parametrize("vmin,vmax", [[-1, 2], [3, 1]])
 def test_lognorm_invalid(vmin, vmax):
     # Check that invalid limits in LogNorm error
diff --git a/tutorials/colors/colormapnorms.py b/tutorials/colors/colormapnorms.py
@@ -47,6 +47,7 @@
 import matplotlib.pyplot as plt
 import matplotlib.colors as colors
 import matplotlib.cbook as cbook
+from matplotlib import cm
 
 N = 100
 X, Y = np.mgrid[-3:3:complex(0, N), -2:2:complex(0, N)]
@@ -69,6 +70,46 @@
 fig.colorbar(pcm, ax=ax[1], extend='max')
 plt.show()
 
+###############################################################################
+# Centered
+# --------
+#
+# In many cases, data is symmetrical around a center, for example, positive and
+# negative anomalies around a center 0. In this case, we would like the center
+# to be mapped to 0.5 and the datapoint with the largest deviation from the
+# center to be mapped to 1.0, if its value is greater than the center, or 0.0
+# otherwise. The norm `.colors.CenteredNorm` creates such a mapping
+# automatically. It is well suited to be combined with a divergent colormap
+# which uses different colors edges that meet in the center at an unsaturated
+# color.
+#
+# If the center of symmetry is different from 0, it can be set with the
+# *vcenter* argument. For logarithmic scaling on both sides of the center, see
+# `.colors.SymLogNorm` below; to apply a different mapping above and below the
+# center, use `.colors.TwoSlopeNorm` below.
+
+delta = 0.1
+x = np.arange(-3.0, 4.001, delta)
+y = np.arange(-4.0, 3.001, delta)
+X, Y = np.meshgrid(x, y)
+Z1 = np.exp(-X**2 - Y**2)
+Z2 = np.exp(-(X - 1)**2 - (Y - 1)**2)
+Z = (0.9*Z1 - 0.5*Z2) * 2
+
+# select a divergent colormap
+cmap = cm.coolwarm
+
+fig, (ax1, ax2) = plt.subplots(ncols=2)
+pc = ax1.pcolormesh(Z, cmap=cmap)
+fig.colorbar(pc, ax=ax1)
+ax1.set_title('Normalize()')
+
+pc = ax2.pcolormesh(Z, norm=colors.CenteredNorm(), cmap=cmap)
+fig.colorbar(pc, ax=ax2)
+ax2.set_title('CenteredNorm()')
+
+plt.show()
+
 ###############################################################################
 # Symmetric logarithmic
 # ---------------------
