Make plot_trisurf support custom triangulations

dmcdougall · dmcdougall · commit 17159dc85a3c · 2012-09-22T11:54:13.000+01:00
This brings plot_trisurf more inline with the 2d tri* method
equivalents. An example has also been added to: a) reflect the new
support; and b) illustrate the additional flexibility this change
enables.
diff --git a/examples/mplot3d/trisurf3d_demo2.py b/examples/mplot3d/trisurf3d_demo2.py
@@ -0,0 +1,19 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.tri as mtri
+
+u = (np.linspace(0, 2.0 * np.pi, endpoint=True, num=50) * np.ones((10, 1))).flatten()
+v = np.repeat(np.linspace(-0.5, 0.5, endpoint=True, num=10), repeats=50).flatten()
+
+x = (1 + 0.5 * v * np.cos(u / 2.0)) * np.cos(u)
+y = (1 + 0.5 * v * np.cos(u / 2.0)) * np.sin(u)
+z = 0.5 * v * np.sin(u / 2.0)
+
+tri = mtri.Triangulation(u, v)
+
+fig = plt.figure()
+ax = fig.add_subplot(1, 1, 1, projection='3d')
+ax.plot_trisurf(x, y, z, triangles=tri.triangles)
+ax.set_zlim(-1, 1)
+plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1562,7 +1562,7 @@ def plot_wireframe(self, X, Y, Z, *args, **kwargs):
 
         return linec
 
-    def plot_trisurf(self, X, Y, Z, *args, **kwargs):
+    def plot_trisurf(self, *args, **kwargs):
         """
         ============= ================================================
         Argument      Description
@@ -1596,15 +1596,14 @@ def plot_trisurf(self, X, Y, Z, *args, **kwargs):
         shade = kwargs.pop('shade', cmap is None)
         lightsource = kwargs.pop('lightsource', None)
 
+        tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs)
+        z = np.asarray(args[0])
+
         # TODO: Support masked triangulations
-        tri = Triangulation(X, Y)
-        x = tri.x
-        y = tri.y
         triangles = tri.triangles
-
-        xt = x[triangles][...,np.newaxis]
-        yt = y[triangles][...,np.newaxis]
-        zt = np.array(Z)[triangles][...,np.newaxis]
+        xt = tri.x[triangles][...,np.newaxis]
+        yt = tri.y[triangles][...,np.newaxis]
+        zt = np.array(z)[triangles][...,np.newaxis]
 
         verts = np.concatenate((xt, yt, zt), axis=2)
 
@@ -1649,7 +1648,7 @@ def plot_trisurf(self, X, Y, Z, *args, **kwargs):
             polyc.set_facecolors(colset)
 
         self.add_collection(polyc)
-        self.auto_scale_xyz(X, Y, Z, had_data)
+        self.auto_scale_xyz(tri.x, tri.y, z, had_data)
 
         return polyc
 
