Feature/time delay fit

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/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

@@ -761,18 +761,23 @@ 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_delay_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
+    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.
 
-        A full description of the calculation is given in the `autolens.lens.tracer.tracer_util.time_delay_from`
+        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_delay_from(
+        return tracer_util.time_delays_from(
             galaxies=ag.Galaxies(self.galaxies_ascending_redshift),
             grid=grid,
             cosmology=self.cosmology,

autolens/lens/tracer_util.py

@@ -251,7 +251,7 @@ def grid_2d_at_redshift_from(
     return traced_grid_list[plane_index_insert]
 
 
-def time_delay_from(
+def time_delays_from(
     galaxies: List[ag.Galaxy],
     grid: aa.type.Grid2DLike,
     cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),

autolens/point/dataset.py

@@ -11,6 +11,10 @@ def __init__(
         positions_noise_map: Union[float, aa.ArrayIrregular, List[float]],
         fluxes: Optional[Union[aa.ArrayIrregular, List[float]]] = None,
         fluxes_noise_map: Optional[Union[float, aa.ArrayIrregular, List[float]]] = None,
+        time_delays: Optional[Union[aa.ArrayIrregular, List[float]]] = None,
+        time_delays_noise_map: Optional[
+            Union[float, aa.ArrayIrregular, List[float]]
+        ] = None,
     ):
         """
         A collection of the data component that can be used for point-source model-fitting, for example fitting the
@@ -33,42 +37,49 @@ def __init__(
         fluxes
             The image-plane flux of each observed point-source of light.
         fluxes_noise_map
-            The noise-value of every observed flux.
+            The noise-value of every observed flux, which is typically measured from the pixel values of the pixel
+            containing the point source after convolution with the PSF.
+        time_delays
+            The time delays of each observed point-source of light in days.
+        time_delays_noise_map
+            The noise-value of every observed time delay, which is typically measured from the time delay analysis.
         """
 
         self.name = name
 
-        if not isinstance(positions, aa.Grid2DIrregular):
-            positions = aa.Grid2DIrregular(values=positions)
-
-        self.positions = positions
+        # Ensure positions is a Grid2DIrregular
+        self.positions = (
+            positions
+            if isinstance(positions, aa.Grid2DIrregular)
+            else aa.Grid2DIrregular(values=positions)
+        )
 
+        # Ensure positions_noise_map is an ArrayIrregular
         if isinstance(positions_noise_map, float):
-            positions_noise_map = aa.ArrayIrregular(
-                values=len(positions) * [positions_noise_map]
-            )
-
-        if not isinstance(positions_noise_map, aa.ArrayIrregular):
-            positions_noise_map = aa.ArrayIrregular(values=positions_noise_map)
-
-        self.positions_noise_map = positions_noise_map
-
-        if fluxes is not None:
-            if not isinstance(fluxes, aa.ArrayIrregular):
-                fluxes = aa.ArrayIrregular(values=fluxes)
-
-        self.fluxes = fluxes
-
-        if isinstance(fluxes_noise_map, float):
-            fluxes_noise_map = aa.ArrayIrregular(
-                values=len(fluxes) * [fluxes_noise_map]
+            positions_noise_map = [positions_noise_map] * len(self.positions)
+
+        self.positions_noise_map = (
+            positions_noise_map
+            if isinstance(positions_noise_map, aa.ArrayIrregular)
+            else aa.ArrayIrregular(values=positions_noise_map)
+        )
+
+        def convert_to_array_irregular(values):
+            """
+            Convert data to ArrayIrregular if it is not already.
+            """
+            return (
+                aa.ArrayIrregular(values=values)
+                if values is not None and not isinstance(values, aa.ArrayIrregular)
+                else values
             )
 
-        if fluxes_noise_map is not None:
-            if not isinstance(fluxes_noise_map, aa.ArrayIrregular):
-                fluxes_noise_map = aa.ArrayIrregular(values=fluxes_noise_map)
+        # Convert fluxes, time delays and their noise maps to ArrayIrregular if provided as values and not already this type
 
-        self.fluxes_noise_map = fluxes_noise_map
+        self.fluxes = convert_to_array_irregular(fluxes)
+        self.fluxes_noise_map = convert_to_array_irregular(fluxes_noise_map)
+        self.time_delays = convert_to_array_irregular(time_delays)
+        self.time_delays_noise_map = convert_to_array_irregular(time_delays_noise_map)
 
     @property
     def info(self) -> str:
@@ -82,6 +93,8 @@ def info(self) -> str:
         info += f"positions_noise_map : {self.positions_noise_map}\n"
         info += f"fluxes : {self.fluxes}\n"
         info += f"fluxes_noise_map : {self.fluxes_noise_map}\n"
+        info += f"time_delays : {self.time_delays}\n"
+        info += f"time_delays_noise_map : {self.time_delays_noise_map}\n"
         return info
 
     def extent_from(self, buffer: float = 0.1):

autolens/point/fit/dataset.py

@@ -4,6 +4,7 @@
 from autolens.point.dataset import PointDataset
 from autolens.point.solver import PointSolver
 from autolens.point.fit.fluxes import FitFluxes
+from autolens.point.fit.times_delays import FitTimeDelays
 from autolens.lens.tracer import Tracer
 
 from autolens.point.fit.positions.image.pair import FitPositionsImagePair
@@ -36,7 +37,8 @@ def __init__(
         - The fluxes of the point source, which use the magnification of the point source to compute the fluxes in the
           image-plane.
 
-        - The time delays of the point source (NOT IMPLEMENTED YET).
+        - The time delays of the point source in delays, which use the tracer to compute the model time delays
+          at the image-plane positions of the point source in the dataset.
 
         The fit may use one or combinations of the above components to compute the log likelihood, depending on what
         components are available in the point source dataset and the model point source profiles input. For example:
@@ -54,7 +56,8 @@ def __init__(
         2) Fit the fluxes of the point source dataset using the `FitFluxes` object, where the object type may be
           extended in the future to support different types of point source profiles.
 
-        3) Time delays are not currently supported but this API will extend to include time delays in the future.
+        3) Fits the time delays of the point source dataset using the `FitTimeDelays` object, which is an image-plane
+           evaluation of the time delays at the image-plane positions of the point source in the dataset.
 
         Point source fitting uses name pairing, whereby the `name` of the `Point` object is paired to the name of the
         point source dataset to ensure that point source datasets are fitted to the correct point source.
@@ -100,17 +103,34 @@ def __init__(
             self.positions = None
 
         try:
-            self.flux = FitFluxes(
-                name=dataset.name,
-                data=dataset.fluxes,
-                noise_map=dataset.fluxes_noise_map,
-                positions=dataset.positions,
-                tracer=tracer,
-            )
+            if dataset.fluxes is not None:
+                self.flux = FitFluxes(
+                    name=dataset.name,
+                    data=dataset.fluxes,
+                    noise_map=dataset.fluxes_noise_map,
+                    positions=dataset.positions,
+                    tracer=tracer,
+                )
+            else:
+                self.flux = None
 
         except exc.PointExtractionException:
             self.flux = None
 
+        try:
+            if dataset.time_delays is not None:
+                self.time_delays = FitTimeDelays(
+                    name=dataset.name,
+                    data=dataset.time_delays,
+                    noise_map=dataset.time_delays_noise_map,
+                    positions=dataset.positions,
+                    tracer=tracer,
+                )
+            else:
+                self.time_delays = None
+        except exc.PointExtractionException:
+            self.time_delays = None
+
     @property
     def model_obj(self):
         return self.tracer
@@ -125,8 +145,13 @@ def log_likelihood(self) -> float:
             self.positions.log_likelihood if self.positions is not None else 0.0
         )
         log_likelihood_flux = self.flux.log_likelihood if self.flux is not None else 0.0
+        log_likelihood_time_delays = (
+            self.time_delays.log_likelihood if self.time_delays is not None else 0.0
+        )
 
-        return log_likelihood_positions + log_likelihood_flux
+        return (
+            log_likelihood_positions + log_likelihood_flux + log_likelihood_time_delays
+        )
 
     @property
     def figure_of_merit(self) -> float:

autolens/point/fit/times_delays.py

@@ -0,0 +1,106 @@
+import numpy as np
+from typing import Optional
+
+import autoarray as aa
+import autogalaxy as ag
+
+from autolens.point.fit.abstract import AbstractFitPoint
+from autolens.lens.tracer import Tracer
+
+
+class FitTimeDelays(AbstractFitPoint):
+    def __init__(
+        self,
+        name: str,
+        data: aa.ArrayIrregular,
+        noise_map: aa.ArrayIrregular,
+        positions: aa.Grid2DIrregular,
+        tracer: Tracer,
+        profile: Optional[ag.ps.Point] = None,
+    ):
+        """
+        Fits the time delays of a point source dataset using a `Tracer` object,
+        where every model time delay of the point-source is compared with its observed time delay.
+
+        The fit performs the following steps:
+
+        1) Compute the model time delays at the input image-plane `positions` using the tracer.
+
+        2) Compute the relative time delays of the dataset time delays and the time delays of the point source at
+           these positions, which are the time delays  relative to the shortest time delay
+
+        2) Subtract the observed relative time delays from the model relative time delays to compute the residuals.
+
+        3) Compute the chi-squared of each time delay residual.
+
+        4) Sum the chi-squared values to compute the overall log likelihood of the fit.
+
+        Time delay fitting uses name pairing similar to flux fitting to ensure
+        the correct point source profile is used.
+
+        Parameters
+        ----------
+        name
+            The name of the point source dataset which is paired to a `Point` profile.
+        data
+            The observed time delays in days of the point source.
+        noise_map
+            The noise-map of the time delays in days used to compute the log likelihood.
+        tracer
+            The tracer of galaxies whose point source profile is used to fit the time delays.
+        positions
+            The image-plane positions of the point source where the time delays are calculated.
+        profile
+            Manually input the profile of the point source, used instead of one extracted from the tracer.
+        """
+        self.positions = positions
+
+        super().__init__(
+            name=name,
+            data=data,
+            noise_map=noise_map,
+            tracer=tracer,
+            solver=None,
+            profile=profile,
+        )
+
+    @property
+    def model_data(self) -> aa.ArrayIrregular:
+        """
+        The model time delays of the tracer at each of the input image-plane positions.
+
+        These values are not subtracted by the shorter time delay of the point source, which would make the shorter
+        delay have a value of zero. However, this subtraction is performed in the `residual_map` property, in order
+        to ensure the residuals are computed relative to the shorter time delay.
+        """
+        return self.tracer.time_delays_from(grid=self.positions)
+
+    @property
+    def model_time_delays(self) -> aa.ArrayIrregular:
+        return self.model_data
+
+    @property
+    def residual_map(self) -> aa.ArrayIrregular:
+        """
+        Returns the difference between the observed and model time delays of the point source,
+        which is the residual time delay of the fit.
+
+        The residuals are computed relative to the shortest time delay of the point source, which is subtracted
+        from the dataset time delays and model time delays before the subtraction.
+        """
+
+        data = self.data - np.min(self.data)
+        model_data = self.model_data - np.min(self.model_data)
+
+        residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
+        return aa.ArrayIrregular(values=residual_map)
+
+    @property
+    def chi_squared(self) -> float:
+        """
+        Returns the chi-squared of the fit of the point source time delays,
+        which is the residual values divided by the RMS noise-map squared.
+        """
+        return ag.util.fit.chi_squared_from(
+            chi_squared_map=self.chi_squared_map,
+        )

test_autolens/lens/test_tracer.py

@@ -561,7 +561,7 @@ def test__potential_2d_from(grid_2d_7x7):
     ).all()
 
 
-def test__time_delay_from():
+def test__time_delays_from():
 
     grid = al.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0)])
 
@@ -574,7 +574,7 @@ def test__time_delay_from():
 
     tracer = al.Tracer(galaxies=[g0, g1])
 
-    time_delay = tracer.time_delay_from(grid=grid)
+    time_delay = tracer.time_delays_from(grid=grid)
 
     assert time_delay == pytest.approx(np.array([8.52966247, -29.0176387]), 1.0e-4)
 

test_autolens/lens/test_tracer_util.py

@@ -175,7 +175,7 @@ def test__grid_2d_at_redshift_from__redshift_between_planes(grid_2d_7x7):
     assert (grid_at_redshift == grid_2d_7x7.mask.derive_grid.all_false).all()
 
 
-def test__time_delay_from():
+def test__time_delays_from():
 
     grid = al.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0)])
 
@@ -186,7 +186,7 @@ def test__time_delay_from():
     lens = al.Galaxy(redshift=0.2, mass=mp)
     source = al.Galaxy(redshift=0.7)
 
-    time_delay = al.util.tracer.time_delay_from(
+    time_delay = al.util.tracer.time_delays_from(
         galaxies=al.Galaxies([lens, source]),
         grid=grid,
         cosmology=al.cosmo.Planck15(),