all unit tests pass

Author: Jammy2211

autogalaxy/galaxy/galaxies.py

@@ -8,10 +8,9 @@
 from autogalaxy.profiles.basis import Basis
 from autogalaxy.profiles.light.linear import LightProfileLinear
 from autogalaxy.operate.image import OperateImageGalaxies
-from autogalaxy.operate.deflections import OperateDeflections
 
 
-class Galaxies(List, OperateImageGalaxies, OperateDeflections):
+class Galaxies(List, OperateImageGalaxies):
     def __init__(
         self,
         galaxies: List[Galaxy],
@@ -244,6 +243,34 @@ def potential_2d_from(
         """
         return sum(map(lambda g: g.potential_2d_from(grid=grid, xp=xp), self))
 
+    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
+        """
+        Returns the Fermat potential of the galaxies for a given grid of image-plane positions.
+
+        This is the sum of the geometric time delay term and the lensing potential, computed
+        using the combined deflections and potential of all galaxies:
+
+        .. math::
+            \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2
+            - \\psi(\\boldsymbol{\\theta})
+
+        Parameters
+        ----------
+        grid
+            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
+        xp
+            The array module (``numpy`` or ``jax.numpy``).
+        """
+        from autogalaxy.operate.deflections import OperateDeflections
+
+        od = OperateDeflections.from_mass_obj(self)
+        time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp)
+        potential = self.potential_2d_from(grid=grid, xp=xp)
+        fermat_potential = time_delay - potential
+        if isinstance(grid, aa.Grid2DIrregular):
+            return aa.ArrayIrregular(values=fermat_potential)
+        return aa.Array2D(values=fermat_potential, mask=grid.mask)
+
     def has(self, cls: Union[Type, Tuple[Type]]) -> bool:
         """
         Returns a bool specifying whether any of the galaxies has a certain class type.

autogalaxy/galaxy/galaxy.py

@@ -8,7 +8,6 @@
 import autofit as af
 
 from autogalaxy import exc
-from autogalaxy.operate.deflections import OperateDeflections
 from autogalaxy.operate.image import OperateImageList
 from autogalaxy.profiles.geometry_profiles import GeometryProfile
 from autogalaxy.profiles.light.abstract import LightProfile
@@ -17,7 +16,7 @@
 from autogalaxy.profiles.mass.abstract.abstract import MassProfile
 
 
-class Galaxy(af.ModelObject, OperateImageList, OperateDeflections):
+class Galaxy(af.ModelObject, OperateImageList):
     """
     @DynamicAttrs
     """
@@ -344,6 +343,34 @@ def potential_2d_from(
             )
         return xp.zeros((grid.shape[0],))
 
+    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
+        """
+        Returns the Fermat potential of the galaxy for a given grid of image-plane positions.
+
+        This is the sum of the geometric time delay term and the lensing potential, computed
+        using the galaxy's total deflections and total potential across all mass profiles:
+
+        .. math::
+            \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2
+            - \\psi(\\boldsymbol{\\theta})
+
+        Parameters
+        ----------
+        grid
+            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
+        xp
+            The array module (``numpy`` or ``jax.numpy``).
+        """
+        from autogalaxy.operate.deflections import OperateDeflections
+
+        od = OperateDeflections.from_mass_obj(self)
+        time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp)
+        potential = self.potential_2d_from(grid=grid, xp=xp)
+        fermat_potential = time_delay - potential
+        if isinstance(grid, aa.Grid2DIrregular):
+            return aa.ArrayIrregular(values=fermat_potential)
+        return aa.Array2D(values=fermat_potential, mask=grid.mask)
+
     @property
     def half_light_radius(self):
         return None

autogalaxy/operate/deflections.py

@@ -86,88 +86,94 @@ class OperateDeflections:
     The majority of methods are those which from the 2D deflection angle map compute lensing quantities like a 2D
     shear field, magnification map or the Einstein Radius.
 
-    The methods in `CalcLens` are passed to the mass object to provide a concise API.
-
     Parameters
     ----------
     deflections_yx_2d_from
-        The function which returns the mass object's 2D deflection angles.
+        A callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D deflection angles on the given
+        grid.  Typically a bound method of a ``MassProfile``, ``Galaxy``, or ``Galaxies`` instance.
     """
 
-    def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
-        raise NotImplementedError
+    def __init__(self, deflections_yx_2d_from):
+        self.deflections_yx_2d_from = deflections_yx_2d_from
 
-    def __eq__(self, other):
-        return self.__dict__ == other.__dict__ and self.__class__ is other.__class__
+    @classmethod
+    def from_mass_obj(cls, mass_obj):
+        """Construct from any object that has a ``deflections_yx_2d_from`` method."""
+        return cls(deflections_yx_2d_from=mass_obj.deflections_yx_2d_from)
 
-    def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D:
+    @classmethod
+    def from_tracer(cls, tracer, use_multi_plane: bool = True, plane_i: int = 0, plane_j: int = -1):
         """
-            Returns the geometric time delay term of the Fermat potential for a given grid of image-plane positions.
-
-            This term is given by:
-
-        .. math::
-                \[\tau_{\text{geom}}(\boldsymbol{\theta}) = \frac{1}{2} |\boldsymbol{\theta} - \boldsymbol{\beta}|^2\]
+        Construct from a PyAutoLens ``Tracer`` object.
 
-            where:
-            - \( \boldsymbol{\theta} \) is the image-plane coordinate,
-            - \( \boldsymbol{\beta} = \boldsymbol{\theta} - \boldsymbol{\alpha}(\boldsymbol{\theta}) \) is the source-plane coordinate,
-            - \( \boldsymbol{\alpha} \) is the deflection angle at each image-plane coordinate.
-
-            Parameters
-            ----------
-            grid
-                The 2D grid of (y,x) arc-second coordinates the deflection angles and time delay geometric term are computed
-                on.
-
-            Returns
-            -------
-            The geometric time delay term at each grid position.
-        """
-        deflections = self.deflections_yx_2d_from(grid=grid, xp=xp)
+        The ``Tracer`` type is deliberately left unannotated: ``autogalaxy`` does not
+        depend on ``autolens``, so no import of ``Tracer`` is performed here.  Callers
+        (which live inside ``autolens``) are responsible for passing the correct object.
 
-        src_y = grid[:, 0] - deflections[:, 0]
-        src_x = grid[:, 1] - deflections[:, 1]
-
-        delay = 0.5 * ((grid[:, 0] - src_y) ** 2 + (grid[:, 1] - src_x) ** 2)
-
-        if isinstance(grid, aa.Grid2DIrregular):
-            return aa.ArrayIrregular(values=delay)
-        return aa.Array2D(values=delay, mask=grid.mask)
+        Parameters
+        ----------
+        tracer
+            A PyAutoLens ``Tracer`` instance.  Must expose ``deflections_yx_2d_from``
+            and, when ``use_multi_plane=True``, ``deflections_between_planes_from``.
+        use_multi_plane
+            If ``True`` the stored callable is
+            ``tracer.deflections_between_planes_from`` with ``plane_i`` and ``plane_j``
+            bound via ``functools.partial``, matching the multi-plane ray-tracing path.
+            If ``False`` the stored callable is ``tracer.deflections_yx_2d_from``,
+            i.e. the standard two-plane path.
+        plane_i
+            Index of the first plane used by ``deflections_between_planes_from``.
+            Ignored when ``use_multi_plane=False``.  Defaults to ``0`` (image plane).
+        plane_j
+            Index of the second plane used by ``deflections_between_planes_from``.
+            Ignored when ``use_multi_plane=False``.  Defaults to ``-1`` (source plane).
+        """
+        if use_multi_plane:
+            from functools import partial
+
+            return cls(
+                deflections_yx_2d_from=partial(
+                    tracer.deflections_between_planes_from,
+                    plane_i=plane_i,
+                    plane_j=plane_j,
+                )
+            )
+        return cls(deflections_yx_2d_from=tracer.deflections_yx_2d_from)
 
-    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
+    def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D:
         """
-        Returns the Fermat potential for a given grid of image-plane positions.
+        Returns the geometric time delay term of the Fermat potential for a given grid of image-plane positions.
 
-        This is the sum of the geometric time delay term and the gravitational (Shapiro) delay term (i.e. the lensing
-        potential), and is given by:
+        This term is given by:
 
         .. math::
-            \[\phi(\boldsymbol{\theta}) = \frac{1}{2} |\boldsymbol{\theta} - \boldsymbol{\beta}|^2 - \psi(\boldsymbol{\theta})\]
+            \[\tau_{\text{geom}}(\boldsymbol{\theta}) = \frac{1}{2} |\boldsymbol{\theta} - \boldsymbol{\beta}|^2\]
 
         where:
         - \( \boldsymbol{\theta} \) is the image-plane coordinate,
         - \( \boldsymbol{\beta} = \boldsymbol{\theta} - \boldsymbol{\alpha}(\boldsymbol{\theta}) \) is the source-plane coordinate,
-        - \( \psi(\boldsymbol{\theta}) \) is the lensing potential,
-        - \( \phi(\boldsymbol{\theta}) \) is the Fermat potential.
+        - \( \boldsymbol{\alpha} \) is the deflection angle at each image-plane coordinate.
 
         Parameters
         ----------
         grid
-            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
+            The 2D grid of (y,x) arc-second coordinates the deflection angles and time delay geometric term are computed
+            on.
 
         Returns
         -------
-        The Fermat potential at each grid position.
+        The geometric time delay term at each grid position.
         """
-        time_delay_geometry_term = self.time_delay_geometry_term_from(grid=grid, xp=xp)
-        potential = self.potential_2d_from(grid=grid, xp=xp)
+        deflections = self.deflections_yx_2d_from(grid=grid, xp=xp)
+
+        src_y = grid[:, 0] - deflections[:, 0]
+        src_x = grid[:, 1] - deflections[:, 1]
 
-        fermat_potential = time_delay_geometry_term - potential
+        delay = 0.5 * ((grid[:, 0] - src_y) ** 2 + (grid[:, 1] - src_x) ** 2)
 
         if isinstance(grid, aa.Grid2DIrregular):
-            return aa.ArrayIrregular(values=fermat_potential)
-        return aa.Array2D(values=fermat_potential, mask=grid.mask)
+            return aa.ArrayIrregular(values=delay)
+        return aa.Array2D(values=delay, mask=grid.mask)
 
     def tangential_eigen_value_from(self, grid, xp=np) -> aa.Array2D:
         """
@@ -188,9 +194,6 @@ def tangential_eigen_value_from(self, grid, xp=np) -> aa.Array2D:
         convergence = self.convergence_2d_via_hessian_from(grid=grid, xp=xp)
         shear_yx = self.shear_yx_2d_via_hessian_from(grid=grid, xp=xp)
 
-        if xp is not np:
-            shear_magnitudes = xp.sqrt(shear_yx[:, 0] ** 2 + shear_yx[:, 1] ** 2)
-            return xp.array(1 - convergence - shear_magnitudes)
         return aa.Array2D(values=1 - convergence - shear_yx.magnitudes, mask=grid.mask)
 
     def radial_eigen_value_from(self, grid, xp=np) -> aa.Array2D:
@@ -211,9 +214,6 @@ def radial_eigen_value_from(self, grid, xp=np) -> aa.Array2D:
         convergence = self.convergence_2d_via_hessian_from(grid=grid, xp=xp)
         shear = self.shear_yx_2d_via_hessian_from(grid=grid, xp=xp)
 
-        if xp is not np:
-            shear_magnitudes = xp.sqrt(shear[:, 0] ** 2 + shear[:, 1] ** 2)
-            return xp.array(1 - convergence + shear_magnitudes)
         return aa.Array2D(values=1 - convergence + shear.magnitudes, mask=grid.mask)
 
     def magnification_2d_from(self, grid, xp=np) -> aa.Array2D:
@@ -235,8 +235,6 @@ def magnification_2d_from(self, grid, xp=np) -> aa.Array2D:
 
         det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx
 
-        if xp is not np:
-            return xp.array(1 / det_A)
         return aa.Array2D(values=1 / det_A, mask=grid.mask)
 
     def deflections_yx_scalar(self, y, x, pixel_scales):
@@ -295,10 +293,9 @@ def hessian_from(self, grid, xp=np) -> Tuple:
             The array module (``numpy`` or ``jax.numpy``). Controls which computational path is
             used and the type of the returned arrays.
         """
-        if xp is not np:
-            return self._hessian_via_jax(grid=grid, xp=xp)
-
-        return self._hessian_via_finite_difference(grid=grid)
+        if xp is np:
+            return self._hessian_via_finite_difference(grid=grid)
+        return self._hessian_via_jax(grid=grid, xp=xp)
 
     def _hessian_via_jax(self, grid, xp) -> Tuple:
         import jax
@@ -417,8 +414,6 @@ def convergence_2d_via_hessian_from(
 
         convergence = 0.5 * (hessian_yy + hessian_xx)
 
-        if xp is not np:
-            return xp.array(convergence)
         return aa.ArrayIrregular(values=convergence)
 
     def shear_yx_2d_via_hessian_from(
@@ -460,13 +455,7 @@ def shear_yx_2d_via_hessian_from(
         gamma_1 = 0.5 * (hessian_xx - hessian_yy)
         gamma_2 = hessian_xy
 
-        if xp is not np:
-            return xp.stack([gamma_2, gamma_1], axis=-1)
-
-        shear_yx_2d = np.zeros(shape=(grid.shape[0], 2))
-
-        shear_yx_2d[:, 0] = gamma_2
-        shear_yx_2d[:, 1] = gamma_1
+        shear_yx_2d = xp.stack([gamma_2, gamma_1], axis=-1)
 
         return ShearYX2DIrregular(values=shear_yx_2d, grid=grid)
 
@@ -498,8 +487,6 @@ def magnification_2d_via_hessian_from(
 
         det_A = (1 - hessian_xx) * (1 - hessian_yy) - hessian_xy * hessian_yx
 
-        if xp is not np:
-            return xp.array(1.0 / det_A)
         return aa.ArrayIrregular(values=1.0 / det_A)
 
     def contour_list_from(self, grid, contour_array):

autogalaxy/plot/mass_plotter.py

@@ -118,8 +118,12 @@ def figures_2d(
             )
 
         if magnification:
+            from autogalaxy.operate.deflections import OperateDeflections
+
             self.mat_plot_2d.plot_array(
-                array=self.mass_obj.magnification_2d_from(grid=self.grid),
+                array=OperateDeflections.from_mass_obj(self.mass_obj).magnification_2d_from(
+                    grid=self.grid
+                ),
                 visuals_2d=self.visuals_2d_with_critical_curves,
                 auto_labels=aplt.AutoLabels(
                     title=f"Magnification{title_suffix}",

autogalaxy/plot/visuals/one_d.py

@@ -170,10 +170,13 @@ def add_einstein_radius(
             The collection of attributes that can be plotted by a `Plotter` object.
         """
 
+        from autogalaxy.operate.deflections import OperateDeflections
+
         einstein_radius = None
 
         try:
-            einstein_radius = mass_obj.einstein_radius_from(grid=grid)
+            od = OperateDeflections.from_mass_obj(mass_obj)
+            einstein_radius = od.einstein_radius_from(grid=grid)
         except (TypeError, AttributeError):
             pass
 
@@ -214,11 +217,14 @@ def add_einstein_radius_errors(
             The mean value and errors of each attribute that are plotted in 1D by a `Plotter` object.
         """
 
+        from autogalaxy.operate.deflections import OperateDeflections
+
         einstein_radius_list = []
 
         for mass_obj in mass_obj_list:
             try:
-                einstein_radius_list.append(mass_obj.einstein_radius_from(grid=grid))
+                od = OperateDeflections.from_mass_obj(mass_obj)
+                einstein_radius_list.append(od.einstein_radius_from(grid=grid))
             except TypeError:
                 einstein_radius_list.append(None)