Make hexbin much faster by creating a PathCollection that is a single polygon and a number of offsets. This allows the Agg backend to be dealing with far less data, and it allows the vector backends to write the polygon once and "use" it multiple times.

mdboom · mdboom · commit 4dd3de1b580a · 2012-05-22T12:00:45.000-04:00
diff --git a/examples/pylab_examples/hexbin_demo.py b/examples/pylab_examples/hexbin_demo.py
@@ -9,6 +9,7 @@
 import matplotlib.cm as cm
 import  matplotlib.pyplot as plt
 
+np.random.seed(0)
 n = 100000
 x = np.random.standard_normal(n)
 y = 2.0 + 3.0 * x + 4.0 * np.random.standard_normal(n)
@@ -33,4 +34,3 @@
 cb.set_label('log10(N)')
 
 plt.show()
-
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -6175,43 +6175,47 @@ def hexbin(self, x, y, C = None, gridsize = 100, bins = None,
                 lattice1.astype(float).ravel(), lattice2.astype(float).ravel()))
             good_idxs = ~np.isnan(accum)
 
-        px = xmin + sx * np.array([ 0.5, 0.5, 0.0, -0.5, -0.5,  0.0])
-        py = ymin + sy * np.array([-0.5, 0.5, 1.0,  0.5, -0.5, -1.0]) / 3.0
-
-        polygons = np.zeros((6, n, 2), float)
-        polygons[:,:nx1*ny1,0] = np.repeat(np.arange(nx1), ny1)
-        polygons[:,:nx1*ny1,1] = np.tile(np.arange(ny1), nx1)
-        polygons[:,nx1*ny1:,0] = np.repeat(np.arange(nx2) + 0.5, ny2)
-        polygons[:,nx1*ny1:,1] = np.tile(np.arange(ny2), nx2) + 0.5
-
+        px = sx * np.array([ 0.5, 0.5, 0.0, -0.5, -0.5,  0.0])
+        py = sy * np.array([-0.5, 0.5, 1.0,  0.5, -0.5, -1.0]) / 3.0
+
+        offsets = np.zeros((n, 2), float)
+        offsets[:nx1*ny1,0] = np.repeat(np.arange(nx1), ny1)
+        offsets[:nx1*ny1,1] = np.tile(np.arange(ny1), nx1)
+        offsets[nx1*ny1:,0] = np.repeat(np.arange(nx2) + 0.5, ny2)
+        offsets[nx1*ny1:,1] = np.tile(np.arange(ny2), nx2) + 0.5
+        offsets[:,0] *= sx
+        offsets[:,1] *= sy
+        offsets[:,0] += xmin
+        offsets[:,1] += ymin
         # remove accumulation bins with no data
-        polygons = polygons[:,good_idxs,:]
+        offsets = offsets[good_idxs,:]
         accum = accum[good_idxs]
 
-        polygons = np.transpose(polygons, axes=[1,0,2])
-        polygons[:,:,0] *= sx
-        polygons[:,:,1] *= sy
-        polygons[:,:,0] += px
-        polygons[:,:,1] += py
-
         if xscale=='log':
-            polygons[:,:,0] = 10**(polygons[:,:,0])
+            offsets[:,0] = 10**(offsets[:,0])
             xmin = 10**xmin
             xmax = 10**xmax
             self.set_xscale('log')
         if yscale=='log':
-            polygons[:,:,1] = 10**(polygons[:,:,1])
+            offsets[:,1] = 10**(offsets[:,1])
             ymin = 10**ymin
             ymax = 10**ymax
             self.set_yscale('log')
 
+        polygons = np.zeros((6, 2), float)
+        polygons[:,0] = px
+        polygons[:,1] = py
+
         if edgecolors=='none':
             edgecolors = 'face'
