Add a generic PatchCollection class.

mdboom · mdboom · commit bf8f8cb058a4 · 2008-06-17T18:02:08.000Z
svn path=/trunk/matplotlib/; revision=5583
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,6 @@
+2008-06-17 Add a generic PatchCollection class that can contain any
+           kind of patch. - MGD
+
 2008-06-13 Change pie chart label alignment to avoid having labels
            overwrite the pie - MGD
 
diff --git a/examples/api/patch_collection.py b/examples/api/patch_collection.py
@@ -0,0 +1,38 @@
+import matplotlib
+from matplotlib.patches import Circle, Wedge, Polygon
+from matplotlib.collections import PatchCollection
+import pylab
+
+fig=pylab.figure()
+ax=fig.add_subplot(111)
+
+resolution = 50 # the number of vertices
+N = 3
+x       = pylab.rand(N)
+y       = pylab.rand(N)
+radii   = 0.1*pylab.rand(N)
+patches = []
+for x1,y1,r in zip(x, y, radii):
+    circle = Circle((x1,y1), r)
+    patches.append(circle)
+
+x       = pylab.rand(N)
+y       = pylab.rand(N)
+radii   = 0.1*pylab.rand(N)
+theta1  = 360.0*pylab.rand(N)
+theta2  = 360.0*pylab.rand(N)
+for x1,y1,r,t1,t2 in zip(x, y, radii, theta1, theta2):
+    wedge = Wedge((x1,y1), r, t1, t2)
+    patches.append(wedge)
+
+for i in range(N):
+    polygon = Polygon(pylab.rand(N,2), True)
+    patches.append(polygon)
+
+colors = 100*pylab.rand(len(patches))
+p = PatchCollection(patches, cmap=matplotlib.cm.jet)
+p.set_array(pylab.array(colors))
+ax.add_collection(p)
+pylab.colorbar(p)
+
+pylab.show()
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -775,6 +775,67 @@ def get_color(self):
         return self._edgecolors
     get_colors = get_color  # for compatibility with old versions
 
+class PatchCollection(Collection):
+    """
+    A generic collection of patches.
+
+    This makes it easier to assign a color map to a heterogeneous
+    collection of patches.
+
+    This also may improve plotting speed, since PatchCollection will
+    draw faster than a large number of patches.
+    """
+
+    def __init__(self, patches, match_original=False, **kwargs):
+        """
+        *patches*
+            a sequence of Patch objects.  This list may include
+            a heterogeneous assortment of different patch types.
+
+        *match_original* If True, use the colors and linewidths of the
+            original patches.  If False, new colors may be assigned by
+            providing the standard collection arguments, facecolor,
+            edgecolor, linewidths, norm or cmap.
+
+        If any of *edgecolors*, *facecolors*, *linewidths*,
+        *antialiaseds* are None, they default to their
+        :data:`matplotlib.rcParams` patch setting, in sequence form.
+
+        The use of :class:`~matplotlib.cm.ScalarMappable` is optional.
+        If the :class:`~matplotlib.cm.ScalarMappable` matrix _A is not
+        None (ie a call to set_array has been made), at draw time a
+        call to scalar mappable will be made to set the face colors.
+        """
+
+        if match_original:
+            def determine_facecolor(patch):
+                if patch.fill():
+                    return patch.get_facecolor()
+                return [0, 0, 0, 0]
+
+            facecolors   = [determine_facecolor(p) for p in patches]
+            edgecolors   = [p.get_edgecolor() for p in patches]
+            linewidths   = [p.get_linewidths() for p in patches]
+            antialiaseds = [p.get_antialiased() for p in patches]
+
+            Collection.__init__(
+                self,
+                edgecolors=edgecolors,
+                facecolors=facecolors,
+                linewidths=linewidths,
+                linestyles='solid',
+                antialiaseds = antialiaseds)
+        else:
+            Collection.__init__(self, **kwargs)
+
+        paths        = [p.get_transform().transform_path(p.get_path())
+                        for p in patches]
+
+        self._paths = paths
+
+    def get_paths(self):
+        return self._paths
+
 
 artist.kwdocd['Collection'] = patchstr = artist.kwdoc(Collection)
 for k in ('QuadMesh', 'PolyCollection', 'BrokenBarHCollection', 'RegularPolyCollection',
