integrated new shape based array triangles

autolens/point/solver/point_solver.py

@@ -1,31 +1,21 @@
 import logging
 
-from typing import Tuple, List, Iterator, Optional
-
-import numpy as np
+from typing import Tuple, Optional
 
 import autoarray as aa
+from autoarray.structures.triangles.shape import Point
 
 from autofit.jax_wrapper import jit, register_pytree_node_class
-from .abstract_solver import AbstractSolver
+from .shape_solver import ShapeSolver
 
 
 from autolens.lens.tracer import Tracer
-from .step import Step
 
 logger = logging.getLogger(__name__)
 
 
 @register_pytree_node_class
-class PointSolver(AbstractSolver):
-    # noinspection PyMethodOverriding
-    def _filter_indexes(
-        self,
-        source_triangles: aa.AbstractTriangles,
-        source_plane_coordinate: Tuple[float, float],
-    ) -> np.ndarray:
-        return source_triangles.containing_indices(point=source_plane_coordinate)
-
+class PointSolver(ShapeSolver):
     @jit
     def solve(
         self,
@@ -58,36 +48,6 @@ def solve(
         """
         return super().solve(
             tracer=tracer,
-            source_plane_coordinate=source_plane_coordinate,
-            source_plane_redshift=source_plane_redshift,
-        )
-
-    # noinspection PyMethodOverriding
-    def steps(
-        self,
-        tracer: Tracer,
-        source_plane_coordinate: Tuple[float, float],
-        source_plane_redshift: Optional[float] = None,
-        **kwargs,
-    ) -> Iterator[Step]:
-        """
-        Iterate over the steps of the triangle solver algorithm.
-
-        Parameters
-        ----------
-        tracer
-            The tracer that traces from the image plane to the source plane.
-        source_plane_coordinate
-        source_plane_redshift
-            The redshift of the source plane.
-
-        Returns
-        -------
-        An iterator over the steps of the triangle solver algorithm.
-        """
-        yield from super().steps(
-            tracer=tracer,
-            source_plane_coordinate=source_plane_coordinate,
+            shape=Point(*source_plane_coordinate),
             source_plane_redshift=source_plane_redshift,
-            **kwargs,
         )

autolens/point/solver/abstract_solver.py → autolens/point/solver/shape_solver.py

@@ -1,12 +1,13 @@
 import logging
 import math
-from abc import ABC, abstractmethod
 
 from typing import Tuple, List, Iterator, Type, Optional
 
 import autoarray as aa
 
 import numpy as np
+
+from autoarray.structures.triangles.shape import Shape
 from autofit.jax_wrapper import jit, use_jax
 
 try:
@@ -23,7 +24,7 @@
 logger = logging.getLogger(__name__)
 
 
-class AbstractSolver(ABC):
+class ShapeSolver:
     # noinspection PyPep8Naming
     def __init__(
         self,
@@ -133,8 +134,8 @@ def _source_plane_grid(
     def solve(
         self,
         tracer: Tracer,
+        shape: Shape,
         source_plane_redshift: Optional[float] = None,
-        **kwargs,
     ) -> aa.Grid2DIrregular:
         """
         Solve for the image plane coordinates that are traced to the source plane coordinate.
@@ -150,6 +151,8 @@ def solve(
         ----------
         tracer
             The tracer to use to trace the image plane coordinates to the source plane.
+        shape
+            The shape in the source plane for which we want to identify the image plane coordinates.
         source_plane_redshift
             The redshift of the source plane.
 
@@ -165,8 +168,8 @@ def solve(
         steps = list(
             self.steps(
                 tracer=tracer,
+                shape=shape,
                 source_plane_redshift=source_plane_redshift,
-                **kwargs,
             )
         )
         final_step = steps[-1]
@@ -233,7 +236,7 @@ def _filtered_triangles(
         tracer: Tracer,
         triangles: aa.AbstractTriangles,
         source_plane_redshift,
-        **kwargs,
+        shape: Shape,
     ):
         """
         Filter the triangles to keep only those that meet the solver condition
@@ -245,23 +248,15 @@ def _filtered_triangles(
         )
         source_triangles = triangles.with_vertices(source_plane_grid.array)
 
-        return triangles.for_indexes(
-            indexes=self._filter_indexes(source_triangles, **kwargs)
-        )
+        indexes = source_triangles.containing_indices(shape=shape)
 
-    @abstractmethod
-    def _filter_indexes(
-        self,
-        source_triangles: aa.AbstractTriangles,
-        **kwargs,
-    ) -> np.ndarray:
-        pass
+        return triangles.for_indexes(indexes=indexes)
 
     def steps(
         self,
         tracer: Tracer,
+        shape: Shape,
         source_plane_redshift: Optional[float] = None,
-        **kwargs,
     ) -> Iterator[Step]:
         """
         Iterate over the steps of the triangle solver algorithm.
@@ -272,8 +267,8 @@ def steps(
             The tracer to use to trace the image plane coordinates to the source plane.
         source_plane_redshift
             The redshift of the source plane.
-        kwargs
-            Additional arguments to pass to the triangle filter.
+        shape
+            The shape in the source plane for which we want to identify the image plane coordinates.
 
         Returns
         -------
@@ -292,7 +287,7 @@ def steps(
                 tracer=tracer,
                 triangles=initial_triangles,
                 source_plane_redshift=source_plane_redshift,
-                **kwargs,
+                shape=shape,
             )
             neighbourhood = kept_triangles.neighborhood()
             up_sampled = neighbourhood.up_sample()

test_autolens/point/triangles/test_extended.py

@@ -1,12 +1,13 @@
 import pytest
 
+from autoarray.structures.triangles.shape import Circle
 from autolens.mock import NullTracer
-from autolens.point.solver.circle_solver import CircleSolver
+from autolens.point.solver.shape_solver import ShapeSolver
 
 
 @pytest.fixture
 def solver(grid):
-    return CircleSolver.for_grid(
+    return ShapeSolver.for_grid(
         grid=grid,
         pixel_scale_precision=0.01,
     )
@@ -15,12 +16,13 @@ def solver(grid):
 def test_solver_basic(solver):
     result = solver.solve(
         tracer=NullTracer(),
-        source_plane_coordinate=(0.0, 0.0),
-        radius=0.01,
+        shape=Circle(
+            0.0,
+            0.0,
+            radius=0.01,
+        ),
     )
     assert list(map(tuple, result)) == [
-        (-0.012003766846269881, 0.0078125),
         (-0.0029826688901819823, -0.0078125),
-        (0.015059527021993818, -0.0078125),
-        (0.010548978043949867, 0.015625),
+        (-0.0029826688901819823, 0.0078125),
     ]