ENH: Add vertex normals for smooth surface

Author: larsoner

surfer/viz.py

@@ -2022,6 +2022,10 @@ def __init__(self, subject_id, hemi, surf, figure, geo, curv, title,
         x, y, z, f = self._geo.x, self._geo.y, self._geo.z, self._geo.faces
         self._geo_mesh = mlab.pipeline.triangular_mesh_source(x, y, z, f,
                                                               **meshargs)
+        # add surface normals
+        nn = _compute_normals(x, y, z, f)
+        self._geo_mesh.data.point_data.normals = nn
+        self._geo_mesh.data.cell_data.normals = None
         self._geo_surf = mlab.pipeline.surface(self._geo_mesh,
                                                figure=self._f, reset_zoom=True,
                                                **kwargs)
@@ -2418,6 +2422,75 @@ def _format_colorbar(self):
             self.neg_bar.scalar_bar_representation.position2 = (0.42, 0.09)
 
 
+def _fast_cross_3d(x, y):
+    """Compute cross product between list of 3D vectors
+
+    Much faster than np.cross() when the number of cross products
+    becomes large (>500). This is because np.cross() methods become
+    less memory efficient at this stage.
+
+    Parameters
+    ----------
+    x : array
+        Input array 1.
+    y : array
+        Input array 2.
+
+    Returns
+    -------
+    z : array
+        Cross product of x and y.
+
+    Notes
+    -----
+    x and y must both be 2D row vectors. One must have length 1, or both
+    lengths must match.
+    """
+    assert x.ndim == 2
+    assert y.ndim == 2
+    assert x.shape[1] == 3
+    assert y.shape[1] == 3
+    assert (x.shape[0] == 1 or y.shape[0] == 1) or x.shape[0] == y.shape[0]
+    if max([x.shape[0], y.shape[0]]) >= 500:
+        return np.c_[x[:, 1] * y[:, 2] - x[:, 2] * y[:, 1],
+                     x[:, 2] * y[:, 0] - x[:, 0] * y[:, 2],
+                     x[:, 0] * y[:, 1] - x[:, 1] * y[:, 0]]
+    else:
+        return np.cross(x, y)
+
+
+def _compute_normals(x, y, z, tris):
+    """Efficiently compute vertex normals for triangulated surface"""
+    # first, compute triangle normals
+    rr = np.array([x, y, z]).T
+    r1 = rr[tris[:, 0], :]
+    r2 = rr[tris[:, 1], :]
+    r3 = rr[tris[:, 2], :]
+    tri_nn = _fast_cross_3d((r2 - r1), (r3 - r1))
+
+    #   Triangle normals and areas
+    size = np.sqrt(np.sum(tri_nn * tri_nn, axis=1))
+    zidx = np.where(size == 0)[0]
+    size[zidx] = 1.0  # prevent ugly divide-by-zero
+    tri_nn /= size[:, np.newaxis]
+
+    npts = len(rr)
+
+    # the following code replaces this, but is faster (vectorized):
+    #
+    # for p, verts in enumerate(tris):
+    #     nn[verts, :] += tri_nn[p, :]
+    #
+    nn = np.zeros((npts, 3))
+    for verts in tris.T:  # note this only loops 3x (number of verts per tri)
+        counts = np.bincount(verts, minlength=npts)
+        reord = np.argsort(verts)
+        vals = np.r_[np.zeros((1, 3)), np.cumsum(tri_nn[reord, :], 0)]
+        idx = np.cumsum(np.r_[0, counts])
+        nn += vals[idx[1:], :] - vals[idx[:-1], :]
+    return nn
+
+
 class TimeViewer(HasTraits):
     """TimeViewer object providing a GUI for visualizing time series