From 7848677988f6a21af8f42a4a10cbcb809c41da00 Mon Sep 17 00:00:00 2001
From: Richard <richard@rghsoftware.co.uk>
Date: Mon, 5 Aug 2024 09:00:39 +0100
Subject: [PATCH] revert jax changes to get working array triangles again

---
 autoarray/structures/triangles/array.py       | 31 ++-----------------
 .../triangles/test_array_representation.py    |  6 ++--
 .../structures/triangles/test_jax.py          | 21 -------------
 3 files changed, 5 insertions(+), 53 deletions(-)
 delete mode 100644 test_autoarray/structures/triangles/test_jax.py

diff --git a/autoarray/structures/triangles/array.py b/autoarray/structures/triangles/array.py
index 7b917c3c2..ddcf369d2 100644
--- a/autoarray/structures/triangles/array.py
+++ b/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],
-        )
diff --git a/test_autoarray/structures/triangles/test_array_representation.py b/test_autoarray/structures/triangles/test_array_representation.py
index 4877f8241..75eb232ad 100644
--- a/test_autoarray/structures/triangles/test_array_representation.py
+++ b/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]),
         ),
     ],
 )
diff --git a/test_autoarray/structures/triangles/test_jax.py b/test_autoarray/structures/triangles/test_jax.py
deleted file mode 100644
index 8fad46d76..000000000
--- a/test_autoarray/structures/triangles/test_jax.py
+++ /dev/null
@@ -1,21 +0,0 @@
-from autoarray.structures.triangles.array import ArrayTriangles
-
-
-def test_flatten(triangles):
-    (indices, vertices), _ = triangles.tree_flatten()
-
-    assert (indices == triangles.indices).all()
-    assert (vertices == triangles.vertices).all()
-
-
-def test_unflatten(triangles):
-    new_triangles = ArrayTriangles.tree_unflatten(
-        (),
-        (
-            triangles.indices,
-            triangles.vertices,
-        ),
-    )
-
-    assert (new_triangles.indices == triangles.indices).all()
-    assert (new_triangles.vertices == triangles.vertices).all()