From 5dbe0cece39f899a0929b1973acf186fb32b62cf Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 27 Jul 2026 21:24:24 +0100 Subject: [PATCH 1/2] feat: analytically-solved point-source likelihood variants (#657) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FitPositionsSourceSolved (tensor-weighted, marginalized centre; Lombardi 2024 arXiv:2406.15280 §5.1), FitPositionsImagePairAllSolved / FitPositionsImagePairRepeatSolved (solved centre driving the existing PointSolver forward solve), FitFluxesSolved (analytic flux, flux-space, magnification-first), FitTimeDelaysSolved (analytic reference time); SolvedCentre mixin; fit_flux_cls / fit_time_delays_cls hooks; pytree registration incl. FitPositionsSource; informative mismatch errors; docstring truth sweep. Co-Authored-By: Claude Fable 5 --- autolens/__init__.py | 6 + autolens/point/fit/abstract.py | 16 + autolens/point/fit/dataset.py | 14 +- autolens/point/fit/fluxes.py | 141 ++++++ autolens/point/fit/positions/abstract.py | 6 +- .../point/fit/positions/image/abstract.py | 7 +- autolens/point/fit/positions/image/pair.py | 10 +- .../point/fit/positions/image/pair_all.py | 26 +- .../point/fit/positions/image/pair_repeat.py | 27 +- .../point/fit/positions/source/separations.py | 114 ++++- autolens/point/fit/solved.py | 240 ++++++++++ autolens/point/fit/times_delays.py | 82 ++++ autolens/point/model/analysis.py | 68 ++- autolens/point/solver/shape_solver.py | 2 +- .../fit/positions/image/test_pair_all.py | 69 +++ .../fit/positions/image/test_pair_repeat.py | 27 ++ .../fit/positions/source/test_separations.py | 29 ++ test_autolens/point/fit/test_fit_dataset.py | 61 +++ test_autolens/point/fit/test_fluxes.py | 70 +++ test_autolens/point/fit/test_solved.py | 442 ++++++++++++++++++ test_autolens/point/fit/test_time_delays.py | 56 +++ .../point/model/test_analysis_point.py | 37 ++ 22 files changed, 1527 insertions(+), 23 deletions(-) create mode 100644 autolens/point/fit/solved.py create mode 100644 test_autolens/point/fit/test_solved.py diff --git a/autolens/__init__.py b/autolens/__init__.py index 9839c51e4..e1d02a2d7 100644 --- a/autolens/__init__.py +++ b/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 diff --git a/autolens/point/fit/abstract.py b/autolens/point/fit/abstract.py index b6559a579..abd97c487 100644 --- a/autolens/point/fit/abstract.py +++ b/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.PointExtractionException( + 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 diff --git a/autolens/point/fit/dataset.py b/autolens/point/fit/dataset.py index 8d2df4c77..f4d653f60 100644 --- a/autolens/point/fit/dataset.py +++ b/autolens/point/fit/dataset.py @@ -32,6 +32,8 @@ def __init__( tracer: Tracer, solver: PointSolver, fit_positions_cls=FitPositionsImagePair, + fit_flux_cls=FitFluxes, + fit_time_delays_cls=FitTimeDelays, xp=np, ): """ @@ -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, diff --git a/autolens/point/fit/fluxes.py b/autolens/point/fit/fluxes.py index 01899611c..d41ce5e54 100644 --- a/autolens/point/fit/fluxes.py +++ b/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.PointExtractionException( + 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 + ) diff --git a/autolens/point/fit/positions/abstract.py b/autolens/point/fit/positions/abstract.py index 1956e5ea3..11c125bdc 100644 --- a/autolens/point/fit/positions/abstract.py +++ b/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. diff --git a/autolens/point/fit/positions/image/abstract.py b/autolens/point/fit/positions/image/abstract.py index bd3e351eb..524b71d30 100644 --- a/autolens/point/fit/positions/image/abstract.py +++ b/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 diff --git a/autolens/point/fit/positions/image/pair.py b/autolens/point/fit/positions/image/pair.py index 8c08e7a65..1a303b0c1 100644 --- a/autolens/point/fit/positions/image/pair.py +++ b/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 diff --git a/autolens/point/fit/positions/image/pair_all.py b/autolens/point/fit/positions/image/pair_all.py index cdcf27775..59c255ca2 100644 --- a/autolens/point/fit/positions/image/pair_all.py +++ b/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" diff --git a/autolens/point/fit/positions/image/pair_repeat.py b/autolens/point/fit/positions/image/pair_repeat.py index 52a1c3787..f37621f0b 100644 --- a/autolens/point/fit/positions/image/pair_repeat.py +++ b/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" diff --git a/autolens/point/fit/positions/source/separations.py b/autolens/point/fit/positions/source/separations.py index dbbedbe5a..3968769db 100644 --- a/autolens/point/fit/positions/source/separations.py +++ b/autolens/point/fit/positions/source/separations.py @@ -6,9 +6,12 @@ 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. +to (approximately) the same source-plane coordinate. The figure of merit is the squared +separation of the back-traced positions from that source-plane coordinate, normalised by +the position noise map. That coordinate is not a computed centroid of the back-traced +positions: it is the paired point-source profile's ``centre`` (a free model parameter, for +``FitPositionsSource``) or, for ``FitPositionsSourceSolved``, the analytically-solved centre +``β*`` (see ``autolens.point.fit.solved.SolvedCentre``). This approach avoids the need for a ``PointSolver`` (no forward-solving is required) and is well-suited to JAX-accelerated model fits. @@ -21,6 +24,7 @@ from autolens.lens.tracer import Tracer from autolens.point.fit.positions.abstract import AbstractFitPositions +from autolens.point.fit.solved import SolvedCentre, precision_tensor_components_from from autolens.point.solver import PointSolver @@ -42,8 +46,10 @@ def __init__( The fit performs the following steps: - 1) Determine the source-plane centre of the source-galaxy, 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 source-galaxy, which is either a free model parameter read + from the profile's `centre` (`ag.ps.Point` / `ag.ps.PointFlux`) or, for `FitPositionsSourceSolved`, + solved for analytically given the current tracer (see `autolens.point.fit.solved.SolvedCentre`), + using name pairing (see below). 2) Ray-trace the positions in the point source to the source-plane via the `Tracer`, including accounting for multi-plane ray-tracing. @@ -116,7 +122,9 @@ def model_data(self) -> aa.Grid2DIrregular: grid=self.data, xp=self._xp, plane_i=0, plane_j=self.plane_index ) - return self.data.grid_2d_via_deflection_grid_from(deflection_grid=deflections) + return self.data.grid_2d_via_deflection_grid_from( + deflection_grid=deflections, xp=self._xp + ) @property def residual_map(self) -> aa.ArrayIrregular: @@ -161,3 +169,97 @@ def log_likelihood(self) -> float: Returns the log likelihood of the point-source source-plane fit, which is the sum of the chi-squared values. """ return -0.5 * (sum(self.chi_squared_map) + self.noise_normalization) + + +class FitPositionsSourceSolved(SolvedCentre, FitPositionsSource): + """ + ``FitPositionsSource`` with the source-plane centre solved for analytically, rather than read from a free + ``centre`` model parameter, following Lombardi 2024 (arXiv:2406.15280) §5.1. + + The lens equation is Taylor-expanded around each observed image position `θ̂ᵢ`, giving a source-plane + position that is linear in the per-image back-traced position `β̂ᵢ` (`SolvedCentre._beta_hat`, computed the + same way as `FitPositionsSource.model_data`) and the per-image precision `Wᵢ` (`weighting = "jacobian"` by + default — the tensor weighting `Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹`; set `weighting = "magnification"` for the scalar + `Wᵢ = (µᵢ²/σᵢ²) I₂` weighting used by `FitPositionsSource`, for Lenstool-style comparisons). This makes the + source-plane centre solvable in closed form (`SolvedCentre.source_plane_coordinate`, `β*`), and the + likelihood analytically marginalizes over it (flat prior): + + `log_likelihood = -0.5*(χ² + noise_norm) - 0.5*log(det(Σᵢ Wᵢ) / (2π)²)` + + where `χ² = Σᵢ (β̂ᵢ−β*)ᵀ Wᵢ (β̂ᵢ−β*)` (`chi_squared_map` / `chi_squared`) and + `noise_norm = Σᵢ log((2π)²/det Wᵢ)` (`noise_normalization`), each a separately-testable property, alongside + the marginalization term itself (`marginalization_term`). + + 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 = "FitPositionsSource" + + @property + def residual_vectors(self) -> np.ndarray: + """ + The (n_positions, 2) array of vector residuals `β̂ᵢ − β*`, i.e. the back-traced source-plane positions + minus the solved source-plane centre. + """ + beta_hat = self._beta_hat.array + beta_star_y, beta_star_x = self.source_plane_coordinate + beta_star = self._xp.array([beta_star_y, beta_star_x]) + return beta_hat - beta_star + + @property + def chi_squared_map(self) -> aa.ArrayIrregular: + """ + The per-image quadratic form `(β̂ᵢ−β*)ᵀ Wᵢ (β̂ᵢ−β*)`. + """ + w11, w12, w21, w22 = precision_tensor_components_from(self, self.weighting) + + delta = self.residual_vectors + dy = delta[:, 0] + dx = delta[:, 1] + + terms = dy * (w11 * dy + w12 * dx) + dx * (w21 * dy + w22 * dx) + + return aa.ArrayIrregular(values=terms) + + @property + def chi_squared(self) -> float: + """ + `χ² = Σᵢ (β̂ᵢ−β*)ᵀ Wᵢ (β̂ᵢ−β*)`. + """ + return self._xp.sum(self.chi_squared_map.array) + + @property + def noise_normalization(self) -> float: + """ + `noise_norm = Σᵢ log((2π)² / det Wᵢ)`. + """ + w11, w12, w21, w22 = precision_tensor_components_from(self, self.weighting) + det_w = w11 * w22 - w12 * w21 + return self._xp.sum(self._xp.log((2.0 * np.pi) ** 2.0 / det_w)) + + @property + def marginalization_term(self) -> float: + """ + The analytic-marginalization contribution to the log likelihood from integrating out the (flat-prior) + source-plane centre: `-0.5 * log(det(Σᵢ Wᵢ) / (2π)²)` — the exact 2-D Gaussian integral over the + centre contributes `(2π)^{d/2} / sqrt(det)` with `d = 2`. + """ + w11, w12, w21, w22 = precision_tensor_components_from(self, self.weighting) + sum_w11 = self._xp.sum(w11) + sum_w12 = self._xp.sum(w12) + sum_w21 = self._xp.sum(w21) + sum_w22 = self._xp.sum(w22) + det_sum_w = sum_w11 * sum_w22 - sum_w12 * sum_w21 + return -0.5 * self._xp.log(det_sum_w / (2.0 * np.pi) ** 2.0) + + @property + def log_likelihood(self) -> float: + """ + `log_likelihood = -0.5*(χ² + noise_norm) - 0.5*log(det(Σᵢ Wᵢ) / (2π)²)`. + """ + return ( + -0.5 * (self.chi_squared + self.noise_normalization) + + self.marginalization_term + ) diff --git a/autolens/point/fit/solved.py b/autolens/point/fit/solved.py new file mode 100644 index 000000000..7636edc1f --- /dev/null +++ b/autolens/point/fit/solved.py @@ -0,0 +1,240 @@ +""" +Shared machinery for the analytically-solved point-source fit variants. + +Rather than sampling the source-plane centre of a point source as a free model parameter +(a `centre` on `ag.ps.Point` / `ag.ps.PointFlux`), the `*Solved` fit classes throughout +`autolens.point.fit` solve for it analytically given the current tracer, using a +zero-parameter `ag.ps.PointSolved` profile. This module provides the shared linear algebra +and the `SolvedCentre` mixin that every solved-centre fit class uses. + +**Source-plane solved likelihood** (`FitPositionsSourceSolved`, `positions/source/separations.py`) +follows Lombardi 2024 (arXiv:2406.15280) §5.1: the lens equation is Taylor-expanded around +each observed image position, giving a source-plane centre that is linear in the data and can +be solved in closed form, with the source-plane likelihood analytically marginalized over +that centre. + +**Solved image-plane variants** (`FitPositionsImagePairAllSolved`, `FitPositionsImagePairRepeatSolved`) +reuse the same analytically-solved centre to drive the existing `PointSolver` forward solve, +but the image-plane pairing chi-squareds themselves are not part of the paper. This is an +extension in the spirit of glafic's source-position optimization, not a result from the paper, +and should not be attributed to it. + +Convention +---------- +`A ≡ ∂β/∂θ` is the standard lensing Jacobian of the lens equation (mapping image-plane +displacements to source-plane displacements), so `A⁻¹` is the magnification tensor and the +scalar magnification is `µ = 1/det(A)`. For an observed image position `θ̂ᵢ` with scalar +position noise `σᵢ` (image-plane precision `Θᵢ = σᵢ⁻² I₂`), the per-image precision on the +back-traced source-plane position is: + + `Wᵢ = Aᵢ⁻ᵀ Θᵢ Aᵢ⁻¹ = σᵢ⁻² Aᵢ⁻ᵀAᵢ⁻¹` + +with `Aᵢ` evaluated at the observed position `θ̂ᵢ` (not the model/solved position). The +eigenvalues of `Wᵢ` are `λ²/σᵢ²` with `λ` the *linear* stretch along each eigendirection, so +the tensor is exact in every regime: isotropically each stretch is `√µ` (giving `µᵢ/σᵢ²`), +while near a critical curve the tangential stretch is `≈ µ` (giving `µᵢ²/σᵢ²` along the arc). + +The `weighting = "magnification"` class-attribute option instead uses the scalar weighting +`Wᵢ = (µᵢ²/σᵢ²) I₂` already used by `FitPositionsSource.chi_squared_map` — the traditional +Lenstool-style approximation, which coincides with the tensor only in the near-critical +tangential limit (NOT isotropically) — retained for comparisons with that convention. + +The solved source-plane centre is then the precision-weighted mean of the per-image +back-traced positions `β̂ᵢ`: + + `β* = (Σᵢ Wᵢ)⁻¹ Σᵢ Wᵢ β̂ᵢ` +""" +import numpy as np +from typing import Tuple + +import autoarray as aa +import autogalaxy as ag + +from autolens import exc + + +def _as_array(x): + """ + Returns the underlying array of `x` if `x` is one of this project's array-structure wrappers (anything + exposing `.array`, e.g. `aa.ArrayIrregular` / `aa.Grid2DIrregular`), otherwise returns `x` unchanged. Some + fit classes are exercised in tests with plain `numpy.ndarray` `data` / `noise_map` inputs (never wrapped), + which this module must tolerate. + """ + return x.array if hasattr(x, "array") else x + + +def _lens_calc_for(fit) -> ag.LensCalc: + """ + Returns the `ag.LensCalc` object used to compute the Hessian / Jacobian / magnification of `fit.tracer` at + the plane containing the point source paired to `fit`, accounting for multi-plane ray-tracing. + + Mirrors `AbstractFitPoint.magnifications_at_positions`. + """ + use_multi_plane = len(fit.tracer.planes) > 2 + plane_j = ( + fit.tracer.extract_plane_index_of_profile(profile_name=fit.name) + if use_multi_plane + else -1 + ) + return ag.LensCalc.from_tracer( + tracer=fit.tracer, + use_multi_plane=use_multi_plane, + plane_j=plane_j, + ) + + +def precision_tensor_components_from(fit, weighting: str) -> Tuple: + """ + Returns the (w11, w12, w21, w22) components of the per-image 2x2 precision matrix `Wᵢ`, evaluated at the + fit's observed image-plane positions (`fit.positions`), for every image. + + Each component is an array of shape `(n_positions,)`. The matrix is symmetric by construction, so + `w12 == w21`. + + Parameters + ---------- + fit + The point-source fit object (any `AbstractFitPositions` subclass) providing `positions`, `tracer`, + `name`, `noise_map` and `_xp`. + weighting + `"jacobian"` — the tensor weighting `Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹`, with `Aᵢ` the lensing Jacobian at the observed + position (see module docstring). + `"magnification"` — the scalar isotropic weighting `Wᵢ = (µᵢ²/σᵢ²) I₂`, matching + `FitPositionsSource.chi_squared_map`. + """ + xp = fit._xp + precision_scalar = _as_array(fit.noise_map) ** -2.0 # Θ = σ⁻² I + + if weighting == "magnification": + mu = _as_array(fit.magnifications_at_positions) + w = (mu**2.0) * precision_scalar + zero = xp.zeros_like(w) + return w, zero, zero, w + + if weighting != "jacobian": + raise exc.PointExtractionException( + f"Unsupported weighting '{weighting}' for the analytically-solved source-plane centre. " + f"Valid options are 'jacobian' (tensor weighting, the default) or 'magnification' " + f"(scalar isotropic weighting)." + ) + + lens_calc = _lens_calc_for(fit) + (a_xx, a_xy), (a_yx, a_yy) = lens_calc.jacobian_from(grid=fit.positions, xp=xp) + + # `jacobian_from` returns its 2x2 list in (x, y) row/col order (`a11`/`a22` there are the xx/yy + # components — see its docstring: `A = [[1-hessian_xx, -hessian_xy], [-hessian_yx, 1-hessian_yy]]`), + # whereas every position/grid array in this module is (y, x) (`positions[:, 0]` is y). Re-index into + # (y, x) row/col order before using these as a 2x2 matrix alongside (y, x)-ordered residual vectors. + a11, a12, a21, a22 = a_yy, a_yx, a_xy, a_xx + + det_a = a11 * a22 - a12 * a21 + + ainv11 = a22 / det_a + ainv12 = -a12 / det_a + ainv21 = -a21 / det_a + ainv22 = a11 / det_a + + # M = Aᵢ⁻ᵀAᵢ⁻¹, symmetric positive-(semi)definite by construction. + m11 = ainv11 * ainv11 + ainv21 * ainv21 + m12 = ainv11 * ainv12 + ainv21 * ainv22 + m22 = ainv12 * ainv12 + ainv22 * ainv22 + + w11 = m11 * precision_scalar + w12 = m12 * precision_scalar + w22 = m22 * precision_scalar + + return w11, w12, w12, w22 + + +class SolvedCentre: + """ + Mixin overriding `AbstractFitPoint.source_plane_coordinate` (the single funnel every position-based fit + class reads the source-plane centre from, `fit/abstract.py:143-153`) to return the analytically-solved + centre `β*`, rather than reading a `centre` model parameter. + + Must be listed before the concrete fit class in the MRO, e.g. `class FooSolved(SolvedCentre, Foo)`, so this + override takes precedence over `AbstractFitPoint.source_plane_coordinate`. + + Requires the paired profile to have **no** `centre` attribute (e.g. `ag.ps.PointSolved`): if a + `centre`-bearing profile (`ag.ps.Point` / `ag.ps.PointFlux`) is used with a `SolvedCentre` fit, its centre + priors would be sampled by the non-linear search but silently ignored (the analytic solve does not read + them) — a 2-parameter waste that is instead raised loudly. + """ + + weighting = "jacobian" + + #: Name of the corresponding free-centre fit class, used in the error message when a `centre`-bearing + #: profile is paired to this fit. Overridden per concrete `*Solved` class. + _non_solved_alternative_name = "the corresponding non-Solved fit class" + + @property + def _beta_hat(self) -> aa.Grid2DIrregular: + """ + The observed image-plane positions (`self.positions`) ray-traced to the source-plane, `β̂ᵢ`. Computed + independently of `self.model_data` (which, for the image-plane pairing fit classes this mixin is also + used with, means the forward-solved model image positions instead), via the same path as + `FitPositionsSource.model_data`. + """ + positions = self.positions + if not hasattr(positions, "grid_2d_via_deflection_grid_from"): + # Some fit classes are exercised in tests with a plain `numpy.ndarray` `data` + # (never wrapped): promote to a `Grid2DIrregular` so the deflection / ray-tracing + # calls below (which expect `.array`) have something to call it on. + positions = aa.Grid2DIrregular(values=_as_array(positions), xp=self._xp) + + if len(self.tracer.planes) <= 2: + deflections = self.tracer.deflections_yx_2d_from( + grid=positions, xp=self._xp + ) + else: + deflections = self.tracer.deflections_between_planes_from( + grid=positions, xp=self._xp, plane_i=0, plane_j=self.plane_index + ) + + return positions.grid_2d_via_deflection_grid_from( + deflection_grid=deflections, xp=self._xp + ) + + @property + def source_plane_coordinate(self) -> Tuple[float, float]: + """ + Returns the analytically-solved source-plane centre: + + `β* = (Σᵢ Wᵢ)⁻¹ Σᵢ Wᵢ β̂ᵢ` + + with `Wᵢ` given by `precision_tensor_components_from` (using `self.weighting`) and `β̂ᵢ` given by + `self._beta_hat`. + """ + if hasattr(self.profile, "centre"): + raise exc.PointExtractionException( + f"The point-source profile paired to dataset '{self.name}' " + f"({self.profile.__class__.__name__}) has a `centre` attribute, so its free-centre priors " + f"would be sampled by the non-linear search but silently ignored by " + f"{self.__class__.__name__}, which solves for the source-plane centre analytically. Use a " + f"parameter-free profile (e.g. `ag.ps.PointSolved`) with {self.__class__.__name__}, or use " + f"{self._non_solved_alternative_name} with a `centre`-bearing profile such as `ag.ps.Point` " + f"/ `ag.ps.PointFlux`." + ) + + xp = self._xp + + w11, w12, w21, w22 = precision_tensor_components_from(self, self.weighting) + + beta_hat = self._beta_hat.array + beta_hat_y = beta_hat[:, 0] + beta_hat_x = beta_hat[:, 1] + + sum_w11 = xp.sum(w11) + sum_w12 = xp.sum(w12) + sum_w21 = xp.sum(w21) + sum_w22 = xp.sum(w22) + + rhs_y = xp.sum(w11 * beta_hat_y + w12 * beta_hat_x) + rhs_x = xp.sum(w21 * beta_hat_y + w22 * beta_hat_x) + + det_sum_w = sum_w11 * sum_w22 - sum_w12 * sum_w21 + + beta_star_y = (sum_w22 * rhs_y - sum_w12 * rhs_x) / det_sum_w + beta_star_x = (-sum_w21 * rhs_y + sum_w11 * rhs_x) / det_sum_w + + return beta_star_y, beta_star_x diff --git a/autolens/point/fit/times_delays.py b/autolens/point/fit/times_delays.py index 3a43ca6a0..305f5185a 100644 --- a/autolens/point/fit/times_delays.py +++ b/autolens/point/fit/times_delays.py @@ -121,3 +121,85 @@ def chi_squared(self) -> float: return ag.util.fit.chi_squared_from( chi_squared_map=self.chi_squared_map.array, ) + + +class FitTimeDelaysSolved(FitTimeDelays): + """ + Fits the time delays of a point source dataset with the reference time solved for analytically, following + Lombardi 2024 (arXiv:2406.15280) §6.1, rather than by subtracting the shortest delay from both the data and + the model (`FitTimeDelays.residual_map`, left unchanged for continuity). + + With model per-image delays `Tᵢ` (`tracer.time_delays_from`, `model_data`, inherited unchanged) and per-image + precision `τᵢ = σᵢ⁻²` (`tau`): + + `T* = Σᵢ τᵢ(t̂ᵢ − Tᵢ) / Σᵢ τᵢ` (`solved_reference_time`) + + with residuals `t̂ᵢ − (Tᵢ + T*)` (`residual_map`), a standard chi-squared and noise normalization, and the + likelihood analytically marginalized over `T*` (flat prior): + + `log_likelihood = -0.5*(χ² + noise_norm) - 0.5*log((Σᵢ τᵢ)/(2π))` + + Does not depend on which profile is paired to the dataset (time delays are computed from the tracer alone), + so works with `ag.ps.Point`, `ag.ps.PointFlux` or `ag.ps.PointSolved` interchangeably. + """ + + @property + def tau(self) -> np.ndarray: + """ + `τᵢ = σᵢ⁻²` — the per-image time-delay precision. + """ + return self.noise_map.array**-2.0 + + @property + def tau_sum(self) -> float: + """ + `Σᵢ τᵢ` — the precision of the solved reference time `T*`, and the marginalization normalization. + """ + return self._xp.sum(self.tau) + + @property + def solved_reference_time(self) -> float: + """ + `T* = Σᵢ τᵢ(t̂ᵢ − Tᵢ) / Σᵢ τᵢ`. + """ + t_hat = self.data.array + model_delays = self.model_data.array + return self._xp.sum(self.tau * (t_hat - model_delays)) / self.tau_sum + + @property + def residual_map(self) -> aa.ArrayIrregular: + """ + Returns the difference between the observed time delays and the model time delays offset by the solved + reference time: `t̂ᵢ − (Tᵢ + T*)`. + """ + residual_map = self.data.array - ( + self.model_data.array + self.solved_reference_time + ) + 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. + """ + 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) + reference time: `-0.5 * log((Σᵢ τᵢ)/(2π))`. + """ + return -0.5 * self._xp.log(self.tau_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 + ) diff --git a/autolens/point/model/analysis.py b/autolens/point/model/analysis.py index 4a7d4a3b2..cad32d387 100644 --- a/autolens/point/model/analysis.py +++ b/autolens/point/model/analysis.py @@ -25,6 +25,8 @@ from autolens.analysis.exceptions import raise_fit_exception from autolens.point.fit.positions.image.pair_repeat import FitPositionsImagePairRepeat from autolens.point.fit.dataset import FitPointDataset +from autolens.point.fit.fluxes import FitFluxes +from autolens.point.fit.times_delays import FitTimeDelays from autolens.point.dataset import PointDataset from autolens.point.model.result import ResultPoint from autolens.point.model.visualizer import VisualizerPoint @@ -40,6 +42,8 @@ def __init__( dataset: PointDataset, solver: PointSolver, fit_positions_cls=FitPositionsImagePairRepeat, + fit_flux_cls=FitFluxes, + fit_time_delays_cls=FitTimeDelays, image=None, cosmology: ag.cosmo.LensingCosmology = None, title_prefix: str = None, @@ -73,6 +77,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. cosmology The Cosmology assumed for this analysis. title_prefix @@ -90,6 +100,8 @@ def __init__( self.solver = solver self.fit_positions_cls = fit_positions_cls + self.fit_flux_cls = fit_flux_cls + self.fit_time_delays_cls = fit_time_delays_cls self.title_prefix = title_prefix def log_likelihood_function(self, instance): @@ -175,6 +187,8 @@ def fit_from( tracer=tracer, solver=self.solver, fit_positions_cls=self.fit_positions_cls, + fit_flux_cls=self.fit_flux_cls, + fit_time_delays_cls=self.fit_time_delays_cls, xp=self._xp, ) @@ -189,14 +203,32 @@ def _register_fit_point_pytrees() -> None: """ from autoarray.abstract_ndarray import register_instance_pytree from autolens.lens.tracer import Tracer - from autolens.point.fit.positions.image.pair_all import FitPositionsImagePairAll - from autolens.point.fit.positions.image.pair_repeat import FitPositionsImagePairRepeat + from autolens.point.fit.positions.image.pair_all import ( + FitPositionsImagePairAll, + FitPositionsImagePairAllSolved, + ) + from autolens.point.fit.positions.image.pair_repeat import ( + FitPositionsImagePairRepeat, + FitPositionsImagePairRepeatSolved, + ) from autolens.point.fit.positions.image.pair import FitPositionsImagePair + from autolens.point.fit.positions.source.separations import ( + FitPositionsSource, + FitPositionsSourceSolved, + ) + from autolens.point.fit.fluxes import FitFluxesSolved + from autolens.point.fit.times_delays import FitTimeDelaysSolved import autogalaxy as ag register_instance_pytree( FitPointDataset, - no_flatten=("dataset", "solver", "fit_positions_cls"), + no_flatten=( + "dataset", + "solver", + "fit_positions_cls", + "fit_flux_cls", + "fit_time_delays_cls", + ), ) register_instance_pytree(Tracer, no_flatten=("cosmology",)) # fit-point-pytree: observed data/noise are per-analysis constants; solver/name/use_jax are non-JAX @@ -214,7 +246,35 @@ def _register_fit_point_pytrees() -> None: FitPositionsImagePair, no_flatten=("solver", "name", "use_jax", "_data", "_noise_map"), ) - # fit-point-pytree: ag.ps.Point / PointFlux are handled by + # fit-point-pytree: no solver is used (source-plane / analytic fits), so only + # name/use_jax/observed data+noise are non-JAX. + register_instance_pytree( + FitPositionsSource, + no_flatten=("name", "use_jax", "_data", "_noise_map"), + ) + register_instance_pytree( + FitPositionsSourceSolved, + no_flatten=("name", "use_jax", "_data", "_noise_map"), + ) + # fit-point-pytree: solved image-plane variants still forward-solve via `solver`. + register_instance_pytree( + FitPositionsImagePairAllSolved, + no_flatten=("solver", "name", "use_jax", "_data", "_noise_map"), + ) + register_instance_pytree( + FitPositionsImagePairRepeatSolved, + no_flatten=("solver", "name", "use_jax", "_data", "_noise_map"), + ) + # fit-point-pytree: flux/time-delay fits carry no solver, only observed positions. + register_instance_pytree( + FitFluxesSolved, + no_flatten=("name", "use_jax", "_data", "_noise_map", "positions"), + ) + register_instance_pytree( + FitTimeDelaysSolved, + no_flatten=("name", "use_jax", "_data", "_noise_map", "positions"), + ) + # fit-point-pytree: ag.ps.Point / PointFlux / PointSolved are handled by # autofit.jax.pytrees.register_model before jit is called; skip here. def save_attributes(self, paths: af.DirectoryPaths): diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py index dc58886cd..7b155b3dc 100644 --- a/autolens/point/solver/shape_solver.py +++ b/autolens/point/solver/shape_solver.py @@ -251,7 +251,7 @@ def _plane_grid( grid=grid, plane_i=0, plane_j=plane_index, xp=xp ) # noinspection PyTypeChecker - return grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections) + return grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections, xp=xp) def solve_triangles( self, diff --git a/test_autolens/point/fit/positions/image/test_pair_all.py b/test_autolens/point/fit/positions/image/test_pair_all.py index 4481abdd9..bd4601169 100644 --- a/test_autolens/point/fit/positions/image/test_pair_all.py +++ b/test_autolens/point/fit/positions/image/test_pair_all.py @@ -103,3 +103,72 @@ def test__fit_positions_image_pair_all__model_has_duplicate_position__duplicate_ [-1.14237812, -0.87193683], ) assert fit.chi_squared == -2.0 * -4.211539531047171 + + +def test__fit_positions_image_pair_all__penalty_regression__n_permutations_and_monotonic_likelihood( + data, noise_map +): + # Test 5: n_permutations must equal n_finite_model_positions ** n_observed, and adding a + # spurious (but finite) extra model position must strictly lower the likelihood -- an + # extra candidate explanation dilutes each permutation's probability without improving + # the fit to any observed position. + two_model_positions = al.Grid2DIrregular( + [(-1.0749, -1.1), (1.19117, 1.175)] + ) + fit_two = al.FitPositionsImagePairAll( + name="point_0", + data=data, + noise_map=noise_map, + tracer=tracer, + solver=al.mock.MockPointSolver(two_model_positions), + ) + + n_non_nan = np.count_nonzero(np.isfinite(two_model_positions.array).any(axis=1)) + assert n_non_nan == 2 + n_permutations = n_non_nan ** len(data) + assert n_permutations == 2 ** len(data) + + three_model_positions = al.Grid2DIrregular( + [(-1.0749, -1.1), (1.19117, 1.175), (5.0, 5.0)] + ) + fit_three = al.FitPositionsImagePairAll( + name="point_0", + data=data, + noise_map=noise_map, + tracer=tracer, + solver=al.mock.MockPointSolver(three_model_positions), + ) + + n_non_nan_three = np.count_nonzero( + np.isfinite(three_model_positions.array).any(axis=1) + ) + assert n_non_nan_three == 3 + assert n_non_nan_three ** len(data) == 3 ** len(data) + + # log_likelihood = -0.5 * chi_squared, so a strictly lower likelihood is a strictly + # higher chi_squared. + assert fit_three.chi_squared > fit_two.chi_squared + + +def test__fit_positions_image_pair_all_solved__source_plane_coordinate_feeds_solver( + data, noise_map +): + galaxy_solved = al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved()) + tracer_solved = al.Tracer(galaxies=[al.Galaxy(redshift=0.5), galaxy_solved]) + + model_positions = al.Grid2DIrregular([(-1.0749, -1.1), (1.19117, 1.175)]) + + fit = al.FitPositionsImagePairAllSolved( + name="point_0", + data=data, + noise_map=noise_map, + tracer=tracer_solved, + solver=al.mock.MockPointSolver(model_positions), + ) + + assert np.isfinite(fit.source_plane_coordinate[0]) + assert np.isfinite(fit.source_plane_coordinate[1]) + # The pairing chi-squared itself is untouched (solver is mocked, so model positions are + # fixed regardless of the solved centre): matches the plain FitPositionsImagePairAll + # value from the fixture-equivalent test above. + assert fit.chi_squared == -2.0 * -4.40375330990644 diff --git a/test_autolens/point/fit/positions/image/test_pair_repeat.py b/test_autolens/point/fit/positions/image/test_pair_repeat.py index 91409c505..aced82f2e 100644 --- a/test_autolens/point/fit/positions/image/test_pair_repeat.py +++ b/test_autolens/point/fit/positions/image/test_pair_repeat.py @@ -168,3 +168,30 @@ def test__under_prediction__no_model_images_hits_finite_floor(): assert fit.residual_map.in_list == [1.0e4, 1.0e4] assert np.isfinite(float(fit.log_likelihood)) + + +def test__fit_positions_image_pair_repeat_solved__source_plane_coordinate_feeds_solver(): + galaxy_solved = al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved()) + tracer_solved = al.Tracer(galaxies=[al.Galaxy(redshift=0.5), galaxy_solved]) + + data = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) + noise_map = al.ArrayIrregular([0.5, 1.0]) + model_data = al.Grid2DIrregular([(3.0, 1.0), (2.0, 3.0)]) + + solver = al.m.MockPointSolver(model_positions=model_data) + + fit = al.FitPositionsImagePairRepeatSolved( + name="point_0", + data=data, + noise_map=noise_map, + tracer=tracer_solved, + solver=solver, + ) + + assert np.isfinite(fit.source_plane_coordinate[0]) + assert np.isfinite(fit.source_plane_coordinate[1]) + # The pairing chi-squared itself is untouched (solver is mocked, so model positions are + # fixed regardless of the solved centre): matches the plain FitPositionsImagePairRepeat + # value from the equivalent free-centre test above. + assert fit.chi_squared == pytest.approx(42.0, 1.0e-4) + assert fit.log_likelihood == pytest.approx(-22.14472, 1.0e-4) diff --git a/test_autolens/point/fit/positions/source/test_separations.py b/test_autolens/point/fit/positions/source/test_separations.py index 1ac96058f..4f61283be 100644 --- a/test_autolens/point/fit/positions/source/test_separations.py +++ b/test_autolens/point/fit/positions/source/test_separations.py @@ -1,3 +1,4 @@ +import numpy as np import pytest import autolens as al @@ -77,3 +78,31 @@ def test__fit_positions_source__multi_plane_tracer__model_data_traces_to_correct ) assert (fit_1.model_data == traced_grids[2]).all() + + +def test__fit_positions_source_solved__source_plane_centre_matches_no_free_centre_prior(): + point_source = al.ps.PointSolved() + galaxy_point_source = al.Galaxy(redshift=1.0, point_0=point_source) + galaxy_mass = al.Galaxy( + redshift=0.5, mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.1) + ) + tracer = al.Tracer(galaxies=[galaxy_mass, galaxy_point_source]) + + positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0)]) + noise_map = al.ArrayIrregular([0.5, 1.0]) + + fit = al.FitPositionsSourceSolved( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + + beta_star = fit.source_plane_coordinate + + assert np.isfinite(beta_star[0]) + assert np.isfinite(beta_star[1]) + assert np.isfinite(fit.chi_squared) + assert np.isfinite(fit.noise_normalization) + assert np.isfinite(fit.marginalization_term) + assert np.isfinite(float(fit.log_likelihood)) + + # Default weighting is the tensor ("jacobian") weighting. + assert fit.weighting == "jacobian" diff --git a/test_autolens/point/fit/test_fit_dataset.py b/test_autolens/point/fit/test_fit_dataset.py index e0bab3106..eec6db580 100644 --- a/test_autolens/point/fit/test_fit_dataset.py +++ b/test_autolens/point/fit/test_fit_dataset.py @@ -1,3 +1,4 @@ +import numpy as np import pytest import autolens as al @@ -115,3 +116,63 @@ def test__fit_dataset__positions_and_flux__both_log_likelihoods_correct_and_sum( assert fit.positions.log_likelihood == pytest.approx(-22.14472, 1.0e-4) assert fit.flux.log_likelihood == pytest.approx(-2.9920449, 1.0e-4) assert fit.log_likelihood == fit.positions.log_likelihood + fit.flux.log_likelihood + + +def test__fit_dataset__fit_flux_cls_and_fit_time_delays_cls_hooks_are_forwarded_and_default_unchanged( + point_source_tracer, positions_and_noise, mock_solver +): + positions, noise_map = positions_and_noise + dataset = al.PointDataset( + name="point_0", positions=positions, positions_noise_map=noise_map + ) + + default_fit = al.FitPointDataset( + dataset=dataset, tracer=point_source_tracer, solver=mock_solver + ) + + # Defaults preserve current behaviour exactly. + assert default_fit.fit_flux_cls is al.FitFluxes + assert default_fit.fit_time_delays_cls is al.FitTimeDelays + + explicit_fit = al.FitPointDataset( + dataset=dataset, + tracer=point_source_tracer, + solver=mock_solver, + fit_flux_cls=al.FitFluxesSolved, + fit_time_delays_cls=al.FitTimeDelaysSolved, + ) + + assert explicit_fit.fit_flux_cls is al.FitFluxesSolved + assert explicit_fit.fit_time_delays_cls is al.FitTimeDelaysSolved + + +def test__fit_dataset__fit_flux_cls_hook_is_actually_used_to_construct_the_flux_fit(): + point_source = al.ps.PointSolved() + galaxy_point_source = al.Galaxy(redshift=1.0, point_0=point_source) + tracer = al.Tracer(galaxies=[al.Galaxy(redshift=0.5), galaxy_point_source]) + + positions = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) + noise_map = al.ArrayIrregular([0.5, 1.0]) + model_positions = al.Grid2DIrregular([(3.0, 1.0), (2.0, 3.0)]) + fluxes = al.ArrayIrregular([1.0, 2.0]) + flux_noise_map = al.ArrayIrregular([3.0, 1.0]) + + solver = al.m.MockPointSolver(model_positions=model_positions) + + dataset = al.PointDataset( + name="point_0", + positions=positions, + positions_noise_map=noise_map, + fluxes=fluxes, + fluxes_noise_map=flux_noise_map, + ) + + fit = al.FitPointDataset( + dataset=dataset, + tracer=tracer, + solver=solver, + fit_flux_cls=al.FitFluxesSolved, + ) + + assert isinstance(fit.flux, al.FitFluxesSolved) + assert np.isfinite(fit.flux.log_likelihood) diff --git a/test_autolens/point/fit/test_fluxes.py b/test_autolens/point/fit/test_fluxes.py index 5141d1fef..5f6efced8 100644 --- a/test_autolens/point/fit/test_fluxes.py +++ b/test_autolens/point/fit/test_fluxes.py @@ -1,3 +1,5 @@ +import numpy as np +from scipy.optimize import minimize_scalar import pytest import autolens as al @@ -50,3 +52,71 @@ def test__fit_fluxes__model_flux_magnified_correctly_with_real_isothermal_tracer assert fit.model_fluxes.in_list[1] == pytest.approx(2.5, 1.0e-4) assert fit.log_likelihood == pytest.approx(-3.11702, 1.0e-4) + + +def test__fit_fluxes_solved__solved_flux_equals_brute_force_scan(): + lens = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0)) + galaxy_point_source = al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved()) + tracer = al.Tracer(galaxies=[lens, galaxy_point_source]) + + positions = al.Grid2DIrregular([(0.0, 1.5), (0.0, -1.3), (1.4, 0.05)]) + data = al.ArrayIrregular([5.0, 3.2, 4.1]) + noise_map = al.ArrayIrregular([0.3, 0.4, 0.2]) + + fit = al.FitFluxesSolved( + name="point_0", + data=data, + noise_map=noise_map, + positions=positions, + tracer=tracer, + ) + + mu = fit.magnifications_at_positions.array + + def chi_squared(flux): + return np.sum((data.array - mu * flux) ** 2 / noise_map.array**2) + + brute_force = minimize_scalar(chi_squared) + + assert fit.solved_flux == pytest.approx(brute_force.x, 1.0e-6) + assert fit.model_data.in_list == pytest.approx((mu * fit.solved_flux).tolist(), 1.0e-8) + + +def test__fit_fluxes_solved__profile_with_flux_attribute__raises_naming_alternative(): + galaxy_point_source = al.Galaxy(redshift=1.0, point_0=al.ps.PointFlux(flux=2.0)) + tracer = al.Tracer(galaxies=[al.Galaxy(redshift=0.5), galaxy_point_source]) + + data = al.ArrayIrregular([1.0, 2.0]) + noise_map = al.ArrayIrregular([3.0, 1.0]) + positions = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) + + with pytest.raises(al.exc.PointExtractionException, match="FitFluxes"): + al.FitFluxesSolved( + name="point_0", + data=data, + noise_map=noise_map, + positions=positions, + tracer=tracer, + ) + + +def test__fit_fluxes_solved__works_with_plain_point_profile_no_flux_attribute(): + galaxy_point_source = al.Galaxy( + redshift=1.0, point_0=al.ps.Point(centre=(0.1, 0.1)) + ) + tracer = al.Tracer(galaxies=[al.Galaxy(redshift=0.5), galaxy_point_source]) + + data = al.ArrayIrregular([1.0, 2.0]) + noise_map = al.ArrayIrregular([3.0, 1.0]) + positions = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) + + fit = al.FitFluxesSolved( + name="point_0", + data=data, + noise_map=noise_map, + positions=positions, + tracer=tracer, + ) + + assert np.isfinite(fit.solved_flux) + assert np.isfinite(float(fit.log_likelihood)) diff --git a/test_autolens/point/fit/test_solved.py b/test_autolens/point/fit/test_solved.py new file mode 100644 index 000000000..2e1caa30d --- /dev/null +++ b/test_autolens/point/fit/test_solved.py @@ -0,0 +1,442 @@ +""" +Tests for the analytically-solved point-source fit variants (`autolens.point.fit.solved` and the concrete +`*Solved` fit classes it underpins). + +Numpy-only, per project convention (unit tests never import jax). +""" +import numpy as np +from scipy.optimize import minimize +import pytest + +import autolens as al +from autolens.point.fit.solved import precision_tensor_components_from + + +def _isothermal_sph_tracer(einstein_radius=1.0, profile=None): + profile = profile or al.ps.PointSolved() + lens = al.Galaxy( + redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=einstein_radius) + ) + source = al.Galaxy(redshift=1.0, point_0=profile) + return al.Tracer(galaxies=[lens, source]) + + +def _elliptical_isothermal_tracer(axis_ratio, angle, einstein_radius=1.0, profile=None): + profile = profile or al.ps.PointSolved() + mass = al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=al.convert.ell_comps_from(axis_ratio=axis_ratio, angle=angle), + einstein_radius=einstein_radius, + ) + lens = al.Galaxy(redshift=0.5, mass=mass) + source = al.Galaxy(redshift=1.0, point_0=profile) + return al.Tracer(galaxies=[lens, source]) + + +def _chi_squared_at(beta_hat, weighting_components, centre): + w11, w12, w21, w22 = weighting_components + dy = beta_hat[:, 0] - centre[0] + dx = beta_hat[:, 1] - centre[1] + return np.sum(dy * (w11 * dy + w12 * dx) + dx * (w21 * dy + w22 * dx)) + + +class TestSourcePlaneSolvedCentre: + def test__beta_star_equals_brute_force_minimizer_of_tensor_chi_squared(self): + # Test 1: beta* (the analytic closed-form solve) must equal a brute-force numerical + # minimization of the tensor-weighted source-plane chi-squared over the centre, on a + # fixed mock tracer. + tracer = _isothermal_sph_tracer(einstein_radius=1.0) + + positions = al.Grid2DIrregular( + [(0.0, 1.5), (0.0, -1.3), (1.4, 0.05), (-0.9, -1.1)] + ) + noise_map = al.ArrayIrregular([0.05, 0.05, 0.05, 0.05]) + + fit = al.FitPositionsSourceSolved( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + + beta_star = fit.source_plane_coordinate + + beta_hat = fit._beta_hat.array + weighting_components = precision_tensor_components_from(fit, "jacobian") + + result = minimize( + lambda centre: _chi_squared_at(beta_hat, weighting_components, centre), + x0=np.zeros(2), + method="Nelder-Mead", + options={"xatol": 1.0e-10, "fatol": 1.0e-12, "maxiter": 20000}, + ) + + assert beta_star[0] == pytest.approx(result.x[0], abs=1.0e-6) + assert beta_star[1] == pytest.approx(result.x[1], abs=1.0e-6) + + # And the analytic chi-squared at beta* is (at least as good as, in practice equal to) + # the brute-force minimum. + assert fit.chi_squared == pytest.approx(result.fun, abs=1.0e-6) + + def test__beta_star_scalar_weighting_also_matches_brute_force_minimizer(self): + tracer = _isothermal_sph_tracer(einstein_radius=1.2) + + positions = al.Grid2DIrregular([(0.0, 1.5), (0.0, -1.3), (1.4, 0.05)]) + noise_map = al.ArrayIrregular([0.05, 0.05, 0.05]) + + fit = al.FitPositionsSourceSolved( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + fit.weighting = "magnification" + + beta_star = fit.source_plane_coordinate + + beta_hat = fit._beta_hat.array + weighting_components = precision_tensor_components_from(fit, "magnification") + + result = minimize( + lambda centre: _chi_squared_at(beta_hat, weighting_components, centre), + x0=np.zeros(2), + method="Nelder-Mead", + options={"xatol": 1.0e-10, "fatol": 1.0e-12, "maxiter": 20000}, + ) + + assert beta_star[0] == pytest.approx(result.x[0], abs=1.0e-6) + assert beta_star[1] == pytest.approx(result.x[1], abs=1.0e-6) + + def test__solved_log_likelihood_matches_tensor_profiled_maximum_plus_marginalization( + self, + ): + # Test 2 (source-plane, S1): the solved-centre log likelihood equals the log + # likelihood of the *same* tensor-weighted chi-squared profiled (maximized) over a + # free centre, plus the analytic marginalization term -- the marginalization term is + # exactly the known correction between "point estimate at the MLE" and "analytically + # integrated over a flat prior". + tracer = _elliptical_isothermal_tracer(axis_ratio=0.7, angle=30.0) + + positions = al.Grid2DIrregular( + [(0.0, 1.4), (0.0, -1.2), (1.3, 0.1), (-0.8, -1.0)] + ) + noise_map = al.ArrayIrregular([0.03, 0.03, 0.03, 0.03]) + + fit = al.FitPositionsSourceSolved( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + + beta_hat = fit._beta_hat.array + weighting_components = precision_tensor_components_from(fit, "jacobian") + + result = minimize( + lambda centre: _chi_squared_at(beta_hat, weighting_components, centre), + x0=np.zeros(2), + method="Nelder-Mead", + options={"xatol": 1.0e-10, "fatol": 1.0e-12, "maxiter": 20000}, + ) + + profiled_max_log_likelihood = -0.5 * (result.fun + fit.noise_normalization) + + assert fit.log_likelihood == pytest.approx( + profiled_max_log_likelihood + fit.marginalization_term, abs=1.0e-4 + ) + + +class TestImagePlaneSolvedCentre: + def _perfect_data_setup(self, axis_ratio=0.7, angle=30.0, true_centre=(0.05, 0.03)): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=al.convert.ell_comps_from(axis_ratio=axis_ratio, angle=angle), + einstein_radius=1.0, + ), + ) + tracer_free = al.Tracer( + galaxies=[lens, al.Galaxy(redshift=1.0, point_0=al.ps.Point(centre=true_centre))] + ) + tracer_solved = al.Tracer( + galaxies=[lens, al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved())] + ) + + grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05) + solver = al.PointSolver.for_grid(grid=grid, pixel_scale_precision=0.01) + observed = solver.solve(tracer=tracer_free, source_plane_coordinate=true_centre) + + noise_map = al.ArrayIrregular([0.02] * len(observed)) + mock_solver = al.m.MockPointSolver(model_positions=observed) + + return tracer_free, tracer_solved, observed, noise_map, mock_solver, true_centre + + def test__pair_repeat_solved_log_likelihood_matches_free_centre_profiled_at_truth( + self, + ): + # Test 2 (I2): with noiseless data generated at a known true centre, the solved + # centre back-traces to (very close to) that true centre, so both the solved and a + # free-centre fit *at* the true centre achieve the same (here: zero) chi-squared and + # hence the same log likelihood -- the pairing chi-squared itself is untouched by the + # mixin, so there is no marginalization offset to account for. + ( + tracer_free, + tracer_solved, + observed, + noise_map, + mock_solver, + true_centre, + ) = self._perfect_data_setup() + + fit_solved = al.FitPositionsImagePairRepeatSolved( + name="point_0", + data=observed, + noise_map=noise_map, + tracer=tracer_solved, + solver=mock_solver, + ) + fit_free_at_truth = al.FitPositionsImagePairRepeat( + name="point_0", + data=observed, + noise_map=noise_map, + tracer=tracer_free, + solver=mock_solver, + ) + + assert fit_solved.source_plane_coordinate[0] == pytest.approx( + true_centre[0], abs=1.0e-3 + ) + assert fit_solved.source_plane_coordinate[1] == pytest.approx( + true_centre[1], abs=1.0e-3 + ) + assert fit_solved.chi_squared == pytest.approx(0.0, abs=1.0e-2) + assert fit_solved.log_likelihood == pytest.approx( + fit_free_at_truth.log_likelihood, abs=1.0e-6 + ) + + def test__pair_all_solved_log_likelihood_matches_free_centre_profiled_at_truth(self): + # Test 2 (I1): as above, for the all-to-all pairing scheme. + ( + tracer_free, + tracer_solved, + observed, + noise_map, + mock_solver, + true_centre, + ) = self._perfect_data_setup() + + fit_solved = al.FitPositionsImagePairAllSolved( + name="point_0", + data=observed, + noise_map=noise_map, + tracer=tracer_solved, + solver=mock_solver, + ) + fit_free_at_truth = al.FitPositionsImagePairAll( + name="point_0", + data=observed, + noise_map=noise_map, + tracer=tracer_free, + solver=mock_solver, + ) + + assert fit_solved.log_likelihood == pytest.approx( + fit_free_at_truth.log_likelihood, abs=1.0e-6 + ) + + +class TestTensorVsScalarWeighting: + def test__anisotropic_case__tensor_ordering_matches_image_plane__scalar_does_not(self): + # Test 3: near a critical curve of an elliptical lens, the local precision tensor W + # is strongly anisotropic (one eigen-direction is far more informative about the + # source-plane centre than the other). Construct two candidate source-plane centres + # A and B, both perturbations of the same back-traced position, chosen so that: + # + # - candidate A is displaced along the *low*-precision eigen-direction (a large + # source-plane displacement barely moves the image), + # - candidate B is displaced by a *smaller* amount, but along the *high*-precision + # eigen-direction (a small source-plane displacement moves the image more). + # + # The scalar (magnification-squared) weighting is isotropic, so it ranks candidates + # purely by source-plane Euclidean distance and prefers B (the smaller displacement). + # The tensor weighting accounts for the anisotropy and prefers A. Ground truth is the + # real (nonlinear) image-plane position each candidate centre ray-traces to: A must be + # the genuinely better candidate (smaller true image-plane residual), matching the + # tensor ordering and contradicting the scalar ordering. + from scipy.optimize import root + + tracer = _elliptical_isothermal_tracer(axis_ratio=0.6, angle=45.0) + + observed_theta = (0.0, 0.85) + positions = al.Grid2DIrregular([observed_theta]) + noise_map = al.ArrayIrregular([0.01]) + + fit = al.FitPositionsSourceSolved( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + beta_hat = fit._beta_hat.array[0] + + w11, w12, w21, w22 = precision_tensor_components_from(fit, "jacobian") + w = np.array([[w11[0], w12[0]], [w21[0], w22[0]]]) + eigvals, eigvecs = np.linalg.eigh(w) + low_vec = eigvecs[:, 0] + high_vec = eigvecs[:, 1] + + delta_a = 0.01 * low_vec + delta_b = 0.001 * high_vec + + def chi_squared(delta, weighting): + components = precision_tensor_components_from(fit, weighting) + a, b, c, d = (float(x[0]) for x in components) + dy, dx = delta + return dy * (a * dy + b * dx) + dx * (c * dy + d * dx) + + tensor_a = chi_squared(delta_a, "jacobian") + tensor_b = chi_squared(delta_b, "jacobian") + scalar_a = chi_squared(delta_a, "magnification") + scalar_b = chi_squared(delta_b, "magnification") + + # The two weightings disagree: + assert tensor_a < tensor_b # tensor prefers candidate A + assert scalar_a > scalar_b # scalar prefers candidate B + + # Ground truth: precisely ray-trace (via root-finding on the exact deflection field, + # not the coarse triangulated PointSolver) the image-plane position each candidate + # source centre corresponds to, near the observed image. + def beta_of_theta(theta): + grid = al.Grid2DIrregular([tuple(theta)]) + deflections = tracer.deflections_yx_2d_from(grid=grid) + beta = grid.grid_2d_via_deflection_grid_from(deflection_grid=deflections) + return beta.array[0] + + def theta_for_beta(target_beta): + solution = root( + lambda theta: beta_of_theta(theta) - target_beta, + x0=np.array(observed_theta), + method="hybr", + tol=1.0e-14, + ) + assert solution.success + return solution.x + + theta_a = theta_for_beta(beta_hat + delta_a) + theta_b = theta_for_beta(beta_hat + delta_b) + + image_distance_a = np.sqrt(np.sum((theta_a - np.array(observed_theta)) ** 2)) + image_distance_b = np.sqrt(np.sum((theta_b - np.array(observed_theta)) ** 2)) + + # The tensor ordering (A better than B) matches the true image-plane ordering; the + # scalar ordering (B better than A) does not. + assert image_distance_a < image_distance_b + + +class TestLoudFailures: + def test__solved_fit_given_centre_bearing_profile__raises_naming_alternative(self): + tracer = _isothermal_sph_tracer(profile=al.ps.Point(centre=(0.0, 0.0))) + positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0)]) + noise_map = al.ArrayIrregular([1.0, 1.0]) + + fit = al.FitPositionsSourceSolved( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + + with pytest.raises(al.exc.PointExtractionException, match="FitPositionsSource"): + fit.source_plane_coordinate + + def test__centre_requiring_fit_given_solved_profile__raises_naming_alternative(self): + tracer = _isothermal_sph_tracer(profile=al.ps.PointSolved()) + positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0)]) + noise_map = al.ArrayIrregular([1.0, 1.0]) + + fit = al.FitPositionsSource( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + + with pytest.raises(al.exc.PointExtractionException, match="PointSolved"): + fit.source_plane_coordinate + + def test__pair_repeat_solved_given_centre_bearing_profile__raises(self): + tracer = _isothermal_sph_tracer(profile=al.ps.Point(centre=(0.0, 0.0))) + positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0)]) + noise_map = al.ArrayIrregular([1.0, 1.0]) + solver = al.m.MockPointSolver(model_positions=positions) + + fit = al.FitPositionsImagePairRepeatSolved( + name="point_0", + data=positions, + noise_map=noise_map, + tracer=tracer, + solver=solver, + ) + + with pytest.raises( + al.exc.PointExtractionException, match="FitPositionsImagePairRepeat" + ): + fit.source_plane_coordinate + + def test__pair_all_solved_given_centre_bearing_profile__raises(self): + tracer = _isothermal_sph_tracer(profile=al.ps.Point(centre=(0.0, 0.0))) + positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0)]) + noise_map = al.ArrayIrregular([1.0, 1.0]) + solver = al.m.MockPointSolver(model_positions=positions) + + fit = al.FitPositionsImagePairAllSolved( + name="point_0", + data=positions, + noise_map=noise_map, + tracer=tracer, + solver=solver, + ) + + with pytest.raises( + al.exc.PointExtractionException, match="FitPositionsImagePairAll" + ): + fit.source_plane_coordinate + + def test__point_flux_given_solved_fit__raises(self): + # PointFlux also has a `centre`, so is just as invalid a pairing for a *Solved fit + # as plain Point (the check is on `centre`, not on the specific profile subclass). + tracer = _isothermal_sph_tracer( + profile=al.ps.PointFlux(centre=(0.0, 0.0), flux=1.0) + ) + positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0)]) + noise_map = al.ArrayIrregular([1.0, 1.0]) + + fit = al.FitPositionsSourceSolved( + name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None + ) + + with pytest.raises(al.exc.PointExtractionException): + fit.source_plane_coordinate + + def test__point_given_fit_fluxes__raises(self): + # `Point` has no `flux` attribute, so the (pre-existing) free-flux `FitFluxes` + # cannot fit it -- completing the matrix column for `FitFluxes`. + galaxy_point_source = al.Galaxy( + redshift=1.0, point_0=al.ps.Point(centre=(0.1, 0.1)) + ) + tracer = al.Tracer(galaxies=[al.Galaxy(redshift=0.5), galaxy_point_source]) + + data = al.ArrayIrregular([1.0, 2.0]) + noise_map = al.ArrayIrregular([3.0, 1.0]) + positions = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) + + with pytest.raises(al.exc.PointExtractionException): + al.FitFluxes( + name="point_0", + data=data, + noise_map=noise_map, + positions=positions, + tracer=tracer, + ) + + def test__point_solved_given_fit_fluxes__raises(self): + # `PointSolved` also has no `flux` attribute, so the free-flux `FitFluxes` cannot + # fit it either. + galaxy_point_source = al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved()) + tracer = al.Tracer(galaxies=[al.Galaxy(redshift=0.5), galaxy_point_source]) + + data = al.ArrayIrregular([1.0, 2.0]) + noise_map = al.ArrayIrregular([3.0, 1.0]) + positions = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) + + with pytest.raises(al.exc.PointExtractionException): + al.FitFluxes( + name="point_0", + data=data, + noise_map=noise_map, + positions=positions, + tracer=tracer, + ) diff --git a/test_autolens/point/fit/test_time_delays.py b/test_autolens/point/fit/test_time_delays.py index 0f5f468b3..62455027b 100644 --- a/test_autolens/point/fit/test_time_delays.py +++ b/test_autolens/point/fit/test_time_delays.py @@ -1,4 +1,5 @@ import numpy as np +from scipy.optimize import minimize_scalar import pytest import autolens as al @@ -54,3 +55,58 @@ def test__fit_time_delays__model_time_delays_correct_with_real_isothermal_tracer assert fit.model_time_delays.in_list[1] == pytest.approx(-573.994580905, 1.0e-4) assert fit.log_likelihood == pytest.approx(-22600.81488747, 1.0e-4) + + +def test__fit_time_delays_solved__solved_reference_time_equals_brute_force_scan(): + tracer = al.m.MockTracerPoint( + profile=al.ps.PointSolved(), + time_delays=al.ArrayIrregular([10.0, 12.0, 15.0]), + ) + + positions = al.Grid2DIrregular([(0.0, 1.5), (0.0, -1.3), (1.4, 0.05)]) + data = al.ArrayIrregular([1.0, 4.0, 6.0]) + noise_map = al.ArrayIrregular([0.5, 0.5, 1.0]) + + fit = al.FitTimeDelaysSolved( + name="point_0", + data=data, + noise_map=noise_map, + positions=positions, + tracer=tracer, + ) + + model_delays = fit.model_data.array + + def chi_squared(reference_time): + return np.sum( + (data.array - (model_delays + reference_time)) ** 2 / noise_map.array**2 + ) + + brute_force = minimize_scalar(chi_squared) + + assert fit.solved_reference_time == pytest.approx(brute_force.x, 1.0e-6) + assert fit.chi_squared == pytest.approx(chi_squared(fit.solved_reference_time), 1.0e-8) + + +def test__fit_time_delays_solved__works_with_point_flux_profile_too(): + # Time-delay fitting does not read any attribute of the paired profile (the model + # delays come from the tracer alone), so FitTimeDelaysSolved is not restricted to + # PointSolved -- it works with any profile that can be name-paired. + tracer = al.m.MockTracerPoint( + profile=al.ps.PointFlux(flux=1.0), + time_delays=al.ArrayIrregular([2.0, 2.0]), + ) + + positions = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) + data = al.ArrayIrregular([1.0, 2.0]) + noise_map = al.ArrayIrregular([3.0, 1.0]) + + fit = al.FitTimeDelaysSolved( + name="point_0", + data=data, + noise_map=noise_map, + positions=positions, + tracer=tracer, + ) + + assert np.isfinite(float(fit.log_likelihood)) diff --git a/test_autolens/point/model/test_analysis_point.py b/test_autolens/point/model/test_analysis_point.py index 218009db3..a9f3dc78c 100644 --- a/test_autolens/point/model/test_analysis_point.py +++ b/test_autolens/point/model/test_analysis_point.py @@ -206,3 +206,40 @@ def test__figure_of_merit__includes_fit_fluxes( fit_positions.log_likelihood + fit_fluxes.log_likelihood == analysis_log_likelihood ) + + +def test__fit_flux_cls_and_fit_time_delays_cls_hooks__forwarded_and_default_unchanged( + point_dataset, +): + solver = al.m.MockPointSolver(model_positions=point_dataset.positions) + + default_analysis = al.AnalysisPoint( + dataset=point_dataset, solver=solver, use_jax=False + ) + + # Defaults preserve current behaviour exactly. + assert default_analysis.fit_flux_cls is al.FitFluxes + assert default_analysis.fit_time_delays_cls is al.FitTimeDelays + + explicit_analysis = al.AnalysisPoint( + dataset=point_dataset, + solver=solver, + fit_flux_cls=al.FitFluxesSolved, + fit_time_delays_cls=al.FitTimeDelaysSolved, + use_jax=False, + ) + + assert explicit_analysis.fit_flux_cls is al.FitFluxesSolved + assert explicit_analysis.fit_time_delays_cls is al.FitTimeDelaysSolved + + model = af.Collection( + galaxies=af.Collection( + lens=al.Galaxy(redshift=0.5, point_0=al.ps.Point(centre=(0.0, 0.0))) + ) + ) + instance = model.instance_from_unit_vector([]) + + fit = explicit_analysis.fit_from(instance=instance) + + assert fit.fit_flux_cls is al.FitFluxesSolved + assert fit.fit_time_delays_cls is al.FitTimeDelaysSolved From 9f6fa7971c3e4a12f4f346e672f7cb9b2b619a13 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 27 Jul 2026 21:54:33 +0100 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20codex-review=20=E2=80=94=20observed-?= =?UTF-8?q?plane=20normalization,=20unswallowable=20mismatches,=20kwarg=20?= =?UTF-8?q?order,=20test=20gaps?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - FitPositionsSourceSolved.noise_normalization now uses the observed-plane det(Theta_i) (model-independent; Lombardi Eq. 46), not det(W_i) which spuriously favoured high-magnification models. - New exc.PointProfileMismatchException (NOT a PointExtractionException subclass) for profile/fit mismatches, so FitPointDataset's component-skip handlers can no longer swallow them; regression test through FitPointDataset added. - fit_flux_cls / fit_time_delays_cls moved to the end of the FitPointDataset / AnalysisPoint signatures — positional callers unbroken. - Tests: RecordingMockSolver asserts beta* is actually forwarded to the solver; flux/time marginalization + normalization constants pinned to closed forms. - Solved source-plane path verified numpy==jax under vmap+jit (-94.70750993 both backends on the workspace_test harness). Co-Authored-By: Claude Fable 5 --- autolens/exc.py | 12 ++++ autolens/point/fit/abstract.py | 2 +- autolens/point/fit/dataset.py | 2 +- autolens/point/fit/fluxes.py | 2 +- .../point/fit/positions/source/separations.py | 20 ++++-- autolens/point/fit/solved.py | 8 +-- autolens/point/model/analysis.py | 4 +- test_autolens/point/fit/test_fit_dataset.py | 42 ++++++++++++ test_autolens/point/fit/test_fluxes.py | 33 +++++++++- test_autolens/point/fit/test_solved.py | 64 +++++++++++++++++-- test_autolens/point/fit/test_time_delays.py | 27 ++++++++ 11 files changed, 195 insertions(+), 21 deletions(-) diff --git a/autolens/exc.py b/autolens/exc.py index ed739a462..5e9c4e831 100644 --- a/autolens/exc.py +++ b/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. + """ diff --git a/autolens/point/fit/abstract.py b/autolens/point/fit/abstract.py index abd97c487..a55d86117 100644 --- a/autolens/point/fit/abstract.py +++ b/autolens/point/fit/abstract.py @@ -157,7 +157,7 @@ def source_plane_coordinate(self) -> Tuple[float, float]: The (y,x) arc-second coordinates of the point-source in the source-plane. """ if not hasattr(self.profile, "centre"): - raise exc.PointExtractionException( + 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. " diff --git a/autolens/point/fit/dataset.py b/autolens/point/fit/dataset.py index f4d653f60..c295787d4 100644 --- a/autolens/point/fit/dataset.py +++ b/autolens/point/fit/dataset.py @@ -32,9 +32,9 @@ def __init__( tracer: Tracer, solver: PointSolver, fit_positions_cls=FitPositionsImagePair, + xp=np, fit_flux_cls=FitFluxes, fit_time_delays_cls=FitTimeDelays, - xp=np, ): """ Fits a point source dataset using a `Tracer` object, where the following components of the point source data diff --git a/autolens/point/fit/fluxes.py b/autolens/point/fit/fluxes.py index d41ce5e54..79deccf33 100644 --- a/autolens/point/fit/fluxes.py +++ b/autolens/point/fit/fluxes.py @@ -210,7 +210,7 @@ def __init__( ) if hasattr(self.profile, "flux"): - raise exc.PointExtractionException( + 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 " diff --git a/autolens/point/fit/positions/source/separations.py b/autolens/point/fit/positions/source/separations.py index 3968769db..15bfa3e65 100644 --- a/autolens/point/fit/positions/source/separations.py +++ b/autolens/point/fit/positions/source/separations.py @@ -187,8 +187,15 @@ class FitPositionsSourceSolved(SolvedCentre, FitPositionsSource): `log_likelihood = -0.5*(χ² + noise_norm) - 0.5*log(det(Σᵢ Wᵢ) / (2π)²)` where `χ² = Σᵢ (β̂ᵢ−β*)ᵀ Wᵢ (β̂ᵢ−β*)` (`chi_squared_map` / `chi_squared`) and - `noise_norm = Σᵢ log((2π)²/det Wᵢ)` (`noise_normalization`), each a separately-testable property, alongside - the marginalization term itself (`marginalization_term`). + `noise_norm = Σᵢ log((2π)²/det Θᵢ) = Σᵢ log((2π)² σᵢ⁴)` (`noise_normalization`), each a separately-testable + property, alongside the marginalization term itself (`marginalization_term`). + + The normalization deliberately uses the observed-plane precision `Θᵢ` (model-independent), NOT `det Wᵢ`: + this is a likelihood of the *observed image-plane positions* under the linearized lens equation (Lombardi + 2024 Eq. 46), so the Gaussian normalization is over the data space. A `det Wᵢ` normalization would add + `-2Σᵢ log|µᵢ|`-like model-dependent terms that spuriously favour high-magnification models. (The + long-standing `FitPositionsSource` uses the magnified-noise source-plane-data convention instead — its + normalization is intentionally unchanged.) 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 @@ -233,11 +240,12 @@ def chi_squared(self) -> float: @property def noise_normalization(self) -> float: """ - `noise_norm = Σᵢ log((2π)² / det Wᵢ)`. + `noise_norm = Σᵢ log((2π)² / det Θᵢ) = Σᵢ log((2π)² σᵢ⁴)` — the observed-plane (model-independent) + Gaussian normalization of the linearized image-plane likelihood (see class docstring; NOT `det Wᵢ`, + which would spuriously favour high-magnification models). """ - w11, w12, w21, w22 = precision_tensor_components_from(self, self.weighting) - det_w = w11 * w22 - w12 * w21 - return self._xp.sum(self._xp.log((2.0 * np.pi) ** 2.0 / det_w)) + sigma_sq = self.noise_map.array**2.0 + return self._xp.sum(self._xp.log((2.0 * np.pi) ** 2.0 * sigma_sq**2.0)) @property def marginalization_term(self) -> float: diff --git a/autolens/point/fit/solved.py b/autolens/point/fit/solved.py index 7636edc1f..14814c3fe 100644 --- a/autolens/point/fit/solved.py +++ b/autolens/point/fit/solved.py @@ -99,8 +99,8 @@ def precision_tensor_components_from(fit, weighting: str) -> Tuple: weighting `"jacobian"` — the tensor weighting `Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹`, with `Aᵢ` the lensing Jacobian at the observed position (see module docstring). - `"magnification"` — the scalar isotropic weighting `Wᵢ = (µᵢ²/σᵢ²) I₂`, matching - `FitPositionsSource.chi_squared_map`. + `"magnification"` — the traditional scalar weighting `Wᵢ = (µᵢ²/σᵢ²) I₂`, matching + `FitPositionsSource.chi_squared_map` (the near-critical tangential limit of the tensor). """ xp = fit._xp precision_scalar = _as_array(fit.noise_map) ** -2.0 # Θ = σ⁻² I @@ -112,7 +112,7 @@ def precision_tensor_components_from(fit, weighting: str) -> Tuple: return w, zero, zero, w if weighting != "jacobian": - raise exc.PointExtractionException( + raise exc.PointProfileMismatchException( f"Unsupported weighting '{weighting}' for the analytically-solved source-plane centre. " f"Valid options are 'jacobian' (tensor weighting, the default) or 'magnification' " f"(scalar isotropic weighting)." @@ -206,7 +206,7 @@ def source_plane_coordinate(self) -> Tuple[float, float]: `self._beta_hat`. """ if hasattr(self.profile, "centre"): - raise exc.PointExtractionException( + raise exc.PointProfileMismatchException( f"The point-source profile paired to dataset '{self.name}' " f"({self.profile.__class__.__name__}) has a `centre` attribute, so its free-centre priors " f"would be sampled by the non-linear search but silently ignored by " diff --git a/autolens/point/model/analysis.py b/autolens/point/model/analysis.py index cad32d387..1221575d1 100644 --- a/autolens/point/model/analysis.py +++ b/autolens/point/model/analysis.py @@ -42,12 +42,12 @@ def __init__( dataset: PointDataset, solver: PointSolver, fit_positions_cls=FitPositionsImagePairRepeat, - fit_flux_cls=FitFluxes, - fit_time_delays_cls=FitTimeDelays, image=None, cosmology: ag.cosmo.LensingCosmology = None, title_prefix: str = None, use_jax: bool = True, + fit_flux_cls=FitFluxes, + fit_time_delays_cls=FitTimeDelays, **kwargs, ): """ diff --git a/test_autolens/point/fit/test_fit_dataset.py b/test_autolens/point/fit/test_fit_dataset.py index eec6db580..395fa313a 100644 --- a/test_autolens/point/fit/test_fit_dataset.py +++ b/test_autolens/point/fit/test_fit_dataset.py @@ -176,3 +176,45 @@ def test__fit_dataset__fit_flux_cls_hook_is_actually_used_to_construct_the_flux_ assert isinstance(fit.flux, al.FitFluxesSolved) assert np.isfinite(fit.flux.log_likelihood) + + +def test__profile_fit_mismatches_are_not_swallowed_by_component_skipping(): + # Codex-review finding 4: FitPointDataset swallows PointExtractionException to skip absent + # dataset components; profile/fit mismatches raise PointProfileMismatchException instead and + # must propagate loudly rather than silently dropping the component. + lens = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0)) + positions = al.Grid2DIrregular([(0.0, 1.5), (0.0, -1.3), (1.4, 0.05)]) + + dataset = al.PointDataset( + name="point_0", + positions=positions, + positions_noise_map=al.ArrayIrregular([0.1] * 3), + fluxes=al.ArrayIrregular([1.0, 2.0, 3.0]), + fluxes_noise_map=al.ArrayIrregular([0.3] * 3), + ) + + # PointFlux profile + solved-flux fit: the flux prior would be sampled but ignored. + tracer_flux = al.Tracer( + galaxies=[ + lens, + al.Galaxy(redshift=1.0, point_0=al.ps.PointFlux(centre=(0.07, 0.07))), + ] + ) + with pytest.raises(al.exc.PointProfileMismatchException): + fit = al.FitPointDataset( + dataset=dataset, + tracer=tracer_flux, + solver=al.m.MockPointSolver(model_positions=positions), + fit_flux_cls=al.FitFluxesSolved, + ) + _ = fit.log_likelihood + + # Centre-bearing profile + solved-centre positions fit: the centre priors would be ignored. + with pytest.raises(al.exc.PointProfileMismatchException): + fit = al.FitPointDataset( + dataset=dataset, + tracer=tracer_flux, + solver=al.m.MockPointSolver(model_positions=positions), + fit_positions_cls=al.FitPositionsSourceSolved, + ) + _ = fit.log_likelihood diff --git a/test_autolens/point/fit/test_fluxes.py b/test_autolens/point/fit/test_fluxes.py index 5f6efced8..d5bdcb265 100644 --- a/test_autolens/point/fit/test_fluxes.py +++ b/test_autolens/point/fit/test_fluxes.py @@ -90,7 +90,7 @@ def test__fit_fluxes_solved__profile_with_flux_attribute__raises_naming_alternat noise_map = al.ArrayIrregular([3.0, 1.0]) positions = al.Grid2DIrregular([(0.0, 0.0), (3.0, 4.0)]) - with pytest.raises(al.exc.PointExtractionException, match="FitFluxes"): + with pytest.raises(al.exc.PointProfileMismatchException, match="FitFluxes"): al.FitFluxesSolved( name="point_0", data=data, @@ -120,3 +120,34 @@ def test__fit_fluxes_solved__works_with_plain_point_profile_no_flux_attribute(): assert np.isfinite(fit.solved_flux) assert np.isfinite(float(fit.log_likelihood)) + + +def test__fit_fluxes_solved__marginalization_and_normalization_constants_pinned(): + # Codex-review finding 9: pin each likelihood term to its closed form so a deleted or + # sign-flipped marginalization term cannot pass silently. + lens = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0)) + tracer = al.Tracer( + galaxies=[lens, al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved())] + ) + positions = al.Grid2DIrregular([(0.0, 1.5), (0.0, -1.3), (1.4, 0.05)]) + data = al.ArrayIrregular([5.0, 3.2, 4.1]) + noise_map = al.ArrayIrregular([0.3, 0.4, 0.2]) + + fit = al.FitFluxesSolved( + name="point_0", data=data, noise_map=noise_map, positions=positions, tracer=tracer + ) + + mu = fit.magnifications_at_positions.array + sigma = noise_map.array + precision_sum = np.sum(mu**2.0 / sigma**2.0) + + assert fit.marginalization_term == pytest.approx( + -0.5 * np.log(precision_sum / (2.0 * np.pi)), rel=1e-10 + ) + assert fit.noise_normalization == pytest.approx( + np.sum(np.log(2.0 * np.pi * sigma**2.0)), rel=1e-10 + ) + assert fit.log_likelihood == pytest.approx( + -0.5 * (fit.chi_squared + fit.noise_normalization) + fit.marginalization_term, + rel=1e-10, + ) diff --git a/test_autolens/point/fit/test_solved.py b/test_autolens/point/fit/test_solved.py index 2e1caa30d..8cd2d5af9 100644 --- a/test_autolens/point/fit/test_solved.py +++ b/test_autolens/point/fit/test_solved.py @@ -332,7 +332,7 @@ def test__solved_fit_given_centre_bearing_profile__raises_naming_alternative(sel name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None ) - with pytest.raises(al.exc.PointExtractionException, match="FitPositionsSource"): + with pytest.raises(al.exc.PointProfileMismatchException, match="FitPositionsSource"): fit.source_plane_coordinate def test__centre_requiring_fit_given_solved_profile__raises_naming_alternative(self): @@ -344,7 +344,7 @@ def test__centre_requiring_fit_given_solved_profile__raises_naming_alternative(s name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None ) - with pytest.raises(al.exc.PointExtractionException, match="PointSolved"): + with pytest.raises(al.exc.PointProfileMismatchException, match="PointSolved"): fit.source_plane_coordinate def test__pair_repeat_solved_given_centre_bearing_profile__raises(self): @@ -362,7 +362,7 @@ def test__pair_repeat_solved_given_centre_bearing_profile__raises(self): ) with pytest.raises( - al.exc.PointExtractionException, match="FitPositionsImagePairRepeat" + al.exc.PointProfileMismatchException, match="FitPositionsImagePairRepeat" ): fit.source_plane_coordinate @@ -381,7 +381,7 @@ def test__pair_all_solved_given_centre_bearing_profile__raises(self): ) with pytest.raises( - al.exc.PointExtractionException, match="FitPositionsImagePairAll" + al.exc.PointProfileMismatchException, match="FitPositionsImagePairAll" ): fit.source_plane_coordinate @@ -398,7 +398,7 @@ def test__point_flux_given_solved_fit__raises(self): name="point_0", data=positions, noise_map=noise_map, tracer=tracer, solver=None ) - with pytest.raises(al.exc.PointExtractionException): + with pytest.raises(al.exc.PointProfileMismatchException): fit.source_plane_coordinate def test__point_given_fit_fluxes__raises(self): @@ -440,3 +440,57 @@ def test__point_solved_given_fit_fluxes__raises(self): positions=positions, tracer=tracer, ) + + +class RecordingMockSolver: + """ + Mock solver that records the `source_plane_coordinate` it was called with — `al.m.MockPointSolver` + ignores it, so the parity tests above cannot detect a stale or unforwarded solved centre + (codex-review finding 8). + """ + + def __init__(self, model_positions): + self.model_positions = model_positions + self.last_source_plane_coordinate = None + + def solve( + self, + tracer, + source_plane_coordinate, + xp=np, + plane_redshift=None, + remove_infinities=True, + ): + self.last_source_plane_coordinate = source_plane_coordinate + return self.model_positions + + +def test__solved_image_plane_fits_forward_beta_star_to_the_solver(): + lens = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0)) + tracer = al.Tracer( + galaxies=[lens, al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved())] + ) + observed = al.Grid2DIrregular([(1.1, 0.2), (-0.9, -0.3), (0.3, 1.0), (-0.2, -1.05)]) + noise_map = al.ArrayIrregular([0.05] * 4) + + for fit_cls in [ + al.FitPositionsImagePairAllSolved, + al.FitPositionsImagePairRepeatSolved, + ]: + solver = RecordingMockSolver(model_positions=observed) + fit = fit_cls( + name="point_0", + data=observed, + noise_map=noise_map, + tracer=tracer, + solver=solver, + ) + _ = fit.model_data + + assert solver.last_source_plane_coordinate is not None + assert solver.last_source_plane_coordinate[0] == pytest.approx( + fit.source_plane_coordinate[0], rel=1e-12 + ) + assert solver.last_source_plane_coordinate[1] == pytest.approx( + fit.source_plane_coordinate[1], rel=1e-12 + ) diff --git a/test_autolens/point/fit/test_time_delays.py b/test_autolens/point/fit/test_time_delays.py index 62455027b..eae4d97da 100644 --- a/test_autolens/point/fit/test_time_delays.py +++ b/test_autolens/point/fit/test_time_delays.py @@ -110,3 +110,30 @@ def test__fit_time_delays_solved__works_with_point_flux_profile_too(): ) assert np.isfinite(float(fit.log_likelihood)) + + +def test__fit_time_delays_solved__marginalization_and_normalization_constants_pinned(): + # Codex-review finding 9: pin each likelihood term to its closed form. + lens = al.Galaxy(redshift=0.5, mass=al.mp.IsothermalSph(einstein_radius=1.0)) + tracer = al.Tracer( + galaxies=[lens, al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved())] + ) + positions = al.Grid2DIrregular([(0.0, 1.5), (0.0, -1.3), (1.4, 0.05)]) + data = al.ArrayIrregular([10.0, 14.5, 21.0]) + noise_map = al.ArrayIrregular([0.5, 0.7, 0.4]) + + fit = al.FitTimeDelaysSolved( + name="point_0", data=data, noise_map=noise_map, positions=positions, tracer=tracer + ) + + tau = 1.0 / noise_map.array**2.0 + assert fit.marginalization_term == pytest.approx( + -0.5 * np.log(np.sum(tau) / (2.0 * np.pi)), rel=1e-10 + ) + assert fit.noise_normalization == pytest.approx( + np.sum(np.log(2.0 * np.pi * noise_map.array**2.0)), rel=1e-10 + ) + assert fit.log_likelihood == pytest.approx( + -0.5 * (fit.chi_squared + fit.noise_normalization) + fit.marginalization_term, + rel=1e-10, + )