+
         collection = mcoll.PolyCollection(
-            polygons,
+            [polygons],
             edgecolors = edgecolors,
             linewidths = linewidths,
-            transOffset = self.transData,
+            offsets = offsets,
+            transOffset = mtransforms.IdentityTransform(),
+            offset_position = "data"
             )
 
         if isinstance(norm, mcolors.LogNorm):
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -186,14 +186,16 @@ def draw_markers(self, gc, marker_path, marker_trans, path, trans, rgbFace=None)
 
     def draw_path_collection(self, gc, master_transform, paths, all_transforms,
                              offsets, offsetTrans, facecolors, edgecolors,
-                             linewidths, linestyles, antialiaseds, urls):
+                             linewidths, linestyles, antialiaseds, urls,
+                             offset_position):
         """
         Draws a collection of paths selecting drawing properties from
         the lists *facecolors*, *edgecolors*, *linewidths*,
         *linestyles* and *antialiaseds*. *offsets* is a list of
         offsets to apply to each of the paths.  The offsets in
         *offsets* are first transformed by *offsetTrans* before being
-        applied.
+        applied.  *offset_position* may be either "screen" or "data"
+        depending on the space that the offsets are in.
 
         This provides a fallback implementation of
         :meth:`draw_path_collection` that makes multiple calls to
@@ -213,8 +215,9 @@ def draw_path_collection(self, gc, master_transform, paths, all_transforms,
             path_ids.append((path, transform))
 
         for xo, yo, path_id, gc0, rgbFace in self._iter_collection(
-            gc, path_ids, offsets, offsetTrans, facecolors, edgecolors,
-            linewidths, linestyles, antialiaseds, urls):
+            gc, master_transform, all_transforms, path_ids, offsets,
+            offsetTrans, facecolors, edgecolors, linewidths, linestyles,
+            antialiaseds, urls, offset_position):
             path, transform = path_id
             transform = transforms.Affine2D(transform.get_matrix()).translate(xo, yo)
             self.draw_path(gc0, path, transform, rgbFace)
@@ -240,7 +243,7 @@ def draw_quad_mesh(self, gc, master_transform, meshWidth, meshHeight,
 
         return self.draw_path_collection(
             gc, master_transform, paths, [], offsets, offsetTrans, facecolors,
-            edgecolors, linewidths, [], [antialiased], [None])
+            edgecolors, linewidths, [], [antialiased], [None], 'screen')
 
     def draw_gouraud_triangle(self, gc, points, colors, transform):
         """
