Refactor shading

oscargus · oscargus · commit 0c248a231bfa · 2022-10-23T12:11:18.000+02:00
diff --git a/doc/users/next_whats_new/shade_poly3dcollection.rst b/doc/users/next_whats_new/shade_poly3dcollection.rst
@@ -0,0 +1,33 @@
+``Poly3DCollection`` supports shading
+-------------------------------------
+
+It is now possible to shade a `.Poly3DCollection`. This is useful if the
+polygons are obtained from e.g. a 3D model.
+
+.. plot::
+    :include-source: true
+
+    import numpy as np
+    import matplotlib.pyplot as plt
+    from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+    # Define 3D shape
+    block = np.array([
+        [[1, 1, 0],
+         [1, 0, 0],
+         [0, 1, 0]],
+        [[1, 1, 0],
+         [1, 1, 1],
+         [1, 0, 0]],
+        [[1, 1, 0],
+         [1, 1, 1],
+         [0, 1, 0]],
+        [[1, 0, 0],
+         [1, 1, 1],
+         [0, 1, 0]]
+    ])
+
+    ax = plt.subplot(projection='3d')
+    pc = Poly3DCollection(block, facecolors='b', shade=True)
+    ax.add_collection(pc)
+    plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -12,7 +12,8 @@
 import numpy as np
 
 from matplotlib import (
-    artist, cbook, colors as mcolors, lines, text as mtext, path as mpath)
+    _api, artist, cbook, colors as mcolors, lines, text as mtext,
+    path as mpath)
 from matplotlib.collections import (
     LineCollection, PolyCollection, PatchCollection, PathCollection)
 from matplotlib.colors import Normalize
@@ -808,7 +809,8 @@ class Poly3DCollection(PolyCollection):
         triangulation and thus generates consistent surfaces.
     """
 
-    def __init__(self, verts, *args, zsort='average', **kwargs):
+    def __init__(self, verts, *args, zsort='average', shade=False,
+                 lightsource=None, **kwargs):
         """
         Parameters
         ----------
@@ -819,6 +821,20 @@ def __init__(self, verts, *args, zsort='average', **kwargs):
         zsort : {'average', 'min', 'max'}, default: 'average'
             The calculation method for the z-order.
             See `~.Poly3DCollection.set_zsort` for details.
+        shade : bool, default: False
+            Whether to shade *facecolors* and *edgecolors*.
+
+            .. versionadded:: 3.7
+
+            .. note::
+                *facecolors* and/or *edgecolors* must be provided for shading
+                to work.
+
+        lightsource : `~matplotlib.colors.LightSource`
+            The lightsource to use when *shade* is True.
+
+            .. versionadded:: 3.7
+
         *args, **kwargs
             All other parameters are forwarded to `.PolyCollection`.
 
