diff --git a/autolens/point/dataset.py b/autolens/point/dataset.py index d3245e797..d54d58782 100644 --- a/autolens/point/dataset.py +++ b/autolens/point/dataset.py @@ -1,3 +1,20 @@ +""" +Data structures for point-source strong lens observations. + +Point-source lensing arises when the background source is compact enough to be treated +as a point (e.g. a quasar, supernova, or compact radio source). Gravitational lensing +splits the source into multiple images whose positions, fluxes, and time delays constrain +the lens mass distribution. + +Two concrete data classes are provided: + +- ``PointDataset`` — holds the image-plane positions, fluxes, and time delays of a + single named point source together with their noise maps. The ``name`` attribute is + used to pair this dataset with the corresponding ``Point`` model component during + fitting. +- ``PointDict`` — a dictionary of ``PointDataset`` objects keyed by name, used when + multiple point sources (e.g. different background quasars) are fitted simultaneously. +""" from typing import List, Tuple, Optional, Union import autoarray as aa diff --git a/autolens/point/fit/abstract.py b/autolens/point/fit/abstract.py index e6903ec75..35f26078c 100644 --- a/autolens/point/fit/abstract.py +++ b/autolens/point/fit/abstract.py @@ -1,3 +1,19 @@ +""" +Abstract base class for all point-source fit components. + +``AbstractFitPoint`` provides shared functionality that is common to all point-source fit +classes (position fits, flux fits, time-delay fits): + +- Solving for the image-plane positions of a source-plane coordinate using a + ``PointSolver``. +- Computing the magnification at each image position from the tracer's mass model. +- Computing deflection angles at each image position. +- Providing the name-based pairing between a ``PointDataset`` and its ``Point`` model + component. + +Concrete subclasses implement ``figure_of_merit`` to compute the chi-squared or +log-likelihood for the specific observable (positions, fluxes, time delays) being fitted. +""" from abc import ABC import numpy as np from typing import Optional, Tuple diff --git a/autolens/point/fit/dataset.py b/autolens/point/fit/dataset.py index 0e3ec7a4d..b86f43169 100644 --- a/autolens/point/fit/dataset.py +++ b/autolens/point/fit/dataset.py @@ -1,3 +1,18 @@ +""" +Top-level fit class for a complete point-source dataset. + +``FitPointDataset`` orchestrates fitting of all observables in a ``PointDataset`` +(image-plane positions, fluxes, and/or time delays) simultaneously. It creates and +stores individual fit objects for each component that is present in the dataset: + +- ``FitPositionsImagePair`` (or another positions fit class) — fits image-plane positions. +- ``FitFluxes`` — fits flux ratios (if fluxes are in the dataset). +- ``FitTimeDelays`` — fits time delays (if time delays are in the dataset). + +The ``log_likelihood`` is the sum of the individual component log likelihoods. This +class is used by ``AnalysisPoint`` as the evaluation engine inside the +``log_likelihood_function``. +""" import numpy as np from autolens.point.dataset import PointDataset diff --git a/autolens/point/fit/fluxes.py b/autolens/point/fit/fluxes.py index 9be71dc07..70da16345 100644 --- a/autolens/point/fit/fluxes.py +++ b/autolens/point/fit/fluxes.py @@ -1,3 +1,14 @@ +""" +Flux-ratio fit component for point-source lensing. + +``FitFluxes`` computes the likelihood of the observed image-plane fluxes given the +predicted magnification ratios from the tracer's mass model. + +The predicted fluxes are proportional to the absolute magnification at each solved image +position, normalised so that the brightest image has flux 1.0 (flux ratios). A chi-squared +is computed against the observed flux values and noise map, contributing to the total +``FitPointDataset`` log likelihood. +""" import numpy as np from typing import Optional diff --git a/autolens/point/fit/positions/abstract.py b/autolens/point/fit/positions/abstract.py index 66c3264dd..77002fefc 100644 --- a/autolens/point/fit/positions/abstract.py +++ b/autolens/point/fit/positions/abstract.py @@ -1,3 +1,20 @@ +""" +Abstract base class for position-based point-source fit components. + +``AbstractFitPositions`` extends ``AbstractFitPoint`` with shared logic for all position +fitting strategies (image-plane pair fits and source-plane separation fits). It provides +common attributes — observed positions, solved image positions, deflection angles — and +leaves ``figure_of_merit`` to be implemented by each concrete subclass. + +The concrete subclasses (in the ``image/`` and ``source/`` sub-packages) implement +different statistical approaches for fitting image positions: + +- **Image-plane pair fits** — compare each observed position to the nearest predicted + image position, computing a chi-squared from the separation in arcseconds. +- **Source-plane separation fits** — trace each observed image back to the source plane + and compute the scatter of the traced positions around their mean, penalising poor + source-plane convergence. +""" from abc import ABC import numpy as np from typing import Optional diff --git a/autolens/point/fit/positions/image/abstract.py b/autolens/point/fit/positions/image/abstract.py index 798e42610..2b84b2875 100644 --- a/autolens/point/fit/positions/image/abstract.py +++ b/autolens/point/fit/positions/image/abstract.py @@ -1,3 +1,20 @@ +""" +Abstract base for image-plane position fitting strategies. + +``AbstractFitPositionsImagePair`` solves for the predicted image positions using a +``PointSolver`` and then pairs each observed position with the closest predicted image, +computing a chi-squared from their separation. + +Three concrete pairing strategies are provided in this sub-package: + +- ``FitPositionsImagePair`` — pairs each observed to its nearest predicted image via + the Hungarian (linear sum assignment) algorithm. +- ``FitPositionsImagePairAll`` — computes the chi-squared for every observed/predicted + pair combination and takes the minimum. +- ``FitPositionsImagePairRepeat`` — allows predicted images to be paired to more than + one observed position, useful for highly magnified systems where some images may be + too close to separate. +""" from abc import ABC import numpy as np from typing import Optional diff --git a/autolens/point/fit/positions/source/separations.py b/autolens/point/fit/positions/source/separations.py index 53102ac1f..9919f3207 100644 --- a/autolens/point/fit/positions/source/separations.py +++ b/autolens/point/fit/positions/source/separations.py @@ -1,3 +1,18 @@ +""" +Source-plane position fitting via traced-position separations. + +Instead of comparing predicted and observed positions in the image plane, +``FitPositionsSourcePlane`` traces each *observed* image position back to the source +plane via the tracer's deflection angles and measures how tightly they converge. + +If the lens model is correct, all observed images of the same source should trace back +to (approximately) the same source-plane coordinate. The figure of merit is the mean +squared separation of the back-traced positions from their common centroid, normalised by +the position noise map. + +This approach avoids the need for a ``PointSolver`` (no forward-solving is required) and +is well-suited to JAX-accelerated model fits. +""" import numpy as np from typing import Optional diff --git a/autolens/point/fit/times_delays.py b/autolens/point/fit/times_delays.py index 598abf01f..ca21e051e 100644 --- a/autolens/point/fit/times_delays.py +++ b/autolens/point/fit/times_delays.py @@ -1,3 +1,18 @@ +""" +Time-delay fit component for point-source lensing. + +``FitTimeDelays`` computes the likelihood of the observed time delays between multiple +images of a point source given the Fermat potential predicted by the tracer's mass and +light models. + +The time delay between images i and j is: + + Δt_{ij} = (D_Δt / c) × [φ(θ_i) − φ(θ_j)] + +where D_Δt is the time-delay distance computed from the cosmology and φ is the Fermat +potential. A chi-squared is computed against the observed delays and their noise map, +contributing to the total ``FitPointDataset`` log likelihood. +""" import numpy as np from typing import Optional diff --git a/autolens/point/max_separation.py b/autolens/point/max_separation.py index 6b0fbb55c..dd87da170 100644 --- a/autolens/point/max_separation.py +++ b/autolens/point/max_separation.py @@ -1,3 +1,15 @@ +""" +Maximum source-plane separation statistic for point-source position fitting. + +``SourceMaxSeparation`` computes the maximum separation between the source-plane +positions obtained by tracing each observed image-plane position through the lens model. +For a perfect lens model all images of the same source trace back to exactly the same +source-plane coordinate, so this maximum separation should be zero. + +This statistic is used as a fast figure-of-merit in settings where a full chi-squared +calculation is not required, or as a hard prior threshold that rejects models for which +the back-traced positions diverge by more than a user-specified amount. +""" import numpy as np from typing import Optional diff --git a/autolens/point/model/analysis.py b/autolens/point/model/analysis.py index d0f702528..54dd1f865 100644 --- a/autolens/point/model/analysis.py +++ b/autolens/point/model/analysis.py @@ -1,3 +1,19 @@ +""" +Analysis class for fitting a ``Tracer`` model to a point-source dataset. + +``AnalysisPoint`` implements the ``log_likelihood_function`` called by a ``PyAutoFit`` +non-linear search at each iteration. It: + +1. Constructs a ``Tracer`` from the current model instance. +2. Calls ``FitPointDataset`` to fit the point-source positions (and optionally fluxes + and time delays) using the ``PointSolver`` to find predicted image positions. +3. Optionally adds a position-based prior via ``PositionsLH`` that penalises models + where image positions are not self-consistent. +4. Returns the total log likelihood as the figure of merit. + +It also manages result output (``ResultPoint``) and on-the-fly visualisation +(``VisualizerPoint``). +""" import numpy as np import autofit as af diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py index cb9521cef..6a5cad4b6 100644 --- a/autolens/point/solver/point_solver.py +++ b/autolens/point/solver/point_solver.py @@ -1,3 +1,21 @@ +""" +Image-plane point-source solver for strong gravitational lensing. + +Finding the multiple images of a point source requires solving the lens equation +θ = β + α(θ) for θ given a fixed source-plane position β. This is an inverse +problem with no analytic solution for general mass distributions. + +``PointSolver`` solves this numerically using a triangle-tiling approach: + +1. The image plane is tiled with triangles. +2. Each triangle is ray-traced to the source plane. +3. Triangles that contain the source-plane coordinate β are refined recursively. +4. The centroids of the final refined triangles give the image-plane positions. + +The output positions array is padded to a fixed size (``MAX_CONTAINING_SIZE``) using the +sentinel value ``inf`` for JAX compatibility — these ``inf`` entries are stripped by +default but can be retained for use inside a ``jax.jit``-traced function. +""" import logging from typing import Tuple, Optional diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py index 92d7fcad8..9d96289a3 100644 --- a/autolens/point/solver/shape_solver.py +++ b/autolens/point/solver/shape_solver.py @@ -1,3 +1,18 @@ +""" +Abstract triangle-tiling solver and shape-based solver for point-source positions. + +``AbstractSolver`` (and its concrete subclass ``ShapeSolver``) implement the hierarchical +triangle-refinement algorithm that underlies ``PointSolver``: + +1. An initial grid of triangles covers the image plane. +2. Each triangle is ray-traced to the source plane; those that contain the target + source coordinate are kept. +3. Kept triangles are sub-divided for the next refinement iteration. +4. After ``n_steps`` levels the centroids of the finest triangles give the image positions. + +``ShapeSolver`` extends this base with support for fitting extended source *shapes* +(e.g. rings, arcs) rather than point coordinates, used for morphological constraints. +""" import numpy as np import logging import math