ENH: optimize Collection non-affine transform to only call transform once

jklymak · jklymak · commit acf4ac46fae0 · 2018-06-20T14:39:07.000-07:00
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -189,7 +189,7 @@ def get_datalim(self, transData):
         paths = self.get_paths()
 
         if not transform.is_affine:
-            paths = [transform.transform_path_non_affine(p) for p in paths]
+            paths = self._non_affine_transform_paths(paths, transform)
             transform = transform.get_affine()
         if not transOffset.is_affine:
             offsets = transOffset.transform_non_affine(offsets)
@@ -213,6 +213,22 @@ def get_window_extent(self, renderer):
         # cases other than scatter plot legend
         return self.get_datalim(transforms.IdentityTransform())
 
+    def _non_affine_transform_paths(self, paths, transform):
+        """
+        Sometimes there will be many paths, and the transform is
+        slow to run.  This helper concatenates the paths, transforms,
+        the points, and then splits them back into the individual paths.
+        """
+        vertices = np.vstack(path.vertices for path in paths)
+        pos = np.cumsum([len(x) for x in paths])
+        vertices = transform.transform_path_non_affine(vertices)
+        # repopulate paths
+        verts = np.split(vertices, pos[:-1])
+        for nn, verti in enumerate(verts):
+            paths[nn].vertices = verti
+
+        return paths
+
     def _prepare_points(self):
         """Point prep for drawing and hit testing"""
 
@@ -236,8 +252,7 @@ def _prepare_points(self):
                 offsets = np.column_stack([xs, ys])
 
         if not transform.is_affine:
-            paths = [transform.transform_path_non_affine(path)
-                     for path in paths]
+            paths = self._non_affine_transform_paths(paths, transform)
             transform = transform.get_affine()
         if not transOffset.is_affine:
             offsets = transOffset.transform_non_affine(offsets)
