Merge pull request #78 from Eric89GXL/vert-norm

ENH: Plot with vertex normals

Author: mwaskom

surfer/tests/test_utils.py

@@ -1,7 +1,7 @@
 from os.path import join as pjoin
 
 import numpy as np
-from numpy.testing import assert_array_almost_equal
+from numpy.testing import assert_array_almost_equal, assert_array_equal
 
 from surfer import utils
 
@@ -10,12 +10,36 @@
 data_path = pjoin(subj_dir, subject_id)
 
 
+def _slow_compute_normals(rr, tris):
+    """Efficiently compute vertex normals for triangulated surface"""
+    # first, compute triangle normals
+    r1 = rr[tris[:, 0], :]
+    r2 = rr[tris[:, 1], :]
+    r3 = rr[tris[:, 2], :]
+    tri_nn = np.cross((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]
+
+    # accumulate the normals
+    nn = np.zeros((len(rr), 3))
+    for p, verts in enumerate(tris):
+        nn[verts] += tri_nn[p, :]
+    size = np.sqrt(np.sum(nn * nn, axis=1))
+    size[size == 0] = 1.0  # prevent ugly divide-by-zero
+    nn /= size[:, np.newaxis]
+    return nn
+
+
 @utils.requires_fsaverage
 def test_surface():
     """Test IO for Surface class"""
     for subjects_dir in [None, subj_dir]:
         surface = utils.Surface('fsaverage', 'lh', 'inflated',
-                             subjects_dir=subjects_dir)
+                                subjects_dir=subjects_dir)
         surface.load_geometry()
         surface.load_label('BA1')
         surface.load_curvature()
@@ -25,3 +49,18 @@ def test_surface():
         surface.apply_xfm(xfm)
         x_ = surface.x
         assert_array_almost_equal(x + 2, x_)
+
+        # normals
+        nn = _slow_compute_normals(surface.coords, surface.faces[:10000])
+        nn_fast = utils._compute_normals(surface.coords, surface.faces[:10000])
+        assert_array_almost_equal(nn, nn_fast)
+
+
+def test_huge_cross():
+    """Test cross product with lots of elements
+    """
+    x = np.random.rand(100000, 3)
+    y = np.random.rand(1, 3)
+    z = np.cross(x, y)
+    zz = utils._fast_cross_3d(x, y)
+    assert_array_equal(z, zz)

surfer/utils.py

@@ -39,6 +39,8 @@ class Surface(object):
         The vertices coordinates
     faces : 2d array
         The faces ie. the triangles
+    nn : 2d array
+        Normalized surface normals for vertices.
     subjects_dir : str | None
         If not None, this directory will be used as the subjects directory
         instead of the value set using the SUBJECTS_DIR environment variable.
@@ -80,6 +82,7 @@ def load_geometry(self):
                 self.coords[:, 0] -= (np.max(self.coords[:, 0]) + self.offset)
             else:
                 self.coords[:, 0] -= (np.min(self.coords[:, 0]) + self.offset)
+        self.nn = _compute_normals(self.coords, self.faces)
 
     def save_geometry(self):
         surf_path = op.join(self.data_path, "surf",
@@ -128,6 +131,77 @@ def apply_xfm(self, mtx):
                                      mtx.T)[:, :3]
 
 
+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(rr, tris):
+    """Efficiently compute vertex normals for triangulated surface"""
+    # first, compute triangle normals
+    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], :]
+    size = np.sqrt(np.sum(nn * nn, axis=1))
+    size[size == 0] = 1.0  # prevent ugly divide-by-zero
+    nn /= size[:, np.newaxis]
+    return nn
+
+
 ###############################################################################
 # LOGGING (courtesy of mne-python)
 

surfer/viz.py

@@ -2022,6 +2022,9 @@ 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
+        self._geo_mesh.data.point_data.normals = self._geo.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)