From 091f562ca2b918b97872de0e80c8b4a459c8624d Mon Sep 17 00:00:00 2001 From: harshanavkis Date: Wed, 11 Sep 2019 21:27:50 +0530 Subject: [PATCH 1/2] Feature: stem3D to generate 3D stem plots This commit introduces the stem3D function. It is similar to the stem3 function from Matab. Resolves Issue #15066. Signed-off-by: harshanavkis --- lib/mpl_toolkits/mplot3d/axes3d.py | 43 ++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) 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. From 5684348db5930f43c4c84c4def8a8cee0db7ff6e Mon Sep 17 00:00:00 2001 From: harshanavkis Date: Fri, 13 Sep 2019 22:19:01 +0530 Subject: [PATCH 2/2] Example: Add stem3d example This commit provides an example on generating 3d stem plots. Resolves Issue #15066. Signed-off-by: harshanavkis --- examples/mplot3d/stem3d.py | 41 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 examples/mplot3d/stem3d.py 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()