Skip to content

Commit 50d898c

Browse files
eric-wiesereric-wieser
committed
MAINT: Unify calculation of normal vectors from polygons
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) ```
1 parent 82a4b24 commit 50d898c

1 file changed

Lines changed: 46 additions & 33 deletions

File tree

‎lib/mpl_toolkits/mplot3d/axes3d.py‎

Lines changed: 46 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1669,27 +1669,14 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
16691669
if fcolors is not None:
16701670
colset.append(fcolors[rs][cs])
16711671

1672-
def get_normals(polygons):
1673-
"""
1674-
Takes a list of polygons and return an array of their normals
1675-
"""
1676-
v1 = np.empty((len(polygons), 3))
1677-
v2 = np.empty((len(polygons), 3))
1678-
for poly_i, ps in enumerate(polygons):
1679-
# pick three points around the polygon at which to find the
1680-
# normal doesn't vectorize because polygons is jagged
1681-
i1, i2, i3 = 0, len(ps)//3, 2*len(ps)//3
1682-
v1[poly_i, :] = ps[i1, :] - ps[i2, :]
1683-
v2[poly_i, :] = ps[i2, :] - ps[i3, :]
1684-
return np.cross(v1, v2)
1685-
16861672
# note that the striding causes some polygons to have more coordinates
16871673
# than others
16881674
polyc = art3d.Poly3DCollection(polys, *args, **kwargs)
16891675

16901676
if fcolors is not None:
16911677
if shade:
1692-
colset = self._shade_colors(colset, get_normals(polys), lightsource)
1678+
colset = self._shade_colors(
1679+
colset, self._generate_normals(polys), lightsource)
16931680
polyc.set_facecolors(colset)
16941681
polyc.set_edgecolors(colset)
16951682
elif cmap:
@@ -1702,7 +1689,8 @@ def get_normals(polygons):
17021689
polyc.set_norm(norm)
17031690
else:
17041691
if shade:
1705-
colset = self._shade_colors(color, get_normals(polys), lightsource)
1692+
colset = self._shade_colors(
1693+
color, self._generate_normals(polys), lightsource)
17061694
else:
17071695
colset = color
17081696
polyc.set_facecolors(colset)
@@ -1713,18 +1701,45 @@ def get_normals(polygons):
17131701
return polyc
17141702

17151703
def _generate_normals(self, polygons):
1716-
'''
1717-
Generate normals for polygons by using the first three points.
1718-
This normal of course might not make sense for polygons with
1719-
more than three points not lying in a plane.
1720-
'''
1704+
"""
1705+
Takes a list of polygons and return an array of their normals.
1706+
1707+
Uses three points equally spaced around the polygon.
1708+
This normal of course might not make sense for polygons with more than
1709+
three points not lying in a plane, but it's a plausible and fast
1710+
approximation.
17211711
1722-
normals = []
1723-
for verts in polygons:
1724-
v1 = np.array(verts[0]) - np.array(verts[1])
1725-
v2 = np.array(verts[2]) - np.array(verts[0])
1726-
normals.append(np.cross(v1, v2))
1727-
return normals
1712+
Parameters
1713+
----------
1714+
polygons: list of (M_i, 3) array_like, or (..., M, 3) array_like
1715+
A sequence of polygons to compute normals for, which can have
1716+
varying numbers of vertices. If the polygons all have the same
1717+
number of vertices and array is passed, then the operation will
1718+
be vectorized.
1719+
1720+
Returns
1721+
-------
1722+
normals: (..., 3) array_like
1723+
A normal vector estimated for the polygon.
1724+
1725+
"""
1726+
if isinstance(polygons, np.ndarray):
1727+
# optimization: polygons all have the same number of points, so can
1728+
# vectorize
1729+
n = polygons.shape[-2]
1730+
i1, i2, i3 = 0, n//3, 2*n//3
1731+
v1 = ps[..., i1, :] - ps[..., i2, :]
1732+
v2 = ps[..., i2, :] - ps[..., i3, :]
1733+
else:
1734+
# The subtraction doesn't vectorize because polygons is jagged.
1735+
v1 = np.empty((len(polygons), 3))
1736+
v2 = np.empty((len(polygons), 3))
1737+
for poly_i, ps in enumerate(polygons):
1738+
n = len(ps)
1739+
i1, i2, i3 = 0, n//3, 2*n//3
1740+
v1[poly_i, :] = ps[i1, :] - ps[i2, :]
1741+
v2[poly_i, :] = ps[i2, :] - ps[i3, :]
1742+
return np.cross(v1, v2)
17281743

17291744
def _shade_colors(self, color, normals, lightsource=None):
17301745
'''
@@ -1971,9 +1986,7 @@ def plot_trisurf(self, *args, color=None, norm=None, vmin=None, vmax=None,
19711986
polyc.set_norm(norm)
19721987
else:
19731988
if shade:
1974-
v1 = verts[:, 0, :] - verts[:, 1, :]
1975-
v2 = verts[:, 1, :] - verts[:, 2, :]
1976-
normals = np.cross(v1, v2)
1989+
normals = self._generate_normals(verts)
19771990
colset = self._shade_colors(color, normals, lightsource)
19781991
else:
19791992
colset = color
@@ -2020,9 +2033,9 @@ def _3d_extend_contour(self, cset, stride=5):
20202033
botverts[0][i2],
20212034
botverts[0][i1]])
20222035

2023-
v1 = np.array(topverts[0][i1]) - np.array(topverts[0][i2])
2024-
v2 = np.array(topverts[0][i1]) - np.array(botverts[0][i1])
2025-
normals.append(np.cross(v1, v2))
2036+
# all polygons have 4 vertices, so vectorize
2037+
polyverts = np.array(polyverts)
2038+
normals = self._generate_normals(polyverts)
20262039

20272040
colors = self._shade_colors(color, normals)
20282041
colors2 = self._shade_colors(color, normals)

0 commit comments

Comments
 (0)