@@ -827,6 +843,28 @@ def __init__(self, verts, *args, zsort='average', **kwargs):
         Note that this class does a bit of magic with the _facecolors
         and _edgecolors properties.
         """
+        if 'facecolor' in kwargs:
+            kwargs['facecolors'] = kwargs.pop('facecolor')
+        if 'edgecolor' in kwargs:
+            kwargs['edgecolors'] = kwargs.pop('edgecolor')
+        if shade:
+            normals = _generate_normals(verts)
+            facecolors = kwargs.get('facecolors', None)
+            if facecolors is not None:
+                kwargs['facecolors'] = _shade_colors(
+                    facecolors, normals, lightsource
+                )
+
+            edgecolors = kwargs.get('edgecolors', None)
+            if edgecolors is not None:
+                kwargs['edgecolors'] = _shade_colors(
+                    edgecolors, normals, lightsource
+                )
+            if facecolors is None and edgecolors in None:
+                _api.warn_external(
+                    "You must provide at least one of facecolors and "
+                    "edgecolors for shade to work as expected.")
+
         super().__init__(verts, *args, **kwargs)
         if isinstance(verts, np.ndarray):
             if verts.ndim != 3:
@@ -1086,3 +1124,84 @@ def _zalpha(colors, zs):
     sats = 1 - norm(zs) * 0.7
     rgba = np.broadcast_to(mcolors.to_rgba_array(colors), (len(zs), 4))
     return np.column_stack([rgba[:, :3], rgba[:, 3] * sats])
+
+
+def _generate_normals(polygons):
+    """
+    Compute the normals of a list of polygons, one normal per polygon.
+
+    Normals point towards the viewer for a face with its vertices in
+    counterclockwise order, following the right hand rule.
+
+    Uses three points equally spaced around the polygon.
+    If the polygon points are not in a plane, this does not make sense, but
+    it is on the other hand impossible to compute a single shade in that case.
+
+    Parameters
+    ----------
+    polygons : list of (M_i, 3) array-like, or (..., M, 3) array-like
+        A sequence of polygons to compute normals for, which can have
+        varying numbers of vertices. If the polygons all have the same
+        number of vertices and array is passed, then the operation will
+        be vectorized.
+
+    Returns
+    -------
+    normals : (..., 3) array
+        A normal vector estimated for the polygon.
+    """
+    if isinstance(polygons, np.ndarray):
+        # optimization: polygons all have the same number of points, so can
+        # vectorize
+        n = polygons.shape[-2]
+        i1, i2, i3 = 0, n//3, 2*n//3
+        v1 = polygons[..., i1, :] - polygons[..., i2, :]
+        v2 = polygons[..., i2, :] - polygons[..., i3, :]
+    else:
+        # The subtraction doesn't vectorize because polygons is jagged.
+        v1 = np.empty((len(polygons), 3))
+        v2 = np.empty((len(polygons), 3))
+        for poly_i, ps in enumerate(polygons):
+            n = len(ps)
+            i1, i2, i3 = 0, n//3, 2*n//3
+            v1[poly_i, :] = ps[i1, :] - ps[i2, :]
+            v2[poly_i, :] = ps[i2, :] - ps[i3, :]
+    return np.cross(v1, v2)
+
+
+def _shade_colors(color, normals, lightsource=None):
+    """
+    Shade *color* using normal vectors given by *normals*,
+    assuming a *lightsource* (using default position if not given).
+    *color* can also be an array of the same length as *normals*.
+    """
+    if lightsource is None:
+        # chosen for backwards-compatibility
+        lightsource = mcolors.LightSource(azdeg=225, altdeg=19.4712)
+
+    with np.errstate(invalid="ignore"):
+        shade = ((normals / np.linalg.norm(normals, axis=1, keepdims=True))
+                 @ lightsource.direction)
+    mask = ~np.isnan(shade)
+
+    if mask.any():
+        # convert dot product to allowed shading fractions
+        in_norm = mcolors.Normalize(-1, 1)
+        out_norm = mcolors.Normalize(0.3, 1).inverse
+
+        def norm(x):
+            return out_norm(in_norm(x))
+
+        shade[~mask] = 0
+
+        color = mcolors.to_rgba_array(color)
+        # shape of color should be (M, 4) (where M is number of faces)
+        # shape of shade should be (M,)
+        # colors should have final shape of (M, 4)
+        alpha = color[:, 3]
+        colors = norm(shade)[:, np.newaxis] * color
+        colors[:, 3] = alpha
+    else:
+        colors = np.asanyarray(color).copy()
+
+    return colors
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1668,15 +1668,13 @@ def plot_surface(self, X, Y, Z, *, norm=None, vmin=None,
 
         # note that the striding causes some polygons to have more coordinates
         # than others
-        polyc = art3d.Poly3DCollection(polys, **kwargs)
 
         if fcolors is not None:
-            if shade:
-                colset = self._shade_colors(
-                    colset, self._generate_normals(polys), lightsource)
-            polyc.set_facecolors(colset)
-            polyc.set_edgecolors(colset)
+            polyc = art3d.Poly3DCollection(
+                polys, edgecolors=colset, facecolors=colset, shade=shade,
+                lightsource=lightsource, **kwargs)
         elif cmap:
+            polyc = art3d.Poly3DCollection(polys, **kwargs)
             # can't always vectorize, because polys might be jagged
             if isinstance(polys, np.ndarray):
                 avg_z = polys[..., 2].mean(axis=-1)
@@ -1688,97 +1686,15 @@ def plot_surface(self, X, Y, Z, *, norm=None, vmin=None,
             if norm is not None:
                 polyc.set_norm(norm)
         else:
-            if shade:
-                colset = self._shade_colors(
-                    color, self._generate_normals(polys), lightsource)
-            else:
-                colset = color
-            polyc.set_facecolors(colset)
+            polyc = art3d.Poly3DCollection(
+                polys, facecolors=color, shade=shade,
+                lightsource=lightsource, **kwargs)
 
         self.add_collection(polyc)
         self.auto_scale_xyz(X, Y, Z, had_data)
 
         return polyc
 
-    def _generate_normals(self, polygons):
-        """
-        Compute the normals of a list of polygons.
-
-        Normals point towards the viewer for a face with its vertices in
-        counterclockwise order, following the right hand rule.
-
-        Uses three points equally spaced around the polygon.
-        This normal of course might not make sense for polygons with more than
-        three points not lying in a plane, but it's a plausible and fast
-        approximation.
-
-        Parameters
-        ----------
-        polygons : list of (M_i, 3) array-like, or (..., M, 3) array-like
-            A sequence of polygons to compute normals for, which can have
-            varying numbers of vertices. If the polygons all have the same
-            number of vertices and array is passed, then the operation will
-            be vectorized.
-
-        Returns
-        -------
-        normals : (..., 3) array
-            A normal vector estimated for the polygon.
-        """
-        if isinstance(polygons, np.ndarray):
-            # optimization: polygons all have the same number of points, so can
-            # vectorize
-            n = polygons.shape[-2]
-            i1, i2, i3 = 0, n//3, 2*n//3
-            v1 = polygons[..., i1, :] - polygons[..., i2, :]
-            v2 = polygons[..., i2, :] - polygons[..., i3, :]
-        else:
-            # The subtraction doesn't vectorize because polygons is jagged.
-            v1 = np.empty((len(polygons), 3))
-            v2 = np.empty((len(polygons), 3))
-            for poly_i, ps in enumerate(polygons):
-                n = len(ps)
-                i1, i2, i3 = 0, n//3, 2*n//3
-                v1[poly_i, :] = ps[i1, :] - ps[i2, :]
-                v2[poly_i, :] = ps[i2, :] - ps[i3, :]
-        return np.cross(v1, v2)
-
-    def _shade_colors(self, color, normals, lightsource=None):
-        """
-        Shade *color* using normal vectors given by *normals*.
-        *color* can also be an array of the same length as *normals*.
-        """
-        if lightsource is None:
-            # chosen for backwards-compatibility
-            lightsource = mcolors.LightSource(azdeg=225, altdeg=19.4712)
-
-        with np.errstate(invalid="ignore"):
-            shade = ((normals / np.linalg.norm(normals, axis=1, keepdims=True))
-                     @ lightsource.direction)
-        mask = ~np.isnan(shade)
-
-        if mask.any():
-            # convert dot product to allowed shading fractions
-            in_norm = mcolors.Normalize(-1, 1)
-            out_norm = mcolors.Normalize(0.3, 1).inverse
-
-            def norm(x):
-                return out_norm(in_norm(x))
-
-            shade[~mask] = 0
-
-            color = mcolors.to_rgba_array(color)
-            # shape of color should be (M, 4) (where M is number of faces)
-            # shape of shade should be (M,)
-            # colors should have final shape of (M, 4)
-            alpha = color[:, 3]
-            colors = norm(shade)[:, np.newaxis] * color
-            colors[:, 3] = alpha
-        else:
-            colors = np.asanyarray(color).copy()
-
-        return colors
-
     def plot_wireframe(self, X, Y, Z, **kwargs):
         """
         Plot a 3D wireframe.
