diff --git a/autolens/point/solver/circle_solver.py b/autolens/point/solver/circle_solver.py deleted file mode 100644 index 437428338..000000000 --- a/autolens/point/solver/circle_solver.py +++ /dev/null @@ -1,101 +0,0 @@ -from typing import Tuple, Iterator, Optional - -import numpy as np - -import autoarray as aa - -from autofit.jax_wrapper import jit, register_pytree_node_class -from .abstract_solver import AbstractSolver - -from autolens.lens.tracer import Tracer -from .step import Step - - -@register_pytree_node_class -class CircleSolver(AbstractSolver): - # noinspection PyMethodOverriding - def _filter_indexes( - self, - source_triangles: aa.AbstractTriangles, - source_plane_coordinate: Tuple[float, float], - radius: float, - ) -> np.ndarray: - return source_triangles.containing_indices_circle( - center=source_plane_coordinate, - radius=radius, - ) - - @jit - def solve( - self, - tracer: Tracer, - source_plane_coordinate: Tuple[float, float], - radius: float, - source_plane_redshift: Optional[float] = None, - ) -> aa.Grid2DIrregular: - """ - Solve for the image plane coordinates that are traced to the a circle in the source plane. - - 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 - 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 - threshold. - - Parameters - ---------- - source_plane_coordinate - The source plane coordinate to trace to the image plane. - radius - The radius of the circle. - tracer - The tracer that traces the image plane coordinates to the source plane - source_plane_redshift - The redshift of the source plane coordinate. - - Returns - ------- - A list of image plane coordinates that are traced to the source plane coordinate. - """ - return super().solve( - tracer=tracer, - source_plane_coordinate=source_plane_coordinate, - radius=radius, - source_plane_redshift=source_plane_redshift, - ) - - # noinspection PyMethodOverriding - def steps( - self, - tracer: Tracer, - source_plane_coordinate: Tuple[float, float], - radius: 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 - The centre of the circle in the source plane. - radius - The radius of the circle - 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, - radius=radius, - **kwargs, - ) diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py index 27156bbaa..25fd804e0 100644 --- a/autolens/point/solver/point_solver.py +++ b/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()] ) diff --git a/autolens/point/solver/abstract_solver.py b/autolens/point/solver/shape_solver.py similarity index 84% rename from autolens/point/solver/abstract_solver.py rename to autolens/point/solver/shape_solver.py index 74ef6af3f..2717e0f52 100644 --- a/autolens/point/solver/abstract_solver.py +++ b/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 diff --git a/test_autolens/point/triangles/test_extended.py b/test_autolens/point/triangles/test_extended.py index b50353cab..1c07b0772 100644 --- a/test_autolens/point/triangles/test_extended.py +++ b/test_autolens/point/triangles/test_extended.py @@ -1,26 +1,24 @@ 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, ) def test_solver_basic(solver): - result = solver.solve( + assert solver.find_magnification( tracer=NullTracer(), - source_plane_coordinate=(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), - ] + shape=Circle( + 0.0, + 0.0, + radius=0.01, + ), + ) == pytest.approx(1.0, abs=0.01)