MAINT: Simplify logic in plot_surface

eric-wieser · eric-wieser · commit 86e1518abb11 · 2017-12-13T23:36:40.000-08:00
Previously:
* "cell" perimeters were clumsily calculated with duplicates, which were then removed at runtime
* code to calculate normals was spread into multiple places
* average z was calculated even if not used
* repeated conversion between stride and count was done

Should have no visible behavior changes
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1671,64 +1671,57 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
             fcolors = self._shade_colors_lightsource(Z, cmap, lightsource)
 
         polys = []
-        # Only need these vectors to shade if there is no cmap
-        if cmap is None and shade :
-            totpts = int(np.ceil(float(rows - 1) / rstride) *
-                         np.ceil(float(cols - 1) / cstride))
-            v1 = np.empty((totpts, 3))
-            v2 = np.empty((totpts, 3))
-            # This indexes the vertex points
-            which_pt = 0
+
+        # evenly spaced, and including both endpoints
+        row_inds = list(xrange(0, rows-1, rstride)) + [rows-1]
+        col_inds = list(xrange(0, cols-1, cstride)) + [cols-1]
 
 
-        #colset contains the data for coloring: either average z or the facecolor
+        #colset contains the sampled facecolors
         colset = []
-        for rs in xrange(0, rows-1, rstride):
-            for cs in xrange(0, cols-1, cstride):
+        for rs, rs_next in zip(row_inds[:-1], row_inds[1:]):
+            for cs, cs_next in zip(col_inds[:-1], col_inds[1:]):
                 ps = []
                 for a in (X, Y, Z):
-                    ztop = a[rs,cs:min(cols, cs+cstride+1)]
-                    zleft = a[rs+1:min(rows, rs+rstride+1),
-                              min(cols-1, cs+cstride)]
-                    zbase = a[min(rows-1, rs+rstride), cs:min(cols, cs+cstride+1):][::-1]
-                    zright = a[rs:min(rows-1, rs+rstride):, cs][::-1]
+                    # the edges of the projected quadrilateral
+                    # note we use pythons half-open ranges to avoid repeating
+                    # the corners
+                    ztop   = a[rs,             cs:cs_next   ]
+                    zleft  = a[rs:rs_next,     cs_next      ]
+                    zbase  = a[rs_next,        cs_next:cs:-1]
+                    zright = a[rs_next:rs:-1,  cs           ]
                     z = np.concatenate((ztop, zleft, zbase, zright))
                     ps.append(z)
 
-                # The construction leaves the array with duplicate points, which
-                # are removed here.
-                ps = list(zip(*ps))
-                lastp = np.array([])
-                ps2 = [ps[0]] + [ps[i] for i in xrange(1, len(ps)) if ps[i] != ps[i-1]]
-                avgzsum = sum(p[2] for p in ps2)
-                polys.append(ps2)
+                # ps = np.stack(ps, axis=-1)
+                ps = np.array(ps).T
+                polys.append(ps)
 
                 if fcolors is not None:
                     colset.append(fcolors[rs][cs])
-                else:
-                    colset.append(avgzsum / len(ps2))
-
-                # Only need vectors to shade if no cmap
-                if cmap is None and shade:
-                    i1, i2, i3 = 0, int(len(ps2)/3), int(2*len(ps2)/3)
-                    v1[which_pt] = np.array(ps2[i1]) - np.array(ps2[i2])
-                    v2[which_pt] = np.array(ps2[i2]) - np.array(ps2[i3])
-                    which_pt += 1
-        if cmap is None and shade:
-            normals = np.cross(v1, v2)
-        else :
-            normals = []
 
         polyc = art3d.Poly3DCollection(polys, *args, **kwargs)
 
+        if shade and cmap is None:
+            v1 = np.empty((len(polys), 3))
+            v2 = np.empty((len(polys), 3))
+            for poly_i, ps in enumerate(polys):
+                # pick three points around the polygon to find the normal at
+                # hard to vectorize, since len(ps) is different at the edges
+                i1, i2, i3 = 0, int(len(ps)/3), int(2*len(ps)/3)
+                v1[poly_i] = ps[i1] - ps[i2]
+                v2[poly_i] = ps[i2] - ps[i3]
+            normals = np.cross(v1, v2)
+
         if fcolors is not None:
             if shade:
                 colset = self._shade_colors(colset, normals)
             polyc.set_facecolors(colset)
             polyc.set_edgecolors(colset)
         elif cmap:
-            colset = np.array(colset)
-            polyc.set_array(colset)
+            # doesn't vectorize because polys is jagged
+            avg_z = np.array([ps[:,2].mean() for ps in polys])
+            polyc.set_array(avg_z)
             if vmin is not None or vmax is not None:
                 polyc.set_clim(vmin, vmax)
             if norm is not None:
