feat: improve StepPatch baseline handling

andrzejnovak · andrzejnovak · commit 855e25b0a4f2 · 2020-09-22T23:08:51.000+02:00
diff --git a/examples/lines_bars_and_markers/stairs_demo.py b/examples/lines_bars_and_markers/stairs_demo.py
@@ -46,7 +46,7 @@
 #############################################################################
 # ``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.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):
diff --git a/lib/matplotlib/patches.py b/lib/matplotlib/patches.py
@@ -996,6 +996,7 @@ class StepPatch(PathPatch):
     The path is unclosed. It starts and stops at baseline.
     """
 
+    @docstring.dedent_interpd
     def __init__(self, values, edges, *,
                  orientation='vertical', baseline=0, **kwargs):
         """
@@ -1020,11 +1021,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)
 
@@ -1039,34 +1039,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.
 
@@ -1075,11 +1076,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
 
@@ -1104,20 +1108,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)
 
 
