Feature/deflections operate jax

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
 

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():

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
 

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],

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

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]))

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

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,

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:

docs/api/source.rst

@@ -31,4 +31,4 @@ Operators
    :recursive:
 
    OperateImage
-   OperateDeflections
+   LensCalc

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)

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
 

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):

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