mplot3d: support contours in directions other than z

heeres · heeres · commit 787755e6b483 · 2009-12-10T23:32:41.000Z
svn path=/trunk/matplotlib/; revision=8019
diff --git a/examples/mplot3d/contour3d_demo3.py b/examples/mplot3d/contour3d_demo3.py
@@ -0,0 +1,20 @@
+from mpl_toolkits.mplot3d import axes3d
+import matplotlib.pyplot as plt
+
+fig = plt.figure()
+ax = axes3d.Axes3D(fig)
+X, Y, Z = axes3d.get_test_data(0.05)
+ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
+cset = ax.contour(X, Y, Z, zdir='z', offset=-100)
+cset = ax.contour(X, Y, Z, zdir='x', offset=-40)
+cset = ax.contour(X, Y, Z, zdir='y', offset=40)
+
+ax.set_xlabel('X')
+ax.set_xlim3d(-40, 40)
+ax.set_ylabel('Y')
+ax.set_ylim3d(-40, 40)
+ax.set_zlabel('Z')
+ax.set_zlim3d(-100, 100)
+
+plt.show()
+
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -437,9 +437,15 @@ def juggle_axes(xs, ys, zs, zdir):
     Reorder coordinates so that zdir
     """
     if zdir == 'x':
+        return ys, zs, xs
+    elif zdir == '-x':
         return zs, xs, ys
+
     elif zdir == 'y':
-        return xs, zs, ys
+        return zs, xs, ys
+    elif zdir == '-y':
+        return ys, zs, xs
+
     else:
         return xs, ys, zs
 
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -839,10 +839,11 @@ def contour(self, X, Y, Z, levels=10, **kwargs):
         ==========  ================================================
         *X*, *Y*,   Data values as numpy.arrays
         *Z*
-        *levels*    Number of levels to use, defaults to 10. Can
-                    also be a tuple of specific levels.
         *extend3d*  Whether to extend contour in 3D (default: False)
         *stride*    Stride (step size) for extending contour
+        *zdir*      The direction to use: x, y or z (default)
+        *offset*    If specified plot a projection of the contour
+                    lines on this position in plane normal to zdir
         ==========  ================================================
 
         Other keyword arguments are passed on to
@@ -851,16 +852,22 @@ def contour(self, X, Y, Z, levels=10, **kwargs):
 
         extend3d = kwargs.pop('extend3d', False)
         stride = kwargs.pop('stride', 5)
-        nlevels = kwargs.pop('nlevels', 15)
+        zdir = kwargs.pop('zdir', 'z')
+        offset = kwargs.pop('offset', None)
 
         had_data = self.has_data()
-        cset = Axes.contour(self, X, Y, Z, levels, **kwargs)
 
+        jX, jY, jZ = art3d.juggle_axes(X, Y, Z, zdir)
+        cset = Axes.contour(self, jX, jY, jZ, **kwargs)
+
+        zdir = '-' + zdir
         if extend3d:
             self._3d_extend_contour(cset, stride)
         else:
             for z, linec in zip(cset.levels, cset.collections):
-                art3d.line_collection_2d_to_3d(linec, z)
+                if offset is not None:
+                    z = offset
+                art3d.line_collection_2d_to_3d(linec, z, zdir=zdir)
 
         self.auto_scale_xyz(X, Y, Z, had_data)
         return cset
diff --git a/lib/mpl_toolkits/mplot3d/axis3d.py b/lib/mpl_toolkits/mplot3d/axis3d.py
@@ -57,11 +57,11 @@ class Axis(maxis.XAxis):
 
     # Some properties for the axes
     _AXINFO = {
-        'x': {'i': 0, 'tickdir': 1,
+        'x': {'i': 0, 'tickdir': 1, 'juggled': (1, 0, 2),
             'color': (0.95, 0.95, 0.95, 0.5)},
-        'y': {'i': 1, 'tickdir': 0,
+        'y': {'i': 1, 'tickdir': 0, 'juggled': (0, 1, 2),
             'color': (0.90, 0.90, 0.90, 0.5)},
-        'z': {'i': 2, 'tickdir': 0,
+        'z': {'i': 2, 'tickdir': 0, 'juggled': (0, 2, 1),
             'color': (0.925, 0.925, 0.925, 0.5)},
     }
 
@@ -191,7 +191,7 @@ def draw(self, renderer):
         minmax = np.where(highs, maxs, mins)
 
         # Draw main axis line
-        juggled = art3d.juggle_axes(0, 2, 1, self.adir)
+        juggled = info['juggled']
         edgep1 = minmax.copy()
         edgep1[juggled[0]] = get_flip_min_max(edgep1, juggled[0], mins, maxs)
 
