Merge branch 'main' into feature/jax_wrapper

Author: Jammy2211

autolens/__init__.py

@@ -101,6 +101,7 @@
 from .point.dataset import PointDataset
 from .point.fit.dataset import FitPointDataset
 from .point.fit.fluxes import FitFluxes
+from .point.fit.times_delays import FitTimeDelays
 from .point.fit.positions.image.abstract import AbstractFitPositionsImagePair
 from .point.fit.positions.image.pair import FitPositionsImagePair
 from .point.fit.positions.image.pair_all import FitPositionsImagePairAll

autolens/aggregator/subplot.py

@@ -50,9 +50,9 @@ class SubplotTracer(Enum):
     image = (0, 0)
     source_image = (1, 0)
     source_plane_image = (2, 0)
-    lens_light_image = (1, 0)
+    lens_light_image = (0, 1)
     convergence = (1, 1)
-    potential = (1, 2)
+    potential = (2, 1)
     magnification = (0, 2)
     deflections_y = (1, 2)
     deflections_x = (2, 2)
@@ -69,14 +69,14 @@ class SubplotFit(Enum):
     data_source_scale = (1, 0)
     signal_to_noise_map = (2, 0)
     model_data = (3, 0)
-    lens_light_model = (1, 0)
+    lens_light_model = (0, 1)
     lens_light_subtracted_image = (1, 1)
-    source_model_image = (1, 2)
-    source_plane_image_zoom = (1, 3)
-    normalized_residual_map = (2, 0)
-    normalized_residual_map_one_sigma = (2, 1)
+    source_model_image = (2, 1)
+    source_plane_image_zoom = (3, 1)
+    normalized_residual_map = (0, 2)
+    normalized_residual_map_one_sigma = (1, 2)
     chi_squared_map = (2, 2)
-    source_plane_image = (2, 3)
+    source_plane_image = (3, 2)
 
 
 class SubplotFitLog10(Enum):
@@ -90,11 +90,11 @@ class SubplotFitLog10(Enum):
     data_source_scale = (1, 0)
     signal_to_noise_map = (2, 0)
     model_data = (3, 0)
-    lens_light_model = (1, 0)
+    lens_light_model = (0, 1)
     lens_light_subtracted_image = (1, 1)
-    source_model_image = (1, 2)
-    source_plane_image_zoom = (1, 3)
-    normalized_residual_map = (2, 0)
-    normalized_residual_map_one_sigma = (2, 1)
+    source_model_image = (2, 1)
+    source_plane_image_zoom = (3, 1)
+    normalized_residual_map = (0, 2)
+    normalized_residual_map_one_sigma = (1, 2)
     chi_squared_map = (2, 2)
-    source_plane_image = (2, 3)
+    source_plane_image = (3, 2)

autolens/analysis/analysis/lens.py

