feat: analytically-solved point-source likelihood variants (#657)

autolens/__init__.py

@@ -115,12 +115,18 @@
 from .point.dataset import output_to_csv
 from .point.fit.dataset import FitPointDataset
 from .point.fit.fluxes import FitFluxes
+from .point.fit.fluxes import FitFluxesSolved
 from .point.fit.times_delays import FitTimeDelays
+from .point.fit.times_delays import FitTimeDelaysSolved
+from .point.fit.solved import SolvedCentre
 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
+from .point.fit.positions.image.pair_all import FitPositionsImagePairAllSolved
 from .point.fit.positions.image.pair_repeat import FitPositionsImagePairRepeat
+from .point.fit.positions.image.pair_repeat import FitPositionsImagePairRepeatSolved
 from .point.fit.positions.source.separations import FitPositionsSource
+from .point.fit.positions.source.separations import FitPositionsSourceSolved
 from .point.max_separation import SourceMaxSeparation
 from .point.model.analysis import AnalysisPoint
 from .point.solver import PointSolver

autolens/exc.py

@@ -52,3 +52,15 @@ class PointExtractionException(Exception):
     """
 
     pass
+
+
+class PointProfileMismatchException(Exception):
+    """
+    Raised when a point-source profile is paired with a fit class that cannot honestly use it — e.g. a
+    centre-bearing `ps.Point` / `ps.PointFlux` with a `*Solved` fit (whose analytic solve would leave the centre
+    or flux priors sampled but silently ignored), or a profile without the attribute a fit class requires.
+
+    Deliberately NOT a subclass of `PointExtractionException`: `FitPointDataset` swallows that exception to skip
+    absent dataset components (its long-standing name-pairing semantics), and profile/fit mismatches must never
+    be silently skipped — they invalidate the composed model.
+    """

autolens/point/fit/abstract.py

@@ -146,10 +146,26 @@ def source_plane_coordinate(self) -> Tuple[float, float]:
         Returns the centre of the point-source in the source-plane, which is used when computing the model
         image-plane positions from the tracer.
 
+        This is the single funnel every position-based fit reads the source-plane centre from. By default it
+        reads the `centre` of the paired point-source profile (a free model parameter on `ag.ps.Point` /
+        `ag.ps.PointFlux`). The `autolens.point.fit.solved.SolvedCentre` mixin overrides this property on the
+        `*Solved` fit classes (e.g. `FitPositionsSourceSolved`) to instead return a centre solved for
+        analytically given the current tracer.
+
         Returns
         -------
         The (y,x) arc-second coordinates of the point-source in the source-plane.
         """
+        if not hasattr(self.profile, "centre"):
+            raise exc.PointProfileMismatchException(
+                f"The point-source profile paired to dataset '{self.name}' "
+                f"({self.profile.__class__.__name__}) has no `centre` attribute, so {self.__class__.__name__} "
+                f"cannot read a source-plane coordinate from it. Use a `centre`-bearing profile (e.g. "
+                f"`ag.ps.Point` / `ag.ps.PointFlux`), or use one of the analytically-solved fit classes (e.g. "
+                f"`FitPositionsSourceSolved`, `FitPositionsImagePairAllSolved`, "
+                f"`FitPositionsImagePairRepeatSolved`) which solve for the source-plane centre analytically "
+                f"and require a parameter-free profile such as `ag.ps.PointSolved`."
+            )
         return self.profile.centre
 
     @property

autolens/point/fit/dataset.py

@@ -33,6 +33,8 @@ def __init__(
         solver: PointSolver,
         fit_positions_cls=FitPositionsImagePair,
         xp=np,
+        fit_flux_cls=FitFluxes,
+        fit_time_delays_cls=FitTimeDelays,
     ):
         """
         Fits a point source dataset using a `Tracer` object, where the following components of the point source data
@@ -84,6 +86,12 @@ def __init__(
         fit_positions_cls
             The class used to fit the positions of the point source dataset, which could be an image-plane or
             source-plane chi-squared.
+        fit_flux_cls
+            The class used to fit the fluxes of the point source dataset, which could be a free-flux
+            (`FitFluxes`) or analytically-solved-flux (`FitFluxesSolved`) fit.
+        fit_time_delays_cls
+            The class used to fit the time delays of the point source dataset, which could be the
+            min-subtraction (`FitTimeDelays`) or analytically-solved-reference-time (`FitTimeDelaysSolved`) fit.
         profile
             Manually input the profile of the point source, which is used instead of the one extracted from the
             tracer via name pairing if that profile is not found.
@@ -95,6 +103,8 @@ def __init__(
         profile = self.tracer.extract_profile(profile_name=dataset.name)
 
         self.fit_positions_cls = fit_positions_cls
+        self.fit_flux_cls = fit_flux_cls
+        self.fit_time_delays_cls = fit_time_delays_cls
 
         try:
             self.positions = self.fit_positions_cls(
@@ -111,7 +121,7 @@ def __init__(
 
         try:
             if dataset.fluxes is not None:
-                self.flux = FitFluxes(
+                self.flux = self.fit_flux_cls(
                     name=dataset.name,
                     data=dataset.fluxes,
                     noise_map=dataset.fluxes_noise_map,
@@ -127,7 +137,7 @@ def __init__(
 
         try:
             if dataset.time_delays is not None:
-                self.time_delays = FitTimeDelays(
+                self.time_delays = self.fit_time_delays_cls(
                     name=dataset.name,
                     data=dataset.time_delays,
                     noise_map=dataset.time_delays_noise_map,

autolens/point/fit/fluxes.py

@@ -146,3 +146,144 @@ def chi_squared(self) -> float:
         return ag.util.fit.chi_squared_from(
             chi_squared_map=self.chi_squared_map.array,
         )
+
+
+class FitFluxesSolved(AbstractFitPoint):
+    """
+    Fits the fluxes of a point source dataset with the source-plane flux solved for analytically (in flux space,
+    magnification-first), following Lombardi 2024 (arXiv:2406.15280) §6.1, rather than read from a free `flux`
+    model parameter.
+
+    With image-plane magnifications `µᵢ` (`magnifications_at_positions`), observed fluxes `f̂ᵢ` and noise `σᵢ`:
+
+        `F* = (Σᵢ µᵢ f̂ᵢ/σᵢ²) / (Σᵢ µᵢ²/σᵢ²)`  (`solved_flux`)
+
+    with model fluxes `µᵢF*` (`model_data`), a standard chi-squared and noise normalization, and the likelihood
+    analytically marginalized over `F*` (flat prior):
+
+        `log_likelihood = -0.5*(χ² + noise_norm) - 0.5*log((Σᵢ µᵢ²/σᵢ²)/(2π))`
+
+    The paper's magnitude-space form is not used here: the flux noise maps in this fit are flux-space Gaussians,
+    and converting to magnitude space would change the error model, not just its parametrization.
+
+    Works with any profile that has **no** `flux` attribute (`ag.ps.Point` or `ag.ps.PointSolved`); a profile
+    with a `flux` attribute (`ag.ps.PointFlux`) raises, since its flux prior would otherwise be sampled by the
+    non-linear search but silently ignored by the analytic solve. Use `FitFluxes` for a free-flux fit.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        data: aa.ArrayIrregular,
+        noise_map: aa.ArrayIrregular,
+        positions: aa.Grid2DIrregular,
+        tracer: Tracer,
+        profile: Optional[ag.ps.Point] = None,
+        xp=np,
+    ):
+        """
+        Parameters
+        ----------
+        name
+            The name of the point source dataset which is paired to a `Point` profile.
+        data
+            The observed fluxes of the point source.
+        noise_map
+            The noise-map of the fluxes which are used to compute the log likelihood.
+        positions
+            The image-plane positions of the point source where the fluxes and magnifications are calculated.
+        tracer
+            The tracer of galaxies whose point source profile is used to fit the fluxes.
+        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,
+            xp=xp,
+        )
+
+        if hasattr(self.profile, "flux"):
+            raise exc.PointProfileMismatchException(
+                f"For the point-source named {name} the extracted point source was the class "
+                f"{self.profile.__class__.__name__}, which has a `flux` attribute. `FitFluxesSolved` solves "
+                f"for the source flux analytically (F*), so a free `flux` prior would be sampled by the "
+                f"non-linear search but silently ignored. Use `FitFluxes` with `ag.ps.PointFlux` for a "
+                f"free-flux fit, or use a profile with no `flux` attribute (e.g. `ag.ps.Point` / "
+                f"`ag.ps.PointSolved`) with `FitFluxesSolved`."
+            )
+
+    @property
+    def flux_precision_sum(self) -> float:
+        """
+        `Σᵢ µᵢ²/σᵢ²` — the precision of the solved flux `F*`, and the marginalization normalization.
+        """
+        mu = self.magnifications_at_positions.array
+        sigma_squared = self.noise_map.array**2.0
+        return self._xp.sum(mu**2.0 / sigma_squared)
+
+    @property
+    def solved_flux(self) -> float:
+        """
+        `F* = (Σᵢ µᵢ f̂ᵢ/σᵢ²) / (Σᵢ µᵢ²/σᵢ²)`.
+        """
+        mu = self.magnifications_at_positions.array
+        f_hat = self.data.array
+        sigma_squared = self.noise_map.array**2.0
+        numerator = self._xp.sum(mu * f_hat / sigma_squared)
+        return numerator / self.flux_precision_sum
+
+    @property
+    def model_data(self) -> aa.ArrayIrregular:
+        """
+        The model fluxes `µᵢF*`.
+        """
+        return aa.ArrayIrregular(
+            values=self.magnifications_at_positions.array * self.solved_flux
+        )
+
+    @property
+    def model_fluxes(self) -> aa.ArrayIrregular:
+        return self.model_data
+
+    @property
+    def residual_map(self) -> aa.ArrayIrregular:
+        """
+        Returns the difference between the observed and model fluxes of the point source.
+        """
+        residual_map = super().residual_map
+
+        return aa.ArrayIrregular(values=residual_map)
+
+    @property
+    def chi_squared(self) -> float:
+        """
+        Returns the chi-squared of the fit of the point source fluxes.
+        """
+        return ag.util.fit.chi_squared_from(
+            chi_squared_map=self.chi_squared_map.array,
+        )
+
+    @property
+    def marginalization_term(self) -> float:
+        """
+        The analytic-marginalization contribution to the log likelihood from integrating out the (flat-prior)
+        source flux: `-0.5 * log((Σᵢ µᵢ²/σᵢ²)/(2π))`.
+        """
+        return -0.5 * self._xp.log(self.flux_precision_sum / (2.0 * np.pi))
+
+    @property
+    def log_likelihood(self) -> float:
+        """
+        `log_likelihood = -0.5*(χ² + noise_norm) - 0.5*log((Σᵢ µᵢ²/σᵢ²)/(2π))`.
+        """
+        return (
+            -0.5 * (self.chi_squared + self.noise_normalization)
+            + self.marginalization_term
+        )

autolens/point/fit/positions/abstract.py

@@ -44,8 +44,10 @@ def __init__(
 
         The fit performs the following steps:
 
-        1) Determine the source-plane centre of the point source, which could be a free model parameter or computed
-           as the barycenter of ray-traced positions in the source-plane, using name pairing (see below).
+        1) Determine the source-plane centre of the point source, which is either a free model parameter read
+           from the profile's `centre` (`ag.ps.Point` / `ag.ps.PointFlux`) or, for the `*Solved` fit classes
+           (e.g. `FitPositionsSourceSolved`), solved for analytically given the current tracer (see
+           `autolens.point.fit.solved.SolvedCentre`), using name pairing (see below).
 
         2) Using the sub-class specific chi-squared, compute the residuals of each image-plane position, chi-squared
            and overall log likelihood of the fit.

autolens/point/fit/positions/image/abstract.py

@@ -45,8 +45,11 @@ def __init__(
 
         The fit performs the following steps:
 
-        1) Determine the source-plane centre of the point source, which could be a free model parameter or computed
-           as the barycenter of ray-traced positions in the source-plane, using name pairing (see below).
+        1) Determine the source-plane centre of the point source, which is either a free model parameter read
+           from the profile's `centre` (`ag.ps.Point` / `ag.ps.PointFlux`) or, for the `*Solved` fit classes
+           (e.g. `FitPositionsImagePairAllSolved`, `FitPositionsImagePairRepeatSolved`), solved for
+           analytically given the current tracer (see `autolens.point.fit.solved.SolvedCentre`), using name
+           pairing (see below).
 
         2) Determine the image-plane model positions using the `PointSolver` and the source-plane centre of the point
            source (e.g. ray tracing triangles to and from  the image and source planes), including accounting for

autolens/point/fit/positions/image/pair.py

@@ -24,10 +24,16 @@ class FitPositionsImagePair(AbstractFitPositionsImagePair):
     contributes the ``no_image_residual`` floor. ``FitPositionsImagePairRepeat`` remains the model-fit default;
     it additionally offers over-prediction policies.
 
+    **No analytically-solved-centre variant**: unlike ``FitPositionsImagePairAll`` /
+    ``FitPositionsImagePairRepeat``, this class has no ``*Solved`` counterpart. Its Hungarian assignment
+    (``scipy.optimize.linear_sum_assignment``) is not JAX-jittable, and its behaviour is superseded by
+    ``FitPositionsImagePairAllSolved`` / ``FitPositionsImagePairRepeatSolved`` for solved-centre fits.
+
     The fit performs the following steps:
 
-    1) Determine the source-plane centre of the point source, which could be a free model parameter or computed
-       as the barycenter of ray-traced positions in the source-plane, using name pairing (see below).
+    1) Determine the source-plane centre of the point source, which is either a free model parameter read from
+       the profile's `centre` (`ag.ps.Point` / `ag.ps.PointFlux`) — this class has no `*Solved` counterpart, see
+       above — using name pairing (see below).
 
     2) Determine the image-plane model positions using the `PointSolver` and the source-plane centre of the point
        source (e.g. ray tracing triangles to and from  the image and source planes), including accounting for

autolens/point/fit/positions/image/pair_all.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from autolens.point.fit.positions.image.abstract import AbstractFitPositionsImagePair
+from autolens.point.fit.solved import SolvedCentre
 
 
 class FitPositionsImagePairAll(AbstractFitPositionsImagePair):
@@ -22,8 +23,10 @@ class FitPositionsImagePairAll(AbstractFitPositionsImagePair):
 
     The fit performs the following steps:
 
-    1) Determine the source-plane centre of the point source, which could be a free model parameter or computed
-       as the barycenter of ray-traced positions in the source-plane, using name pairing (see below).
+    1) Determine the source-plane centre of the point source, which is either a free model parameter read from
+       the profile's `centre` (`ag.ps.Point` / `ag.ps.PointFlux`) or, for `FitPositionsImagePairAllSolved`,
+       solved for analytically given the current tracer (see `autolens.point.fit.solved.SolvedCentre`), using
+       name pairing (see below).
 
     2) Determine the image-plane model positions using the `PointSolver` and the source-plane centre of the point
        source (e.g. ray tracing triangles to and from  the image and source planes), including accounting for
@@ -152,3 +155,22 @@ def chi_squared(self) -> float:
             -self._xp.log(n_permutations)
             + self._xp.sum(self.all_permutations_log_likelihoods())
         )
+
+
+class FitPositionsImagePairAllSolved(SolvedCentre, FitPositionsImagePairAll):
+    """
+    ``FitPositionsImagePairAll`` with the source-plane centre fed into the `PointSolver` forward solve
+    (`model_data`, inherited unchanged from `AbstractFitPositionsImagePair`) solved for analytically
+    (`SolvedCentre.source_plane_coordinate`, `β*`) rather than read from a free `centre` model parameter.
+
+    This is **not** a result from Lombardi 2024 (arXiv:2406.15280) — the paper never substitutes a solved
+    source-plane centre into an image-plane likelihood. It is an extension in the spirit of glafic's
+    source-position-optimized image-plane chi-squared: the all-to-all pairing chi-squared itself
+    (`chi_squared`, `all_permutations_log_likelihoods`) is completely unchanged from `FitPositionsImagePairAll`.
+
+    Must be paired (by name) with a parameter-free profile such as `ag.ps.PointSolved`: a `centre`-bearing
+    profile (`ag.ps.Point` / `ag.ps.PointFlux`) raises (see `SolvedCentre.source_plane_coordinate`), since its
+    centre priors would otherwise be sampled but silently ignored.
+    """
+
+    _non_solved_alternative_name = "FitPositionsImagePairAll"

autolens/point/fit/positions/image/pair_repeat.py

@@ -4,6 +4,7 @@
 import autogalaxy as ag
 
 from autolens.point.fit.positions.image.abstract import AbstractFitPositionsImagePair
+from autolens.point.fit.solved import SolvedCentre
 
 
 class FitPositionsImagePairRepeat(AbstractFitPositionsImagePair):
@@ -14,8 +15,10 @@ class FitPositionsImagePairRepeat(AbstractFitPositionsImagePair):
 
     The fit performs the following steps:
 
-    1) Determine the source-plane centre of the point source, which could be a free model parameter or computed
-       as the barycenter of ray-traced positions in the source-plane, using name pairing (see below).
+    1) Determine the source-plane centre of the point source, which is either a free model parameter read from
+       the profile's `centre` (`ag.ps.Point` / `ag.ps.PointFlux`) or, for `FitPositionsImagePairRepeatSolved`,
+       solved for analytically given the current tracer (see `autolens.point.fit.solved.SolvedCentre`), using
+       name pairing (see below).
 
     2) Determine the image-plane model positions using the `PointSolver` and the source-plane centre of the point
        source (e.g. ray tracing triangles to and from  the image and source planes), including accounting for
@@ -217,3 +220,23 @@ def chi_squared(self) -> float:
         return chi_squared + self._xp.sum(
             (self.unmatched_model_penalty_map / noise_mean) ** 2.0
         )
+
+
+class FitPositionsImagePairRepeatSolved(SolvedCentre, FitPositionsImagePairRepeat):
+    """
+    ``FitPositionsImagePairRepeat`` with the source-plane centre fed into the `PointSolver` forward solve
+    (`model_data`, inherited unchanged from `AbstractFitPositionsImagePair`) solved for analytically
+    (`SolvedCentre.source_plane_coordinate`, `β*`) rather than read from a free `centre` model parameter.
+
+    This is **not** a result from Lombardi 2024 (arXiv:2406.15280) — the paper never substitutes a solved
+    source-plane centre into an image-plane likelihood. It is an extension in the spirit of glafic's
+    source-position-optimized image-plane chi-squared: the pairing chi-squared itself (`chi_squared`,
+    `residual_map`, the over-/under-prediction policies) is completely unchanged from
+    `FitPositionsImagePairRepeat`.
+
+    Must be paired (by name) with a parameter-free profile such as `ag.ps.PointSolved`: a `centre`-bearing
+    profile (`ag.ps.Point` / `ag.ps.PointFlux`) raises (see `SolvedCentre.source_plane_coordinate`), since its
+    centre priors would otherwise be sampled but silently ignored.
+    """
+
+    _non_solved_alternative_name = "FitPositionsImagePairRepeat"