From 730ac4c02572218a3e9aae9350ef494303e9faf6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 18 May 2026 21:53:04 +0100 Subject: [PATCH] feat(weak): FitWeak class + plotters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a `FitWeak` class (`autolens/weak/fit.py`) that compares a model shear field — derived from a Tracer's mass profiles via the same `LensCalc.shear_yx_2d_via_hessian_from` primitive that `SimulatorShearYX` uses — against an observed `WeakDataset`. Exposes the standard fit interface: `model_shear`, `residual_map`, `chi_squared_map`, `chi_squared`, `noise_normalization`, `log_likelihood`, `figure_of_merit`. Each background galaxy contributes two independent measurements (gamma_1 and gamma_2 carry the same per-galaxy noise but are independent Gaussian draws), so chi-squared sums over N*2 elements and `noise_normalization` carries a factor of 2 to match. Add a matching plotter set (`autolens/weak/plot/fit_weak_plots.py`) with four helpers re-exported into the `aplt` namespace: - plot_data_vs_model -> data + model quivers overlaid (black + red alpha 0.6) - plot_residuals -> headless-quiver of the residual map, RdBu_r colormap - plot_chi_squared_map -> scalar scatter of per-galaxy chi-squared - subplot_fit_weak -> 2x2 mosaic `FitWeak` is standalone — not inheriting from `aa.AbstractFit`, matching the `FitPoint` precedent. The chi-squared / log-likelihood formulas verified against a direct numpy hand-computation in the test suite. Step 3 of the weak-lensing series (issue #524). 10 new tests, all green; existing 16 weak tests still pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- autolens/__init__.py | 1 + autolens/plot/__init__.py | 6 + autolens/weak/__init__.py | 1 + autolens/weak/fit.py | 88 ++++++++ autolens/weak/plot/fit_weak_plots.py | 202 ++++++++++++++++++ .../weak/plot/test_fit_weak_plots.py | 88 ++++++++ test_autolens/weak/test_fit.py | 122 +++++++++++ 7 files changed, 508 insertions(+) create mode 100644 autolens/weak/fit.py create mode 100644 autolens/weak/plot/fit_weak_plots.py create mode 100644 test_autolens/weak/plot/test_fit_weak_plots.py create mode 100644 test_autolens/weak/test_fit.py diff --git a/autolens/__init__.py b/autolens/__init__.py index a5f5fbf6b..409da88d3 100644 --- a/autolens/__init__.py +++ b/autolens/__init__.py @@ -122,6 +122,7 @@ from .quantity.fit_quantity import FitQuantity from .quantity.model.analysis import AnalysisQuantity from .weak.dataset import WeakDataset +from .weak.fit import FitWeak from .weak.simulator import SimulatorShearYX from . import exc diff --git a/autolens/plot/__init__.py b/autolens/plot/__init__.py index 3e6e1a991..283985209 100644 --- a/autolens/plot/__init__.py +++ b/autolens/plot/__init__.py @@ -72,6 +72,12 @@ plot_noise_map, subplot_weak_dataset, ) +from autolens.weak.plot.fit_weak_plots import ( + plot_data_vs_model, + plot_residuals, + plot_chi_squared_map, + subplot_fit_weak, +) from autolens.lens.plot.subhalo_plots import ( subplot_detection_imaging, diff --git a/autolens/weak/__init__.py b/autolens/weak/__init__.py index 07552c3d9..aeed5bdad 100644 --- a/autolens/weak/__init__.py +++ b/autolens/weak/__init__.py @@ -1,2 +1,3 @@ from autolens.weak.dataset import WeakDataset +from autolens.weak.fit import FitWeak from autolens.weak.simulator import SimulatorShearYX diff --git a/autolens/weak/fit.py b/autolens/weak/fit.py new file mode 100644 index 000000000..2389278e1 --- /dev/null +++ b/autolens/weak/fit.py @@ -0,0 +1,88 @@ +""" +Weak-lensing fit class. + +``FitWeak`` compares a model shear field (derived from a ``Tracer``'s mass profiles via +``LensCalc.shear_yx_2d_via_hessian_from`` — the same primitive ``SimulatorShearYX`` uses) against an observed +``WeakDataset`` and reports per-galaxy residuals, chi-squared and the log-likelihood. It is the weak-lensing +analogue of :class:`autolens.imaging.fit_imaging.FitImaging` and the input to a future ``AnalysisWeak``. + +Each background source galaxy contributes **two** independent measurements (:math:`\\gamma_1` and +:math:`\\gamma_2` carry the same per-galaxy noise but are independent Gaussian draws), so the chi-squared sum +and ``noise_normalization`` count :math:`N \\times 2` elements rather than just :math:`N`. + +The class is deliberately standalone — it does not inherit from ``autoarray.fit.fit_dataset.AbstractFit``, +which is shaped for "data + noise_map + mask" pixel-grid fits. ``FitPoint`` (in ``autolens.point``) follows +the same standalone pattern. +""" +import math +from functools import cached_property + +import numpy as np + +from autogalaxy.operate.lens_calc import LensCalc +from autogalaxy.util.shear_field import ShearYX2DIrregular + +from autolens.weak.dataset import WeakDataset + + +class FitWeak: + def __init__(self, dataset: WeakDataset, tracer): + """ + Fit a ``Tracer`` lens model to a ``WeakDataset`` shear catalogue. + + Parameters + ---------- + dataset + The observed weak-lensing shear catalogue. + tracer + The PyAutoLens ``Tracer`` whose mass profiles generate the model shear field. + """ + self.dataset = dataset + self.tracer = tracer + + @cached_property + def model_shear(self) -> ShearYX2DIrregular: + """The model shear field evaluated at the galaxy positions, via ``LensCalc``.""" + return LensCalc.from_tracer(self.tracer).shear_yx_2d_via_hessian_from( + grid=self.dataset.positions + ) + + @property + def residual_map(self) -> np.ndarray: + """``(N, 2)`` residuals ``data - model`` for each galaxy's ``(gamma_2, gamma_1)`` components.""" + return np.asarray(self.dataset.shear_yx) - np.asarray(self.model_shear) + + @property + def normalized_residual_map(self) -> np.ndarray: + """``(N, 2)`` residuals divided by the per-galaxy noise broadcast across both shear components.""" + noise = np.asarray(self.dataset.noise_map)[:, None] + return self.residual_map / noise + + @property + def chi_squared_map(self) -> np.ndarray: + """``(N, 2)`` per-component chi-squared contributions.""" + return self.normalized_residual_map**2 + + @property + def chi_squared(self) -> float: + """Scalar chi-squared summed over all ``N x 2`` shear measurements.""" + return float(np.sum(self.chi_squared_map)) + + @property + def noise_normalization(self) -> float: + r""" + Gaussian likelihood normalisation :math:`\sum \log(2 \pi \sigma^2)` summed over all ``N x 2`` shear + measurements — the factor of 2 reflects that each galaxy contributes two independent components. + """ + noise = np.asarray(self.dataset.noise_map) + return float(2.0 * np.sum(np.log(2.0 * math.pi * noise**2))) + + @property + def log_likelihood(self) -> float: + r"""Standard Gaussian log-likelihood :math:`-\tfrac{1}{2}(\chi^2 + \text{noise normalization})`.""" + return -0.5 * (self.chi_squared + self.noise_normalization) + + @property + def figure_of_merit(self) -> float: + """Quantity returned to non-linear searches; same as ``log_likelihood`` (no inversion / evidence).""" + return self.log_likelihood diff --git a/autolens/weak/plot/fit_weak_plots.py b/autolens/weak/plot/fit_weak_plots.py new file mode 100644 index 000000000..875b6684f --- /dev/null +++ b/autolens/weak/plot/fit_weak_plots.py @@ -0,0 +1,202 @@ +""" +Module-level matplotlib helpers for visualising a ``FitWeak``. + +Same headless-quiver convention as ``weak_dataset_plots`` (`pivot="middle", headwidth=0, headlength=0, +headaxislength=0`): shear is a spin-2 quantity, so the segments are drawn without arrowheads. The data is +plotted in black, the model in red with low alpha — overlaying them on a single axes makes deviations +visible at a glance. +""" +from typing import Optional + +import numpy as np + +from autoarray.plot.grid import plot_grid +from autoarray.plot.utils import ( + subplots, + save_figure, + conf_subplot_figsize, + tight_layout, +) + + +def _positions_yx(shear_yx) -> np.ndarray: + grid = shear_yx.grid + return np.array(grid.array if hasattr(grid, "array") else grid) + + +def _quiver_components(shear_yx): + """Return ``(y, x, u, v, mag)`` for a quiver of a ``ShearYX2DIrregular``.""" + positions = _positions_yx(shear_yx) + y, x = positions[:, 0], positions[:, 1] + mag = np.asarray(shear_yx.ellipticities) + phi_rad = np.deg2rad(np.asarray(shear_yx.phis)) + u = mag * np.cos(phi_rad) + v = mag * np.sin(phi_rad) + return y, x, u, v, mag + + +def plot_data_vs_model( + fit, + ax=None, + title: str = "Data vs Model", + output_path: Optional[str] = None, + output_filename: str = "data_vs_model", + output_format: Optional[str] = None, +): + """ + Overlay the data and model shear fields as headless quivers on a single axes. + + Data is drawn in black, model in red (alpha=0.6). Deviations are visible where the two segments + disagree in length or orientation. + """ + standalone = ax is None + if standalone: + fig, ax = subplots(1, 1) + + y_d, x_d, u_d, v_d, _ = _quiver_components(fit.dataset.shear_yx) + y_m, x_m, u_m, v_m, _ = _quiver_components(fit.model_shear) + + quiver_kwargs = dict( + pivot="middle", headwidth=0, headlength=0, headaxislength=0 + ) + ax.quiver(x_d, y_d, u_d, v_d, color="black", label="data", **quiver_kwargs) + ax.quiver( + x_m, y_m, u_m, v_m, color="red", alpha=0.6, label="model", **quiver_kwargs + ) + ax.set_xlabel('x (")') + ax.set_ylabel('y (")') + ax.set_title(title) + ax.set_aspect("equal") + ax.legend(loc="best", fontsize="small") + + if standalone: + tight_layout() + save_figure( + fig, + path=output_path, + filename=output_filename, + format=output_format, + ) + + +def plot_residuals( + fit, + ax=None, + title: str = "Residuals", + output_path: Optional[str] = None, + output_filename: str = "residuals", + output_format: Optional[str] = None, +): + """ + Quiver of the per-galaxy residual shear ``data - model``, colour-coded by ``|residual|``. + + Uses the diverging ``RdBu_r`` colormap centred on the median residual magnitude so over- and + under-shears are visually distinguishable. + """ + standalone = ax is None + if standalone: + fig, ax = subplots(1, 1) + + positions = _positions_yx(fit.dataset.shear_yx) + y, x = positions[:, 0], positions[:, 1] + + residual = fit.residual_map + mag = np.sqrt(residual[:, 0] ** 2 + residual[:, 1] ** 2) + phi_rad = 0.5 * np.arctan2(residual[:, 0], residual[:, 1]) + u = mag * np.cos(phi_rad) + v = mag * np.sin(phi_rad) + + ax.quiver( + x, + y, + u, + v, + mag, + pivot="middle", + headwidth=0, + headlength=0, + headaxislength=0, + cmap="RdBu_r", + ) + ax.set_xlabel('x (")') + ax.set_ylabel('y (")') + ax.set_title(title) + ax.set_aspect("equal") + + if standalone: + tight_layout() + save_figure( + fig, + path=output_path, + filename=output_filename, + format=output_format, + ) + + +def plot_chi_squared_map( + fit, + ax=None, + title: str = r"$\chi^2$ Map", + output_path: Optional[str] = None, + output_filename: str = "chi_squared_map", + output_format: Optional[str] = None, +): + """ + Colour-coded scatter of per-galaxy chi-squared (sum over the two shear components). + + Uses ``plot_grid`` from ``autoarray.plot.grid`` with ``color_array`` set to ``chi_squared_map.sum(axis=-1)``. + """ + per_galaxy_chi_squared = np.asarray(fit.chi_squared_map).sum(axis=-1) + + plot_grid( + grid=_positions_yx(fit.dataset.shear_yx), + ax=ax, + color_array=per_galaxy_chi_squared, + colormap="magma", + title=title, + output_path=output_path if ax is None else None, + output_filename=output_filename, + output_format=output_format, + ) + + +def subplot_fit_weak( + fit, + output_path: Optional[str] = None, + output_filename: str = "subplot_fit_weak", + output_format: Optional[str] = None, + title_prefix: Optional[str] = None, +): + """ + Produce a 2x2 subplot mosaic visualising a ``FitWeak``. + + Panels: data shear field, model shear field, data-vs-model overlay, chi-squared map. + """ + from autolens.weak.plot.weak_dataset_plots import plot_shear_yx_2d + + fig, axes = subplots(2, 2, figsize=conf_subplot_figsize(2, 2)) + ax_data, ax_model, ax_overlay, ax_chi = ( + axes[0, 0], + axes[0, 1], + axes[1, 0], + axes[1, 1], + ) + + _prefix = f"{title_prefix.rstrip()} " if title_prefix else "" + + plot_shear_yx_2d( + shear_yx=fit.dataset.shear_yx, ax=ax_data, title=f"{_prefix}Data" + ) + plot_shear_yx_2d( + shear_yx=fit.model_shear, ax=ax_model, title=f"{_prefix}Model" + ) + plot_data_vs_model(fit=fit, ax=ax_overlay, title=f"{_prefix}Data vs Model") + plot_chi_squared_map(fit=fit, ax=ax_chi, title=f"{_prefix}$\\chi^2$ Map") + + tight_layout() + save_figure( + fig, + path=output_path, + filename=output_filename, + format=output_format, + ) diff --git a/test_autolens/weak/plot/test_fit_weak_plots.py b/test_autolens/weak/plot/test_fit_weak_plots.py new file mode 100644 index 000000000..43cd35e1f --- /dev/null +++ b/test_autolens/weak/plot/test_fit_weak_plots.py @@ -0,0 +1,88 @@ +from pathlib import Path + +import autoarray as aa +import autolens as al + +import pytest + +from autolens.weak.plot.fit_weak_plots import ( + plot_data_vs_model, + plot_residuals, + plot_chi_squared_map, + subplot_fit_weak, +) + +directory = Path(__file__).resolve().parent + + +def _isothermal_tracer(einstein_radius=1.6, ell_comps=(0.0, 0.05)): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=ell_comps, + einstein_radius=einstein_radius, + ), + ) + source = al.Galaxy(redshift=1.0) + return al.Tracer(galaxies=[lens, source]) + + +@pytest.fixture(name="fit_weak") +def make_fit_weak(): + """Deterministic 4-galaxy FitWeak with non-zero residuals (model einstein_radius is slightly wrong).""" + grid = aa.Grid2DIrregular( + values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6), (-1.1, -0.8)] + ) + truth = _isothermal_tracer(einstein_radius=1.6) + dataset = al.SimulatorShearYX(noise_sigma=0.0, seed=0).via_tracer_from( + tracer=truth, grid=grid, name="test" + ) + dataset.noise_map = aa.ArrayIrregular(values=[0.3, 0.3, 0.3, 0.3]) + model = _isothermal_tracer(einstein_radius=1.5) + return al.FitWeak(dataset=dataset, tracer=model) + + +@pytest.fixture(name="plot_path") +def make_plot_path(): + return directory / "files" / "plots" / "fit_weak" + + +def test__plot_data_vs_model(fit_weak, plot_path, plot_patch): + plot_data_vs_model( + fit=fit_weak, + output_path=plot_path, + output_filename="data_vs_model", + output_format="png", + ) + assert str(plot_path / "data_vs_model.png") in plot_patch.paths + + +def test__plot_residuals(fit_weak, plot_path, plot_patch): + plot_residuals( + fit=fit_weak, + output_path=plot_path, + output_filename="residuals", + output_format="png", + ) + assert str(plot_path / "residuals.png") in plot_patch.paths + + +def test__plot_chi_squared_map(fit_weak, plot_path, plot_patch): + plot_chi_squared_map( + fit=fit_weak, + output_path=plot_path, + output_filename="chi_squared_map", + output_format="png", + ) + assert str(plot_path / "chi_squared_map.png") in plot_patch.paths + + +def test__subplot_fit_weak(fit_weak, plot_path, plot_patch): + subplot_fit_weak( + fit=fit_weak, + output_path=plot_path, + output_filename="subplot_fit_weak", + output_format="png", + ) + assert str(plot_path / "subplot_fit_weak.png") in plot_patch.paths diff --git a/test_autolens/weak/test_fit.py b/test_autolens/weak/test_fit.py new file mode 100644 index 000000000..c34f2976c --- /dev/null +++ b/test_autolens/weak/test_fit.py @@ -0,0 +1,122 @@ +import math + +import numpy as np +import pytest + +import autoarray as aa +import autolens as al + +from autogalaxy.operate.lens_calc import LensCalc + +from autolens.weak.fit import FitWeak + + +def _isothermal_tracer(einstein_radius=1.6, ell_comps=(0.0, 0.05)): + lens = al.Galaxy( + redshift=0.5, + mass=al.mp.Isothermal( + centre=(0.0, 0.0), + ell_comps=ell_comps, + einstein_radius=einstein_radius, + ), + ) + source = al.Galaxy(redshift=1.0) + return al.Tracer(galaxies=[lens, source]) + + +def _make_dataset(noise_sigma=0.3, seed=0, einstein_radius=1.6): + grid = aa.Grid2DIrregular( + values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6), (-1.1, -0.8)] + ) + tracer = _isothermal_tracer(einstein_radius=einstein_radius) + simulator = al.SimulatorShearYX(noise_sigma=noise_sigma, seed=seed) + return simulator.via_tracer_from(tracer=tracer, grid=grid, name="test") + + +def test__zero_noise_round_trip_zero_residuals(): + """ + The fit's `model_shear` is derived from the exact same `LensCalc` primitive that `SimulatorShearYX` uses, + so a noise-free round-trip through `FitWeak` with the truth tracer must produce a residual map of zeros. + """ + truth = _isothermal_tracer() + dataset = _make_dataset(noise_sigma=0.0) + + fit = FitWeak(dataset=dataset, tracer=truth) + + np.testing.assert_allclose(fit.residual_map, 0.0, atol=1e-12) + + +def test__chi_squared_zero_for_perfect_fit(): + """ + A noise-free dataset fit by the truth tracer has zero residuals; well-defined ``noise_sigma > 0`` then + guarantees zero chi-squared (any-finite-σ divided into zero residuals gives zero). + """ + truth = _isothermal_tracer() + dataset = al.SimulatorShearYX(noise_sigma=0.0, seed=0).via_tracer_from( + tracer=truth, + grid=aa.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6)]), + ) + # Replace the zero noise_map with a finite σ so the chi-squared division is well-defined. + dataset.noise_map = aa.ArrayIrregular(values=[0.3, 0.3, 0.3]) + + fit = FitWeak(dataset=dataset, tracer=truth) + + assert fit.chi_squared == pytest.approx(0.0, abs=1e-12) + + +def test__log_likelihood_against_hand_computed(): + """ + Verify the log-likelihood formula directly: compute chi-squared and noise normalisation from + `dataset.shear_yx`, `model_shear` and `noise_map` using plain numpy, and assert the class agrees to 1e-9. + """ + dataset = _make_dataset(noise_sigma=0.3, seed=42) + perturbed = _isothermal_tracer(einstein_radius=1.5) + + fit = FitWeak(dataset=dataset, tracer=perturbed) + + data = np.asarray(dataset.shear_yx) + model = np.asarray( + LensCalc.from_tracer(perturbed).shear_yx_2d_via_hessian_from( + grid=dataset.positions + ) + ) + sigma = np.asarray(dataset.noise_map) + + expected_chi_squared = np.sum(((data - model) / sigma[:, None]) ** 2) + expected_noise_norm = 2.0 * np.sum(np.log(2.0 * math.pi * sigma**2)) + expected_log_likelihood = -0.5 * (expected_chi_squared + expected_noise_norm) + + assert fit.chi_squared == pytest.approx(expected_chi_squared, abs=1e-9) + assert fit.noise_normalization == pytest.approx(expected_noise_norm, abs=1e-9) + assert fit.log_likelihood == pytest.approx(expected_log_likelihood, abs=1e-9) + + +def test__residual_map_shape_is_n_galaxies_by_2(): + dataset = _make_dataset(noise_sigma=0.3, seed=7) + fit = FitWeak(dataset=dataset, tracer=_isothermal_tracer(einstein_radius=1.5)) + + assert fit.residual_map.shape == (dataset.n_galaxies, 2) + assert fit.chi_squared_map.shape == (dataset.n_galaxies, 2) + + +def test__log_likelihood_drops_for_wrong_model(): + """ + A 0.1 perturbation of the lens einstein_radius away from truth should drop the log-likelihood by at + least 1 nat — sanity-check that wrong models are penalised. + """ + truth = _isothermal_tracer(einstein_radius=1.6) + dataset = _make_dataset(noise_sigma=0.05, seed=11, einstein_radius=1.6) + + fit_truth = FitWeak(dataset=dataset, tracer=truth) + fit_wrong = FitWeak( + dataset=dataset, tracer=_isothermal_tracer(einstein_radius=1.7) + ) + + assert fit_wrong.log_likelihood < fit_truth.log_likelihood - 1.0 + + +def test__figure_of_merit_equals_log_likelihood(): + dataset = _make_dataset(noise_sigma=0.3, seed=3) + fit = FitWeak(dataset=dataset, tracer=_isothermal_tracer()) + + assert fit.figure_of_merit == fit.log_likelihood