@@ -302,9 +305,10 @@ def _iter_collection_raw_paths(self, master_transform, paths,
                 transform = all_transforms[i % Ntransforms]
             yield path, transform + master_transform
 
-    def _iter_collection(self, gc, path_ids, offsets, offsetTrans, facecolors,
-                         edgecolors, linewidths, linestyles, antialiaseds,
-                         urls):
+    def _iter_collection(self, gc, master_transform, all_transforms,
+                         path_ids, offsets, offsetTrans, facecolors,
+                         edgecolors, linewidths, linestyles,
+                         antialiaseds, urls, offset_position):
         """
         This is a helper method (along with
         :meth:`_iter_collection_raw_paths`) to make it easier to write
@@ -330,6 +334,7 @@ def _iter_collection(self, gc, path_ids, offsets, offsetTrans, facecolors,
         *path_ids*; *gc* is a graphics context and *rgbFace* is a color to
         use for filling the path.
         """
+        Ntransforms = len(all_transforms)
         Npaths      = len(path_ids)
         Noffsets    = len(offsets)
         N           = max(Npaths, Noffsets)
@@ -359,6 +364,15 @@ def _iter_collection(self, gc, path_ids, offsets, offsetTrans, facecolors,
             path_id = path_ids[i % Npaths]
             if Noffsets:
                 xo, yo = toffsets[i % Noffsets]
+                if offset_position == 'data':
+                    if Ntransforms:
+                        transform = all_transforms[i % Ntransforms] + master_transform
+                    else:
+                        transform = master_transform
+                    xo, yo = transform.transform_point((xo, yo))
+                    xp, yp = transform.transform_point((0, 0))
+                    xo = -(xp - xo)
+                    yo = -(yp - yo)
             if Nfacecolors:
                 rgbFace = facecolors[i % Nfacecolors]
             if Nedgecolors:
diff --git a/lib/matplotlib/backends/backend_macosx.py b/lib/matplotlib/backends/backend_macosx.py
@@ -60,7 +60,8 @@ def draw_markers(self, gc, marker_path, marker_trans, path, trans, rgbFace=None)
 
     def draw_path_collection(self, gc, master_transform, paths, all_transforms,
                              offsets, offsetTrans, facecolors, edgecolors,
-                             linewidths, linestyles, antialiaseds, urls):
+                             linewidths, linestyles, antialiaseds, urls,
+                             offset_position):
         cliprect = gc.get_clip_rectangle()
         clippath, clippath_transform = gc.get_clip_path()
         if all_transforms:
diff --git a/lib/matplotlib/backends/backend_pdf.py b/lib/matplotlib/backends/backend_pdf.py
@@ -452,6 +452,8 @@ def __init__(self, filename):
         self.markers = {}
         self.multi_byte_charprocs = {}
 
+        self.paths = []
+
         # The PDF spec recommends to include every procset
         procsets = [ Name(x)
                      for x in "PDF Text ImageB ImageC ImageI".split() ]
@@ -505,9 +507,12 @@ def close(self):
             xobjects[tup[0]] = tup[1]
         for name, value in self.multi_byte_charprocs.iteritems():
             xobjects[name] = value
+        for name, path, trans, ob, join, cap, padding in self.paths:
+            xobjects[name] = ob
         self.writeObject(self.XObjectObject, xobjects)
         self.writeImages()
         self.writeMarkers()
+        self.writePathCollectionTemplates()
         self.writeObject(self.pagesObject,
                          { 'Type': Name('Pages'),
                            'Kids': self.pageList,
@@ -1259,6 +1264,28 @@ def writeMarkers(self):
             self.output(Op.paint_path(False, fillp, strokep))
             self.endStream()
 
+    def pathCollectionObject(self, gc, path, trans, padding):
+        name = Name('P%d' % len(self.paths))
+        ob = self.reserveObject('path %d' % len(self.paths))
+        self.paths.append(
+            (name, path, trans, ob, gc.get_joinstyle(), gc.get_capstyle(), padding))
+        return name
+
+    def writePathCollectionTemplates(self):
+        for (name, path, trans, ob, joinstyle, capstyle, padding) in self.paths:
+            pathops = self.pathOperations(path, trans, simplify=False)
+            bbox = path.get_extents(trans)
+            bbox = bbox.padded(padding)
+            self.beginStream(
+                ob.id, None,
+                {'Type': Name('XObject'), 'Subtype': Name('Form'),
+                 'BBox': list(bbox.extents)})
+            self.output(GraphicsContextPdf.joinstyles[joinstyle], Op.setlinejoin)
+            self.output(GraphicsContextPdf.capstyles[capstyle], Op.setlinecap)
+            self.output(*pathops)
+            self.output(Op.paint_path(False, True, True))
+            self.endStream()
+
     @staticmethod
     def pathOperations(path, transform, clip=None, simplify=None):
         cmds = []
@@ -1466,6 +1493,32 @@ def draw_path(self, gc, path, transform, rgbFace=None):
             rgbFace is None and gc.get_hatch_path() is None)
         self.file.output(self.gc.paint())
 
+    def draw_path_collection(self, gc, master_transform, paths, all_transforms,
+                             offsets, offsetTrans, facecolors, edgecolors,
+                             linewidths, linestyles, antialiaseds, urls,
+                             offset_position):
+
+        padding = np.max(linewidths)
+        path_codes = []
+        for i, (path, transform) in enumerate(self._iter_collection_raw_paths(
+            master_transform, paths, all_transforms)):
+            name = self.file.pathCollectionObject(gc, path, transform, padding)
+            path_codes.append(name)
+
+        output = self.file.output
+        output(Op.gsave)
+        lastx, lasty = 0, 0
+        for xo, yo, path_id, gc0, rgbFace in self._iter_collection(
+            gc, master_transform, all_transforms, path_codes, offsets,
+            offsetTrans, facecolors, edgecolors, linewidths, linestyles,
+            antialiaseds, urls, offset_position):
+
+            self.check_gc(gc0, rgbFace)
+            dx, dy = xo - lastx, yo - lasty
+            output(1, 0, 0, 1, dx, dy, Op.concat_matrix, path_id, Op.use_xobject)
+            lastx, lasty = xo, yo
+        output(Op.grestore)
+
     def draw_markers(self, gc, marker_path, marker_trans, path, trans, rgbFace=None):
         # For simple paths or small numbers of markers, don't bother
         # making an XObject
diff --git a/lib/matplotlib/backends/backend_ps.py b/lib/matplotlib/backends/backend_ps.py
@@ -625,7 +625,8 @@ def draw_markers(self, gc, marker_path, marker_trans, path, trans, rgbFace=None)
 
     def draw_path_collection(self, gc, master_transform, paths, all_transforms,
                              offsets, offsetTrans, facecolors, edgecolors,
-                             linewidths, linestyles, antialiaseds, urls):
+                             linewidths, linestyles, antialiaseds, urls,
+                             offset_position):
         write = self._pswriter.write
 
         path_codes = []
@@ -640,8 +641,9 @@ def draw_path_collection(self, gc, master_transform, paths, all_transforms,
             path_codes.append(name)
 
         for xo, yo, path_id, gc0, rgbFace in self._iter_collection(
-            gc, path_codes, offsets, offsetTrans, facecolors, edgecolors,
-            linewidths, linestyles, antialiaseds, urls):
+            gc, master_transform, all_transforms, path_codes, offsets,
+            offsetTrans, facecolors, edgecolors, linewidths, linestyles,
+            antialiaseds, urls, offset_position):
             ps = "%g %g %s" % (xo, yo, path_id)
             self._draw_ps(ps, gc0, rgbFace)
 
diff --git a/lib/matplotlib/backends/backend_svg.py b/lib/matplotlib/backends/backend_svg.py
@@ -577,7 +577,8 @@ def draw_markers(self, gc, marker_path, marker_trans, path, trans, rgbFace=None)
 
     def draw_path_collection(self, gc, master_transform, paths, all_transforms,
                              offsets, offsetTrans, facecolors, edgecolors,
-                             linewidths, linestyles, antialiaseds, urls):
+                             linewidths, linestyles, antialiaseds, urls,
+                             offset_position):
         writer = self.writer
         path_codes = []
         writer.start(u'defs')
