Feature/time delays

autogalaxy/operate/deflections.py

@@ -112,6 +112,92 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs):
     def __eq__(self, other):
         return self.__dict__ == other.__dict__ and self.__class__ is other.__class__
 
+    def time_delay_geometry_term_from(self, grid) -> aa.Array2D:
+        """
+        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\]
+
+        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)
+
+        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)
+
+    def fermat_potential_from(self, grid) -> aa.Array2D:
+        """
+        Returns 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:
+
+        .. math::
+            \[\phi(\boldsymbol{\theta}) = \frac{1}{2} |\boldsymbol{\theta} - \boldsymbol{\beta}|^2 - \psi(\boldsymbol{\theta})\]
+
+        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.
+
+        Parameters
+        ----------
+        grid
+            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
+
+        Returns
+        -------
+        The Fermat potential at each grid position.
+        """
+        time_delay_geometry_term = self.time_delay_geometry_term_from(grid=grid)
+        potential = self.potential_2d_from(grid=grid)
+
+        fermat_potential = time_delay_geometry_term - potential
+
+        if isinstance(grid, aa.Grid2DIrregular):
+            return aa.ArrayIrregular(values=fermat_potential)
+        return aa.Array2D(values=fermat_potential, mask=grid.mask)
+
+    def time_delays_from(self, grid) -> aa.Array2D:
+        """
+        Returns the 2D time delay map of lensing object, which is computed as the deflection angles in the y and x
+        directions multiplied by the y and x coordinates of the grid.
+
+        Parameters
+        ----------
+        grid
+            The 2D grid of (y,x) arc-second coordinates the deflection angles and time delay are computed on.
+        """
+        deflections_yx = self.deflections_yx_2d_from(grid=grid)
+
+        return aa.Array2D(
+            values=deflections_yx[:, 0] * grid[:, 0] + deflections_yx[:, 1] * grid[:, 1],
+            mask=grid.mask,
+        )
+
     @precompute_jacobian
     def tangential_eigen_value_from(self, grid, jacobian=None) -> aa.Array2D:
         """

test_autogalaxy/operate/test_deflections.py

@@ -55,6 +55,31 @@ def caustics_via_magnification_from(mass_profile, grid):
 
     return caustics
 
+def test__time_delay_geometry_term_from():
+
+    grid = ag.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0)])
+
+    mp = ag.mp.Isothermal(
+        centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0
+    )
+
+    time_delay_geometry_term = mp.time_delay_geometry_term_from(grid=grid)
+
+    assert time_delay_geometry_term == pytest.approx(np.array([1.92815688, 1.97625436]), 1.0e-4)
+
+
+def test__fermat_potential_from():
+
+    grid = ag.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0)])
+
+    mp = ag.mp.Isothermal(
+        centre=(0.0, 0.0), ell_comps=(0.0, -0.111111), einstein_radius=2.0
+    )
+
+    fermat_potential = mp.fermat_potential_from(grid=grid)
+
+    assert fermat_potential == pytest.approx(np.array([0.24329033, -0.82766592]), 1.0e-4)
+
 
 def test__hessian_from():
     grid = ag.Grid2DIrregular(values=[(0.5, 0.5), (1.0, 1.0)])