vectorizing some trisurf functions for performance improvement

Author: choldgraf

plotly/tools.py

@@ -1464,11 +1464,10 @@ def _find_intermediate_color(lowcolor, highcolor, intermed):
         diff_1 = float(highcolor[1] - lowcolor[1])
         diff_2 = float(highcolor[2] - lowcolor[2])
 
-        new_tuple = (lowcolor[0] + intermed*diff_0,
-                     lowcolor[1] + intermed*diff_1,
-                     lowcolor[2] + intermed*diff_2)
-
-        return new_tuple
+        inter_colors = np.array([lowcolor[0] + intermed * diff_0,
+                                 lowcolor[1] + intermed * diff_1,
+                                 lowcolor[2] + intermed * diff_2])
+        return inter_colors
 
     @staticmethod
     def _unconvert_from_RGB_255(colors):
@@ -1491,7 +1490,7 @@ def _unconvert_from_RGB_255(colors):
         return un_rgb_colors
 
     @staticmethod
-    def _map_z2color(zval, colormap, vmin, vmax):
+    def _map_z2color(zvals, colormap, vmin, vmax):
         """
         Returns the color corresponding zval's place between vmin and vmax
 
@@ -1508,21 +1507,14 @@ def _map_z2color(zval, colormap, vmin, vmax):
                                          "of vmax.")
         # find distance t of zval from vmin to vmax where the distance
         # is normalized to be between 0 and 1
-        t = (zval - vmin)/float((vmax - vmin))
-        t_color = FigureFactory._find_intermediate_color(colormap[0],
-                                                         colormap[1],
-                                                         t)
-        t_color = (t_color[0]*255.0, t_color[1]*255.0, t_color[2]*255.0)
-        labelled_color = 'rgb{}'.format(t_color)
-
-        return labelled_color
-
-    @staticmethod
-    def _tri_indices(simplices):
-        """
-        Returns a triplet of lists containing simplex coordinates
-        """
-        return ([triplet[c] for triplet in simplices] for c in range(3))
+        t = (zvals - vmin) / float((vmax - vmin))
+        t_colors = FigureFactory._find_intermediate_color(colormap[0],
+                                                          colormap[1],
+                                                          t)
+        t_colors = t_colors * 255.
+        labelled_colors = ['rgb(%s, %s, %s)' % (i, j, k)
+                           for i, j, k in t_colors.T]
+        return labelled_colors
 
     @staticmethod
     def _trisurf(x, y, z, simplices, colormap=None, dist_func=None,
@@ -1539,11 +1531,11 @@ def _trisurf(x, y, z, simplices, colormap=None, dist_func=None,
         points3D = np.vstack((x, y, z)).T
 
         # vertices of the surface triangles
-        tri_vertices = list(map(lambda index: points3D[index], simplices))
+        tri_vertices = points3D[simplices]
 
         if not dist_func:
             # mean values of z-coordinates of triangle vertices
-            mean_dists = [np.mean(tri[:, 2]) for tri in tri_vertices]
+            mean_dists = tri_vertices[:, :, 2].mean(-1)
         else:
             # apply user inputted function to calculate
             # custom coloring for triangle vertices
@@ -1559,9 +1551,9 @@ def _trisurf(x, y, z, simplices, colormap=None, dist_func=None,
 
         min_mean_dists = np.min(mean_dists)
         max_mean_dists = np.max(mean_dists)
-        facecolor = ([FigureFactory._map_z2color(zz, colormap, min_mean_dists,
-                      max_mean_dists) for zz in mean_dists])
-        ii, jj, kk = FigureFactory._tri_indices(simplices)
+        facecolor = FigureFactory._map_z2color(mean_dists, colormap,
+                                               min_mean_dists, max_mean_dists)
+        ii, jj, kk = zip(*simplices)
 
         triangles = graph_objs.Mesh3d(x=x, y=y, z=z, facecolor=facecolor,
                                       i=ii, j=jj, k=kk, name='')
@@ -5621,4 +5613,3 @@ def make_table_annotations(self):
                         font=dict(color=font_color),
                         showarrow=False))
         return annotations
-