Merge pull request matplotlib#18511 from andrzejnovak/baselinearr

feat: StepPatch to take array as baseline

Author: timhoffm

examples/lines_bars_and_markers/stairs_demo.py

@@ -43,6 +43,16 @@
     ax.legend()
 plt.show()
 
+#############################################################################
+# *baseline* can take an array to allow for stacked histogram plots
+A = [[0, 0, 0],
+     [1, 2, 3],
+     [2, 4, 6],
+     [3, 6, 9]]
+
+for i in range(len(A) - 1):
+    plt.stairs(A[i+1], baseline=A[i], fill=True)
+
 #############################################################################
 # Comparison of `.pyplot.step` and `.pyplot.stairs`
 # -------------------------------------------------

lib/matplotlib/axes/_axes.py

@@ -6730,7 +6730,8 @@ def hist(self, x, bins=None, range=None, density=False, weights=None,
     def stairs(self, values, edges=None, *,
                orientation='vertical', baseline=0, fill=False, **kwargs):
         """
-        A stepwise constant line or filled plot.
+        A stepwise constant function as a line with bounding edges
+        or a filled plot.
 
         Parameters
         ----------
@@ -6745,9 +6746,11 @@ def stairs(self, values, edges=None, *,
             The direction of the steps. Vertical means that *values* are along
             the y-axis, and edges are along the x-axis.
 
-        baseline : float or None, default: 0
-            Determines starting value of the bounding edges or when
-            ``fill=True``, position of lower edge.
+        baseline : float, array-like or None, default: 0
+            The bottom value of the bounding edges or when
+            ``fill=True``, position of lower edge. If *fill* is
+            True or an array is passed to *baseline*, a closed
+            path is drawn.
 
         fill : bool, default: False
             Whether the area under the step curve should be filled.
@@ -6776,8 +6779,8 @@ def stairs(self, values, edges=None, *,
         if edges is None:
             edges = np.arange(len(values) + 1)
 
-        edges, values = self._process_unit_info(
-            [("x", edges), ("y", values)], kwargs)
+        edges, values, baseline = self._process_unit_info(
+            [("x", edges), ("y", values), ("y", baseline)], kwargs)
 
         patch = mpatches.StepPatch(values,
                                    edges,
@@ -6789,9 +6792,9 @@ def stairs(self, values, edges=None, *,
         if baseline is None:
             baseline = 0
         if orientation == 'vertical':
-            patch.sticky_edges.y.append(baseline)
+            patch.sticky_edges.y.append(np.min(baseline))
         else:
-            patch.sticky_edges.x.append(baseline)
+            patch.sticky_edges.x.append(np.min(baseline))
         self._request_autoscale_view()
         return patch
 

lib/matplotlib/patches.py

@@ -4,6 +4,7 @@
 import math
 from numbers import Number
 import textwrap
+from collections import namedtuple
 
 import numpy as np
 
@@ -926,7 +927,8 @@ class StepPatch(PathPatch):
     """
     A path patch describing a stepwise constant function.
 
-    The path is unclosed. It starts and stops at baseline.
+    By default the path is not closed and starts and stops at
+    baseline value.
     """
 
     _edge_default = False
@@ -945,21 +947,23 @@ def __init__(self, values, edges, *,
             between which the curve takes on vals values.
 
         orientation : {'vertical', 'horizontal'}, default: 'vertical'
-            The direction of the steps. Vertical means that *values* are along
-            the y-axis, and edges are along the x-axis.
+            The direction of the steps. Vertical means that *values* are
+            along the y-axis, and edges are along the x-axis.
 
-        baseline : float or None, default: 0
-            Determines starting value of the bounding edges or when
-            ``fill=True``, position of lower edge.
+        baseline : float, array-like or None, default: 0
+            The bottom value of the bounding edges or when
+            ``fill=True``, position of lower edge. If *fill* is
+            True or an array is passed to *baseline*, a closed
+            path is drawn.
 
         Other valid keyword arguments are:
 
         %(Patch)s
         """
-        self.baseline = baseline
         self.orientation = orientation
         self._edges = np.asarray(edges)
         self._values = np.asarray(values)
+        self._baseline = np.asarray(baseline) if baseline is not None else None
         self._update_path()
         super().__init__(self._path, **kwargs)
 
@@ -972,13 +976,24 @@ def _update_path(self):
                              f"`len(values) = {self._values.size}` and "
                              f"`len(edges) = {self._edges.size}`.")
         verts, codes = [], []
-        for idx0, idx1 in cbook.contiguous_regions(~np.isnan(self._values)):
+
+        _nan_mask = np.isnan(self._values)
+        if self._baseline is not None:
+            _nan_mask |= np.isnan(self._baseline)
+        for idx0, idx1 in cbook.contiguous_regions(~_nan_mask):
             x = np.repeat(self._edges[idx0:idx1+1], 2)
             y = np.repeat(self._values[idx0:idx1], 2)
-            if self.baseline is not None:
-                y = np.hstack((self.baseline, y, self.baseline))
-            else:
+            if self._baseline is None:
                 y = np.hstack((y[0], y, y[-1]))
