Refactor shading

oscargus · oscargus · commit 5bee4ad12640 · 2022-10-29T14:03:26.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,23 @@ def __init__(self, verts, *args, zsort='average', **kwargs):
         Note that this class does a bit of magic with the _facecolors
         and _edgecolors properties.
         """
+        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 facecolors, edgecolors, or both for "
+                    "shade to work as expected.")
         super().__init__(verts, *args, **kwargs)
         if isinstance(verts, np.ndarray):
             if verts.ndim != 3:
@@ -1086,3 +1119,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. This method assumes
+    that the points are in a plane. Otherwise, more than one shade is required,
+    which is not supported.
+
+    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
