Merge pull request #252 from PyAutoLabs/feature/unit-test-profiling

perf: speed up unit tests 63% by removing JAX from triangle tests

Author: Jammy2211

test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py

@@ -260,13 +260,13 @@ def test__data_vector_via_weighted_data_two_methods_agree():
 
 def test__curvature_matrix_via_psf_weighted_noise_two_methods_agree():
 
-    mask = aa.Mask2D.circular(shape_native=(51, 51), pixel_scales=0.1, radius=2.0)
+    mask = aa.Mask2D.circular(shape_native=(21, 21), pixel_scales=0.1, radius=0.8)
 
     noise_map = np.random.uniform(size=mask.shape_native)
     noise_map = aa.Array2D(values=noise_map, mask=mask)
 
     kernel = aa.Convolver.from_gaussian(
-        shape_native=(7, 7), pixel_scales=mask.pixel_scales, sigma=1.0, normalize=True
+        shape_native=(5, 5), pixel_scales=mask.pixel_scales, sigma=1.0, normalize=True
     )
 
     psf = kernel
@@ -277,7 +277,7 @@ def test__curvature_matrix_via_psf_weighted_noise_two_methods_agree():
         psf=psf.kernel.native,
     )
 
-    mesh = aa.mesh.RectangularAdaptDensity(shape=(20, 20))
+    mesh = aa.mesh.RectangularAdaptDensity(shape=(8, 8))
 
     interpolator = mesh.interpolator_from(
         source_plane_data_grid=mask.derive_grid.unmasked,

test_autoarray/structures/triangles/conftest.py

@@ -1,9 +1,9 @@
-import jax.numpy as jnp
+import numpy as np
 from matplotlib import pyplot as plt
 import pytest
 
-from autoarray.structures.triangles.array import ArrayTriangles
-from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
+from autoarray.structures.triangles.array_np import ArrayTrianglesNp
+from autoarray.structures.triangles.coordinate_array_np import CoordinateArrayTrianglesNp
 
 
 @pytest.fixture
@@ -12,8 +12,8 @@ def plot():
 
     def plot(triangles, color="black"):
         for triangle in triangles:
-            triangle = jnp.array(triangle)
-            triangle = jnp.append(triangle, jnp.array([triangle[0]]), axis=0)
+            triangle = np.array(triangle)
+            triangle = np.append(triangle, np.array([triangle[0]]), axis=0)
             plt.plot(triangle[:, 0], triangle[:, 1], "o-", color=color)
 
     yield plot
@@ -24,27 +24,27 @@ def plot(triangles, color="black"):
 @pytest.fixture
 def compare_with_nans():
     def compare_with_nans_(arr1, arr2):
-        nan_mask1 = jnp.isnan(arr1)
-        nan_mask2 = jnp.isnan(arr2)
+        nan_mask1 = np.isnan(arr1)
+        nan_mask2 = np.isnan(arr2)
 
         arr1 = arr1[~nan_mask1]
         arr2 = arr2[~nan_mask2]
 
-        return jnp.all(arr1 == arr2)
+        return np.all(arr1 == arr2)
 
     return compare_with_nans_
 
 
 @pytest.fixture
 def triangles():
-    return ArrayTriangles(
-        indices=jnp.array(
+    return ArrayTrianglesNp(
+        indices=np.array(
             [
                 [0, 1, 2],
                 [1, 2, 3],
             ]
         ),
-        vertices=jnp.array(
+        vertices=np.array(
             [
                 [0.0, 0.0],
                 [1.0, 0.0],
@@ -57,15 +57,15 @@ def triangles():
 
 @pytest.fixture
 def one_triangle():
-    return CoordinateArrayTriangles(
-        coordinates=jnp.array([[0, 0]]),
+    return CoordinateArrayTrianglesNp(
+        coordinates=np.array([[0, 0]]),
         side_length=1.0,
     )
 
 
 @pytest.fixture
 def two_triangles():
-    return CoordinateArrayTriangles(
-        coordinates=jnp.array([[0, 0], [1, 0]]),
+    return CoordinateArrayTrianglesNp(
+        coordinates=np.array([[0, 0], [1, 0]]),
         side_length=1.0,
     )

test_autoarray/structures/triangles/test_area.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from autoarray.structures.triangles.array import ArrayTriangles
+from autoarray.structures.triangles.array_np import ArrayTrianglesNp as ArrayTriangles
 from autoarray.structures.triangles.shape import Triangle, Circle, Square, Polygon
 
 

test_autoarray/structures/triangles/test_coordinate.py

@@ -1,16 +1,13 @@
 import numpy as np
-from jax.tree_util import register_pytree_node_class
 
 import pytest
 
 from autoarray.structures.triangles.abstract import HEIGHT_FACTOR
 
-from autoarray.structures.triangles.coordinate_array import (
-    CoordinateArrayTriangles,
+from autoarray.structures.triangles.coordinate_array_np import (
+    CoordinateArrayTrianglesNp,
 )
 
-CoordinateArrayTriangles = register_pytree_node_class(CoordinateArrayTriangles)
-
 
 def test__two(two_triangles):
 
@@ -52,7 +49,7 @@ def test__trivial_triangles(one_triangle):
 
 
 def test__above():
-    triangles = CoordinateArrayTriangles(
+    triangles = CoordinateArrayTrianglesNp(
         coordinates=np.array([[0, 1]]),
         side_length=1.0,
     )
@@ -87,7 +84,7 @@ def test__above():
 
 @pytest.fixture
 def upside_down():
-    return CoordinateArrayTriangles(
+    return CoordinateArrayTrianglesNp(
         coordinates=np.array([[1, 0]]),
         side_length=1.0,
     )
@@ -279,104 +276,37 @@ def test_means(one_triangle):
 
 
 def test_triangles_touch():
-    triangles = CoordinateArrayTriangles(
+    triangles = CoordinateArrayTrianglesNp(
         np.array([[0, 0], [2, 0]]),
     )
 
     assert max(triangles.triangles[0][:, 0]) == min(triangles.triangles[1][:, 0])
 
-    triangles = CoordinateArrayTriangles(
+    triangles = CoordinateArrayTrianglesNp(
         np.array([[0, 0], [0, 1]]),
     )
     assert max(triangles.triangles[0][:, 1]) == min(triangles.triangles[1][:, 1])
 
 
 def test_from_grid_regression():
-    triangles = CoordinateArrayTriangles.for_limits_and_scale(
-        x_min=-4.75,
-        x_max=4.75,
-        y_min=-4.75,
-        y_max=4.75,
-        scale=0.5,
+    triangles = CoordinateArrayTrianglesNp.for_limits_and_scale(
+        x_min=-2.0,
+        x_max=2.0,
+        y_min=-2.0,
+        y_max=2.0,
+        scale=1.5,
     )
 
     x = triangles.vertices[:, 0]
-    assert min(x) <= -4.75
-    assert max(x) >= 4.75
+    assert min(x) <= -2.0
+    assert max(x) >= 2.0
 
     y = triangles.vertices[:, 1]
-    assert min(y) <= -4.75
-    assert max(y) >= 4.75
-
+    assert min(y) <= -2.0
+    assert max(y) >= 2.0
 
-@pytest.fixture
-def one_triangle():
-    return CoordinateArrayTriangles(
-        coordinates=np.array([[0, 0]]),
-        side_length=1.0,
-    )
 
-
-def test_neighborhood(one_triangle):
-    import jax
-
-    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):
-    import jax
-
-    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):
+def test_means_up_sampled(one_triangle):
     assert len(one_triangle.means) == 1
 
     up_sampled = one_triangle.up_sample()

test_autoarray/structures/triangles/test_extended_source.py

@@ -1,13 +1,13 @@
 import pytest
 import numpy as np
 
-from autoarray.structures.triangles.array import ArrayTriangles
+from autoarray.structures.triangles.array_np import ArrayTrianglesNp
 from autoarray.structures.triangles.shape import Circle
 
 
 @pytest.fixture
 def triangles():
-    return ArrayTriangles(
+    return ArrayTrianglesNp(
         indices=np.array(
             [
                 [0, 1, 2],
@@ -49,7 +49,7 @@ def test_small_point(triangles, point, indices):
             radius=0.001,
         )
     )
-    assert [i for i in containing_triangles.tolist() if i != -1] == indices
+    assert containing_triangles.tolist() == indices
 
 
 @pytest.mark.parametrize(
@@ -72,4 +72,4 @@ def test_large_circle(
             radius=radius,
         )
     )
-    assert [i for i in containing_triangles.tolist() if i != -1] == indices
+    assert containing_triangles.tolist() == indices