Merge pull request #12380 from dstansby/stem-speedup2

timhoffm · web-flow · commit 7ca934f8bd36 · 2019-02-15T07:29:10.000+01:00
Stem speedup2
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
@@ -2625,7 +2625,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.
 
@@ -2684,6 +2684,13 @@ 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 the old
+            behavior of using a list of `.Line2D` objects.
+
 
         Returns
         -------
@@ -2715,6 +2722,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:
@@ -2763,24 +2773,39 @@ 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 = []
+            stemlines = [((xi, bottom), (xi, yi)) for xi, yi in zip(x, y)]
+            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:
+            cbook._warn_external(
+                'In Matplotlib 3.3 individual lines on a stem plot will be '
+                'added as a LineCollection instead of individual lines. '
+                'This significantly improves the performance of a stem plot. '
+                'To remove this warning and switch to the new behaviour, '
+                'set the "use_line_collection" keyword argument to True.')
+            stemlines = []
+            for xi, yi in zip(x, y):
+                l, = self.plot([xi, xi], [bottom, yi],
+                               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
@@ -175,10 +175,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,43 @@ 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:
+            # change the function used by update_prop() from the default
+            # to one that handles LineCollection
+            orig_update_func = self._update_prop_func
+            self._update_prop_func = self._copy_collection_props
+
+            for line in leg_stemlines:
+                self.update_prop(line, stemlines, legend)
+
+        else:
+            for lm, m in zip(leg_stemlines, stemlines):
+                self.update_prop(lm, m, legend)
+
+        if using_linecoll:
+            self._update_prop_func = orig_update_func
 
         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)
+        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/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2877,11 +2877,12 @@ def stackplot(
 @docstring.copy(Axes.stem)
 def stem(
         *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0,
-        label=None, data=None):
+        label=None, use_line_collection=False, data=None):
     return gca().stem(
         *args, linefmt=linefmt, markerfmt=markerfmt, basefmt=basefmt,
-        bottom=bottom, label=label, **({"data": data} if data is not
-        None else {}))
+        bottom=bottom, label=label,
+        use_line_collection=use_line_collection, **({"data": data} if
+        data is not None else {}))
 
 
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
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
@@ -3143,6 +3143,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)
