diff --git a/doc/users/next_whats_new/stem3d.rst b/doc/users/next_whats_new/stem3d.rst new file mode 100644 index 000000000000..0d1939a8885d --- /dev/null +++ b/doc/users/next_whats_new/stem3d.rst @@ -0,0 +1,24 @@ +Stem plots in 3D Axes +--------------------- + +Stem plots are now supported on 3D Axes. Much like 2D stems, +`~.axes3d.Axes3D.stem3D` supports plotting the stems in various orientations: + +.. plot:: + + theta = np.linspace(0, 2*np.pi) + x = np.cos(theta - np.pi/2) + y = np.sin(theta - np.pi/2) + z = theta + directions = ['z', 'x', 'y'] + names = [r'$\theta$', r'$\cos\theta$', r'$\sin\theta$'] + + fig, axs = plt.subplots(1, 3, figsize=(8, 4), + constrained_layout=True, + subplot_kw={'projection': '3d'}) + for ax, zdir, name in zip(axs, directions, names): + ax.stem(x, y, z, orientation=zdir) + ax.set_title(name) + fig.suptitle(r'A parametric circle: $(x, y) = (\cos\theta, \sin\theta)$') + +See also the :doc:`/gallery/mplot3d/stem3d_demo` demo. diff --git a/examples/mplot3d/stem3d_demo.py b/examples/mplot3d/stem3d_demo.py new file mode 100644 index 000000000000..712f93227ae2 --- /dev/null +++ b/examples/mplot3d/stem3d_demo.py @@ -0,0 +1,51 @@ +""" +======= +3D stem +======= + +Demonstration of a stem plot in 3D, which plots vertical lines from a baseline +to the *z*-coordinate and places a marker at the tip. +""" + +import matplotlib.pyplot as plt +import numpy as np + +theta = np.linspace(0, 2*np.pi) +x = np.cos(theta - np.pi/2) +y = np.sin(theta - np.pi/2) +z = theta + +fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) +ax.stem(x, y, z) + +plt.show() + +############################################################################# +# +# The position of the baseline can be adapted using *bottom*. The parameters +# *linefmt*, *markerfmt*, and *basefmt* control basic format properties of the +# plot. However, in contrast to `~.axes3d.Axes3D.plot` not all properties are +# configurable via keyword arguments. For more advanced control adapt the line +# objects returned by `~.stem3D`. + +fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) +markerline, stemlines, baseline = ax.stem( + x, y, z, linefmt='grey', markerfmt='D', bottom=np.pi) +markerline.set_markerfacecolor('none') + +plt.show() + +############################################################################# +# +# The orientation of the stems and baseline can be changed using *orientation*. +# This determines in which direction the stems are projected from the head +# points, towards the *bottom* baseline. +# +# For examples, by setting ``orientation='x'``, the stems are projected along +# the *x*-direction, and the baseline is in the *yz*-plane. + +fig, ax = plt.subplots(subplot_kw=dict(projection='3d')) +markerline, stemlines, baseline = ax.stem(x, y, z, bottom=-1, orientation='x') +ax.set(xlabel='x', ylabel='y', zlabel='z') + +plt.show() diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py index 15858f59cc5f..31b3b610789b 100644 --- a/lib/mpl_toolkits/mplot3d/axes3d.py +++ b/lib/mpl_toolkits/mplot3d/axes3d.py @@ -3288,6 +3288,121 @@ def get_tightbbox(self, renderer, call_axes_locator=True, batch.append(axis_bb) return mtransforms.Bbox.union(batch) + def stem(self, x, y, z, *, linefmt='C0-', markerfmt='C0o', basefmt='C3-', + bottom=0, label=None, orientation='z'): + """ + Create a 3D stem plot. + + A stem plot draws lines perpendicular to a baseline, and places markers + at the heads. By default, the baseline is defined by *x* and *y*, and + stems are drawn vertically from *bottom* to *z*. + + Parameters + ---------- + x, y, z : array-like + The positions of the heads of the stems. The stems are drawn along + the *orientation*-direction from the baseline at *bottom* (in the + *orientation*-coordinate) to the heads. By default, the *x* and *y* + positions are used for the baseline and *z* for the head position, + but this can be changed by *orientation*. + + linefmt : str, default: 'C0-' + A string defining the properties of the vertical lines. Usually, + this will be a color or a color and a linestyle: + + ========= ============= + Character Line Style + ========= ============= + ``'-'`` solid line + ``'--'`` dashed line + ``'-.'`` dash-dot line + ``':'`` dotted line + ========= ============= + + Note: While it is technically possible to specify valid formats + other than color or color and linestyle (e.g. 'rx' or '-.'), this + is beyond the intention of the method and will most likely not + result in a reasonable plot. + + markerfmt : str, default: 'C0o' + A string defining the properties of the markers at the stem heads. + + basefmt : str, default: 'C3-' + A format string defining the properties of the baseline. + + bottom : float, default: 0 + The position of the baseline, in *orientation*-coordinates. + + label : str, default: None + The label to use for the stems in legends. + + orientation : {'x', 'y', 'z'}, default: 'z' + The direction along which stems are drawn. + + Returns + ------- + `.StemContainer` + The container may be treated like a tuple + (*markerline*, *stemlines*, *baseline*) + + Examples + -------- + .. plot:: gallery/mplot3d/stem3d_demo.py + """ + + from matplotlib.container import StemContainer + + had_data = self.has_data() + + _api.check_in_list(['x', 'y', 'z'], orientation=orientation) + + xlim = (np.min(x), np.max(x)) + ylim = (np.min(y), np.max(y)) + zlim = (np.min(z), np.max(z)) + + # Determine the appropriate plane for the baseline and the direction of + # stemlines based on the value of orientation. + if orientation == 'x': + basex, basexlim = y, ylim + basey, baseylim = z, zlim + lines = [[(bottom, thisy, thisz), (thisx, thisy, thisz)] + for thisx, thisy, thisz in zip(x, y, z)] + elif orientation == 'y': + basex, basexlim = x, xlim + basey, baseylim = z, zlim + lines = [[(thisx, bottom, thisz), (thisx, thisy, thisz)] + for thisx, thisy, thisz in zip(x, y, z)] + else: + basex, basexlim = x, xlim + basey, baseylim = y, ylim + lines = [[(thisx, thisy, bottom), (thisx, thisy, thisz)] + for thisx, thisy, thisz in zip(x, y, z)] + + # Determine style for stem lines. + linestyle, linemarker, linecolor = _process_plot_format(linefmt) + if linestyle is None: + linestyle = rcParams['lines.linestyle'] + + # Plot everything in required order. + baseline, = self.plot(basex, basey, basefmt, zs=bottom, + zdir=orientation, label='_nolegend_') + stemlines = art3d.Line3DCollection( + lines, linestyles=linestyle, colors=linecolor, label='_nolegend_') + self.add_collection(stemlines) + markerline, = self.plot(x, y, z, markerfmt, label='_nolegend_') + + stem_container = StemContainer((markerline, stemlines, baseline), + label=label) + self.add_container(stem_container) + + jx, jy, jz = art3d.juggle_axes(basexlim, baseylim, [bottom, bottom], + orientation) + self.auto_scale_xyz([*jx, *xlim], [*jy, *ylim], [*jz, *zlim], had_data) + + return stem_container + + stem3D = stem + docstring.interpd.update(Axes3D_kwdoc=artist.kwdoc(Axes3D)) docstring.dedent_interpd(Axes3D.__init__) diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/stem3d.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/stem3d.png new file mode 100644 index 000000000000..cdb5fbdd1b42 Binary files /dev/null and b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/stem3d.png differ diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py index 60436bd6bb93..8d1ea9704acb 100644 --- a/lib/mpl_toolkits/tests/test_mplot3d.py +++ b/lib/mpl_toolkits/tests/test_mplot3d.py @@ -1203,6 +1203,35 @@ def test_errorbar3d(): ax.legend() +@image_comparison(['stem3d.png'], style='mpl20') +def test_stem3d(): + fig, axs = plt.subplots(2, 3, figsize=(8, 6), + constrained_layout=True, + subplot_kw={'projection': '3d'}) + + theta = np.linspace(0, 2*np.pi) + x = np.cos(theta - np.pi/2) + y = np.sin(theta - np.pi/2) + z = theta + + for ax, zdir in zip(axs[0], ['x', 'y', 'z']): + ax.stem(x, y, z, orientation=zdir) + ax.set_title(f'orientation={zdir}') + + x = np.linspace(-np.pi/2, np.pi/2, 20) + y = np.ones_like(x) + z = np.cos(x) + + for ax, zdir in zip(axs[1], ['x', 'y', 'z']): + markerline, stemlines, baseline = ax.stem( + x, y, z, + linefmt='C4-.', markerfmt='C1D', basefmt='C2', + orientation=zdir) + ax.set_title(f'orientation={zdir}') + markerline.set(markerfacecolor='none', markeredgewidth=2) + baseline.set_linewidth(3) + + @image_comparison(["equal_box_aspect.png"], style="mpl20") def test_equal_box_aspect(): from itertools import product, combinations