feature/coordinate array integration

autolens/point/solver/point_solver.py

@@ -28,7 +28,7 @@ def solve(
         Solve for the image plane coordinates that are traced to the source plane coordinate.
 
         This is done by tiling the image plane with triangles and checking if the source plane coordinate is contained
-        within the triangle. The triangles are subsampled to increase the resolution with only the triangles that
+        within the triangle. The triangles are sub-sampled to increase the resolution with only the triangles that
         contain the source plane coordinate and their neighbours being kept.
 
         The means of the triangles  are then filtered to keep only those with an absolute magnification above the
@@ -52,6 +52,7 @@ def solve(
             shape=Point(*source_plane_coordinate),
             source_plane_redshift=source_plane_redshift,
         )
+
         filtered_means = self._filter_low_magnification(
             tracer=tracer, points=kept_triangles.means
         )
@@ -77,7 +78,7 @@ def solve(
         for mean in filtered_means:
             if any(
                 np.linalg.norm(np.array(mean) - np.array(other))
-                <= self.pixel_scale_precision
+                <= self.pixel_scale_precision / 2
                 for other in filtered_close
             ):
                 continue

autolens/point/solver/shape_solver.py

@@ -4,6 +4,7 @@
 from typing import Tuple, List, Iterator, Type, Optional
 
 import autoarray as aa
+from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
 
 from autoarray.structures.triangles.shape import Shape
 from autofit.jax_wrapper import jit, use_jax, numpy as np, register_pytree_node_class
@@ -40,6 +41,7 @@ def __init__(
         initial_triangles: AbstractTriangles,
         pixel_scale_precision: float,
         magnification_threshold=0.1,
+        neighbor_degree: int = 1,
     ):
         """
         Determine the image plane coordinates that are traced to be a source plane coordinate.
@@ -50,12 +52,16 @@ def __init__(
 
         Parameters
         ----------
+        neighbor_degree
+            The number of times recursively add neighbors for the triangles that contain
+            the source plane coordinate.
         pixel_scale_precision
             The target pixel scale of the image grid.
         """
         self.scale = scale
         self.pixel_scale_precision = pixel_scale_precision
         self.magnification_threshold = magnification_threshold
+        self.neighbor_degree = neighbor_degree
 
         self.initial_triangles = initial_triangles
 
@@ -66,7 +72,7 @@ def for_grid(
         grid: aa.Grid2D,
         pixel_scale_precision: float,
         magnification_threshold=0.1,
-        array_triangles_cls: Type[AbstractTriangles] = ArrayTriangles,
+        array_triangles_cls: Type[AbstractTriangles] = CoordinateArrayTriangles,
         max_containing_size=MAX_CONTAINING_SIZE,
     ):
         """
@@ -278,7 +284,10 @@ def steps(
                 source_plane_redshift=source_plane_redshift,
                 shape=shape,
             )
-            neighbourhood = kept_triangles.neighborhood()
+            neighbourhood = kept_triangles
+            for _ in range(self.neighbor_degree):
+                neighbourhood = neighbourhood.neighborhood()
+
             up_sampled = neighbourhood.up_sample()
 
             yield Step(

autolens/point/visualise.py

@@ -22,3 +22,16 @@ def visualise(step: Step):
     plt.title(f"Step {step.number}")
     plt.gca().set_aspect("equal", adjustable="box")
     plt.show()
+
+
+def plot_triangles(triangles, color="black"):
+    plt.figure(figsize=(8, 8))
+    for triangle in triangles:
+        triangle = np.append(triangle, [triangle[0]], axis=0)
+        plt.plot(triangle[:, 0], triangle[:, 1], "o-", color=color)
+
+    plt.xlabel("X")
+    plt.ylabel("Y")
+    plt.title(f"Triangles")
+    plt.gca().set_aspect("equal", adjustable="box")
+    plt.show()

test_autolens/point/triangles/test_extended.py

@@ -1,5 +1,6 @@
 import pytest
 
+from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles
 from autoarray.structures.triangles.shape import Circle
 from autolens.mock import NullTracer
 from autolens.point.solver.shape_solver import ShapeSolver
@@ -9,7 +10,8 @@
 def solver(grid):
     return ShapeSolver.for_grid(
         grid=grid,
-        pixel_scale_precision=0.01,
+        pixel_scale_precision=0.001,
+        array_triangles_cls=CoordinateArrayTriangles,
     )
 
 
@@ -19,6 +21,6 @@ def test_solver_basic(solver):
         shape=Circle(
             0.0,
             0.0,
-            radius=0.01,
+            radius=0.1,
         ),
-    ) == pytest.approx(1.0, abs=0.01)
+    ) == pytest.approx(1.0, abs=0.1)

test_autolens/point/triangles/test_solver.py

@@ -4,7 +4,6 @@
 
 import autolens as al
 import autogalaxy as ag
-from autoarray.structures.triangles.array import ArrayTriangles
 from autolens.mock import NullTracer
 from autolens.point.solver import PointSolver
 
@@ -62,21 +61,21 @@ def test_trivial(
     solver = PointSolver.for_grid(
         grid=grid,
         pixel_scale_precision=0.01,
-        array_triangles_cls=ArrayTriangles,
     )
-    (coordinates,) = solver.solve(
+    coordinates = solver.solve(
         tracer=NullTracer(),
         source_plane_coordinate=source_plane_coordinate,
     )
-    assert coordinates == pytest.approx(source_plane_coordinate, abs=1.0e-1)
+
+    assert coordinates[0] == pytest.approx(source_plane_coordinate, abs=1.0e-1)
 
 
 def test_real_example(grid, tracer):
     solver = PointSolver.for_grid(
         grid=grid,
         pixel_scale_precision=0.001,
-        array_triangles_cls=ArrayTriangles,
     )
+
     result = solver.solve(tracer=tracer, source_plane_coordinate=(0.07, 0.07))
 
     assert len(result) == 5

test_autolens/point/triangles/test_solver_jax.py

@@ -6,7 +6,7 @@
 import autogalaxy as ag
 import autofit as af
 import numpy as np
-from autolens import PointSolver
+from autolens import PointSolver, Tracer
 
 try:
     from autoarray.structures.triangles.jax_array import ArrayTriangles
@@ -33,14 +33,19 @@ def solver(grid):
 
 
 def test_solver(solver):
-    tracer = ag.mp.Isothermal(
+    mass_profile = ag.mp.Isothermal(
         centre=(0.0, 0.0),
         einstein_radius=1.0,
     )
-    assert solver.solve(
+    tracer = Tracer(
+        galaxies=[ag.Galaxy(redshift=0.5, mass=mass_profile)],
+    )
+    result = solver.solve(
         tracer,
         source_plane_coordinate=(0.0, 0.0),
     )
+    print(result)
+    assert result
 
 
 @pytest.mark.parametrize(