From 05f972a1b63cf66c541032659190e71fbf4c98d9 Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 6 Sep 2024 16:22:44 +0100 Subject: [PATCH 1/4] integrated new shape based array triangles --- autolens/point/solver/abstract_solver.py | 33 +++--- autolens/point/solver/circle_solver.py | 101 ------------------ autolens/point/solver/point_solver.py | 50 +-------- .../point/triangles/test_extended.py | 12 ++- 4 files changed, 27 insertions(+), 169 deletions(-) delete mode 100644 autolens/point/solver/circle_solver.py diff --git a/autolens/point/solver/abstract_solver.py b/autolens/point/solver/abstract_solver.py index 74ef6af3f..79bce4a7f 100644 --- a/autolens/point/solver/abstract_solver.py +++ b/autolens/point/solver/abstract_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/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..4d6a51560 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 .abstract_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/test_autolens/point/triangles/test_extended.py b/test_autolens/point/triangles/test_extended.py index b50353cab..763b5076b 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.abstract_solver import ShapeSolver @pytest.fixture def solver(grid): - return CircleSolver.for_grid( + return ShapeSolver.for_grid( grid=grid, pixel_scale_precision=0.01, ) @@ -15,8 +16,11 @@ 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), From 59e7cfc960f23843d9ca34b37bef5ad01f4d9081 Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 13 Sep 2024 13:00:58 +0100 Subject: [PATCH 2/4] rename and fix --- autolens/point/solver/point_solver.py | 2 +- .../point/solver/{abstract_solver.py => shape_solver.py} | 0 test_autolens/point/triangles/test_extended.py | 6 ++---- 3 files changed, 3 insertions(+), 5 deletions(-) rename autolens/point/solver/{abstract_solver.py => shape_solver.py} (100%) diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py index 4d6a51560..808dae6cc 100644 --- a/autolens/point/solver/point_solver.py +++ b/autolens/point/solver/point_solver.py @@ -6,7 +6,7 @@ from autoarray.structures.triangles.shape import Point from autofit.jax_wrapper import jit, register_pytree_node_class -from .abstract_solver import ShapeSolver +from .shape_solver import ShapeSolver from autolens.lens.tracer import Tracer diff --git a/autolens/point/solver/abstract_solver.py b/autolens/point/solver/shape_solver.py similarity index 100% rename from autolens/point/solver/abstract_solver.py rename to autolens/point/solver/shape_solver.py diff --git a/test_autolens/point/triangles/test_extended.py b/test_autolens/point/triangles/test_extended.py index 763b5076b..c40b8fe8c 100644 --- a/test_autolens/point/triangles/test_extended.py +++ b/test_autolens/point/triangles/test_extended.py @@ -2,7 +2,7 @@ from autoarray.structures.triangles.shape import Circle from autolens.mock import NullTracer -from autolens.point.solver.abstract_solver import ShapeSolver +from autolens.point.solver.shape_solver import ShapeSolver @pytest.fixture @@ -23,8 +23,6 @@ def test_solver_basic(solver): ), ) 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), ] From 80a9056fc356423a3c34e6bd263da43bd7c3e978 Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 13 Sep 2024 16:00:25 +0100 Subject: [PATCH 3/4] implementation to find magnification --- autolens/point/solver/point_solver.py | 37 +++++++++++-- autolens/point/solver/shape_solver.py | 52 +++++++------------ .../point/triangles/test_extended.py | 8 +-- 3 files changed, 54 insertions(+), 43 deletions(-) diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py index 808dae6cc..25fd804e0 100644 --- a/autolens/point/solver/point_solver.py +++ b/autolens/point/solver/point_solver.py @@ -1,12 +1,13 @@ import logging - 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 .shape_solver import ShapeSolver +from .shape_solver import AbstractSolver from autolens.lens.tracer import Tracer @@ -15,7 +16,7 @@ @register_pytree_node_class -class PointSolver(ShapeSolver): +class PointSolver(AbstractSolver): @jit def solve( self, @@ -46,8 +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, shape=Point(*source_plane_coordinate), source_plane_redshift=source_plane_redshift, ) + 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()] + ) diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py index 79bce4a7f..4caedf1a6 100644 --- a/autolens/point/solver/shape_solver.py +++ b/autolens/point/solver/shape_solver.py @@ -24,7 +24,7 @@ logger = logging.getLogger(__name__) -class ShapeSolver: +class AbstractSolver: # noinspection PyPep8Naming def __init__( self, @@ -131,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, - ) -> aa.Grid2DIrregular: + ) -> AbstractTriangles: """ Solve for the image plane coordinates that are traced to the source plane coordinate. @@ -173,36 +173,7 @@ def solve( ) ) 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]] @@ -326,3 +297,18 @@ 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: + 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 c40b8fe8c..1c07b0772 100644 --- a/test_autolens/point/triangles/test_extended.py +++ b/test_autolens/point/triangles/test_extended.py @@ -14,15 +14,11 @@ def solver(grid): def test_solver_basic(solver): - result = solver.solve( + assert solver.find_magnification( tracer=NullTracer(), shape=Circle( 0.0, 0.0, radius=0.01, ), - ) - assert list(map(tuple, result)) == [ - (-0.0029826688901819823, -0.0078125), - (-0.0029826688901819823, 0.0078125), - ] + ) == pytest.approx(1.0, abs=0.01) From 91e3028b8a90e3e8daa801b3a017c4777705fc8a Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 13 Sep 2024 16:01:12 +0100 Subject: [PATCH 4/4] docs --- autolens/point/solver/shape_solver.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py index 4caedf1a6..2717e0f52 100644 --- a/autolens/point/solver/shape_solver.py +++ b/autolens/point/solver/shape_solver.py @@ -306,6 +306,22 @@ def find_magnification( 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,