Improve step performance

dstansby · dstansby · commit 2098d1aa72a1 · 2019-01-01T16:53:33.000Z
Do deprecation better

Fix linestyle
diff --git a/doc/api/next_api_changes/2017-10-24-DS.rst b/doc/api/next_api_changes/2017-10-24-DS.rst
@@ -0,0 +1,14 @@
+`StemContainer` now stores `LineCollection`
+-------------------------------------------
+
+`StemContainer` objects can now store a `LineCollection` object instead of a
+list of `Line2D` objects for stem lines plotted using `ax.stem`. This gives a
+very large performance boost to displaying and moving `ax.stem` plots.
+
+This will become the default behaviour in Matplotlib 3.3. To use it now, the
+``use_line_collection`` keyword argument to ~`.axes.stem` can be set to
+``True``.
+
+Individual line segments can be extracted from the `LineCollection` using
+`LineCollection.get_segements()`. See the `LineCollection` documentation for
+other methods to retrieve the collection properties.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2480,7 +2480,7 @@ def broken_barh(self, xranges, yrange, **kwargs):
 
     @_preprocess_data()
     def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
-             label=None):
+             label=None, use_line_collection=False):
         """
         Create a stem plot.
 
@@ -2539,6 +2539,12 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
         label : str, optional, default: None
             The label to use for the stems in legends.
 
+        use_line_collection : bool, optional, default: False
+            If ``True``, store and plot the stem lines as a
+            ~`.collections.LineCollection` instead of individual lines. This
+            significantly increases performance, and will become the default
+            option in Matplotlib 3.3. If ``False``, defaults to old behaviour.
+
 
         Returns
         -------
@@ -2570,6 +2576,9 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
             x = y
             y = np.asarray(args[0], dtype=float)
             args = args[1:]
+        self._process_unit_info(xdata=x, ydata=y)
+        x = self.convert_xunits(x)
+        y = self.convert_yunits(y)
 
         # defaults for formats
         if linefmt is None:
@@ -2618,24 +2627,40 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
         else:
             basestyle, basemarker, basecolor = _process_plot_format(basefmt)
 
+        # New behaviour in 3.1 is to use a LineCollection for the stemlines
+        if use_line_collection:
+            stemlines = []
+            for thisx, thisy in zip(x, y):
+                stemlines.append(((thisx, bottom), (thisx, thisy)))
+            stemlines = mcoll.LineCollection(stemlines, linestyles=linestyle,
+                                             colors=linecolor,
+                                             label='_nolegend_')
+            self.add_collection(stemlines)
+        # Old behaviour is to plot each of the lines individually
+        else:
+            warnings.warn(
+                'In Matplotlib 3.3 individual lines on a stem plot will be '
+                'added as a LineCollection instead of individual lines.\n'
+                'This significantly improves the performance of a stem plot.\n'
+                'To remove this warning and switch to the new behaviour, '
+                'set the "use_line_collection" keyword argument to True.')
+            stemlines = []
+            for thisx, thisy in zip(x, y):
+                l, = self.plot([thisx, thisx], [bottom, thisy],
+                               color=linecolor, linestyle=linestyle,
+                               marker=linemarker, label="_nolegend_")
+            stemlines.append(l)
+
         markerline, = self.plot(x, y, color=markercolor, linestyle=markerstyle,
                                 marker=markermarker, label="_nolegend_")
 
-        stemlines = []
-        for thisx, thisy in zip(x, y):
-            l, = self.plot([thisx, thisx], [bottom, thisy],
-                           color=linecolor, linestyle=linestyle,
-                           marker=linemarker, label="_nolegend_")
-            stemlines.append(l)
-
         baseline, = self.plot([np.min(x), np.max(x)], [bottom, bottom],
                               color=basecolor, linestyle=basestyle,
                               marker=basemarker, label="_nolegend_")
 
         stem_container = StemContainer((markerline, stemlines, baseline),
                                        label=label)
         self.add_container(stem_container)
-
         return stem_container
 
     @_preprocess_data(replace_names=["x", "explode", "labels", "colors"])
diff --git a/lib/matplotlib/container.py b/lib/matplotlib/container.py
@@ -174,10 +174,17 @@ class StemContainer(Container):
 
     baseline : :class:`~matplotlib.lines.Line2D`
         The artist of the horizontal baseline.
-
     """
-
     def __init__(self, markerline_stemlines_baseline, **kwargs):