+            elif self._baseline.ndim == 0:  # single baseline value
+                y = np.hstack((self._baseline, y, self._baseline))
+            elif self._baseline.ndim == 1:  # baseline array
+                base = np.repeat(self._baseline[idx0:idx1], 2)[::-1]
+                x = np.concatenate([x, x[::-1]])
+                y = np.concatenate([np.hstack((base[-1], y, base[0],
+                                               base[0], base, base[-1]))])
+            else:  # no baseline
+                raise ValueError('Invalid `baseline` specified')
             if self.orientation == 'vertical':
                 xy = np.column_stack([x, y])
             else:
@@ -988,59 +1003,29 @@ def _update_path(self):
             self._path = Path(np.vstack(verts), np.hstack(codes))
 
     def get_data(self):
-        """Get `.StepPatch` values and edges."""
-        return self._values, self._edges
+        """Get `.StepPatch` values, edges and baseline as namedtuple."""
+        StairData = namedtuple('StairData', 'values edges baseline')
+        return StairData(self._values, self._edges, self._baseline)
 
-    def set_data(self, values, edges=None):
+    def set_data(self, values=None, edges=None, baseline=None):
         """
-        Set `.StepPatch` values and optionally edges.
+        Set `.StepPatch` values, edges and baseline.
 
         Parameters
         ----------
         values : 1D array-like or None
             Will not update values, if passing None
         edges : 1D array-like, optional
+        baseline : float, 1D array-like or None
         """
+        if values is None and edges is None and baseline is None:
+            raise ValueError("Must set *values*, *edges* or *baseline*.")
         if values is not None:
             self._values = np.asarray(values)
         if edges is not None:
             self._edges = np.asarray(edges)
-        self._update_path()
-        self.stale = True
-
-    def set_values(self, values):
-        """
-        Set `.StepPatch` values.
-
-        Parameters
-        ----------
-        values : 1D array-like
-        """
-        self.set_data(values, edges=None)
-
-    def set_edges(self, edges):
-        """
-        Set `.StepPatch` edges.
-
-        Parameters
-        ----------
-        edges : 1D array-like
-        """
-        self.set_data(None, edges=edges)
-
-    def get_baseline(self):
-        """Get `.StepPatch` baseline value."""
-        return self.baseline
-
-    def set_baseline(self, baseline):
-        """
-        Set `.StepPatch` baseline value.
-
-        Parameters
-        ----------
-        baseline : float or None
-        """
-        self.baseline = baseline
+        if baseline is not None:
+            self._baseline = np.asarray(baseline)
         self._update_path()
         self.stale = True
 

lib/matplotlib/tests/baseline_images/test_axes/test_stairs_options.png

@@ -1892,15 +1892,15 @@ def test_stairs_update(fig_test, fig_ref):
     test_ax = fig_test.add_subplot()
     h = test_ax.stairs([1, 2, 3])
     test_ax.set_ylim(ylim)
-    h.set_values([3, 2, 1])
-    h.set_edges(np.arange(4)+2)
+    h.set_data([3, 2, 1])
+    h.set_data(edges=np.arange(4)+2)
     h.set_data([1, 2, 1], np.arange(4)/2)
     h.set_data([1, 2, 3])
     h.set_data(None, np.arange(4))
     assert np.allclose(h.get_data()[0], np.arange(1, 4))
     assert np.allclose(h.get_data()[1], np.arange(4))
-    h.set_baseline(-2)
-    assert h.get_baseline() == -2
+    h.set_data(baseline=-2)
+    assert h.get_data().baseline == -2
 
     # Ref
     ref_ax = fig_ref.add_subplot()
@@ -1921,13 +1921,13 @@ def test_stairs_invalid_mismatch():
 def test_stairs_invalid_update():
     h = plt.stairs([1, 2], [0, 1, 2])
     with pytest.raises(ValueError, match='Nan values in "edges"'):
-        h.set_edges([1, np.nan, 2])
+        h.set_data(edges=[1, np.nan, 2])
 
 
 def test_stairs_invalid_update2():
     h = plt.stairs([1, 2], [0, 1, 2])
     with pytest.raises(ValueError, match='Size mismatch'):
-        h.set_edges(np.arange(5))
+        h.set_data(edges=np.arange(5))
 
 
 @image_comparison(['test_stairs_options.png'], remove_text=True)
@@ -1943,10 +1943,14 @@ def test_stairs_options():
     ax.stairs(yn, x, color='orange', ls='--', lw=2, label="C")
     ax.stairs(yn/3, x*3-2, ls='--', lw=2, baseline=0.5,
               orientation='horizontal', label="D")
-    ax.stairs(y[::-1]*3+12, x, color='red', ls='--', lw=2, baseline=None,
+    ax.stairs(y[::-1]*3+13, x-1, color='red', ls='--', lw=2, baseline=None,
               label="E")
+    ax.stairs(y[::-1]*3+14, x, baseline=26,
+              color='purple', ls='--', lw=2, label="F")
+    ax.stairs(yn[::-1]*3+15, x+1, baseline=np.linspace(27, 25, len(y)),
+              color='blue', ls='--', lw=2, label="G", fill=True)
     ax.stairs(y[:-1][::-1]*2+11, x[:-1]+0.5, color='black', ls='--', lw=2,
-              baseline=12, hatch='//', label="F")
+              baseline=12, hatch='//', label="H")
     ax.legend(loc=0)