feat: improve StepPatch baseline handling

andrzejnovak · andrzejnovak · commit 40c2435f43e9 · 2020-09-18T12:27:19.000+02:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -6787,7 +6787,8 @@ def stairs(self, values, edges=None, *,
         if edges is None:
             edges = np.arange(len(values) + 1)
 
-        self._process_unit_info(xdata=edges, ydata=[values, baseline], 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)
diff --git a/lib/matplotlib/patches.py b/lib/matplotlib/patches.py
@@ -992,6 +992,7 @@ 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):
         """
@@ -1014,11 +1015,10 @@ def __init__(self, values, edges, *,
 
         %(Patch)s
         """
-        self._edge_default = False
         self.orientation = orientation
         self._edges = np.asarray(edges)
         self._values = np.asarray(values)
-        self.baseline = baseline
+        self._baseline = np.asarray(baseline) if baseline is not None else None
         self._update_path()
         super().__init__(self._path, **kwargs)
 
@@ -1033,34 +1033,35 @@ def _update_path(self):
         verts, codes = [], []
 
         _nan_vals = ~np.isnan(self._values)
-        if self.baseline is not None:
-            _nan_vals = _nan_vals & ~np.isnan(self.baseline)
+        if self._baseline is not None:
+            _nan_vals = _nan_vals & ~np.isnan(self._baseline)
         for idx0, idx1 in cbook.contiguous_regions(_nan_vals):
             x = np.repeat(self._edges[idx0:idx1+1], 2)
-            y = np.repeat(self._values[idx0:idx1], 2)     
-            if np.isscalar(self.baseline):  # baseline values
-                y = np.hstack((self.baseline, y, self.baseline))
-            elif hasattr(self.baseline, '__len__'):  # baseline array
-                base = np.repeat(self.baseline[idx0:idx1], 2)[::-1]
+            y = np.repeat(self._values[idx0:idx1], 2)
+            if self._baseline is None:
+                y = np.hstack((y[0], y, y[-1]))
+            elif self._baseline.ndim == 0:  # baseline values
+                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
-                y = np.hstack((y[0], y, y[-1]))
+                raise ValueError('Invalid `baseline` specified')
             if self.orientation == 'vertical':
                 xy = np.column_stack([x, y])
             else:
                 xy = np.column_stack([y, x])
             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."""
-        return self._values, self._edges
+        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.
 
@@ -1069,11 +1070,14 @@ def set_data(self, values, edges=None):
         values : 1D array-like or None
             Will not update values, if passing None
         edges : 1D array-like, optional
+        baseline : 1D array-like, optional
         """
         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
 
@@ -1098,20 +1102,18 @@ def set_edges(self, edges):
         self.set_data(None, edges=edges)
 
     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
         """
-        self.baseline = baseline
-        self._update_path()
-        self.stale = True
+        self.set_data(None, baseline=baseline)
 
 
 class Polygon(Patch):
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1891,6 +1891,9 @@ def test_stairs_update(fig_test, fig_ref):
     assert np.allclose(h.get_data()[1], np.arange(4))
     h.set_baseline(-2)
     assert h.get_baseline() == -2
+    h.set_baseline([-2, -2, -2])
+    assert np.array_equal(h.get_baseline(), np.array([-2, -2, -2]))
+    h.set_baseline(-2)
 
     # # Ref
     fig_ref, ref_ax = plt.subplots()
@@ -1928,18 +1931,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")
+              orientation='horizontal', label="D")
     ax.stairs(y[::-1]*3+13, x-1, color='red', ls='--', lw=2, baseline=None,
-            label="E")
+              label="E")
     ax.stairs(y[::-1]*3+14, x, baseline=26,
-            color='purple', ls='--', lw=2, label="F")
+              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)
+              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="H")
+              baseline=12, hatch='//', label="H")
     ax.legend(loc=0)
 
 
