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..808dae6cc 100644 --- a/autolens/point/solver/point_solver.py +++ b/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, ) diff --git a/autolens/point/solver/abstract_solver.py b/autolens/point/solver/shape_solver.py similarity index 94% rename from autolens/point/solver/abstract_solver.py rename to autolens/point/solver/shape_solver.py index 74ef6af3f..79bce4a7f 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 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() diff --git a/test_autolens/point/triangles/test_extended.py b/test_autolens/point/triangles/test_extended.py index b50353cab..c40b8fe8c 100644 --- a/test_autolens/point/triangles/test_extended.py +++ b/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), ]