@@ -1622,15 +1622,15 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
16221622 # Strides have priority over counts in classic mode.
16231623 # So, only compute strides from counts
16241624 # if counts were explicitly given
1625- if has_count :
1626- rstride = int (max (np .ceil (rows / rcount ), 1 ))
1627- cstride = int (max (np .ceil (cols / ccount ), 1 ))
1625+ compute_strides = has_count
16281626 else :
16291627 # If the strides are provided then it has priority.
16301628 # Otherwise, compute the strides from the counts.
1631- if not has_stride :
1632- rstride = int (max (np .ceil (rows / rcount ), 1 ))
1633- cstride = int (max (np .ceil (cols / ccount ), 1 ))
1629+ compute_strides = not has_stride
1630+
1631+ if compute_strides :
1632+ rstride = int (max (np .ceil (rows / rcount ), 1 ))
1633+ cstride = int (max (np .ceil (cols / ccount ), 1 ))
16341634
16351635 if 'facecolors' in kwargs :
16361636 fcolors = kwargs .pop ('facecolors' )
@@ -1648,71 +1648,60 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
16481648 if shade and cmap is not None and fcolors is not None :
16491649 fcolors = self ._shade_colors_lightsource (Z , cmap , lightsource )
16501650
1651+ # evenly spaced, and including both endpoints
1652+ row_inds = list (range (0 , rows - 1 , rstride )) + [rows - 1 ]
1653+ col_inds = list (range (0 , cols - 1 , cstride )) + [cols - 1 ]
1654+
1655+ colset = [] # the sampled facecolor
16511656 polys = []
1652- # Only need these vectors to shade if there is no cmap
1653- if cmap is None and shade :
1654- totpts = int (np .ceil ((rows - 1 ) / rstride ) *
1655- np .ceil ((cols - 1 ) / cstride ))
1656- v1 = np .empty ((totpts , 3 ))
1657- v2 = np .empty ((totpts , 3 ))
1658- # This indexes the vertex points
1659- which_pt = 0
1660-
1661-
1662- #colset contains the data for coloring: either average z or the facecolor
1663- colset = []
1664- for rs in range (0 , rows - 1 , rstride ):
1665- for cs in range (0 , cols - 1 , cstride ):
1666- ps = []
1667- for a in (X , Y , Z ):
1668- ztop = a [rs ,cs :min (cols , cs + cstride + 1 )]
1669- zleft = a [rs + 1 :min (rows , rs + rstride + 1 ),
1670- min (cols - 1 , cs + cstride )]
1671- zbase = a [min (rows - 1 , rs + rstride ), cs :min (cols , cs + cstride + 1 ):][::- 1 ]
1672- zright = a [rs :min (rows - 1 , rs + rstride ):, cs ][::- 1 ]
1673- z = np .concatenate ((ztop , zleft , zbase , zright ))
1674- ps .append (z )
1675-
1676- # The construction leaves the array with duplicate points, which
1677- # are removed here.
1678- ps = list (zip (* ps ))
1679- ps2 = [ps [0 ]] + [ps [i ] for i in range (1 , len (ps )) if ps [i ] != ps [i - 1 ]]
1680- avgzsum = sum (p [2 ] for p in ps2 )
1681- polys .append (ps2 )
1657+ for rs , rs_next in zip (row_inds [:- 1 ], row_inds [1 :]):
1658+ for cs , cs_next in zip (col_inds [:- 1 ], col_inds [1 :]):
1659+ ps = [
1660+ # +1 ensures we share edges between polygons
1661+ cbook ._array_perimeter (a [rs :rs_next + 1 , cs :cs_next + 1 ])
1662+ for a in (X , Y , Z )
1663+ ]
1664+ # ps = np.stack(ps, axis=-1)
1665+ ps = np .array (ps ).T
1666+ polys .append (ps )
16821667
16831668 if fcolors is not None :
16841669 colset .append (fcolors [rs ][cs ])
1685- else :
1686- colset .append (avgzsum / len (ps2 ))
1687-
1688- # Only need vectors to shade if no cmap
1689- if cmap is None and shade :
1690- i1 , i2 , i3 = 0 , int (len (ps2 )/ 3 ), int (2 * len (ps2 )/ 3 )
1691- v1 [which_pt ] = np .array (ps2 [i1 ]) - np .array (ps2 [i2 ])
1692- v2 [which_pt ] = np .array (ps2 [i2 ]) - np .array (ps2 [i3 ])
1693- which_pt += 1
1694- if cmap is None and shade :
1695- normals = np .cross (v1 , v2 )
1696- else :
1697- normals = []
16981670
1671+ def get_normals (polygons ):
1672+ """
1673+ Takes a list of polygons and return an array of their normals
1674+ """
1675+ v1 = np .empty ((len (polygons ), 3 ))
1676+ v2 = np .empty ((len (polygons ), 3 ))
1677+ for poly_i , ps in enumerate (polygons ):
1678+ # pick three points around the polygon at which to find the normal
1679+ # doesn't vectorize because polygons is jagged
1680+ i1 , i2 , i3 = 0 , len (ps )// 3 , 2 * len (ps )// 3
1681+ v1 [poly_i , :] = ps [i1 , :] - ps [i2 , :]
1682+ v2 [poly_i , :] = ps [i2 , :] - ps [i3 , :]
1683+ return np .cross (v1 , v2 )
1684+
1685+ # note that the striding causes some polygons to have more coordinates
1686+ # than others
16991687 polyc = art3d .Poly3DCollection (polys , * args , ** kwargs )
17001688
17011689 if fcolors is not None :
17021690 if shade :
1703- colset = self ._shade_colors (colset , normals )
1691+ colset = self ._shade_colors (colset , get_normals ( polys ) )
17041692 polyc .set_facecolors (colset )
17051693 polyc .set_edgecolors (colset )
17061694 elif cmap :
1707- colset = np .array (colset )
1708- polyc .set_array (colset )
1695+ # doesn't vectorize because polys is jagged
1696+ avg_z = np .array ([ps [:,2 ].mean () for ps in polys ])
1697+ polyc .set_array (avg_z )
17091698 if vmin is not None or vmax is not None :
17101699 polyc .set_clim (vmin , vmax )
17111700 if norm is not None :
17121701 polyc .set_norm (norm )
17131702 else :
17141703 if shade :
1715- colset = self ._shade_colors (color , normals )
1704+ colset = self ._shade_colors (color , get_normals ( polys ) )
17161705 else :
17171706 colset = color
17181707 polyc .set_facecolors (colset )
0 commit comments