Merge pull request matplotlib#18310 from QuLogic/stem3d

Add 3D stem plot

Author: efiring

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.

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()

lib/mpl_toolkits/mplot3d/axes3d.py

@@ -3351,6 +3351,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__)
 

lib/mpl_toolkits/tests/baseline_images/test_mplot3d/stem3d.png

@@ -1225,6 +1225,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