diff --git a/doc/users/whats_new.rst b/doc/users/whats_new.rst index ab6ab4f39b71..c7ed163c946c 100644 --- a/doc/users/whats_new.rst +++ b/doc/users/whats_new.rst @@ -598,6 +598,14 @@ Example:: figure('figure') fignum_exists('figure') #true +Addition of 3D Stem Plots for mplot3d toolkit +````````````````````````````````````````````` +A team of students in the *Engineering Large Software Systems* course, taught by +Professor Anya Tafliovich at the University of Toronto (Scarborough) during the Winter 2016 term, +implemented a basic version of 3D Stem Plots as a part of a term project. +This feature is documented in :func:`~mpl_toolkits.mplot3d.Axes3D.stem`. +The team consists of 5 students: Kanwar Gill, Alvee Jamal, Samriddhi Kaushik, Jessy Liang, and Srihitha Maryada. + ToolManager ----------- diff --git a/examples/mplot3d/stem3d_demo.py b/examples/mplot3d/stem3d_demo.py new file mode 100644 index 000000000000..68f8a3f29579 --- /dev/null +++ b/examples/mplot3d/stem3d_demo.py @@ -0,0 +1,13 @@ +from mpl_toolkits.mplot3d import Axes3D +import matplotlib.pyplot as plt +import numpy as np + +fig = plt.figure() +ax = fig.gca(projection='3d') +theta = np.linspace(0, 2*np.pi) +x = np.cos(theta) +y = np.sin(theta) +z = theta +markerline, stemlines, baseline = ax.stem(x, y, z) + +plt.show() diff --git a/examples/mplot3d/stem3d_demo2.py b/examples/mplot3d/stem3d_demo2.py new file mode 100644 index 000000000000..f24452aed79a --- /dev/null +++ b/examples/mplot3d/stem3d_demo2.py @@ -0,0 +1,14 @@ +from mpl_toolkits.mplot3d import Axes3D +import matplotlib.pyplot as plt +import numpy as np + +fig = plt.figure() +ax = fig.gca(projection='3d') +x = np.linspace(-np.pi/2, np.pi/2, 40) +y = [1]*len(x) +z = np.cos(x) +markerline, stemlines, baseline = ax.stem(x, y, z, '-.', zdir='-y') +plt.setp(markerline, 'markerfacecolor', 'b') +plt.setp(baseline, 'color', 'r', 'linewidth', 1) + +plt.show() diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_default.png b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_default.png new file mode 100644 index 000000000000..6df6686f2a09 Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_default.png differ diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_x.png b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_x.png new file mode 100644 index 000000000000..e8125628a162 Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_x.png differ diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_y.png b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_y.png new file mode 100644 index 000000000000..228cf19d065f Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_y.png differ diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_z.png b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_z.png new file mode 100644 index 000000000000..6df6686f2a09 Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_zdir_z.png differ diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_default.png b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_default.png new file mode 100644 index 000000000000..5afc81e19ed6 Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_default.png differ diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_x.png b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_x.png new file mode 100644 index 000000000000..bd770c44cbbf Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_x.png differ diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_y.png b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_y.png new file mode 100644 index 000000000000..b33a99eb6daf Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_y.png differ diff --git a/lib/matplotlib/tests/test_3dstem.py b/lib/matplotlib/tests/test_3dstem.py new file mode 100644 index 000000000000..232fa126c0f5 --- /dev/null +++ b/lib/matplotlib/tests/test_3dstem.py @@ -0,0 +1,93 @@ +import numpy as np +import matplotlib as mpl +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import Axes3D +from matplotlib.testing.decorators import image_comparison + + +@image_comparison(baseline_images=['3dstemplot_default'], + extensions=['png']) +def test_3dstemplot_default(): + fig = plt.figure() + ax = fig.gca(projection='3d') + theta = np.linspace(0, 2*np.pi) + x = np.cos(theta) + y = np.sin(theta) + z = np.power(x, 2) + markerline, stemlines, baseline = ax.stem(x, y, z) + + +@image_comparison(baseline_images=['3dstemplot_rotate_along_x'], + extensions=['png']) +def test_3dstemplot_rotate_along_x(): + fig2 = plt.figure() + ax2 = fig2.gca(projection='3d') + theta = np.linspace(0, 2*np.pi) + x = np.cos(theta) + y = np.sin(theta) + z = np.power(x, 2) + markerline, stemlines, baseline = ax2.stem(x, y, z, zdir='-x') + + +@image_comparison(baseline_images=['3dstemplot_rotate_along_y'], + extensions=['png']) +def test_3dstemplot_rotate_along_y(): + fig3 = plt.figure() + ax3 = fig3.gca(projection='3d') + theta = np.linspace(0, 2*np.pi) + x = np.cos(theta) + y = np.sin(theta) + z = np.power(x, 2) + markerline, stemlines, baseline = ax3.stem(x, y, z, zdir='-y') + + +@image_comparison(baseline_images=['2plane_default'], + extensions=['png']) +def test_2plane_default(): + fig4 = plt.figure() + ax4 = fig4.gca(projection='3d') + x = np.linspace(-np.pi/2, np.pi/2, 40) + y = [1]*len(x) + z = np.cos(x) + markerline, stemlines, baseline = ax4.stem(x, y, z, '-.') + plt.setp(markerline, 'markerfacecolor', 'b') + plt.setp(baseline, 'color', 'r', 'linewidth', 1) + + +@image_comparison(baseline_images=['2plane_zdir_x'], + extensions=['png']) +def test_2plane_zdir_x(): + fig5 = plt.figure() + ax5 = fig5.gca(projection='3d') + x = np.linspace(-np.pi/2, np.pi/2, 40) + y = [1]*len(x) + z = np.cos(x) + markerline, stemlines, baseline = ax5.stem(x, y, z, '-.', zdir='x') + plt.setp(markerline, 'markerfacecolor', 'b') + plt.setp(baseline, 'color', 'r', 'linewidth', 1) + + +@image_comparison(baseline_images=['2plane_zdir_y'], + extensions=['png']) +def test_3dstemplot_zdir_y(): + fig6 = plt.figure() + ax6 = fig6.gca(projection='3d') + x = np.linspace(-np.pi/2, np.pi/2, 40) + y = [1]*len(x) + z = np.cos(x) + markerline, stemlines, baseline = ax6.stem(x, y, z, '-.', zdir='y') + plt.setp(markerline, 'markerfacecolor', 'b') + plt.setp(baseline, 'color', 'r', 'linewidth', 1) + + +@image_comparison(baseline_images=['2plane_zdir_z'], + extensions=['png']) +def test_3dstemplot_zdir_z(): + fig7 = plt.figure() + ax7 = fig7.gca(projection='3d') + x = np.linspace(-np.pi/2, np.pi/2, 40) + y = [1]*len(x) + z = np.cos(x) + markerline, stemlines, baseline = ax7.stem(x, y, z, '-.', zdir='z') + plt.setp(markerline, 'markerfacecolor', 'b') + plt.setp(baseline, 'color', 'r', 'linewidth', 1) diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py index 1e2fc5b19829..99e37e704944 100755 --- a/lib/mpl_toolkits/mplot3d/axes3d.py +++ b/lib/mpl_toolkits/mplot3d/axes3d.py @@ -2641,6 +2641,103 @@ def calc_arrow(uvw, angle=15): quiver3D = quiver + def stem(self, *args, **kwargs): + """ + Create a 3D stem plot. + + Call signatures:: + stem3(x, y, z, linefmt='b-', markerfmt='bo', basefmt='r-', + zdir={'x'|'y'|'z'|'-x'|'-y'}) + + By default, stem plots vertical lines (using *linefmt*) in the z direction + at each *x* and *y* location from the baseline to *z*, and places a marker there + using *markerfmt*. By default, the baseline at the xy-plane is plotted using + *basefmt*. + + *x*, *y, and *z* values are required and must be arrays of the same length. + + If no zdir value is provided, then it is set to default and no rotation occurs. + + Return value is a tuple (*markerline*, *stemlines*, *baseline*). + + .. seealso:: + This `document `_ + for details. + + **Example:** + .. plot:: /mplot3d/stem3d_demo.py + """ + + from matplotlib.container import StemContainer + + had_data = self.has_data() + + remember_hold = self._hold + if not self._hold: + self.cla() + self.hold(True) + + # Extract the required values + x, y, z = [np.asarray(i) for i in args[:3]] + args = args[3:] + + # Popping some defaults + try: + linefmt = kwargs.pop('linefmt', args[0]) + except IndexError: + linefmt = kwargs.pop('linefmt', 'b-') + try: + markerfmt = kwargs.pop('markerfmt', args[1]) + except IndexError: + markerfmt = kwargs.pop('markerfmt', 'bo') + try: + basefmt = kwargs.pop('basefmt', args[2]) + except IndexError: + basefmt = kwargs.pop('basefmt', 'r-') + try: + zdir = kwargs.pop('zdir', args[3]) + except IndexError: + zdir = kwargs.pop('zdir', 'z') + + bottom = kwargs.pop('bottom', None) + label = kwargs.pop('label', None) + + if bottom is None: + bottom = 0 + + stemlines = [] + jx, jy, jz = art3d.juggle_axes(x, y, z, zdir) + + #plot the baseline in the appropriate plane + for i in range(len(x)-1): + + baseline, = Axes.plot(self, [x[i], x[i+1]], [y[i], y[i+1]], + basefmt, label="_nolegend_") + art3d.line_2d_to_3d(baseline, [bottom, bottom], zdir) + + #plot the stemlines based on the value of rotate + for thisx, thisy, thisz in zip(x, y, z): + + l, = Axes.plot(self, [thisx, thisx], [thisy, thisy], linefmt, + label="_nolegend_") + art3d.line_2d_to_3d(l, [bottom, thisz], zdir) + + stemlines.append(l) + + markerline, = Axes.plot(self, x, y, markerfmt, label="_nolegend_") + art3d.line_2d_to_3d(markerline, z, zdir) + self.hold(remember_hold) + + stem_container = StemContainer((markerline, stemlines, baseline), + label=label) + self.add_container(stem_container) + + self.auto_scale_xyz(jx, jy, jz, had_data) + + return stem_container + + stem3D = stem + def get_test_data(delta=0.05): '''