mplot3d: add examples, fix NaN bug

heeres · heeres · commit 678713596351 · 2009-05-28T15:59:12.000Z
svn path=/trunk/matplotlib/; revision=7154
diff --git a/examples/mplot3d/contour.py b/examples/mplot3d/contour.py
@@ -0,0 +1,12 @@
+from mpl_toolkits.mplot3d import axes3d
+import pylab
+import random
+
+fig = pylab.figure()
+ax = axes3d.Axes3D(fig)
+X, Y, Z = axes3d.get_test_data(0.05)
+cset = ax.contour3D(X, Y, Z)
+ax.clabel(cset, fontsize=9, inline=1)
+
+pylab.show()
+
diff --git a/examples/mplot3d/contourf.py b/examples/mplot3d/contourf.py
@@ -0,0 +1,12 @@
+from mpl_toolkits.mplot3d import axes3d
+import pylab
+import random
+
+fig = pylab.figure()
+ax = axes3d.Axes3D(fig)
+X, Y, Z = axes3d.get_test_data(0.05)
+cset = ax.contourf3D(X, Y, Z)
+ax.clabel(cset, fontsize=9, inline=1)
+
+pylab.show()
+
diff --git a/examples/mplot3d/polys.py b/examples/mplot3d/polys.py
@@ -0,0 +1,31 @@
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.collections import PolyCollection
+from matplotlib.colors import colorConverter
+import pylab
+import random
+import numpy as np
+
+fig = pylab.figure()
+ax = Axes3D(fig)
+
+cc = lambda arg: colorConverter.to_rgba(arg, alpha=0.6)
+
+xs = np.arange(0, 10, 0.4)
+verts = []
+zs = [0.0, 1.0, 2.0, 3.0]
+for z in zs:
+    ys = [random.random() for x in xs]
+    ys[0], ys[-1] = 0, 0
+    verts.append(zip(xs, ys))
+
+poly = PolyCollection(verts, facecolors = [cc('r'), cc('g'), cc('b'),
+                                           cc('y')])
+poly.set_alpha(0.7)
+ax.add_collection(poly, zs=zs, dir='y')
+
+ax.set_xlim(0, 10)
+ax.set_ylim(-1, 4)
+ax.set_zlim(0, 1)
+
+pylab.show()
+
diff --git a/examples/mplot3d/scatter.py b/examples/mplot3d/scatter.py
@@ -0,0 +1,21 @@
+from mpl_toolkits.mplot3d import Axes3D
+import pylab
+import random
+
+fig = pylab.figure()
+ax = Axes3D(fig)
+n = 100
+for c, zl, zh in [('r', -50, -25), ('b', -30, -5)]:
+    xs, ys, zs = zip(*
+                   [(random.randrange(23, 32),
+                     random.randrange(100),
+                     random.randrange(zl, zh)
+                     ) for i in range(n)])
+    ax.scatter3D(xs, ys, zs, c=c)
+
+ax.set_xlabel('X Label')
+ax.set_ylabel('Y Label')
+ax.set_zlabel('Z Label')
+
+pylab.show()
+
diff --git a/examples/mplot3d/surface.py b/examples/mplot3d/surface.py
@@ -0,0 +1,16 @@
+from mpl_toolkits.mplot3d import Axes3D
+import pylab
+import random
+import numpy as np
+
+fig = pylab.figure()
+ax = Axes3D(fig)
+X = np.arange(-5, 5, 0.5)
+Y = np.arange(-5, 5, 0.5)
+X, Y = np.meshgrid(X, Y)
+R = np.sqrt(X**2 + Y**2)
+Z = np.sin(R)
+ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color='forestgreen')
+
+pylab.show()
+
diff --git a/examples/mplot3d/wire.py b/examples/mplot3d/wire.py
@@ -0,0 +1,12 @@
+from mpl_toolkits.mplot3d import axes3d
+import pylab
+import random
+import numpy as np
+
+fig = pylab.figure()
+ax = axes3d.Axes3D(fig)
+X, Y, Z = axes3d.get_test_data(0.05)
+ax.plot_wireframe(X, Y, Z, rstride=10, cstride=10)
+
+pylab.show()
+
diff --git a/lib/mpl_toolkits/mplot3d/__init__.py b/lib/mpl_toolkits/mplot3d/__init__.py
@@ -0,0 +1 @@
+from axes3d import Axes3D
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -539,6 +539,9 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
         rstride = kwargs.pop('rstride', 10)
         cstride = kwargs.pop('cstride', 10)
 
+        color = kwargs.pop('color', 'b')
+        color = np.array(colorConverter.to_rgba(color))
+
         polys = []
         boxes = []
         for rs in np.arange(0,rows-1,rstride):
@@ -567,8 +570,10 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
             shade.append(np.dot(n,[-1,-1,0.5]))
             lines.append((box[0],n+box[0]))
 
-        color = np.array([0,0,1,1])
-        norm = Normalize(min(shade),max(shade))
+        shade = np.array(shade)
+        mask = ~np.isnan(shade)
+        norm = Normalize(min(shade[mask]), max(shade[mask]))
+
         colors = [color * (0.5+norm(v)*0.5) for v in shade]
         for c in colors: c[3] = 1
         polyc = art3d.Poly3DCollection(polys, facecolors=colors, *args, **kwargs)
