Speedup RegularPolyCollection.get_transformed_patches

efiring · efiring · commit dac9540301df · 2007-05-10T07:20:45.000Z
svn path=/trunk/matplotlib/; revision=3263
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -17,7 +17,7 @@
 from colors import colorConverter
 from cm import ScalarMappable
 from numerix import arange, sin, cos, pi, asarray, sqrt, array, newaxis, ones
-from numerix import isnan, any
+from numerix import isnan, any, resize
 from transforms import identity_transform
 
 import matplotlib.nxutils as nxutils
@@ -141,7 +141,7 @@ def __init__(self,
         #self._offsets = offsets
         self._offsets = offsets
         self._transOffset = transOffset
-        self._verts = []        
+        self._verts = []
 
     __init__.__doc__ = dedent(__init__.__doc__) % kwdocd
 
@@ -153,13 +153,13 @@ def pick(self, mouseevent):
         if not self.pickable(): return
         ind = []
         x, y = mouseevent.x, mouseevent.y
-        for i, thispoly in enumerate(self.get_transformed_patches()):            
+        for i, thispoly in enumerate(self.get_transformed_patches()):
             inside = nxutils.pnpoly(x, y, thispoly)
             if inside: ind.append(i)
         if len(ind):
             self.figure.canvas.pick_event(mouseevent, self, ind=ind)
-        
-        
+
+
     def get_transformed_patches(self):
         """
         get a sequence of the polygons in the collection in display (transformed) space
@@ -348,7 +348,7 @@ def draw(self, renderer):
         transform = self.get_transform()
         transoffset = self.get_transoffset()
 
-        
+
         transform.freeze()
         transoffset.freeze()
         self.update_scalarmappable()
@@ -364,14 +364,14 @@ def draw(self, renderer):
         transoffset.thaw()
         renderer.close_group('polycollection')
 
-        
+
     def get_verts(self, dataTrans=None):
         '''Return vertices in data coordinates.
         The calculation is incomplete in general; it is based
         on the vertices or the offsets, whichever is using
         dataTrans as its transformation, so it does not take
         into account the combined effect of segments and offsets.
-        '''        
+        '''
         verts = []
         if self._offsets is None:
             for seg in self._verts:
@@ -451,24 +451,23 @@ def __init__(self,
     __init__.__doc__ = dedent(__init__.__doc__) % kwdocd
 
     def get_transformed_patches(self):
-        
-        xverts, yverts = zip(*self._verts)
-        xverts = asarray(xverts)
-        yverts = asarray(yverts)
-        sizes = sqrt(asarray(self._sizes)*self._dpi.get()/72.0)
-        Nsizes = len(sizes)
+        # Shouldn't need all these calls to asarray;
+        # the variables should be converted when stored.
+        # Similar speedups with numerix should be attainable
+        # in many other places.
+        verts = asarray(self._verts)
+        offsets = asarray(self._offsets)
+        Npoly = len(offsets)
+        scales = sqrt(asarray(self._sizes)*self._dpi.get()/72.0)
+        Nscales = len(scales)
+        if Nscales >1:
+            scales = resize(scales, (Npoly, 1, 1))
         transOffset = self.get_transoffset()
-        polys = []
-        for i, loc in enumerate(self._offsets):
-            xo,yo = transOffset.xy_tup(loc)
-            #print 'xo, yo', loc, (xo, yo)
-            scale = sizes[i % Nsizes]
-
-            thisxverts = scale*xverts + xo
-            thisyverts = scale*yverts + yo
-            polys.append(zip(thisxverts, thisyverts))
+        xyo = transOffset.numerix_xy(offsets)
+        polys = scales * verts + xyo[:, newaxis, :]
         return polys
 
+
     def _update_verts(self):
         r = 1.0/math.sqrt(math.pi)  # unit area
         theta = (2*math.pi/self.numsides)*arange(self.numsides) + self.rotation
@@ -486,7 +485,7 @@ def draw(self, renderer):
         self._update_verts()
         scales = sqrt(asarray(self._sizes)*self._dpi.get()/72.0)
 
-        
+
         offsets = self._offsets
         if self._offsets is not None:
             xs, ys = zip(*offsets)
@@ -642,7 +641,7 @@ def __init__(self, segments,     # Can be None.
         self._transOffset = transOffset
         self.set_segments(segments)
         self.update(kwargs)
-        
+
     def get_transoffset(self):
         if self._transOffset is None:
             self._transOffset = identity_transform()
@@ -678,7 +677,7 @@ def draw(self, renderer):
 
         transform.freeze()
         transoffset.freeze()
-        
+
         segments = self._segments
         offsets = self._offsets
 
@@ -693,7 +692,7 @@ def draw(self, renderer):
                 xs = self.convert_xunits(self._offsets[:0])
                 ys = self.convert_yunits(self._offsets[:1])
                 offsets = zip(xs, ys)
-            
+
         self.update_scalarmappable()
         renderer.draw_line_collection(
             segments, transform, self.clipbox,
