added Surface.select_vertices()

freesurfer_surface/__init__.py

@@ -515,10 +515,14 @@ class Surface:
     def _triangles_vertex_indices(self) -> numpy.ndarray:
         return numpy.vstack([t.vertex_indices for t in self.triangles])
 
+    def select_vertices(self, vertex_indices: typing.Iterable[int]
+                        ) -> numpy.ndarray:
+        if not hasattr(vertex_indices, '__getitem__'):
+            vertex_indices = list(vertex_indices)
+        return numpy.take(self.vertices, indices=vertex_indices, axis=0)
+
     def _triangles_vertex_coords(self) -> numpy.ndarray:
-        return numpy.take(self.vertices,
-                          indices=self._triangles_vertex_indices(),
-                          axis=0)
+        return self.select_vertices(self._triangles_vertex_indices())
 
     def _triangles_point_normal(self) -> typing.Tuple[numpy.ndarray,
                                                       numpy.ndarray]:

tests/test_surface.py

@@ -524,6 +524,20 @@ def test__triangles_vertex_indices():
             == [[0, 1, 2], [3, 1, 2], [3, 1, 4], [5, 1, 4]]).all()
 
 
+def test_select_vertices():
+    surface = Surface()
+    for i in range(4):
+        surface.add_vertex(Vertex(i, i, i))
+    assert (surface.select_vertices([2, 1]) \
+            == [surface.vertices[2], surface.vertices[1]]).all()
+    assert (surface.select_vertices((3, 2)) \
+            == [surface.vertices[3], surface.vertices[2]]).all()
+    assert (surface.select_vertices((3, 2)) \
+            == [[3, 3, 3], [2, 2, 2]]).all()
+    assert (surface.select_vertices(filter(lambda i: i % 2 == 1, range(4))) \
+            == [[1, 1, 1], [3, 3, 3]]).all()
+
+
 def test__triangles_vertex_coords():
     surface = Surface()
     for i in range(6):