Add 'gouraud' shading option to tripcolor.

bfroehle · bfroehle · commit 2f137020ced7 · 2011-10-27T14:49:01.000-07:00
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1129,6 +1129,74 @@ def set_paths(self, patches):
                         for p in patches]
         self._paths = paths
 
+class TriMesh(Collection):
+    """
+    Class for the efficient drawing of a triangular mesh using
+    Gouraud shading.
+
+    A triangular mesh is a :class:`~matplotlib.tri.Triangulation`
+    object.
+    """
+    def __init__(self, triangulation, **kwargs):
+        Collection.__init__(self, **kwargs)
+        self._triangulation = triangulation;
+        self._shading = 'gouraud'
+        self._is_filled = True
+
+        self._bbox = transforms.Bbox.unit()
+
+        # Unfortunately this requires a copy, unless Triangulation
+        # was rewritten.
+        xy = np.hstack((triangulation.x.reshape(-1,1),
+                        triangulation.y.reshape(-1,1)))
+        self._bbox.update_from_data_xy(xy)
+
+    def get_paths(self):
+        if self._paths is None:
+            self.set_paths()
+        return self._paths
+
+    def set_paths(self):
+        self._paths = self.convert_mesh_to_paths(self._triangulation)
+
+    @staticmethod
+    def convert_mesh_to_paths(tri):
+        """
+        Converts a given mesh into a sequence of
+        :class:`matplotlib.path.Path` objects for easier rendering by
+        backends that do not directly support meshes.
+
+        This function is primarily of use to backend implementers.
+        """
+        Path = mpath.Path
+        triangles = tri.get_masked_triangles()
+        verts = np.concatenate((tri.x[triangles][...,np.newaxis],
+                                tri.y[triangles][...,np.newaxis]), axis=2)
+        return [Path(x) for x in verts]
+
+    @allow_rasterization
+    def draw(self, renderer):
+        if not self.get_visible(): return
+        renderer.open_group(self.__class__.__name__)
+        transform = self.get_transform()
+
+        # Get a list of triangles and the color at each vertex.
+        tri = self._triangulation
+        triangles = tri.get_masked_triangles()
+
+        verts = np.concatenate((tri.x[triangles][...,np.newaxis],
+                                tri.y[triangles][...,np.newaxis]), axis=2)
+
+        self.update_scalarmappable()
+        colors = self._facecolors[triangles];
+
+        gc = renderer.new_gc()
+        self._set_gc_clip(gc)
+        gc.set_linewidth(self.get_linewidth()[0])
+        renderer.draw_gouraud_triangles(gc, verts, colors, transform.frozen())
+        gc.restore()
+        renderer.close_group(self.__class__.__name__)
+
 
 class QuadMesh(Collection):
     """
@@ -1315,7 +1383,7 @@ def draw(self, renderer):
 
 
 patchstr = artist.kwdoc(Collection)
-for k in ('QuadMesh', 'PolyCollection', 'BrokenBarHCollection',
+for k in ('QuadMesh', 'TriMesh', 'PolyCollection', 'BrokenBarHCollection',
            'RegularPolyCollection', 'PathCollection',
           'StarPolygonCollection', 'PatchCollection',
           'CircleCollection', 'Collection',):
diff --git a/lib/matplotlib/tri/tripcolor.py b/lib/matplotlib/tri/tripcolor.py
@@ -1,4 +1,4 @@
-from matplotlib.collections import PolyCollection
+from matplotlib.collections import PolyCollection, TriMesh
 from matplotlib.colors import Normalize
 from matplotlib.tri.triangulation import Triangulation
 import numpy as np
@@ -28,8 +28,11 @@ def tripcolor(ax, *args, **kwargs):
     possibilities.
 
     The next argument must be *C*, the array of color values, one per
-    point in the triangulation.  The colors used for each triangle
-    are from the mean C of the triangle's three points.
+    point in the triangulation.
+
+    *shading* may be 'flat', 'faceted' or 'gouraud'. If *shading* is
+    'flat' or 'faceted', the colors used for each triangle are from
+    the mean C of the triangle's three points.
 
     The remaining kwargs are the same as for
     :meth:`~matplotlib.axes.Axes.pcolor`.
@@ -52,29 +55,30 @@ def tripcolor(ax, *args, **kwargs):
     y = tri.y
     triangles = tri.get_masked_triangles()
 
-    # Vertices of triangles.
-    verts = np.concatenate((x[triangles][...,np.newaxis],
-                            y[triangles][...,np.newaxis]), axis=2)
-
     C = np.asarray(args[0])
     if C.shape != x.shape:
         raise ValueError('C array must have same length as triangulation x and'
                          ' y arrays')
 
-    # Color values, one per triangle, mean of the 3 vertex color values.
-    C = C[triangles].mean(axis=1)
-
-    if shading == 'faceted':
-        edgecolors = (0,0,0,1),
-        linewidths = (0.25,)
+    if shading == 'gouraud':
+        collection = TriMesh(tri, **kwargs)
     else:
-        edgecolors = 'face'
-        linewidths = (1.0,)
-    kwargs.setdefault('edgecolors', edgecolors)
-    kwargs.setdefault('antialiaseds', (0,))
-    kwargs.setdefault('linewidths', linewidths)
-
-    collection = PolyCollection(verts, **kwargs)
+        if shading == 'faceted':
+            edgecolors = (0,0,0,1),
+            linewidths = (0.25,)
+        else:
+            edgecolors = 'face'
+            linewidths = (1.0,)
+        kwargs.setdefault('edgecolors', edgecolors)
+        kwargs.setdefault('antialiaseds', (0,))
+        kwargs.setdefault('linewidths', linewidths)
+
+        # Vertices of triangles.
+        verts = np.concatenate((x[triangles][...,np.newaxis],
+                                y[triangles][...,np.newaxis]), axis=2)
+        # Color values, one per triangle, mean of the 3 vertex color values.
+        C = C[triangles].mean(axis=1)
+        collection = PolyCollection(verts, **kwargs)
 
     collection.set_alpha(alpha)
     collection.set_array(C)