@@ -121,12 +121,20 @@ def log_likelihood_positions_overwrite_from(
         else a None is returned to indicate there is no penalty.
         """
         if self.positions_likelihood_list is not None:
+
+            log_likelihood_overwrite = None
+
             try:
-                log_likelihood = 0.0
                 for positions_likelihood in self.positions_likelihood_list:
-                    log_likelihood += positions_likelihood.log_likelihood_function_positions_overwrite(
+                    log_likelihood_with_penalty = positions_likelihood.log_likelihood_function_positions_overwrite(
                         instance=instance, analysis=self
                     )
-                return log_likelihood
+                    if log_likelihood_with_penalty is not None:
+                        try:
+                            log_likelihood_overwrite += log_likelihood_with_penalty
+                        except TypeError:
+                            log_likelihood_overwrite = log_likelihood_with_penalty
+
+                return log_likelihood_overwrite
             except (ValueError, np.linalg.LinAlgError) as e:
                 raise exc.FitException from e

autolens/analysis/positions.py

@@ -84,7 +84,9 @@ def __init__(
 
         self.log_likelihood_penalty_factor = log_likelihood_penalty_factor
 
-    def output_positions_info(self, output_path: str, tracer: Tracer):
+    def output_positions_info(
+        self, output_path: str, tracer: Tracer, overwrite_file: bool = True
+    ):
         """
         Outputs a `positions.info` file which summarises the positions penalty term for a model fit, including:
 
@@ -102,7 +104,10 @@ def output_positions_info(self, output_path: str, tracer: Tracer):
         -------
 
         """
-        with open_(path.join(output_path, "positions.info"), "a+") as f:
+
+        flag = "w+" if overwrite_file else "a+"
+
+        with open_(path.join(output_path, "positions.info"), flag) as f:
 
             positions_fit = SourceMaxSeparation(
                 data=self.positions,
@@ -202,8 +207,6 @@ def log_likelihood_penalty_from(self, tracer: Tracer) -> Optional[float]:
         if not tracer.has(cls=ag.mp.MassProfile) or len(tracer.planes) == 1:
             return
 
-        log_likelihood_penalty = 0.0
-
         positions_fit = SourceMaxSeparation(
             data=self.positions,
             noise_map=None,
@@ -213,12 +216,10 @@ def log_likelihood_penalty_from(self, tracer: Tracer) -> Optional[float]:
 
         if not positions_fit.max_separation_within_threshold(self.threshold):
 
-            log_likelihood_penalty += self.log_likelihood_penalty_factor * (
+            return self.log_likelihood_penalty_factor * (
                 positions_fit.max_separation_of_plane_positions - self.threshold
             )
 
-        return log_likelihood_penalty
-
     def log_likelihood_function_positions_overwrite(
         self, instance: af.ModelInstance, analysis: AnalysisDataset
     ) -> Optional[float]:

autolens/imaging/model/visualizer.py

@@ -93,12 +93,16 @@ def visualize(
 
         if analysis.positions_likelihood_list is not None:
 
+            overwrite_file = True
+
             for positions_likelihood in analysis.positions_likelihood_list:
 
                 positions_likelihood.output_positions_info(
-                    output_path=paths.output_path, tracer=fit.tracer
+                    output_path=paths.output_path, tracer=fit.tracer, overwrite_file=overwrite_file
                 )
 
+                overwrite_file = False
+
         if fit.inversion is not None:
             try:
                 fit.inversion.reconstruction

autolens/interferometer/model/visualizer.py

@@ -91,12 +91,18 @@ def visualize(
 
         if analysis.positions_likelihood_list is not None:
 
+            overwrite_file = True
+
             for positions_likelihood in analysis.positions_likelihood_list:
 
                 positions_likelihood.output_positions_info(
-                    output_path=paths.output_path, tracer=fit.tracer
+                    output_path=paths.output_path,
+                    tracer=fit.tracer,
+                    overwrite_file=overwrite_file,
                 )
 
+                overwrite_file = False
+
         if fit.inversion is not None:
             try:
                 fit.inversion.reconstruction

autolens/lens/mock/mock_tracer.py

@@ -22,6 +22,7 @@ def __init__(
         magnification=None,
         einstein_radius=None,
         einstein_mass=None,
+        time_delays=None,
     ):
         super().__init__(image_plane_mesh_grid_pg_list=image_plane_mesh_grid_pg_list)
 
@@ -33,6 +34,7 @@ def __init__(
         self.magnification = magnification
         self.einstein_radius = einstein_radius
         self.einstein_mass = einstein_mass
+        self.time_delays = time_delays
 
     @property
     def planes(self):
@@ -58,3 +60,6 @@ def einstein_radius_from(self, grid):
 
     def einstein_mass_angular_from(self, grid):
         return self.einstein_mass
+
+    def time_delays_from(self, grid):
+        return self.time_delays

autolens/lens/tracer.py

@@ -759,6 +759,29 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
         """
         return sum([galaxy.potential_2d_from(grid=grid) for galaxy in self.galaxies])
 
+    @aa.grid_dec.to_array
+    def time_delays_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
+        """
+        Returns the gravitational lensing time delay in days, for a grid of 2D (y, x) coordinates.
+
+        This function calculates the time delay at each image-plane position due to both geometric and gravitational
+        (Shapiro) effects, as described by the Fermat potential, which are computed using the deflection angles of the
+        galaxies in the lens system.
+
+        Time dleays are computed from a reference point, which is the delay one would compute without a mass model.
+        This means a time delay could be negative, which means the light travel time is faster when lensing is
+        accounted for. When performing fitting, all time-delays are subtracted by the time-delay of the
+        shortest time-delay, such that its value is zero and all other time-delays are positive.
+
+        A full description of the calculation is given in the `autolens.lens.tracer.tracer_util.time_delays_from`
+        function, which performs the calculation and has full latex documentation of the equations used.
+        """
+        return tracer_util.time_delays_from(
+            galaxies=ag.Galaxies(self.galaxies_ascending_redshift),
+            grid=grid,
+            cosmology=self.cosmology,
+        )
+
     def has(self, cls: Type) -> bool:
         """
         Returns a bool specifying whether this tracer has a galaxy with a certain class type.

autolens/lens/tracer_util.py

@@ -4,6 +4,8 @@
 import autoarray as aa
 import autogalaxy as ag
 
+from autolens import exc
+
 
 def plane_redshifts_from(galaxies: List[ag.Galaxy]) -> List[float]:
     """
@@ -249,6 +251,92 @@ def grid_2d_at_redshift_from(
     return traced_grid_list[plane_index_insert]
 
 
+def time_delays_from(
+    galaxies: List[ag.Galaxy],
+    grid: aa.type.Grid2DLike,
+    cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
+) -> aa.type.Grid2DLike:
+    """
+    Returns the gravitational lensing time delay in days for a grid of 2D (y, x) coordinates.
+
+    This function calculates the time delay at each image-plane position due to both geometric and gravitational
+    (Shapiro) effects, as described by the Fermat potential, which are computed using the deflection angles of the
+    galaxies in the lens system.
+
+    It requires a two-plane system (lens and source), and does not currently support multi-plane time delay
+    calculations involving more than two planes, but it could be extended to do so in the future.
+
+    The time delay is computed as:
+
+    .. math::
+        \Delta t(\boldsymbol{\theta}) = \frac{D_{\Delta t}}{c} \, \phi(\boldsymbol{\theta})
+
+    where:
+
+    - \( \boldsymbol{\theta} \): image-plane coordinate
+    - \( \phi(\boldsymbol{\theta}) \): Fermat potential at each coordinate
+    - \( c \): speed of light
+    - \( D_{\Delta t} \): time-delay distance
+
+    The time-delay distance is given by:
+
+    .. math::
+        D_{\Delta t} = (1 + z_l) \frac{D_d D_s}{D_{ds}}
+
+    with \( D_d, D_s, D_{ds} \) the angular diameter distances to the lens, to the source, and from lens to source.
+
+    The time delay is computed using the Fermat potential,
+
+    An input `AstroPy` cosmology object can change the cosmological model, which is used to compute the scaling
+    factors between planes (which are derived from their redshifts and angular diameter distances). It is these
+    scaling factors that account for multi-plane ray tracing effects.
+
+    Parameters
+    ----------
+    galaxies
+        List of galaxies whose mass profiles define the lens and source planes. Must contain exactly two redshifts.
+    grid
+        The 2D (y, x) image-plane coordinates where the time delay is computed.
+    cosmology
+        The cosmological model used to calculate angular diameter distances. Defaults to Planck15.
+
+    Returns
+    -------
+    The time delay at each (y, x) coordinate in the input grid, in units of days.
+    """
+    plane_redshifts = plane_redshifts_from(galaxies=galaxies)
+
+    if len(plane_redshifts) != 2:
+        raise exc.RayTracingException(
+            "The time delay calculation requires exactly two planes, but the input galaxies have "
+            f"{len(plane_redshifts)} planes with redshifts {plane_redshifts}."
+        )
+
+    # Constants
+    mpc_in_m = 3.08567758e22  # Mpc in meters
+    arcsec_to_rad = np.deg2rad(1.0 / 3600.0)  # arcsec to radians
+    seconds_per_day = 86400
+    c = 299792458  # speed of light in m/s
+
+    factor = arcsec_to_rad**2 / seconds_per_day
+
+    # Angular diameter distances
+    Dd = cosmology.angular_diameter_distance(plane_redshifts[0]).value  # [Mpc]
+    Ds = cosmology.angular_diameter_distance(plane_redshifts[1]).value  # [Mpc]
+    Dds = cosmology.angular_diameter_distance_z1z2(
+        z1=plane_redshifts[0], z2=plane_redshifts[1]
+    ).value  # [Mpc]
+
+    # Time-delay distance in meters
+    D_dt = (1 + plane_redshifts[0]) * Dd * Ds / Dds * mpc_in_m
+
+    # Fermat potential
+    fermat_potential = galaxies.fermat_potential_from(grid=grid)
+
+    # Final time delay in days
+    return D_dt / c * fermat_potential * factor
+
+
 def ordered_plane_redshifts_with_slicing_from(
     lens_redshifts, planes_between_lenses, source_plane_redshift
 ):