+        '''
+        Parameters
+        ----------
+        markerline_stemlines_baseline : tuple
+            Tuple of ``(markerline, stemlines, baseline)``.
+            ``markerline`` contains the `LineCollection` of the markers,
+            ``stemlines`` is a `LineCollection` of the main lines,
+            ``baseline`` is the `Line2D` of the baseline.
+        '''
         markerline, stemlines, baseline = markerline_stemlines_baseline
         self.markerline = markerline
         self.stemlines = stemlines
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -589,8 +589,11 @@ def get_ydata(self, legend, xdescent, ydescent, width, height, fontsize):
     def create_artists(self, legend, orig_handle,
                        xdescent, ydescent, width, height, fontsize,
                        trans):
-
         markerline, stemlines, baseline = orig_handle
+        # Check to see if the stemcontainer is storing lines as a list or a
+        # LineCollection. Eventually using a list will be removed, and this
+        # logic can also be removed.
+        using_linecoll = isinstance(stemlines, mcoll.LineCollection)
 
         xdata, xdata_marker = self.get_xdata(legend, xdescent, ydescent,
                                              width, height, fontsize)
@@ -609,23 +612,50 @@ def create_artists(self, legend, orig_handle,
         leg_stemlines = [Line2D([x, x], [bottom, y])
                          for x, y in zip(xdata_marker, ydata)]
 
-        for lm, m in zip(leg_stemlines, stemlines):
-            self.update_prop(lm, m, legend)
+        if using_linecoll:
+            # update_prop() usually takes two Line2D collections;
+            # override temporarily to copy properties from a LineCollection
+            orig_update_func = self._update_prop_func
+            self._update_prop_func = self._copy_collection_props
+
+            for thisx, thisy in zip(xdata_marker, ydata):
+                thisline = Line2D([thisx, thisx], [bottom, thisy])
+                leg_stemlines.append(thisline)
+                self.update_prop(thisline, stemlines, legend)
+
+            self._update_prop_func = orig_update_func
+        else:
+            for thisx, thisy in zip(xdata_marker, ydata):
+                thisline = Line2D([thisx, thisx], [bottom, thisy])
+                leg_stemlines.append(thisline)
+            for lm, m in zip(leg_stemlines, stemlines):
+                self.update_prop(lm, m, legend)
 
         leg_baseline = Line2D([np.min(xdata), np.max(xdata)],
                               [bottom, bottom])
 
         self.update_prop(leg_baseline, baseline, legend)
 
-        artists = [leg_markerline]
-        artists.extend(leg_stemlines)
+        leg_markerline = Line2D(xdata_marker, ydata[:len(xdata_marker)])
+        self.update_prop(leg_markerline, markerline, legend)
+
+        artists = leg_stemlines
         artists.append(leg_baseline)
+        artists.append(leg_markerline)
 
         for artist in artists:
             artist.set_transform(trans)
 
         return artists
 
+    def _copy_collection_props(self, legend_handle, orig_handle):
+        '''
+        Method to copy properties from a LineCollection (orig_handle) to a
+        Line2D (legend_handle).
+        '''
+        legend_handle.set_color(orig_handle.get_color()[0])
+        legend_handle.set_linestyle(orig_handle.get_linestyle()[0])
+
 
 class HandlerTuple(HandlerBase):
     """
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/stem.png b/lib/matplotlib/tests/baseline_images/test_axes/stem.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3072,6 +3072,28 @@ def test_hist_stacked_weighted():
     ax.hist((d1, d2), weights=(w1, w2), histtype="stepfilled", stacked=True)
 
 
+@pytest.mark.parametrize("use_line_collection", [True, False],
+                         ids=['w/ line collection', 'w/o line collection'])
+@image_comparison(baseline_images=['stem'], extensions=['png'], style='mpl20',
+                  remove_text=True)
+def test_stem(use_line_collection):
+    x = np.linspace(0.1, 2 * np.pi, 100)
+    args = (x, np.cos(x))
+    # Label is a single space to force a legend to be drawn, but to avoid any
+    # text being drawn
+    kwargs = dict(linefmt='C2-.', markerfmt='k+', basefmt='C1-.',
+                  label=' ', use_line_collection=use_line_collection)
+
+    fig, ax = plt.subplots()
+    if use_line_collection:
+        ax.stem(*args, **kwargs)
+    else:
+        with pytest.warns(UserWarning):
+            ax.stem(*args, **kwargs)
+
+    ax.legend()
+
+
 def test_stem_args():
     fig = plt.figure()
     ax = fig.add_subplot(1, 1, 1)
