Update Axes3D.tricontour for custom triangulations

dmcdougall · dmcdougall · commit 76ac5c945c51 · 2013-06-21T23:58:04.000-05:00
An example has also been added to illustrate the additional flexibility
possible with providing a user-defined triangulation.
diff --git a/examples/mplot3d/tricontour3d_demo.py b/examples/mplot3d/tricontour3d_demo.py
@@ -0,0 +1,36 @@
+"""
+Contour plots of unstructured triangular grids.
+"""
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.tri as tri
+import numpy as np
+import math
+
+# First create the x and y coordinates of the points.
+n_angles = 48
+n_radii = 8
+min_radius = 0.25
+radii = np.linspace(min_radius, 0.95, n_radii)
+
+angles = np.linspace(0, 2*math.pi, n_angles, endpoint=False)
+angles = np.repeat(angles[...,np.newaxis], n_radii, axis=1)
+angles[:,1::2] += math.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 a custom triangulation
+triang = tri.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)
+
+plt.figure()
+plt.gca(projection='3d')
+plt.tricontour(triang, z)
+plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1980,8 +1980,8 @@ def contour(self, X, Y, Z, *args, **kwargs):
 
     contour3D = contour
 
-    def tricontour(self, X, Y, Z, *args, **kwargs):
-        '''
+    def tricontour(self, *args, **kwargs):
+        """
         Create a 3D contour plot.
 
         ==========  ================================================
@@ -2001,8 +2001,9 @@ def tricontour(self, X, Y, Z, *args, **kwargs):
 
         Returns a :class:`~matplotlib.axes.Axes.contour`
 
-        .. versionadded:: 1.1.0
-        '''
+        .. versionchanged:: 1.3.0
+            Added support for custom triangulations
+        """
 
         extend3d = kwargs.pop('extend3d', False)
         stride = kwargs.pop('stride', 5)
@@ -2011,8 +2012,19 @@ def tricontour(self, X, Y, Z, *args, **kwargs):
 
         had_data = self.has_data()
 
+        tri, args, kwargs = Triangulation.get_from_args_and_kwargs(
+                *args, **kwargs)
+        X = tri.x
+        Y = tri.y
+        Z = args[0]
+
+        # We do this so Z doesn't get passed as an arg to Axes.tricontour
+        args = args[1:]
+
         jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir)
-        cset = Axes.tricontour(self, jX, jY, jZ, *args, **kwargs)
+        tri = Triangulation(jX, jY, tri.triangles, tri.mask)
+
+        cset = Axes.tricontour(self, tri, jZ, *args, **kwargs)
         self.add_contour_set(cset, extend3d, stride, zdir, offset)
 
         self.auto_scale_xyz(X, Y, Z, had_data)
