feat: StepPatch to take array as baseline

andrzejnovak · andrzejnovak · commit 088d67237b36 · 2020-09-29T20:09:04.000+02:00
diff --git a/examples/lines_bars_and_markers/stairs_demo.py b/examples/lines_bars_and_markers/stairs_demo.py
@@ -43,6 +43,15 @@
     ax.legend()
 plt.show()
 
+#############################################################################
+# ``baseline`` can take an array to allow for stacked histogram plots
+
+A = np.tile([1, 2, 3], (3, 1)) * np.arange(1, 4).reshape(-1, 1)
+A = np.vstack([np.zeros(3), A])
+
+for i in range(3):
+    plt.stairs(A[i+1], baseline=A[i], fill=True)
+
 #############################################################################
 # Comparison of `.pyplot.step` and `.pyplot.stairs`
 # -------------------------------------------------
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -6729,8 +6729,8 @@ 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
+        baseline : float, array-like or None, default: 0
+            Determines bottom value of the bounding edges or when
             ``fill=True``, position of lower edge.
 
         fill : bool, default: False
@@ -6760,8 +6760,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,
@@ -6773,9 +6773,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
 
diff --git a/lib/matplotlib/patches.py b/lib/matplotlib/patches.py
@@ -1013,18 +1013,18 @@ def __init__(self, values, edges, *,
             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
+        baseline : float, 1D array-like or None, default: 0
+            Determines bottom value of the bounding edges or when
             ``fill=True``, position of lower edge.
 
         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)
 
@@ -1037,13 +1037,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:
@@ -1053,23 +1064,26 @@ 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."""
+        return self._values, self._edges, self._baseline
 
-    def set_data(self, values, edges=None):
+    def set_data(self, values, edges=None, baseline=None):
         """
-        Set `.StepPatch` values and optionally edges.
+        Set `.StepPatch` values and optionally 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 not None:
             self._values = np.asarray(values)
         if edges is not None:
             self._edges = np.asarray(edges)
+        if baseline is not None:
+            self._baseline = np.asarray(baseline)
         self._update_path()
         self.stale = True
 
@@ -1081,7 +1095,7 @@ def set_values(self, values):
         ----------
         values : 1D array-like
         """
-        self.set_data(values, edges=None)
+        self.set_data(values, edges=None, baseline=None)
 
     def set_edges(self, edges):
         """
@@ -1091,23 +1105,21 @@ def set_edges(self, edges):
         ----------
         edges : 1D array-like
         """
-        self.set_data(None, edges=edges)
+        self.set_data(None, edges=edges, baseline=None)
 
     def get_baseline(self):
-        """Get `.StepPatch` baseline value."""
-        return self.baseline
+        """Get `.StepPatch` baseline."""
+        return self._baseline
 
     def set_baseline(self, baseline):
         """
-        Set `.StepPatch` baseline value.
+        Set `.StepPatch` baseline.
 
         Parameters
         ----------
-        baseline : float or None
+        baseline : float, array-like or None, default: 0
         """
-        self.baseline = baseline
-        self._update_path()
-        self.stale = True
+        self.set_data(None, edges=None, baseline=baseline)
 
 
 class Polygon(Patch):
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/test_stairs_options.png b/lib/matplotlib/tests/baseline_images/test_axes/test_stairs_options.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1932,10 +1932,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)
 
 
