simplfiied coordinate_array and removed support for numpy

Author: Jammy2211

…coordinate_array/jax_coordinate_array.py …structures/triangles/coordinate_array.pyautoarray/structures/triangles/coordinate_array/jax_coordinate_array.py renamed to autoarray/structures/triangles/coordinate_array.py autoarray/structures/triangles/coordinate_array/jax_coordinate_array.py renamed to autoarray/structures/triangles/coordinate_array.py

@@ -1,3 +1,4 @@
+import numpy as np
 import jax.numpy as jnp
 import jax
 
@@ -12,6 +13,38 @@
 @register_pytree_node_class
 class CoordinateArrayTriangles:
 
+    def __init__(
+        self,
+        coordinates: np.ndarray,
+        side_length: float = 1.0,
+        x_offset: float = 0.0,
+        y_offset: float = 0.0,
+        flipped: bool = False,
+    ):
+        """
+        Represents a set of triangles by integer coordinates.
+
+        Parameters
+        ----------
+        coordinates
+            Integer x y coordinates for each triangle.
+        side_length
+            The side length of the triangles.
+        flipped
+            Whether the triangles are flipped upside down.
+        y_offset
+            An y_offset to apply to the y coordinates so that up-sampled triangles align.
+        """
+        self.coordinates = coordinates
+        self.side_length = side_length
+        self.flipped = flipped
+
+        self.scaling_factors = jnp.array(
+            [0.5 * side_length, HEIGHT_FACTOR * side_length]
+        )
+        self.x_offset = x_offset
+        self.y_offset = y_offset
+
     @classmethod
     def for_limits_and_scale(
         cls,
@@ -63,6 +96,12 @@ def tree_unflatten(cls, aux_data, children):
         """
         return cls(*children, flipped=aux_data[0])
 
+    def __len__(self):
+        return jnp.count_nonzero(~jnp.isnan(self.coordinates).any(axis=1))
+
+    def __iter__(self):
+        return iter(self.triangles)
+
     @property
     def centres(self) -> jnp.ndarray:
         """
@@ -73,6 +112,48 @@ def centres(self) -> jnp.ndarray:
         )
         return centres
 
+    @cached_property
+    def vertex_coordinates(self) -> np.ndarray:
+        """
+        The vertices of the triangles as an Nx3x2 array.
+        """
+        coordinates = self.coordinates
+        return jnp.concatenate(
+            [
+                coordinates + self.flip_array * np.array([0, 1], dtype=np.int32),
+                coordinates + self.flip_array * np.array([1, -1], dtype=np.int32),
+                coordinates + self.flip_array * np.array([-1, -1], dtype=np.int32),
+            ],
+            dtype=np.int32,
+        )
+
+    @cached_property
+    def triangles(self) -> np.ndarray:
+        """
+        The vertices of the triangles as an Nx3x2 array.
+        """
+        centres = self.centres
+        return jnp.stack(
+            (
+                centres
+                + self.flip_array
+                * jnp.array(
+                    [0.0, 0.5 * self.side_length * HEIGHT_FACTOR],
+                ),
+                centres
+                + self.flip_array
+                * jnp.array(
+                    [0.5 * self.side_length, -0.5 * self.side_length * HEIGHT_FACTOR]
+                ),
+                centres
+                + self.flip_array
+                * jnp.array(
+                    [-0.5 * self.side_length, -0.5 * self.side_length * HEIGHT_FACTOR]
+                ),
+            ),
+            axis=1,
+        )
+
     @cached_property
     def flip_mask(self) -> jnp.ndarray:
         """
@@ -93,9 +174,6 @@ def flip_array(self) -> jnp.ndarray:
         array = jnp.where(self.flip_mask, -1, 1)
         return array[:, None]
 
-    def __iter__(self):
-        return iter(self.triangles)
-
     def up_sample(self) -> "CoordinateArrayTriangles":
         """
         Up-sample the triangles by adding a new vertex at the midpoint of each edge.
@@ -226,4 +304,26 @@ def for_indexes(self, indexes: jnp.ndarray) -> "CoordinateArrayTriangles":
             y_offset=self.y_offset,
             x_offset=self.x_offset,
             flipped=self.flipped,
-        )
+        )
+
+    @property
+    def vertices(self) -> np.ndarray:
+        """
+        The unique vertices of the triangles.
+        """
+        return self._vertices_and_indices[0]
+
+    @property
+    def indices(self) -> np.ndarray:
+        """
+        The indices of the vertices of the triangles.
+        """
+        return self._vertices_and_indices[1]
+
+    @property
+    def means(self):
+        return jnp.mean(self.triangles, axis=1)
+
+    @property
+    def area(self):
+        return (3**0.5 / 4 * self.side_length**2) * len(self)

autoarray/structures/triangles/coordinate_array/__init__.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from autoarray.structures.triangles.coordinate_array import JAXCoordinateArrayTriangles as CoordinateArrayTriangles
+from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
 
 
 @pytest.fixture

autoarray/structures/triangles/coordinate_array/abstract_coordinate_array.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from autoarray.structures.triangles.abstract import HEIGHT_FACTOR
-from autoarray.structures.triangles.coordinate_array import JAXCoordinateArrayTriangles as CoordinateArrayTriangles
+from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
 from autoarray.structures.triangles.shape import Point