Merge pull request #1298 from dmcdougall/update_trisurf

ianthomas23 · ianthomas23 · commit 89482b21c858 · 2012-10-02T05:12:23.000-07:00
Update mplot3d trisurf to support custom triangulations.
diff --git a/examples/mplot3d/trisurf3d_demo2.py b/examples/mplot3d/trisurf3d_demo2.py
@@ -0,0 +1,55 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.tri as mtri
+
+# u, v are parameterisation variables
+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()
+
+# This is the Mobius mapping, taking a u, v pair and returning an x, y, z
+# triple
+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)
+
+# Triangulate parameter space to determine the triangles
+tri = mtri.Triangulation(u, v)
+
+fig = plt.figure()
+ax = fig.add_subplot(1, 1, 1, projection='3d')
+
+# The triangles in parameter space determine which x, y, z points are
+# connected by an edge
+ax.plot_trisurf(x, y, z, triangles=tri.triangles, cmap=plt.cm.Spectral)
+
+ax.set_zlim(-1, 1)
+
+# First create the x and y coordinates of the points.
+n_angles = 36
+n_radii = 8
+min_radius = 0.25
+radii = np.linspace(min_radius, 0.95, n_radii)
+
+angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)
+angles = np.repeat(angles[...,np.newaxis], n_radii, axis=1)
+angles[:,1::2] += np.pi/n_angles
+
+x = (radii*np.cos(angles)).flatten()
+y = (radii*np.sin(angles)).flatten()
+z = (np.cos(radii)*np.cos(angles*3.0)).flatten()
+
+# Create the Triangulation; no triangles so Delaunay triangulation created.
+triang = mtri.Triangulation(x, y)
+
+# Mask off unwanted triangles.
+xmid = x[triang.triangles].mean(axis=1)
+ymid = y[triang.triangles].mean(axis=1)
+mask = np.where(xmid*xmid + ymid*ymid < min_radius*min_radius, 1, 0)
+triang.set_mask(mask)
+
+# tripcolor plot.
+fig = plt.figure()
+ax = fig.add_subplot(1, 1, 1, projection='3d')
+ax.plot_trisurf(triang, z, cmap=plt.cm.CMRmap)
+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
@@ -1576,9 +1576,37 @@ def plot_trisurf(self, X, Y, Z, *args, **kwargs):
         *shade*       Whether to shade the facecolors
         ============= ================================================
 
+        The (optional) triangulation can be specified in one of two ways;
+        either::
+
+          plot_trisurf(triangulation, ...)
+
+        where triangulation is a :class:`~matplotlib.tri.Triangulation`
+        object, or::
+
+          plot_trisurf(X, Y, ...)
+          plot_trisurf(X, Y, triangles, ...)
+          plot_trisurf(X, Y, triangles=triangles, ...)
+
+        in which case a Triangulation object will be created.  See
+        :class:`~matplotlib.tri.Triangulation` for a explanation of
+        these possibilities.
+
+        The remaining arguments are::
+
+          plot_trisurf(..., Z)
+
+        where *Z* is the array of values to contour, one per point
+        in the triangulation.
+
         Other arguments are passed on to
         :class:`~mpl_toolkits.mplot3d.art3d.Poly3DCollection`
 
+        **Examples:**
+
+        .. plot:: mpl_examples/mplot3d/trisurf3d_demo.py
+        .. plot:: mpl_examples/mplot3d/trisurf3d_demo2.py
+
         .. versionadded:: 1.2.0
             This plotting function was added for the v1.2.0 release.
         """
@@ -1596,15 +1624,13 @@ def plot_trisurf(self, X, Y, Z, *args, **kwargs):
         shade = kwargs.pop('shade', cmap is None)
         lightsource = kwargs.pop('lightsource', None)
 
-        # TODO: Support masked triangulations
-        tri = Triangulation(X, Y)
-        x = tri.x
-        y = tri.y
-        triangles = tri.triangles
+        tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs)
+        z = np.asarray(args[0])
 
-        xt = x[triangles][...,np.newaxis]
-        yt = y[triangles][...,np.newaxis]
-        zt = np.array(Z)[triangles][...,np.newaxis]
+        triangles = tri.get_masked_triangles()
+        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 +1675,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
 
