for indexes

Jammy2211 · Nov 1, 2024 · a323a36 · a323a36
1 parent ff7028b
commit a323a36
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/autoarray/structures/triangles/coordinate_array.py b/autoarray/structures/triangles/coordinate_array.py
@@ -176,3 +176,11 @@ def with_vertices(self, vertices: np.ndarray) -> ArrayTriangles:
             indices=self.indices,
             vertices=vertices,
         )
+
+    def for_indexes(self, indexes: np.ndarray) -> "CoordinateArrayTriangles":
+        return CoordinateArrayTriangles(
+            coordinates=self.coordinates[indexes],
+            side_length=self.side_length,
+            flipped=self.flipped,
+            offset=self.offset,
+        )
diff --git a/test_autoarray/structures/triangles/test_coordinate_implementation.py b/test_autoarray/structures/triangles/test_coordinate_implementation.py
@@ -255,3 +255,16 @@ def test_up_sampled_vertices(one_triangle):
 def test_with_vertices(one_triangle):
     triangle = one_triangle.with_vertices(np.array([[0, 0], [1, 0], [0.5, 1]]))
     assert np.all(triangle.triangles == [[[1.0, 0.0], [0.5, 1.0], [0.0, 0.0]]])
+
+
+def test_for_indexes(two_triangles):
+    assert np.all(
+        two_triangles.for_indexes(np.array([0])).triangles
+        == [
+            [
+                [0.0, 0.4330127018922193],
+                [0.5, -0.4330127018922193],
+                [-0.5, -0.4330127018922193],
+            ]
+        ]
+    )