feature/magnification

autolens/point/solver/point_solver.py

@@ -1,31 +1,22 @@
 import logging
-
-from typing import Tuple, List, Iterator, Optional
+from typing import Tuple, Optional
 
 import numpy as np
 
 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 AbstractSolver
 
 
 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)
-
     @jit
     def solve(
         self,
@@ -56,38 +47,36 @@ def solve(
         -------
         A list of image plane coordinates that are traced to the source plane coordinate.
         """
-        return super().solve(
+        kept_triangles = super().solve_triangles(
             tracer=tracer,
-            source_plane_coordinate=source_plane_coordinate,
+            shape=Point(*source_plane_coordinate),
             source_plane_redshift=source_plane_redshift,
         )
+        filtered_means = self._filter_low_magnification(
+            tracer=tracer, points=kept_triangles.means
+        )
 
-    # 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,
-            source_plane_redshift=source_plane_redshift,
-            **kwargs,
+        difference = len(kept_triangles.means) - len(filtered_means)
+        if difference > 0:
+            logger.debug(
+                f"Filtered one multiple-image with magnification below threshold."
+            )
+        elif difference > 1:
+            logger.warning(
+                f"Filtered {difference} multiple-images with magnification below threshold."
+            )
+
+        filtered_close = []
+
+        for mean in filtered_means:
+            if any(
+                np.linalg.norm(np.array(mean) - np.array(other))
+                <= self.pixel_scale_precision
+                for other in filtered_close
+            ):
+                continue
+            filtered_close.append(mean)
+
+        return aa.Grid2DIrregular(
+            [pair for pair in filtered_close if not np.isnan(pair).all()]
         )

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 AbstractSolver:
     # noinspection PyPep8Naming
     def __init__(
         self,
@@ -130,12 +131,12 @@ def _source_plane_grid(
         return grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections)
 
     @jit
-    def solve(
+    def solve_triangles(
         self,
         tracer: Tracer,
+        shape: Shape,
         source_plane_redshift: Optional[float] = None,
-        **kwargs,
-    ) -> aa.Grid2DIrregular:
+    ) -> AbstractTriangles:
         """
         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,41 +168,12 @@ def solve(
         steps = list(
             self.steps(
                 tracer=tracer,
+                shape=shape,
                 source_plane_redshift=source_plane_redshift,
-                **kwargs,
             )
         )
         final_step = steps[-1]
-        kept_triangles = final_step.filtered_triangles
-
-        filtered_means = self._filter_low_magnification(
-            tracer=tracer, points=kept_triangles.means
-        )
-
-        difference = len(kept_triangles.means) - len(filtered_means)
-        if difference > 0:
-            logger.debug(
-                f"Filtered one multiple-image with magnification below threshold."
-            )
-        elif difference > 1:
-            logger.warning(
-                f"Filtered {difference} multiple-images with magnification below threshold."
-            )
-
-        filtered_close = []
-
-        for mean in filtered_means:
-            if any(
-                np.linalg.norm(np.array(mean) - np.array(other))
-                <= self.pixel_scale_precision
-                for other in filtered_close
-            ):
-                continue
-            filtered_close.append(mean)
-
-        return aa.Grid2DIrregular(
-            [pair for pair in filtered_close if not np.isnan(pair).all()]
-        )
+        return final_step.filtered_triangles
 
     def _filter_low_magnification(
         self, tracer: Tracer, points: List[Tuple[float, float]]
@@ -233,7 +207,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 +219,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 +238,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 +258,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()
@@ -331,3 +297,34 @@ def tree_unflatten(cls, aux_data, children):
             magnification_threshold=aux_data[6],
             array_triangles_cls=aux_data[7],
         )
+
+
+class ShapeSolver(AbstractSolver):
+    def find_magnification(
+        self,
+        tracer: Tracer,
+        shape: Shape,
+        source_plane_redshift: Optional[float] = None,
+    ) -> float:
+        """
+        Find the magnification of the shape in the source plane.
+
+        Parameters
+        ----------
+        tracer
+            A tracer that traces the image plane to the source plane.
+        shape
+            The shape of an image plane pixel.
+        source_plane_redshift
+            The redshift of the source plane.
+
+        Returns
+        -------
+        The magnification of the shape in the source plane.
+        """
+        kept_triangles = super().solve_triangles(
+            tracer=tracer,
+            shape=shape,
+            source_plane_redshift=source_plane_redshift,
+        )
+        return kept_triangles.area / shape.area