diff --git a/autolens/__init__.py b/autolens/__init__.py index 44da6dc43..a7ad85ac3 100644 --- a/autolens/__init__.py +++ b/autolens/__init__.py @@ -81,7 +81,7 @@ ) from autogalaxy.profiles.light.linear import LightProfileLinearObjFuncList from autogalaxy.operate.image import OperateImage -from autogalaxy.operate.deflections import OperateDeflections +from autogalaxy.operate.lens_calc import LensCalc from autogalaxy.quantity.dataset_quantity import DatasetQuantity from autogalaxy import convert diff --git a/autolens/fixtures.py b/autolens/fixtures.py index 281a09290..9df614484 100644 --- a/autolens/fixtures.py +++ b/autolens/fixtures.py @@ -35,7 +35,9 @@ def make_point_dataset(): def make_solver(): grid = al.Grid2D.uniform(shape_native=(10, 10), pixel_scales=0.5) - return al.PointSolver.for_grid(grid=grid, pixel_scale_precision=0.25) + return al.PointSolver.for_grid( + grid=grid, pixel_scale_precision=0.25, magnification_threshold=1e-8 + ) def make_tracer_x1_plane_7x7(): diff --git a/autolens/lens/mock/mock_tracer.py b/autolens/lens/mock/mock_tracer.py index c6412028a..7cdc4c2f7 100644 --- a/autolens/lens/mock/mock_tracer.py +++ b/autolens/lens/mock/mock_tracer.py @@ -43,8 +43,8 @@ def __init__( def planes(self): return [0, 1] - def deflections_yx_2d_from(self): - pass + def deflections_yx_2d_from(self, grid, xp=np): + return xp.zeros_like(grid.array) def extract_attribute(self, cls, attr_name): return [self.attribute] @@ -55,9 +55,6 @@ def extract_profile(self, profile_name): def traced_grid_2d_list_from(self, grid, plane_index_limit=None): return [self.positions] - def magnification_2d_via_hessian_from(self, grid, deflections_func=None, xp=np): - return self.magnification - def einstein_radius_from(self, grid): return self.einstein_radius diff --git a/autolens/lens/tracer.py b/autolens/lens/tracer.py index b6bc12eb1..2535ab5b9 100644 --- a/autolens/lens/tracer.py +++ b/autolens/lens/tracer.py @@ -13,7 +13,7 @@ from autolens.lens import tracer_util -class Tracer(ABC, ag.OperateImageGalaxies, ag.OperateDeflections): +class Tracer(ABC, ag.OperateImageGalaxies): def __init__( self, galaxies: Union[List[ag.Galaxy], af.ModelInstance], diff --git a/autolens/lens/tracer_util.py b/autolens/lens/tracer_util.py index 461ef6059..b2fe6556f 100644 --- a/autolens/lens/tracer_util.py +++ b/autolens/lens/tracer_util.py @@ -354,7 +354,11 @@ def time_delays_from( D_dt_m = (1.0 + z_l) * (Dd_kpc * Ds_kpc / Dds_kpc) * kpc_in_m # Fermat potential (should be in arcsec^2 for this formula) - fermat_potential = galaxies.fermat_potential_from(grid=grid, xp=xp) + import autogalaxy as ag + + fermat_potential = ag.LensCalc.from_mass_obj(galaxies).fermat_potential_from( + grid=grid, xp=xp + ) # Final time delay in days return (D_dt_m / c) * fermat_potential * factor diff --git a/autolens/mock.py b/autolens/mock.py index e20872d4a..b591f211c 100644 --- a/autolens/mock.py +++ b/autolens/mock.py @@ -3,7 +3,6 @@ from autofit.mock import * # noqa from autoarray.mock import * # noqa from autogalaxy.mock import * # noqa -import autoarray as aa from autolens import Tracer from autolens.imaging.mock.mock_fit_imaging import MockFitImaging # noqa @@ -21,12 +20,3 @@ def deflections_yx_2d_from(self, grid, xp=np): def deflections_between_planes_from(self, grid, xp=np, plane_i=0, plane_j=-1): return xp.zeros_like(grid.array) - - def magnification_2d_via_hessian_from( - self, - grid, - buffer: float = 0.01, - deflections_func=None, - xp=np, - ) -> aa.ArrayIrregular: - return aa.ArrayIrregular(values=xp.ones(grid.shape[0])) diff --git a/autolens/point/fit/abstract.py b/autolens/point/fit/abstract.py index d59387009..e6903ec75 100644 --- a/autolens/point/fit/abstract.py +++ b/autolens/point/fit/abstract.py @@ -1,5 +1,4 @@ from abc import ABC -from functools import partial import numpy as np from typing import Optional, Tuple @@ -90,40 +89,6 @@ def data(self): def noise_map(self): return self._noise_map - @property - def deflections_func(self): - """ - Returns the deflection angle function, which for example given input image-plane positions computes their - deflection angles. - - The use of this specific `deflections_func` property is not typical, using the `partial` function to wrap - a deflections method of the tracer. This is essentially a trick so that, depending on whether multi-plane - ray-tracing is to be performed, a different deflection function is used. This function is then - used in `magnifications_at_positions` to compute the magnification of the point source account for - multi-plane ray-tracing. - - For multi-plane ray-tracing with more than 2 planes, the deflection function determines the index of the - plane with the last mass profile such that the deflection function does not perform unnecessary computations - beyond this plane. - - TODO: Simplify this property and calculation of the deflection angles, as this property is confusing. - TODO: This could be done by allowing for the Hessian to receive planes as input. - """ - - if len(self.tracer.planes) > 2: - upper_plane_index = self.tracer.extract_plane_index_of_profile( - profile_name=self.name - ) - - return partial( - self.tracer.deflections_between_planes_from, - xp=self._xp, - plane_i=0, - plane_j=upper_plane_index, - ) - - return self.tracer.deflections_yx_2d_from - @property def magnifications_at_positions(self) -> aa.ArrayIrregular: """ @@ -139,10 +104,19 @@ def magnifications_at_positions(self) -> aa.ArrayIrregular: 2) For fitting the fluxes of point sources, the magnification is used to scale the flux of the point source in the source-plane to the image-plane, thus computing the model image-plane fluxes. """ + use_multi_plane = len(self.tracer.planes) > 2 + plane_j = ( + self.tracer.extract_plane_index_of_profile(profile_name=self.name) + if use_multi_plane + else -1 + ) + od = ag.LensCalc.from_tracer( + tracer=self.tracer, + use_multi_plane=use_multi_plane, + plane_j=plane_j, + ) return abs( - self.tracer.magnification_2d_via_hessian_from( - grid=self.positions, deflections_func=self.deflections_func, xp=self._xp - ) + od.magnification_2d_via_hessian_from(grid=self.positions, xp=self._xp) ) @property diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py index fabc3fe1b..cb9521cef 100644 --- a/autolens/point/solver/point_solver.py +++ b/autolens/point/solver/point_solver.py @@ -4,7 +4,7 @@ import autoarray as aa from autoarray.structures.triangles.shape import Point -from autogalaxy import OperateDeflections +from autolens.lens.tracer import Tracer from .shape_solver import AbstractSolver @@ -15,7 +15,7 @@ class PointSolver(AbstractSolver): def solve( self, - tracer: OperateDeflections, + tracer: Tracer, source_plane_coordinate: Tuple[float, float], plane_redshift: Optional[float] = None, remove_infinities: bool = True, diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py index 4cb4f7ba2..11bd4652a 100644 --- a/autolens/point/solver/shape_solver.py +++ b/autolens/point/solver/shape_solver.py @@ -8,7 +8,8 @@ from autoarray.structures.triangles.shape import Shape -from autogalaxy import OperateDeflections +import autogalaxy as ag +from autolens.lens.tracer import Tracer from .step import Step logger = logging.getLogger(__name__) @@ -186,7 +187,7 @@ def n_steps(self) -> int: def _plane_grid( self, - tracer: OperateDeflections, + tracer: Tracer, grid: aa.type.Grid2DLike, plane_redshift: Optional[float] = None, ) -> aa.type.Grid2DLike: @@ -215,7 +216,7 @@ def _plane_grid( def solve_triangles( self, - tracer: OperateDeflections, + tracer: Tracer, shape: Shape, plane_redshift: Optional[float] = None, ): @@ -258,7 +259,7 @@ def solve_triangles( return final_step.filtered_triangles def _filter_low_magnification( - self, tracer: OperateDeflections, points: List[Tuple[float, float]] + self, tracer: Tracer, points: List[Tuple[float, float]] ) -> List[Tuple[float, float]]: """ Filter the points to keep only those with an absolute magnification above the threshold. @@ -273,15 +274,17 @@ def _filter_low_magnification( The points with an absolute magnification above the threshold. """ points = self._xp.array(points) - magnifications = tracer.magnification_2d_via_hessian_from( - grid=aa.Grid2DIrregular(points).array, buffer=self.scale, xp=self._xp + magnifications = ag.LensCalc.from_mass_obj( + tracer + ).magnification_2d_via_hessian_from( + grid=aa.Grid2DIrregular(points).array, xp=self._xp ) mask = self._xp.abs(magnifications.array) > self.magnification_threshold return self._xp.where(mask[:, None], points, self._xp.nan) def _plane_triangles( self, - tracer: OperateDeflections, + tracer: Tracer, triangles: aa.AbstractTriangles, plane_redshift, ): @@ -298,7 +301,7 @@ def _plane_triangles( def steps( self, - tracer: OperateDeflections, + tracer: Tracer, shape: Shape, plane_redshift: Optional[float] = None, ) -> Iterator[Step]: @@ -368,7 +371,7 @@ def tree_unflatten(cls, aux_data, children): class ShapeSolver(AbstractSolver): def find_magnification( self, - tracer: OperateDeflections, + tracer: Tracer, shape: Shape, plane_redshift: Optional[float] = None, ) -> float: diff --git a/docs/api/source.rst b/docs/api/source.rst index 69a24a747..0e75be86e 100644 --- a/docs/api/source.rst +++ b/docs/api/source.rst @@ -31,4 +31,4 @@ Operators :recursive: OperateImage - OperateDeflections \ No newline at end of file + LensCalc \ No newline at end of file diff --git a/test_autolens/lens/test_operate.py b/test_autolens/lens/test_operate.py index ef9936535..f19706720 100644 --- a/test_autolens/lens/test_operate.py +++ b/test_autolens/lens/test_operate.py @@ -2,6 +2,7 @@ import pytest from os import path +import autogalaxy as ag import autolens as al test_path = path.join("{}".format(path.dirname(path.realpath(__file__))), "files") @@ -193,6 +194,8 @@ def test__operate_lens__sums_individual_quantities(): cosmology=al.cosmo.Planck15(), ) - einstein_mass = tracer.einstein_mass_angular_from(grid=grid) + einstein_mass = ag.LensCalc.from_mass_obj( + tracer + ).einstein_mass_angular_from(grid=grid) assert einstein_mass == pytest.approx(np.pi * 2.0**2.0, 1.0e-1) diff --git a/test_autolens/point/fit/test_abstract.py b/test_autolens/point/fit/test_abstract.py index 888d2f983..ddceeb56d 100644 --- a/test_autolens/point/fit/test_abstract.py +++ b/test_autolens/point/fit/test_abstract.py @@ -1,6 +1,6 @@ -from functools import partial import pytest +import autogalaxy as ag import autolens as al @@ -23,13 +23,8 @@ def test__magnifications_at_positions__multi_plane_calculation(gal_x1_mp): tracer=tracer, ) - deflections_func = partial( - tracer.deflections_between_planes_from, plane_i=0, plane_j=1 - ) - - magnification_0 = tracer.magnification_2d_via_hessian_from( - grid=positions, deflections_func=deflections_func - ) + od_0 = ag.LensCalc.from_tracer(tracer, use_multi_plane=True, plane_j=1) + magnification_0 = abs(od_0.magnification_2d_via_hessian_from(grid=positions)) assert fit_0.magnifications_at_positions[0] == magnification_0 @@ -41,13 +36,8 @@ def test__magnifications_at_positions__multi_plane_calculation(gal_x1_mp): tracer=tracer, ) - deflections_func = partial( - tracer.deflections_between_planes_from, plane_i=0, plane_j=2 - ) - - magnification_1 = tracer.magnification_2d_via_hessian_from( - grid=positions, deflections_func=deflections_func - ) + od_1 = ag.LensCalc.from_tracer(tracer, use_multi_plane=True, plane_j=2) + magnification_1 = abs(od_1.magnification_2d_via_hessian_from(grid=positions)) assert fit_1.magnifications_at_positions[0] == magnification_1 diff --git a/test_autolens/point/fit/test_fluxes.py b/test_autolens/point/fit/test_fluxes.py index ef1920cd9..703b8c340 100644 --- a/test_autolens/point/fit/test_fluxes.py +++ b/test_autolens/point/fit/test_fluxes.py @@ -1,13 +1,10 @@ -from functools import partial import pytest import autolens as al def test__one_set_of_fluxes__residuals_likelihood_correct(): - tracer = al.m.MockTracerPoint( - profile=al.ps.PointFlux(flux=2.0), magnification=al.ArrayIrregular([2.0, 2.0]) - ) + tracer = al.m.MockTracerPoint(profile=al.ps.PointFlux(flux=2.0)) data = al.ArrayIrregular([1.0, 2.0]) noise_map = al.ArrayIrregular([3.0, 1.0]) @@ -23,13 +20,13 @@ def test__one_set_of_fluxes__residuals_likelihood_correct(): assert fit.data.in_list == [1.0, 2.0] assert fit.noise_map.in_list == [3.0, 1.0] - assert fit.model_fluxes.in_list == [4.0, 4.0] - assert fit.residual_map.in_list == [-3.0, -2.0] - assert fit.normalized_residual_map.in_list == [-1.0, -2.0] - assert fit.chi_squared_map.in_list == [1.0, 4.0] - assert fit.chi_squared == pytest.approx(5.0, 1.0e-4) + assert fit.model_fluxes.in_list == [2.0, 2.0] + assert fit.residual_map.in_list == [-1.0, -0.0] + assert fit.normalized_residual_map.in_list == [-1.0 / 3.0, -0.0] + assert fit.chi_squared_map.in_list == [1.0 / 9.0, 0.0] + assert fit.chi_squared == pytest.approx(1.0 / 9.0, 1.0e-4) assert fit.noise_normalization == pytest.approx(5.87297, 1.0e-4) - assert fit.log_likelihood == pytest.approx(-5.43648, 1.0e-4) + assert fit.log_likelihood == pytest.approx(-2.992044910633, 1.0e-4) def test__use_real_tracer(gal_x1_mp): diff --git a/test_autolens/point/triangles/test_solver.py b/test_autolens/point/triangles/test_solver.py index 413781ef7..2264afce5 100644 --- a/test_autolens/point/triangles/test_solver.py +++ b/test_autolens/point/triangles/test_solver.py @@ -22,8 +22,7 @@ def register(tracer): @pytest.fixture def solver(grid): return PointSolver.for_grid( - grid=grid, - pixel_scale_precision=0.01, + grid=grid, pixel_scale_precision=0.01, magnification_threshold=1e-8 ) @@ -35,10 +34,7 @@ def test_solver(solver): tracer = Tracer( galaxies=[ag.Galaxy(redshift=0.5, mass=mass_profile)], ) - result = solver.solve( - tracer, - source_plane_coordinate=(0.0, 0.0), - ) + result = solver.solve(tracer, source_plane_coordinate=(0.0, 0.0)) assert result @@ -71,7 +67,9 @@ def test_real_example_jax(grid, tracer): import jax.numpy as jnp - jax_solver = PointSolver.for_grid(grid=grid, pixel_scale_precision=0.001, xp=jnp) + jax_solver = PointSolver.for_grid( + grid=grid, pixel_scale_precision=0.001, xp=jnp, magnification_threshold=1e-8 + ) result = jax_solver.solve( tracer=tracer, source_plane_coordinate=(0.07, 0.07), remove_infinities=True