Working version of ArrayTriangles

autoarray/structures/triangles/array.py

@@ -1,12 +1,9 @@
 from typing import Tuple
 
 from autoarray import Grid2D
-from jax import numpy as np
-from jax.tree_util import register_pytree_node_class
-from jax import jit
+import numpy as np
 
 
-@register_pytree_node_class
 class ArrayTriangles:
     def __init__(
         self,
@@ -28,16 +25,13 @@ def __init__(
         self.vertices = vertices
 
     @property
-    @jit
     def triangles(self):
         return self.vertices[self.indices]
 
     @property
-    @jit
     def means(self):
         return np.mean(self.triangles, axis=1)
 
-    @jit
     def containing_indices(self, point: Tuple[float, float]) -> np.ndarray:
         """
         Find the triangles that contain a given point.
@@ -67,9 +61,8 @@ def containing_indices(self, point: Tuple[float, float]) -> np.ndarray:
 
         inside = (0 <= a) & (a <= 1) & (0 <= b) & (b <= 1) & (0 <= c) & (c <= 1)
 
-        return np.where(inside, size=5, fill_value=-1)[0]
+        return np.where(inside)[0]
 
-    @jit
     def for_indexes(self, indexes: np.ndarray) -> "ArrayTriangles":
         """
         Create a new ArrayTriangles containing indices and vertices corresponding to the given indexes
@@ -95,7 +88,6 @@ def for_indexes(self, indexes: np.ndarray) -> "ArrayTriangles":
 
         return ArrayTriangles(indices=new_indices, vertices=unique_vertices)
 
-    @jit
     def up_sample(self) -> "ArrayTriangles":
         """
         Up-sample the triangles by adding a new vertex at the midpoint of each edge.
@@ -128,7 +120,6 @@ def up_sample(self) -> "ArrayTriangles":
             vertices=unique_vertices,
         )
 
-    @jit
     def neighborhood(self) -> "ArrayTriangles":
         """
         Create a new set of triangles that are the neighborhood of the current triangles.
@@ -166,7 +157,6 @@ def neighborhood(self) -> "ArrayTriangles":
             vertices=unique_vertices,
         )
 
-    @jit
     def with_vertices(self, vertices: np.ndarray) -> "ArrayTriangles":
         """
         Create a new set of triangles with the vertices replaced.
@@ -186,7 +176,6 @@ def with_vertices(self, vertices: np.ndarray) -> "ArrayTriangles":
         )
 
     @classmethod
-    @jit
     def for_grid(cls, grid: Grid2D) -> "ArrayTriangles":
         """
         Create a grid of equilateral triangles from a regular grid.
@@ -278,19 +267,3 @@ def add_vertex(v):
 
     def __iter__(self):
         return iter(self.triangles)
-
-    def tree_flatten(self):
-        """
-        Flatten this model as a PyTree.
-        """
-        return (self.indices, self.vertices), ()
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        """
-        Unflatten a PyTree into a model.
-        """
-        return cls(
-            indices=children[0],
-            vertices=children[1],
-        )

test_autoarray/structures/triangles/test_array_representation.py

@@ -16,7 +16,7 @@
                     [1.0, 0.0],
                 ]
             ),
-            np.array([0, -1, -1, -1, -1]),
+            np.array([0]),
         ),
         (
             (0.6, 0.6),
@@ -27,7 +27,7 @@
                     [1.0, 1.0],
                 ]
             ),
-            np.array([1, -1, -1, -1, -1]),
+            np.array([1]),
         ),
         (
             (0.5, 0.5),
@@ -39,7 +39,7 @@
                     [1.0, 1.0],
                 ]
             ),
-            np.array([0, 1, -1, -1, -1]),
+            np.array([0, 1]),
         ),
     ],
 )