diff --git a/doc/api/next_api_changes/2017-10-24-DS.rst b/doc/api/next_api_changes/2017-10-24-DS.rst new file mode 100644 index 000000000000..d44b4a22c36f --- /dev/null +++ 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 index 686ef9ee76f5..6a1c104c1741 100644 --- 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,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 ------- @@ -2570,6 +2577,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,16 +2628,32 @@ 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_") @@ -2635,7 +2661,6 @@ def stem(self, *args, linefmt=None, markerfmt=None, basefmt=None, bottom=0, 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 index 552d6da5a761..194ae401bcef 100644 --- 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 index 5d4e55a351da..cea1b0a5b2a2 100644 --- 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 index 73c43302cbce..440c1ea05c72 100644 --- a/lib/matplotlib/pyplot.py +++ b/lib/matplotlib/pyplot.py @@ -2899,11 +2899,12 @@ def stackplot(x, *args, data=None, **kwargs): @docstring.copy_dedent(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 new file mode 100644 index 000000000000..4c6b3af3205b Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_axes/stem.png differ diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py index ed7e0ac8168e..c7230dd6b592 100644 --- 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)