MAINT: Unify calculation of normal vectors from polygons

eric-wieser · eric-wieser · commit 5e6e837167a7 · 2018-11-11T10:29:26.000-08:00
This combines `get_normals` and `_generate_normals`, and eliminates all other calls to np.cross.

`get_normals` and `_generate_normals` were profiled, and it was found that vectorizing `np.cross`  like in `get_normals` was faster:

```python
import numpy as np

def get_normals(polygons):
    v1 = np.empty((len(polygons), 3))
    v2 = np.empty((len(polygons), 3))
    for poly_i, ps in enumerate(polygons):
        # pick three points around the polygon at which to find the
        # normal doesn't vectorize because polygons 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, :]
    return np.cross(v1, v2)

def _generate_normals(self, polygons):
    normals = []
    for verts in polygons:
        v1 = np.array(verts[0]) - np.array(verts[1])
        v2 = np.array(verts[2]) - np.array(verts[0])
        normals.append(np.cross(v1, v2))
    return np.array(normals)

polygons = [
    np.random.rand(np.random.randint(10, 1000), 3)
    for i in range(100)
]

%timeit _generate_normals(polygons)
# 3.14 ms ± 255 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit get_normals(polygons)
# 452 µs ± 4.33 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1686,28 +1686,14 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
                 if fcolors is not None:
                     colset.append(fcolors[rs][cs])
 
-        def get_normals(polygons):
-            """
-            Takes a list of polygons and return an array of their normals
-            """
-            v1 = np.empty((len(polygons), 3))
-            v2 = np.empty((len(polygons), 3))
-            for poly_i, ps in enumerate(polygons):
-                # pick three points around the polygon at which to find the
-                # normal doesn't vectorize because polygons 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, :]
-            return np.cross(v1, v2)
-
         # note that the striding causes some polygons to have more coordinates
         # than others
         polyc = art3d.Poly3DCollection(polys, *args, **kwargs)
 
         if fcolors is not None:
             if shade:
                 colset = self._shade_colors(
-                    colset, get_normals(polys), lightsource)
+                    colset, self._generate_normals(polys), lightsource)
             polyc.set_facecolors(colset)
             polyc.set_edgecolors(colset)
         elif cmap:
@@ -1721,7 +1707,7 @@ def get_normals(polygons):
         else:
             if shade:
                 colset = self._shade_colors(
-                    color, get_normals(polys), lightsource)
+                    color, self._generate_normals(polys), lightsource)
             else:
                 colset = color
             polyc.set_facecolors(colset)
@@ -1732,21 +1718,48 @@ def get_normals(polygons):
         return polyc
 
     def _generate_normals(self, polygons):
-        '''
-        Generate normals for polygons by using the first three points.
-        This normal of course might not make sense for polygons with
-        more than three points not lying in a plane.
+        """
+        Takes a list of polygons and return an array of their normals.
 
         Normals point towards the viewer for a face with its vertices in
         counterclockwise order, following the right hand rule.
-        '''
 
-        normals = []
-        for verts in polygons:
-            v1 = np.array(verts[1]) - np.array(verts[0])
-            v2 = np.array(verts[2]) - np.array(verts[0])
-            normals.append(np.cross(v1, v2))
-        return normals
+        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_like
+            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):
         '''
@@ -1993,9 +2006,7 @@ def plot_trisurf(self, *args, color=None, norm=None, vmin=None, vmax=None,
                 polyc.set_norm(norm)
         else:
             if shade:
-                v1 = verts[:, 0, :] - verts[:, 1, :]
-                v2 = verts[:, 1, :] - verts[:, 2, :]
-                normals = np.cross(v1, v2)
+                normals = self._generate_normals(verts)
                 colset = self._shade_colors(color, normals, lightsource)
             else:
                 colset = color
@@ -2042,9 +2053,9 @@ def _3d_extend_contour(self, cset, stride=5):
                     botverts[0][i2],
                     botverts[0][i1]])
 
-                v1 = np.array(topverts[0][i1]) - np.array(topverts[0][i2])
-                v2 = np.array(topverts[0][i1]) - np.array(botverts[0][i1])
-                normals.append(np.cross(v1, v2))
+            # 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)
