matplotlib
diff --git a/‎lib/mpl_toolkits/mplot3d/axes3d.py
Lines changed: 47 additions & 45 deletions b/‎lib/mpl_toolkits/mplot3d/axes3d.py
Lines changed: 47 additions & 45 deletions
diff --git a/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/mixedsubplot.pdf
-9.08 KB b/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/mixedsubplot.pdf
-9.08 KB
diff --git a/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/mixedsubplot.png
195 Bytes b/‎lib/mpl_toolkits/tests/baseline_images/test_mplot3d/mixedsubplot.png
195 Bytes
@@ -1670,65 +1670,67 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
         if shade and cmap is not None and fcolors is not None:
             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]
 
+        def boundary_edge(arr, i0, j0, i1, j1):
+            """
+            Get the boundary elements of the rectangle ``arr[i0:i1+1, j0:j1+1]``
 
-        #colset contains the data for coloring: either average z or the facecolor
-        colset = []
-        for rs in xrange(0, rows-1, rstride):
-            for cs in xrange(0, cols-1, cstride):
-                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]
-                    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)
+            The bounds are inclusive, and the order of the results is:
+
+                [arr[i0,j0] ... arr[i0,j1] ... arr[i1, j1] ... arr[i1,j0] ...]
+
+            With length ``2*((i1 - i0) + (j1 - j0))``
+            """
+            # note we use Python's half-open ranges to avoid repeating
+            # the corners
+            ztop   = a[i0,        j0:j1   ]  # noqa: E201, E221
+            zleft  = a[i0:i1,     j1      ]  # noqa: E201, E221
+            zbase  = a[i1,        j1:j0:-1]  # noqa: E201, E221
+            zright = a[i1:i0:-1,  j0      ]  # noqa: E201, E221
+            return np.concatenate((ztop, zleft, zbase, zright))
+
+        colset = []  # the sampled facecolor
+        polys = []
+        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 = [
+                    boundary_edge(a, rs, cs, rs_next, cs_next)
+                    for a in (X, Y, Z)
+                ]
+                # 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 = []
 
+        # note that the striding causes some polygons to have more coordinates
+        # than others
         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 at which to find the normal
+                # doesn't vectorize because polys is jagged
+                i1, i2, i3 = 0, len(ps)//3, 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: