From 5e6e837167a7f1654fbeb84dd7604bbdcd11dfcc Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Sun, 16 Sep 2018 15:05:22 -0700 Subject: [PATCH] MAINT: Unify calculation of normal vectors from polygons MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) ``` --- lib/mpl_toolkits/mplot3d/axes3d.py | 77 +++++++++++++++++------------- 1 file changed, 44 insertions(+), 33 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py index b8ad90f22b13..5fc47ba55069 100644 --- a/lib/mpl_toolkits/mplot3d/axes3d.py +++ b/lib/mpl_toolkits/mplot3d/axes3d.py @@ -1686,20 +1686,6 @@ 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) @@ -1707,7 +1693,7 @@ def get_normals(polygons): 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)