Implemented 3D stem plots with image comparison tests.

matplotlib · srihitha09 · Apr 5, 2016 · Apr 22, 2016 · Apr 22, 2016 · Apr 22, 2016
commit 855f28a98f1aaa001d2e5cf3da9cdbd0fca7ba74
diff --git a/examples/mplot3d/stem3_demo.py b/examples/mplot3d/stem3_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.stem3(x,y,z)
+
+plt.show()
diff --git a/examples/mplot3d/stem3_demo2.py b/examples/mplot3d/stem3_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.stem3(x,y,z,'-.',rotate='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
diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_rotate_along_x.png b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_rotate_along_x.png
diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_rotate_along_y.png b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_rotate_along_y.png
diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_rotate_along_z.png b/lib/matplotlib/tests/baseline_images/test_3dstem/2plane_rotate_along_z.png
diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_default.png b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_default.png
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
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
diff --git a/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_z.png b/lib/matplotlib/tests/baseline_images/test_3dstem/3dstemplot_rotate_along_z.png
diff --git a/lib/matplotlib/tests/test_3dstem.py b/lib/matplotlib/tests/test_3dstem.py
@@ -0,0 +1,104 @@
+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.stem3(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.stem3(x,y,z,rotate='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.stem3(x,y,z,rotate='y')
+
+@image_comparison(baseline_images=['3dstemplot_rotate_along_z'],
+                  extensions=['png'])
+
+def test_3dstemplot_rotate_along_z():
+    fig4 = plt.figure()
+    ax4 = fig4.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 = ax4.stem3(x,y,z,rotate='z')
+
+
+@image_comparison(baseline_images=['2plane_default'],
+                  extensions=['png'])    
+def test_2plane_default():
+    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.stem3(x,y,z,'-.',rotate='')
+    plt.setp(markerline, 'markerfacecolor', 'b')
+    plt.setp(baseline, 'color', 'r', 'linewidth', 1)
+
+@image_comparison(baseline_images=['2plane_rotate_along_x'],
+                  extensions=['png'])    
+def test_2plane_rotate_along_x():
+    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.stem3(x,y,z,'-.',rotate='x')
+    plt.setp(markerline, 'markerfacecolor', 'b')
+    plt.setp(baseline, 'color', 'r', 'linewidth', 1)
+
+@image_comparison(baseline_images=['2plane_rotate_along_y'],
+                  extensions=['png'])    
+def test_2plane_rotate_along_y():
+    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.stem3(x,y,z,'-.',rotate='y')
+    plt.setp(markerline, 'markerfacecolor', 'b')
+    plt.setp(baseline, 'color', 'r', 'linewidth', 1)
+
+@image_comparison(baseline_images=['2plane_rotate_along_z'],
+                  extensions=['png'])    
+def test_2plane_rotate_along_z():
+    fig8 = plt.figure()
+    ax8 = fig8.gca(projection='3d')
+    x = np.linspace(-np.pi/2, np.pi/2, 40)
+    y = [1]*len(x)
+    z = np.cos(x)
+    markerline, stemlines, baseline = ax8.stem3(x,y,z,'-.',rotate='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
@@ -2641,6 +2641,137 @@ def calc_arrow(uvw, angle=15):
 
     quiver3D = quiver
 
+    def stem3(self, *args, **kwargs):
+        """
+        Create a 3D stem plot.
+
+        Call signatures::
+          stem3(x, y, z, linefmt='b-', markerfmt='bo', basefmt='r-', 
+                rotate={'x'|'y'|'z'})
+
+        By default, stem3 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 rotate value is provided, then it is set to default and no rotation occurs.
+
+        Return value is a tuple (*markerline*, *stemlines*, *baseline*).
+
+        .. seealso::
+            This `document <http://www.mathworks.com/help/techdoc/ref/stem3.html>`_
+            for details.
+
+        **Example:**
+        .. plot:: /mplot3d/stem3_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, zs = [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:
+            rotate = kwargs.pop('rotate', args[3])
+        except IndexError:
+            rotate = kwargs.pop('rotate', 'default')
+
+        bottom = kwargs.pop('bottom', None)
+        label = kwargs.pop('label', None)
+        zdir = kwargs.pop('zdir','z')
+
+        if bottom is None:
+            bottom = 0
+
+        stemlines = []
+        #plot the stemlines based on the value of rotate
+        for thisx, thisy, thisz in zip(x, y, zs):
+
+            if rotate=='y':
+                l, = Axes.plot(self, [bottom, thisz], [thisy, thisy], linefmt,
+                                label="_nolegend_")
+                art3d.line_2d_to_3d(l, [thisx, thisx], zdir=zdir)
+            elif rotate=='x':
+                l, = Axes.plot(self, [thisx, thisx], [bottom, thisz], linefmt,
+                                    label="_nolegend_")
+                art3d.line_2d_to_3d(l, [thisy, thisy], zdir=zdir)
+            elif rotate=='z':
+                l, = Axes.plot(self, [thisy, thisy], [thisx, thisx], linefmt,
+                                    label="_nolegend_")
+                art3d.line_2d_to_3d(l, [bottom, thisz], zdir=zdir)
+            else:
+                l, = Axes.plot(self, [thisx, thisx], [thisy, thisy], linefmt,
+                                label="_nolegend_")
+                art3d.line_2d_to_3d(l, [bottom, thisz], zdir=zdir)            
+
+            stemlines.append(l)
+
+        #plot the baseline in the appropriate plane
+        for i in range(len(x)-1):
+
+            if rotate=='y':
+                baseline, = Axes.plot(self, [bottom, bottom], [y[i], y[i+1]],
+                                      basefmt, label="_nolegend_")
+                art3d.line_2d_to_3d(baseline, [x[i], x[i+1]], zdir=zdir)
+
+            elif rotate=='x':
+                baseline, = Axes.plot(self, [x[i], x[i+1]], [bottom, bottom],
+                                      basefmt, label="_nolegend_")
+                art3d.line_2d_to_3d(baseline, [y[i], y[i+1]], zdir=zdir)
+
+            elif rotate=='z':
+                baseline, = Axes.plot(self, [y[i], y[i+1]],[x[i], x[i+1]],
+                                      basefmt, label="_nolegend_")
+                art3d.line_2d_to_3d(baseline, [bottom, bottom], zdir=zdir)  
+
+            else:
+                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=zdir)            
+
+        #swap x,y,zs arrays to plot markerlines and auto scale        
+        if rotate=='y':  
+            x, y, zs = zs, y, x
+        elif rotate=='x':
+            x, y, zs = x, zs, y
+        elif rotate=='z':
+            x, y, zs = y, x, zs
+
+        markerline, = Axes.plot(self, x, y, markerfmt, label="_nolegend_")
+        art3d.line_2d_to_3d(markerline, zs, zdir=zdir)                
+        self.hold(remember_hold)
+
+        stem_container = StemContainer((markerline, stemlines, baseline),
+                                       label=label)
+        self.add_container(stem_container)
+
+        self.auto_scale_xyz(x, y, zs, had_data)
+
+        return stem_container
+
 
 def get_test_data(delta=0.05):
     '''