Merge pull request vispy#2434 from asnt/vectorize_meshdata_vertex_normals

Author: djhoese

vispy/geometry/meshdata.py

@@ -327,21 +327,25 @@ def get_vertex_normals(self, indexed=None):
             The normals.
         """
         if self._vertex_normals is None:
-            faceNorms = self.get_face_normals()
-            vertFaces = self.get_vertex_faces()
-            self._vertex_normals = np.empty(self._vertices.shape,
-                                            dtype=np.float32)
-            for vindex in range(self._vertices.shape[0]):
-                faces = vertFaces[vindex]
-                if len(faces) == 0:
-                    self._vertex_normals[vindex] = (0, 0, 0)
-                    continue
-                norms = faceNorms[faces]  # get all face normals
-                norm = norms.sum(axis=0)  # sum normals
-                renorm = (norm**2).sum()**0.5
-                if renorm > 0:
-                    norm /= renorm
-                self._vertex_normals[vindex] = norm
+            face_normals = self.get_face_normals().astype(np.float32)
+            vertex_normals = np.zeros(self._vertices.shape, dtype=np.float32)
+
+            face_normals_repeated_on_face_vertices = np.repeat(face_normals,
+                                                               3,
+                                                               axis=0)
+            # NOTE: Cannot use the following intuitive code as it does not
+            # accumulate the values from the right hand side at repeated
+            # indices on the left hand side (instead, the value is overwritten
+            # at each occurence of an index):
+            #  vertex_normals[self._faces.ravel()] += face_normals_repeated_on_face_vertices
+            # The following does the desired accumulation at repeated indices:
+            np.add.at(vertex_normals, self._faces.ravel(),
+                      face_normals_repeated_on_face_vertices)
+
+            norms = (vertex_normals**2).sum(axis=1)**0.5
+            nonzero_norms = norms > 0
+            vertex_normals[nonzero_norms] /= norms[nonzero_norms][:, None]
+            self._vertex_normals = vertex_normals
 
         if indexed is None:
             return self._vertex_normals

vispy/geometry/tests/test_meshdata.py

@@ -31,4 +31,43 @@ def test_meshdata():
     assert_array_equal(square_edges, mesh.get_edges())
 
 
+def test_vertex_normals_indexed_none():
+    dtype_float = np.float32
+    dtype_int = np.int64
+    vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                        dtype=dtype_float)
+    faces = np.array([[0, 2, 1], [0, 3, 2], [0, 1, 3]], dtype=dtype_int)
+    mesh = MeshData(vertices=vertices, faces=faces)
+    vertex_normals_unnormalized = np.array(
+        [[-1, -1, -1], [0, -1, -1], [-1, 0, -1], [-1, -1, 0]],
+        dtype=dtype_float)
+    norms = np.sqrt((vertex_normals_unnormalized**2).sum(axis=1,
+                                                         keepdims=True))
+    expected_vertex_normals = vertex_normals_unnormalized / norms
+
+    computed_vertex_normals = mesh.get_vertex_normals(indexed=None)
+
+    assert_array_equal(expected_vertex_normals, computed_vertex_normals)
+
+
+def test_vertex_normals_indexed_faces():
+    dtype_float = np.float32
+    dtype_int = np.int64
+    vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                        dtype=dtype_float)
+    faces = np.array([[0, 2, 1], [0, 3, 2], [0, 1, 3]], dtype=dtype_int)
+    mesh = MeshData(vertices=vertices, faces=faces)
+    vertex_normals_unnormalized = np.array(
+        [[-1, -1, -1], [0, -1, -1], [-1, 0, -1], [-1, -1, 0]],
+        dtype=dtype_float)
+    norms = np.sqrt((vertex_normals_unnormalized**2).sum(axis=1,
+                                                         keepdims=True))
+    vertex_normals = vertex_normals_unnormalized / norms
+    expected_vertex_normals = vertex_normals[faces]
+
+    computed_vertex_normals = mesh.get_vertex_normals(indexed="faces")
+
+    assert_array_equal(expected_vertex_normals, computed_vertex_normals)
+
+
 run_tests_if_main()