feat: StepPatch to take array as baseline

andrzejnovak · andrzejnovak · commit 973072d68bba · 2020-09-18T10:36:15.000+02:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -6757,8 +6757,8 @@ def stairs(self, values, edges=None, *,
 
         orientation : {'vertical', 'horizontal'}, default: 'vertical'
 
-        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
@@ -6787,9 +6787,10 @@ def stairs(self, values, edges=None, *,
         if edges is None:
             edges = np.arange(len(values) + 1)
 
-        self._process_unit_info(xdata=edges, ydata=values, kwargs=kwargs)
+        self._process_unit_info(xdata=edges, ydata=[values, baseline], kwargs=kwargs)
         edges = self.convert_xunits(edges)
         values = self.convert_yunits(values)
+        baseline = self.convert_yunits(baseline)
 
         patch = mpatches.StepPatch(values,
                                    edges,
@@ -6801,9 +6802,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
@@ -992,7 +992,6 @@ def set_path(self, path):
 class StepPatch(PathPatch):
     """An unclosed step path patch."""
 
-    @docstring.dedent_interpd
     def __init__(self, values, edges, *,
                  orientation='vertical', baseline=0, **kwargs):
         """
@@ -1007,18 +1006,19 @@ def __init__(self, values, edges, *,
 
         orientation : {'vertical', 'horizontal'}, default: 'vertical'
 
-        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.
 
         Other valid keyword arguments are:
 
         %(Patch)s
         """
-        self.baseline = baseline
+        self._edge_default = False
         self.orientation = orientation
         self._edges = np.asarray(edges)
         self._values = np.asarray(values)
+        self.baseline = baseline
         self._update_path()
         super().__init__(self._path, **kwargs)
 
@@ -1033,10 +1033,14 @@ def _update_path(self):
         verts, codes = [], []
         for idx0, idx1 in cbook.contiguous_regions(~np.isnan(self._values)):
             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.repeat(self._values[idx0:idx1], 2)     
+            if np.isscalar(self.baseline):  # baseline values
                 y = np.hstack((self.baseline, y, self.baseline))
-            else:
+            elif hasattr(self.baseline, '__len__'):  # 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
                 y = np.hstack((y[0], y, y[-1]))
             if self.orientation == 'vertical':
                 xy = np.column_stack([x, y])
@@ -1045,6 +1049,7 @@ def _update_path(self):
             verts.append(xy)
             codes.append(np.array([Path.MOVETO] + [Path.LINETO]*(len(xy)-1)))
             self._path = Path(np.vstack(verts), np.hstack(codes))
+            
 
     def get_data(self):
         """Get `.StepPatch` values and edges."""
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
@@ -1928,14 +1928,18 @@ def test_stairs_options():
     fig, ax = plt.subplots()
     ax.stairs(y*3, x, color='green', fill=True, label="A")
     ax.stairs(y, x*3-3, color='red', fill=True,
-              orientation='horizontal', label="B")
+            orientation='horizontal', label="B")
     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,
-              label="E")
+            orientation='horizontal', label="D")
+    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)
 
 
