@@ -1670,65 +1670,67 @@ def plot_surface(self, X, Y, Z, *args, **kwargs):
16701670 if shade and cmap is not None and fcolors is not None :
16711671 fcolors = self ._shade_colors_lightsource (Z , cmap , lightsource )
16721672
1673- polys = []
1674- # Only need these vectors to shade if there is no cmap
1675- if cmap is None and shade :
1676- totpts = int (np .ceil (float (rows - 1 ) / rstride ) *
1677- np .ceil (float (cols - 1 ) / cstride ))
1678- v1 = np .empty ((totpts , 3 ))
1679- v2 = np .empty ((totpts , 3 ))
1680- # This indexes the vertex points
1681- which_pt = 0
1673+ # evenly spaced, and including both endpoints
1674+ row_inds = list (xrange (0 , rows - 1 , rstride )) + [rows - 1 ]
1675+ col_inds = list (xrange (0 , cols - 1 , cstride )) + [cols - 1 ]
16821676
1677+ def boundary_edge (arr , i0 , j0 , i1 , j1 ):
1678+ """
1679+ Get the boundary elements of the rectangle ``arr[i0:i1+1, j0:j1+1]``
16831680
1684- #colset contains the data for coloring: either average z or the facecolor
1685- colset = []
1686- for rs in xrange (0 , rows - 1 , rstride ):
1687- for cs in xrange (0 , cols - 1 , cstride ):
1688- ps = []
1689- for a in (X , Y , Z ):
1690- ztop = a [rs ,cs :min (cols , cs + cstride + 1 )]
1691- zleft = a [rs + 1 :min (rows , rs + rstride + 1 ),
1692- min (cols - 1 , cs + cstride )]
1693- zbase = a [min (rows - 1 , rs + rstride ), cs :min (cols , cs + cstride + 1 ):][::- 1 ]
1694- zright = a [rs :min (rows - 1 , rs + rstride ):, cs ][::- 1 ]
1695- z = np .concatenate ((ztop , zleft , zbase , zright ))
1696- ps .append (z )
1697- 1698- # The construction leaves the array with duplicate points, which
1699- # are removed here.
1700- ps = list (zip (* ps ))
1701- lastp = np .array ([])
1702- ps2 = [ps [0 ]] + [ps [i ] for i in xrange (1 , len (ps )) if ps [i ] != ps [i - 1 ]]
1703- avgzsum = sum (p [2 ] for p in ps2 )
1704- polys .append (ps2 )
1681+ The bounds are inclusive, and the order of the results is:
1682+
1683+ [arr[i0,j0] ... arr[i0,j1] ... arr[i1, j1] ... arr[i1,j0] ...]
1684+
1685+ With length ``2*((i1 - i0) + (j1 - j0))``
1686+ """
1687+ # note we use Python's half-open ranges to avoid repeating
1688+ # the corners
1689+ ztop = a [i0 , j0 :j1 ] # noqa: E201, E221
1690+ zleft = a [i0 :i1 , j1 ] # noqa: E201, E221
1691+ zbase = a [i1 , j1 :j0 :- 1 ] # noqa: E201, E221
1692+ zright = a [i1 :i0 :- 1 , j0 ] # noqa: E201, E221
1693+ return np .concatenate ((ztop , zleft , zbase , zright ))
1694+ 1695+ colset = [] # the sampled facecolor
1696+ polys = []
1697+ for rs , rs_next in zip (row_inds [:- 1 ], row_inds [1 :]):
1698+ for cs , cs_next in zip (col_inds [:- 1 ], col_inds [1 :]):
1699+ ps = [
1700+ boundary_edge (a , rs , cs , rs_next , cs_next )
1701+ for a in (X , Y , Z )
1702+ ]
1703+ # ps = np.stack(ps, axis=-1)
1704+ ps = np .array (ps ).T
1705+ polys .append (ps )
17051706
17061707 if fcolors is not None :
17071708 colset .append (fcolors [rs ][cs ])
1708- else :
1709- colset .append (avgzsum / len (ps2 ))
1710- 1711- # Only need vectors to shade if no cmap
1712- if cmap is None and shade :
1713- i1 , i2 , i3 = 0 , int (len (ps2 )/ 3 ), int (2 * len (ps2 )/ 3 )
1714- v1 [which_pt ] = np .array (ps2 [i1 ]) - np .array (ps2 [i2 ])
1715- v2 [which_pt ] = np .array (ps2 [i2 ]) - np .array (ps2 [i3 ])
1716- which_pt += 1
1717- if cmap is None and shade :
1718- normals = np .cross (v1 , v2 )
1719- else :
1720- normals = []
17211709
1710+ # note that the striding causes some polygons to have more coordinates
1711+ # than others
17221712 polyc = art3d .Poly3DCollection (polys , * args , ** kwargs )
17231713
1714+ if shade and cmap is None :
1715+ v1 = np .empty ((len (polys ), 3 ))
1716+ v2 = np .empty ((len (polys ), 3 ))
1717+ for poly_i , ps in enumerate (polys ):
1718+ # pick three points around the polygon at which to find the normal
1719+ # doesn't vectorize because polys is jagged
1720+ i1 , i2 , i3 = 0 , len (ps )// 3 , 2 * len (ps )// 3
1721+ v1 [poly_i , :] = ps [i1 , :] - ps [i2 , :]
1722+ v2 [poly_i , :] = ps [i2 , :] - ps [i3 , :]
1723+ normals = np .cross (v1 , v2 )
1724+ 17241725 if fcolors is not None :
17251726 if shade :
17261727 colset = self ._shade_colors (colset , normals )
17271728 polyc .set_facecolors (colset )
17281729 polyc .set_edgecolors (colset )
17291730 elif cmap :
1730- colset = np .array (colset )
1731- polyc .set_array (colset )
1731+ # doesn't vectorize because polys is jagged
1732+ avg_z = np .array ([ps [:,2 ].mean () for ps in polys ])
1733+ polyc .set_array (avg_z )
17321734 if vmin is not None or vmax is not None :
17331735 polyc .set_clim (vmin , vmax )
17341736 if norm is not None :
0 commit comments