@@ -1975,9 +1891,8 @@ def plot_trisurf(self, *args, color=None, norm=None, vmin=None, vmax=None,
         zt = z[triangles]
         verts = np.stack((xt, yt, zt), axis=-1)
 
-        polyc = art3d.Poly3DCollection(verts, *args, **kwargs)
-
         if cmap:
+            polyc = art3d.Poly3DCollection(verts, *args, **kwargs)
             # average over the three points of each triangle
             avg_z = verts[:, :, 2].mean(axis=1)
             polyc.set_array(avg_z)
@@ -1986,12 +1901,9 @@ def plot_trisurf(self, *args, color=None, norm=None, vmin=None, vmax=None,
             if norm is not None:
                 polyc.set_norm(norm)
         else:
-            if shade:
-                normals = self._generate_normals(verts)
-                colset = self._shade_colors(color, normals, lightsource)
-            else:
-                colset = color
-            polyc.set_facecolors(colset)
+            polyc = art3d.Poly3DCollection(
+                verts, *args, shade=shade, lightsource=lightsource,
+                facecolors=color, **kwargs)
 
         self.add_collection(polyc)
         self.auto_scale_xyz(tri.x, tri.y, z, had_data)
@@ -2036,13 +1948,10 @@ def _3d_extend_contour(self, cset, stride=5):
 
             # all polygons have 4 vertices, so vectorize
             polyverts = np.array(polyverts)
-            normals = self._generate_normals(polyverts)
-
-            colors = self._shade_colors(color, normals)
-            colors2 = self._shade_colors(color, normals)
             polycol = art3d.Poly3DCollection(polyverts,
-                                             facecolors=colors,
-                                             edgecolors=colors2)
+                                             facecolors=color,
+                                             edgecolors=color,
+                                             shade=True)
             polycol.set_sort_zpos(z)
             self.add_collection3d(polycol)
 
@@ -2587,15 +2496,11 @@ def bar3d(self, x, y, z, dx, dy, dz, color=None,
             if len(facecolors) < len(x):
                 facecolors *= (6 * len(x))
 
-        if shade:
-            normals = self._generate_normals(polys)
-            sfacecolors = self._shade_colors(facecolors, normals, lightsource)
-        else:
-            sfacecolors = facecolors
-
         col = art3d.Poly3DCollection(polys,
                                      zsort=zsort,
-                                     facecolor=sfacecolors,
+                                     facecolors=facecolors,
+                                     shade=shade,
+                                     lightsource=lightsource,
                                      *args, **kwargs)
         self.add_collection(col)
 
@@ -2964,16 +2869,10 @@ def permutation_matrices(n):
             # shade the faces
             facecolor = facecolors[coord]
             edgecolor = edgecolors[coord]
-            if shade:
-                normals = self._generate_normals(faces)
-                facecolor = self._shade_colors(facecolor, normals, lightsource)
-                if edgecolor is not None:
-                    edgecolor = self._shade_colors(
-                        edgecolor, normals, lightsource
-                    )
 
             poly = art3d.Poly3DCollection(
-                faces, facecolors=facecolor, edgecolors=edgecolor, **kwargs)
+                faces, facecolors=facecolor, edgecolors=edgecolor,
+                shade=shade, lightsource=lightsource, **kwargs)
             self.add_collection3d(poly)
             polygons[coord] = poly
 
