浏览代码

added Surface._triangles_vertex_coords()

Fabian Peter Hammerle 6 年之前
父节点
当前提交
13a770b5bb
共有 2 个文件被更改,包括 22 次插入1 次删除
  1. 7 1
      freesurfer_surface/__init__.py
  2. 15 0
      tests/test_surface.py

+ 7 - 1
freesurfer_surface/__init__.py

@@ -181,7 +181,8 @@ class PolygonalChainsNotOverlapingError(ValueError):
 class PolygonalChain:
 
     def __init__(self, vertex_indices: typing.Iterable[int]):
-        self.vertex_indices = collections.deque(vertex_indices) # type: Deque[int]
+        self.vertex_indices \
+            = collections.deque(vertex_indices)  # type: Deque[int]
 
     def __eq__(self, other: 'PolygonalChain') -> bool:
         return self.vertex_indices == other.vertex_indices
@@ -513,3 +514,8 @@ class Surface:
 
     def _triangles_vertex_indices(self) -> numpy.ndarray:
         return numpy.vstack([t.vertex_indices for t in self.triangles])
+
+    def _triangles_vertex_coords(self) -> numpy.ndarray:
+        return numpy.take(self.vertices,
+                          indices=self._triangles_vertex_indices(),
+                          axis=0)

+ 15 - 0
tests/test_surface.py

@@ -522,3 +522,18 @@ def test__triangles_vertex_indices():
     surface.triangles.append(Triangle((5, 1, 4)))
     assert (surface._triangles_vertex_indices()
             == [[0, 1, 2], [3, 1, 2], [3, 1, 4], [5, 1, 4]]).all()
+
+
+def test__triangles_vertex_coords():
+    surface = Surface()
+    for i in range(6):
+        surface.add_vertex(Vertex(i, i, i))
+    surface.triangles.append(Triangle((0, 1, 2)))
+    surface.triangles.append(Triangle((3, 1, 2)))
+    surface.triangles.append(Triangle((3, 1, 4)))
+    surface.triangles.append(Triangle((5, 1, 4)))
+    assert (surface._triangles_vertex_coords()
+            == [[[0, 0, 0], [1, 1, 1], [2, 2, 2]],
+                [[3, 3, 3], [1, 1, 1], [2, 2, 2]],
+                [[3, 3, 3], [1, 1, 1], [4, 4, 4]],
+                [[5, 5, 5], [1, 1, 1], [4, 4, 4]]]).all()