@@ -592,8 +593,9 @@ def draw_path_collection(self, gc, master_transform, paths, all_transforms,
         writer.end(u'defs')
 
         for xo, yo, path_id, gc0, rgbFace in self._iter_collection(
-            gc, path_codes, offsets, offsetTrans, facecolors, edgecolors,
-            linewidths, linestyles, antialiaseds, urls):
+            gc, master_transform, all_transforms, path_codes, offsets,
+            offsetTrans, facecolors, edgecolors, linewidths, linestyles,
+            antialiaseds, urls, offset_position):
             clipid = self._get_clip(gc0)
             url = gc0.get_url()
             if url is not None:
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -42,14 +42,19 @@ class Collection(artist.Artist, cm.ScalarMappable):
         * *antialiaseds*: None
         * *offsets*: None
         * *transOffset*: transforms.IdentityTransform()
+        * *offset_position*: 'screen' (default) or 'data'
         * *norm*: None (optional for
           :class:`matplotlib.cm.ScalarMappable`)
         * *cmap*: None (optional for
           :class:`matplotlib.cm.ScalarMappable`)
         * *hatch*: None
 
     *offsets* and *transOffset* are used to translate the patch after
-    rendering (default no offsets).
+    rendering (default no offsets).  If offset_position is 'screen'
+    (default) the offset is applied after the master transform has
+    been applied, that is, the offsets are in screen coordinates.  If
+    offset_position is 'data', the offset is applied before the master
+    transform, i.e., the offsets are in data coordinates.
 
     If any of *edgecolors*, *facecolors*, *linewidths*, *antialiaseds*
     are None, they default to their :data:`matplotlib.rcParams` patch
@@ -80,6 +85,7 @@ def __init__(self,
                  pickradius = 5.0,
                  hatch=None,
                  urls = None,
+                 offset_position='screen',
                  **kwargs
                  ):
         """
@@ -98,7 +104,7 @@ def __init__(self,
         self.set_pickradius(pickradius)
         self.set_urls(urls)
         self.set_hatch(hatch)
-
+        self.set_offset_position(offset_position)
 
         self._uniform_offsets = None
         self._offsets = np.array([], np.float_)
@@ -242,7 +248,8 @@ def draw(self, renderer):
         renderer.draw_path_collection(
             gc, transform.frozen(), paths, self.get_transforms(),
             offsets, transOffset, self.get_facecolor(), self.get_edgecolor(),
-            self._linewidths, self._linestyles, self._antialiaseds, self._urls)
+            self._linewidths, self._linestyles, self._antialiaseds, self._urls,
+            self._offset_position)
 
         gc.restore()
         renderer.close_group(self.__class__.__name__)
@@ -359,6 +366,24 @@ def get_offsets(self):
         else:
             return self._uniform_offsets
 
+    def set_offset_position(self, offset_position):
+        """
+        The how offsets are applied.  If *offset_position* is 'screen'
+        (default) the offset is applied after the master transform has
+        been applied, that is, the offsets are in screen coordinates.
+        If offset_position is 'data', the offset is applied before the
+        master transform, i.e., the offsets are in data coordinates.
+        """
+        if offset_position not in ('screen', 'data'):
+            raise ValueError("offset_position must be 'screen' or 'data'")
+        self._offset_position = offset_position
+
+    def get_offset_position(self):
+        """
+        Returns the offset position of the collection.
+        """
+        return self._offset_position
+
     def set_linewidth(self, lw):
         """
         Set the linewidth(s) for the collection.  *lw* can be a scalar
diff --git a/src/_backend_agg.cpp b/src/_backend_agg.cpp
diff --git a/src/_backend_agg.h b/src/_backend_agg.h
