diff --git a/autolens/mock.py b/autolens/mock.py index 52f9ec62d..bef94bcbe 100644 --- a/autolens/mock.py +++ b/autolens/mock.py @@ -1,8 +1,19 @@ -from autofit.mock import * -from autoarray.mock import * -from autogalaxy.mock import * - -from autolens.imaging.mock.mock_fit_imaging import MockFitImaging -from autolens.lens.mock.mock_tracer import MockTracer -from autolens.lens.mock.mock_tracer import MockTracerPoint -from autolens.point.mock.mock_point_solver import MockPointSolver +import numpy as np + +from autofit.mock import * # noqa +from autoarray.mock import * # noqa +from autogalaxy.mock import * # noqa +from autolens import Tracer + +from autolens.imaging.mock.mock_fit_imaging import MockFitImaging # noqa +from autolens.lens.mock.mock_tracer import MockTracer # noqa +from autolens.lens.mock.mock_tracer import MockTracerPoint # noqa +from autolens.point.mock.mock_point_solver import MockPointSolver # noqa + + +class NullTracer(Tracer): + def __init__(self): + super().__init__([]) + + def deflections_yx_2d_from(self, grid): + return np.zeros_like(grid.array) diff --git a/autolens/point/analysis.py b/autolens/point/analysis.py index 256fbbaf1..002801638 100644 --- a/autolens/point/analysis.py +++ b/autolens/point/analysis.py @@ -56,7 +56,7 @@ def log_likelihood_function(self, instance): """ lens = instance.lens - solver = TriangleSolver( + solver = TriangleSolver.for_grid( lensing_obj=lens, grid=self.grid, pixel_scale_precision=self.pixel_scale_precision, @@ -147,7 +147,7 @@ def _log_likelihood_for_coordinates( "The number of predicted coordinates must be equal to the number of observed coordinates." ) - predicted_coordinates = set(predicted_coordinates) + predicted_coordinates = set(map(tuple, predicted_coordinates)) observed_coordinates = set(self.observed_coordinates) log_likelihood = 0.0 diff --git a/autolens/point/triangles/triangle_solver.py b/autolens/point/triangles/triangle_solver.py index d9b3887d7..d26f429eb 100644 --- a/autolens/point/triangles/triangle_solver.py +++ b/autolens/point/triangles/triangle_solver.py @@ -1,24 +1,56 @@ import logging import math -from typing import Tuple, List +from dataclasses import dataclass +from typing import Tuple, List, Iterator, Type from autoarray import Grid2D, Grid2DIrregular -from autoarray.structures.triangles.subsample_triangles import SubsampleTriangles -from autoarray.structures.triangles.triangles import Triangles +from autoarray.structures.triangles import array +from autoarray.structures.triangles.abstract import AbstractTriangles from autoarray.type import Grid2DLike from autogalaxy import OperateDeflections - logger = logging.getLogger(__name__) +@dataclass +class Step: + """ + A step in the triangle solver algorithm. + + Attributes + ---------- + number + The number of the step. + initial_triangles + The triangles at the start of the step. + filtered_triangles + The triangles trace to triangles that contain the source plane coordinate. + neighbourhood + The neighbourhood of the filtered triangles. + up_sampled + The neighbourhood up-sampled to increase the resolution. + """ + + number: int + initial_triangles: AbstractTriangles + filtered_triangles: AbstractTriangles + neighbourhood: AbstractTriangles + up_sampled: AbstractTriangles + + class TriangleSolver: + # noinspection PyPep8Naming def __init__( self, lensing_obj: OperateDeflections, - grid: Grid2D, + scale: float, + y_min: float, + y_max: float, + x_min: float, + x_max: float, pixel_scale_precision: float, magnification_threshold=0.1, + ArrayTriangles: Type[AbstractTriangles] = array.ArrayTriangles, ): """ Determine the image plane coordinates that are traced to be a source plane coordinate. @@ -31,22 +63,59 @@ def __init__( ---------- lensing_obj A tracer describing the lensing system. - grid - The grid of image plane coordinates. pixel_scale_precision The target pixel scale of the image grid. + ArrayTriangles + The class to use for the triangles. """ self.lensing_obj = lensing_obj - self.grid = grid + self.scale = scale + self.y_min = y_min + self.y_max = y_max + self.x_min = x_min + self.x_max = x_max self.pixel_scale_precision = pixel_scale_precision self.magnification_threshold = magnification_threshold + self.ArrayTriangles = ArrayTriangles + + # noinspection PyPep8Naming + @classmethod + def for_grid( + cls, + lensing_obj: OperateDeflections, + grid: Grid2D, + pixel_scale_precision: float, + magnification_threshold=0.1, + ArrayTriangles: Type[AbstractTriangles] = array.ArrayTriangles, + ): + scale = grid.pixel_scale + + y = grid[:, 0] + x = grid[:, 1] + + y_min = y.min() + y_max = y.max() + x_min = x.min() + x_max = x.max() + + return cls( + lensing_obj=lensing_obj, + scale=scale, + y_min=y_min, + y_max=y_max, + x_min=x_min, + x_max=x_max, + pixel_scale_precision=pixel_scale_precision, + magnification_threshold=magnification_threshold, + ArrayTriangles=ArrayTriangles, + ) @property def n_steps(self) -> int: """ How many times should triangles be subdivided? """ - return math.ceil(math.log2(self.grid.pixel_scale / self.pixel_scale_precision)) + return math.ceil(math.log2(self.scale / self.pixel_scale_precision)) def _source_plane_grid(self, grid: Grid2DLike) -> Grid2DLike: """ @@ -87,31 +156,18 @@ def solve( ------- A list of image plane coordinates that are traced to the source plane coordinate. """ - triangles = Triangles.for_grid(grid=self.grid) - if self.n_steps == 0: raise ValueError( "The target pixel scale is too large to subdivide the triangles." ) - kept_triangles = [] - - for _ in range(self.n_steps): - kept_triangles = self._filter_triangles( - triangles=triangles, - source_plane_coordinate=source_plane_coordinate, - ) - with_neighbourhood = { - triangle - for kept_triangle in kept_triangles - for triangle in kept_triangle.neighbourhood - } - triangles = SubsampleTriangles(parent_triangles=list(with_neighbourhood)) + steps = list(self.steps(source_plane_coordinate=source_plane_coordinate)) + final_step = steps[-1] + kept_triangles = final_step.filtered_triangles - means = [triangle.mean for triangle in kept_triangles] - filtered_means = self._filter_low_magnification(points=means) + filtered_means = self._filter_low_magnification(points=kept_triangles.means) - difference = len(means) - len(filtered_means) + difference = len(kept_triangles.means) - len(filtered_means) if difference > 0: logger.debug( f"Filtered one multiple-image with magnification below threshold." @@ -144,7 +200,7 @@ def _filter_low_magnification( points, self.lensing_obj.magnification_2d_via_hessian_from( grid=Grid2DIrregular(points), - buffer=self.grid.pixel_scale, + buffer=self.scale, ), ) if abs(magnification) > self.magnification_threshold @@ -152,7 +208,7 @@ def _filter_low_magnification( def _filter_triangles( self, - triangles: Triangles, + triangles: AbstractTriangles, source_plane_coordinate: Tuple[float, float], ): """ @@ -169,16 +225,51 @@ def _filter_triangles( ------- The triangles that contain the source plane coordinate. """ - source_plane_grid = self._source_plane_grid(grid=triangles.grid_2d) - - kept_triangles = [] - for image_triangle, source_triangle in zip( - triangles.triangles, - triangles.with_updated_grid(source_plane_grid), - ): - if source_triangle.contains( - point=source_plane_coordinate, - ): - kept_triangles.append(image_triangle) - - return kept_triangles + source_plane_grid = self._source_plane_grid( + grid=Grid2DIrregular(triangles.vertices) + ) + source_triangles = triangles.with_vertices(source_plane_grid.array) + indexes = source_triangles.containing_indices(point=source_plane_coordinate) + return triangles.for_indexes(indexes=indexes) + + def steps( + self, + source_plane_coordinate: Tuple[float, float], + ) -> Iterator[Step]: + """ + Iterate over the steps of the triangle solver algorithm. + + Parameters + ---------- + source_plane_coordinate + The source plane coordinate to trace to the image plane. + + Returns + ------- + An iterator over the steps of the triangle solver algorithm. + """ + initial_triangles = self.ArrayTriangles.for_limits_and_scale( + y_min=self.y_min, + y_max=self.y_max, + x_min=self.x_min, + x_max=self.x_max, + scale=self.scale, + ) + + for number in range(self.n_steps): + kept_triangles = self._filter_triangles( + initial_triangles, + source_plane_coordinate, + ) + neighbourhood = kept_triangles.neighborhood() + up_sampled = neighbourhood.up_sample() + + yield Step( + number=number, + initial_triangles=initial_triangles, + filtered_triangles=kept_triangles, + neighbourhood=neighbourhood, + up_sampled=up_sampled, + ) + + initial_triangles = up_sampled diff --git a/autolens/point/triangles/visualise.py b/autolens/point/triangles/visualise.py new file mode 100644 index 000000000..3dd5ab402 --- /dev/null +++ b/autolens/point/triangles/visualise.py @@ -0,0 +1,23 @@ +from .triangle_solver import Step +from matplotlib import pyplot as plt +import numpy as np + + +def add_triangles(triangles, color): + for triangle in triangles: + triangle = np.append(triangle, [triangle[0]], axis=0) + plt.plot(triangle[:, 0], triangle[:, 1], "o-", color=color) + + +def visualise(step: Step): + plt.figure(figsize=(8, 8)) + add_triangles(step.initial_triangles, color="black") + add_triangles(step.up_sampled, color="green") + add_triangles(step.neighbourhood, color="red") + add_triangles(step.filtered_triangles, color="blue") + + plt.xlabel("X") + plt.ylabel("Y") + plt.title(f"Step {step.number}") + plt.gca().set_aspect("equal", adjustable="box") + plt.show() diff --git a/test_autolens/point/triangles/conftest.py b/test_autolens/point/triangles/conftest.py index 321efc191..c596799ca 100644 --- a/test_autolens/point/triangles/conftest.py +++ b/test_autolens/point/triangles/conftest.py @@ -5,6 +5,26 @@ @pytest.fixture def grid(): return al.Grid2D.uniform( - shape_native=(100, 100), - pixel_scales=0.05, + shape_native=(10, 10), + pixel_scales=1.0, ) + + +@pytest.fixture +def tracer(): + isothermal_mass_profile = al.mp.Isothermal( + centre=(0.0, 0.0), + einstein_radius=1.6, + ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0), + ) + + lens_galaxy = al.Galaxy( + redshift=0.5, + mass=isothermal_mass_profile, + ) + + point_source = al.ps.PointSourceChi(centre=(0.07, 0.07)) + + source_galaxy = al.Galaxy(redshift=1.0, point_0=point_source) + + return al.Tracer(galaxies=[lens_galaxy, source_galaxy]) diff --git a/test_autolens/point/triangles/test_regressions.py b/test_autolens/point/triangles/test_regressions.py index 5c94adac4..c712b2e5b 100644 --- a/test_autolens/point/triangles/test_regressions.py +++ b/test_autolens/point/triangles/test_regressions.py @@ -66,7 +66,7 @@ def test_missing_multiple_image(grid): tracer = al.Tracer(galaxies=[instance.lens_galaxy, instance.source_galaxy]) - solver = TriangleSolver( + solver = TriangleSolver.for_grid( grid=grid, lensing_obj=tracer, pixel_scale_precision=0.001, diff --git a/test_autolens/point/triangles/test_solver.py b/test_autolens/point/triangles/test_solver.py index 55032fcbe..429c2631e 100644 --- a/test_autolens/point/triangles/test_solver.py +++ b/test_autolens/point/triangles/test_solver.py @@ -1,10 +1,10 @@ from typing import Tuple -import numpy as np import pytest import autolens as al import autogalaxy as ag +from autolens.mock import NullTracer from autolens.point.triangles.triangle_solver import TriangleSolver @@ -22,7 +22,7 @@ def solver(grid): ] ) - return TriangleSolver( + return TriangleSolver.for_grid( lensing_obj=tracer, grid=grid, pixel_scale_precision=0.01, @@ -36,15 +36,7 @@ def test_solver(solver): def test_steps(solver): - assert solver.n_steps == 3 - - -class NullTracer(al.Tracer): - def __init__(self): - super().__init__([]) - - def deflections_yx_2d_from(self, grid): - return np.zeros_like(grid) + assert solver.n_steps == 7 @pytest.mark.parametrize( @@ -52,6 +44,7 @@ def deflections_yx_2d_from(self, grid): [ (0.0, 0.0), (0.0, 1.0), + (1.0, 0.0), (1.0, 1.0), (0.5, 0.5), (0.1, 0.1), @@ -62,7 +55,7 @@ def test_trivial( source_plane_coordinate: Tuple[float, float], grid, ): - solver = TriangleSolver( + solver = TriangleSolver.for_grid( lensing_obj=NullTracer(), grid=grid, pixel_scale_precision=0.01, @@ -70,31 +63,16 @@ def test_trivial( (coordinates,) = solver.solve( source_plane_coordinate=source_plane_coordinate, ) - assert coordinates == pytest.approx(source_plane_coordinate, abs=1.0e-2) - - -def test_real_example(grid): - isothermal_mass_profile = al.mp.Isothermal( - centre=(0.0, 0.0), - einstein_radius=1.6, - ell_comps=al.convert.ell_comps_from(axis_ratio=0.9, angle=45.0), - ) - - lens_galaxy = al.Galaxy( - redshift=0.5, - mass=isothermal_mass_profile, - ) - - point_source = al.ps.PointSourceChi(centre=(0.07, 0.07)) - - source_galaxy = al.Galaxy(redshift=1.0, point_0=point_source) + assert coordinates == pytest.approx(source_plane_coordinate, abs=1.0e-1) - tracer = al.Tracer(galaxies=[lens_galaxy, source_galaxy]) - solver = TriangleSolver( +def test_real_example(grid, tracer): + solver = TriangleSolver.for_grid( grid=grid, lensing_obj=tracer, pixel_scale_precision=0.001, ) result = solver.solve((0.07, 0.07)) - assert len(result) == 4 + for r in result: + print(r) + assert len(result) == 5 diff --git a/test_autolens/point/triangles/test_solver_jax.py b/test_autolens/point/triangles/test_solver_jax.py new file mode 100644 index 000000000..2f2dcd1f0 --- /dev/null +++ b/test_autolens/point/triangles/test_solver_jax.py @@ -0,0 +1,79 @@ +from typing import Tuple + +import pytest + +import autolens as al +import autogalaxy as ag +from autoarray.structures.triangles.jax_array import ArrayTriangles +from autolens.mock import NullTracer +from autolens.point.triangles.triangle_solver import TriangleSolver + + +@pytest.fixture +def solver(grid): + tracer = al.Tracer( + galaxies=[ + al.Galaxy( + redshift=0.5, + mass=ag.mp.Isothermal( + centre=(0.0, 0.0), + einstein_radius=1.0, + ), + ) + ] + ) + + return TriangleSolver.for_grid( + lensing_obj=tracer, + grid=grid, + pixel_scale_precision=0.01, + ArrayTriangles=ArrayTriangles, + ) + + +def test_solver(solver): + assert solver.solve( + source_plane_coordinate=(0.0, 0.0), + ) + + +@pytest.mark.parametrize( + "source_plane_coordinate", + [ + (0.0, 0.0), + (0.0, 1.0), + (1.0, 0.0), + (1.0, 1.0), + (0.5, 0.5), + (0.1, 0.1), + (-1.0, -1.0), + ], +) +def test_trivial( + source_plane_coordinate: Tuple[float, float], + grid, +): + solver = TriangleSolver.for_grid( + lensing_obj=NullTracer(), + grid=grid, + pixel_scale_precision=0.01, + ArrayTriangles=ArrayTriangles, + ) + coordinates = solver.solve( + source_plane_coordinate=source_plane_coordinate, + ) + print(coordinates) + assert coordinates[0] == pytest.approx(source_plane_coordinate, abs=1.0e-1) + + +def test_real_example(grid, tracer): + solver = TriangleSolver.for_grid( + grid=grid, + lensing_obj=tracer, + pixel_scale_precision=0.001, + ArrayTriangles=ArrayTriangles, + ) + result = solver.solve((0.07, 0.07)) + for pair in result: + print(pair) + assert len(result) == 5