diff --git a/examples/mplot3d/stem3d.py b/examples/mplot3d/stem3d.py new file mode 100644 index 000000000000..8f68c1682192 --- /dev/null +++ b/examples/mplot3d/stem3d.py @@ -0,0 +1,41 @@ +''' +============== +3D stemplot +============== + +Demonstration of a basic stemplot in 3D. +''' + +import matplotlib.pyplot as plt +import numpy as np + +# Fixing random state for reproducibility +np.random.seed(19680801) + + +def randrange(n, vmin, vmax): + ''' + Helper function to make an array of random numbers having shape (n, ) + with each number distributed Uniform(vmin, vmax). + ''' + return (vmax - vmin)*np.random.rand(n) + vmin + +fig = plt.figure() +ax = fig.add_subplot(111, projection='3d') + +n = 100 + +# For each set of style and range settings, plot n random points in the box +# defined by x in [23, 32], y in [0, 100], z in [zlow, zhigh]. +for zlow, zhigh in [(-50, -25), (-30, -5)]: + xs = randrange(n, 23, 32) + ys = randrange(n, 0, 100) + zs = randrange(n, zlow, zhigh) + ax.stem(xs, ys, zs, linefmt='grey', markerfmt='D', bottom=1.1, + use_line_collection=True) + +ax.set_xlabel('X Label') +ax.set_ylabel('Y Label') +ax.set_zlabel('Z Label') + +plt.show() diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py index 77586cd58ff1..c1bb60bc7207 100644 --- a/lib/mpl_toolkits/mplot3d/axes3d.py +++ b/lib/mpl_toolkits/mplot3d/axes3d.py @@ -1418,6 +1418,49 @@ def text(self, x, y, z, s, zdir=None, **kwargs): text3D = text text2D = Axes.text + def stem(self, xs, ys, zs, *args, zdir='z', **kwargs): + """ + Plot a 3D stem plot. + + Parameters + ---------- + xs : scalar or 1D array-like + x coordinates of vertices; either one for all points or one + for each point + ys : scalar or 1D array-like + y coordinates of vertices; either one for all points or one + for each point + zs : scalar or 1D array-like + z coordinates of vertices; either one for all points or one + for each point. + zdir : {'x', 'y', 'z'} + When plotting 2D data, the direction to use as z ('x', 'y' + or 'z'); defaults to 'z'. + **kwargs + Other arguments are forwarded to `matplotlib.axes.Axes.stem`. + """ + had_data = self.has_data() + + if 'xs' in kwargs: + raise TypeError( + "stem() for multiple values for argument 'x'") + if 'ys' in kwargs: + raise TypeError( + "stem() for multiple values for argument 'y'") + + xs = np.broadcast_to(xs, len(zs)) + ys = np.broadcast_to(ys, len(zs)) + + lines = super().stem(xs, ys, *args, **kwargs) + for line in lines: + art3d.line_2d_to_3d(line, zs=zs, zdir=zdir) + + xs, ys, zs = art3d.juggle_axes(xs, ys, zs, zdir) + self.auto_scale_xyz(xs, ys, zs, had_data) + return lines + + stem3D = stem + def plot(self, xs, ys, *args, zdir='z', **kwargs): """ Plot 2D or 3D data.