simplify unit tests

Author: Jammy2211

test_autoarray/structures/triangles/conftest.py

@@ -1,5 +1,6 @@
 from autoarray.numpy_wrapper import np
 from autoarray.structures.triangles.array import ArrayTriangles
+from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
 
 from matplotlib import pyplot as plt
 
@@ -54,3 +55,18 @@ def triangles():
             ]
         ),
     )
+
+@pytest.fixture
+def one_triangle():
+    return CoordinateArrayTriangles(
+        coordinates=np.array([[0, 0]]),
+        side_length=1.0,
+    )
+
+
+@pytest.fixture
+def two_triangles():
+    return CoordinateArrayTriangles(
+        coordinates=np.array([[0, 0], [1, 0]]),
+        side_length=1.0,
+    )

test_autoarray/structures/triangles/coordinate/__init__.py

@@ -1,11 +1,19 @@
-import pytest
-
+from jax import numpy as np
+import jax
 import numpy as np
 
+jax.config.update("jax_log_compiles", True)
+
+import pytest
+
 from autoarray.structures.triangles.abstract import HEIGHT_FACTOR
-from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
 from autoarray.structures.triangles.shape import Point
 
+from autoarray.structures.triangles.coordinate_array import (
+    CoordinateArrayTriangles,
+)
+
+
 
 def test__two(two_triangles):
 
@@ -253,3 +261,113 @@ def test_from_grid_regression():
     y = triangles.vertices[:, 1]
     assert min(y) <= -4.75
     assert max(y) >= 4.75
+
+
+@pytest.fixture
+def one_triangle():
+    return CoordinateArrayTriangles(
+        coordinates=np.array([[0, 0]]),
+        side_length=1.0,
+    )
+
+
+@jax.jit
+def full_routine(triangles):
+    neighborhood = triangles.neighborhood()
+    up_sampled = neighborhood.up_sample()
+    with_vertices = up_sampled.with_vertices(up_sampled.vertices)
+    indexes = with_vertices.containing_indices(Point(0.1, 0.1))
+    return up_sampled.for_indexes(indexes)
+
+
+# def test_full_routine(one_triangle, compare_with_nans):
+#     result = full_routine(one_triangle)
+#
+#     assert compare_with_nans(
+#         result.triangles,
+#         np.array(
+#             [
+#                 [
+#                     [0.0, 0.4330126941204071],
+#                     [0.25, 0.0],
+#                     [-0.25, 0.0],
+#                 ]
+#             ]
+#         ),
+#     )
+
+
+def test_neighborhood(one_triangle):
+    assert np.allclose(
+        np.array(jax.jit(one_triangle.neighborhood)().triangles),
+        np.array(
+            [
+                [
+                    [-0.5, -0.4330126941204071],
+                    [-1.0, 0.4330126941204071],
+                    [0.0, 0.4330126941204071],
+                ],
+                [
+                    [0.0, -1.299038052558899],
+                    [-0.5, -0.4330126941204071],
+                    [0.5, -0.4330126941204071],
+                ],
+                [
+                    [0.0, 0.4330126941204071],
+                    [0.5, -0.4330126941204071],
+                    [-0.5, -0.4330126941204071],
+                ],
+                [
+                    [0.5, -0.4330126941204071],
+                    [0.0, 0.4330126941204071],
+                    [1.0, 0.4330126941204071],
+                ],
+            ]
+        ),
+    )
+
+
+def test_up_sample(one_triangle):
+    up_sampled = jax.jit(one_triangle.up_sample)()
+    assert np.allclose(
+        np.array(up_sampled.triangles),
+        np.array(
+            [
+                [
+                    [[0.0, -0.4330126941204071], [-0.25, 0.0], [0.25, 0.0]],
+                    [
+                        [0.25, 0.0],
+                        [0.5, -0.4330126941204071],
+                        [0.0, -0.4330126941204071],
+                    ],
+                    [
+                        [-0.25, 0.0],
+                        [0.0, -0.4330126941204071],
+                        [-0.5, -0.4330126941204071],
+                    ],
+                    [[0.0, 0.4330126941204071], [0.25, 0.0], [-0.25, 0.0]],
+                ]
+            ]
+        ),
+    )
+
+
+def test_means(one_triangle):
+    assert len(one_triangle.means) == 1
+
+    up_sampled = one_triangle.up_sample()
+    neighborhood = up_sampled.neighborhood()
+    assert np.count_nonzero(~np.isnan(neighborhood.means).any(axis=1)) == 10
+
+
+ONE_TRIANGLE_AREA = HEIGHT_FACTOR * 0.5
+
+
+def test_area(one_triangle):
+    assert one_triangle.area == ONE_TRIANGLE_AREA
+    assert one_triangle.up_sample().area == ONE_TRIANGLE_AREA
+
+    neighborhood = one_triangle.neighborhood()
+    assert neighborhood.area == 4 * ONE_TRIANGLE_AREA
+    assert neighborhood.up_sample().area == 4 * ONE_TRIANGLE_AREA
+    assert neighborhood.neighborhood().area == 10 * ONE_TRIANGLE_AREA

test_autoarray/structures/triangles/coordinate/conftest.py

@@ -1,19 +1,13 @@
-from autoarray.structures.triangles.shape import Point
-
-try:
-    from jax import numpy as np
-    import jax
+from jax import numpy as np
+import jax
 
-    jax.config.update("jax_log_compiles", True)
-    from autoarray.structures.triangles.array import ArrayTriangles
-except ImportError:
-    import numpy as np
-    from autoarray.structures.triangles.array import ArrayTriangles
+jax.config.update("jax_log_compiles", True)
 
 import pytest
 
 
-pytest.importorskip("jax")
+from autoarray.structures.triangles.shape import Point
+from autoarray.structures.triangles.array import ArrayTriangles
 
 
 @pytest.fixture