docs(point): add module-level docstrings to point package

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

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

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

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
 

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

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

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
 

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
 

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
 

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

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
 

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