From dba5da473680dfb70cd51ed3b0d47736e55ba577 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 9 Jul 2026 15:17:22 +0100 Subject: [PATCH] Add FitWeak per-galaxy sigma_crit scaling + JAX support - When WeakDataset.redshifts is present, FitWeak scales the model signal per galaxy by beta_i/beta_ref via LensingCosmology.scaling_factor_between_ redshifts_from (unity at the tracer's source plane, zero at/below the lens plane); reduced datasets scale both gamma and kappa: g_i = s*g/(1-s*k). No redshifts -> bit-identical to the previous single-plane behaviour. Factors are concrete NumPy constants (plane/catalogue redshifts are not sampled parameters), keeping them outside any JAX trace. - xp threaded through the fit statistics (LensCalc hessian methods already take xp); model_shear returns a raw array on the jnp path (LensCalc guard pattern); AnalysisWeak._register_fit_weak_pytrees mirrors AnalysisPoint (no_flatten: dataset, _xp, _redshift_scale_factors). use_jax default stays False. - 4 new NumPy-only tests (factor endpoints, per-galaxy scaling, unchanged no-redshift behaviour, reduced+scaled composition); JAX validated by the companion autolens_workspace_test weak vmap-parity script (plain + scaled variants match eager NumPy at rtol 1e-6). Weak series step 10 (PyAutoLabs/PyAutoLens#590). Co-Authored-By: Claude Fable 5 --- autolens/weak/fit.py | 137 ++++++++++++++---- autolens/weak/model/analysis.py | 29 +++- test_autolens/weak/test_sigma_crit_scaling.py | 98 +++++++++++++ 3 files changed, 235 insertions(+), 29 deletions(-) create mode 100644 test_autolens/weak/test_sigma_crit_scaling.py diff --git a/autolens/weak/fit.py b/autolens/weak/fit.py index 2fe111e72..5652e6810 100644 --- a/autolens/weak/fit.py +++ b/autolens/weak/fit.py @@ -4,15 +4,31 @@ ``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``. +analogue of :class:`autolens.imaging.fit_imaging.FitImaging` and the input to ``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 model quantity adapts to what the dataset declares: + +- ``dataset.is_reduced`` — the model is the *reduced* shear :math:`g = \\gamma / (1 - \\kappa)`, the quantity + real surveys measure from galaxy ellipticities. +- ``dataset.redshifts`` — the model signal is scaled per galaxy by the lensing-efficiency ratio + :math:`\\beta_i / \\beta_{\\rm ref}` (``LensingCosmology.scaling_factor_between_redshifts_from``), so a + catalogue spanning a range of source redshifts is fitted self-consistently. The reference plane is the + tracer's outermost (source) plane; galaxies at or below the lens redshift carry zero signal. Without + redshifts the fit assumes the tracer's single effective source plane — exactly the pre-scaling behaviour. + The scale factors are computed eagerly from concrete plane redshifts (galaxy redshifts are fixed + constants, not sampled parameters), which keeps them outside any JAX trace. + 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. + +JAX support follows the ``LensCalc`` guard pattern: with ``xp=jnp`` the fit statistics are traceable and +``model_shear`` returns a raw ``(N, 2)`` array (``ShearYX2DIrregular`` is not a registered pytree); +``AnalysisWeak`` registers the ``FitWeak`` pytree so ``jax.jit(fit_from)`` round-trips a real fit object. """ import math from functools import cached_property @@ -26,7 +42,7 @@ class FitWeak: - def __init__(self, dataset: WeakDataset, tracer): + def __init__(self, dataset: WeakDataset, tracer, xp=np): """ Fit a ``Tracer`` lens model to a ``WeakDataset`` shear catalogue. @@ -36,71 +52,140 @@ def __init__(self, dataset: WeakDataset, tracer): The observed weak-lensing shear catalogue. tracer The PyAutoLens ``Tracer`` whose mass profiles generate the model shear field. + xp + The array module (``numpy`` or ``jax.numpy``). With ``jax.numpy`` the fit statistics are + traceable and ``model_shear`` returns a raw array rather than a ``ShearYX2DIrregular``. """ self.dataset = dataset self.tracer = tracer + self._xp = xp @cached_property - def model_shear(self) -> ShearYX2DIrregular: + def _redshift_scale_factors(self): + """ + Per-galaxy lensing-efficiency ratios ``beta_i / beta_ref``, or ``None`` when the dataset carries no + redshifts. + + Unity for galaxies at the tracer's source-plane redshift, zero at or below the lens redshift (such + galaxies are not lensed), and ``LensingCosmology.scaling_factor_between_redshifts_from`` in between — + the same factor multi-plane ray-tracing applies to deflections. Always a concrete NumPy array: + plane and catalogue redshifts are fixed constants, so this never participates in a JAX trace. """ - The model shear field evaluated at the galaxy positions, via ``LensCalc``. + redshifts = getattr(self.dataset, "redshifts", None) + if redshifts is None: + return None + + plane_redshifts = sorted( + float(galaxy.redshift) for galaxy in self.tracer.galaxies + ) + redshift_lens = plane_redshifts[0] + redshift_ref = plane_redshifts[-1] + + cosmology = self.tracer.cosmology + + factors = [ + 0.0 + if float(redshift_i) <= redshift_lens + else float( + cosmology.scaling_factor_between_redshifts_from( + redshift_0=redshift_lens, + redshift_1=float(redshift_i), + redshift_final=redshift_ref, + ) + ) + for redshift_i in np.asarray(redshifts) + ] + return np.asarray(factors) - When the dataset is marked ``is_reduced`` (real catalogues measure galaxy ellipticities, - i.e. the reduced shear) the model quantity is ``g = gamma / (1 - kappa)``, with the - convergence from the same Hessian primitive — so data and model always live in the same - space. + @cached_property + def model_shear(self): """ + The model signal evaluated at the galaxy positions, via ``LensCalc``. + + This is the (optionally per-galaxy-scaled) shear ``gamma``, or the reduced shear + ``g = gamma / (1 - kappa)`` when the dataset is marked ``is_reduced`` — data and model always live + in the same space. On the NumPy path the return is a ``ShearYX2DIrregular``; with ``xp=jax.numpy`` + it is a raw ``(N, 2)`` array (the ``LensCalc`` guard pattern). + """ + xp = self._xp + lens_calc = LensCalc.from_tracer(self.tracer) - shear = lens_calc.shear_yx_2d_via_hessian_from(grid=self.dataset.positions) + shear = lens_calc.shear_yx_2d_via_hessian_from( + grid=self.dataset.positions, xp=xp + ) - if not getattr(self.dataset, "is_reduced", False): + scale = self._redshift_scale_factors + is_reduced = getattr(self.dataset, "is_reduced", False) + + if scale is None and not is_reduced: return shear - convergence = lens_calc.convergence_2d_via_hessian_from( - grid=self.dataset.positions - ) - return ShearYX2DIrregular( - values=np.asarray(shear) / (1.0 - np.asarray(convergence))[:, None], - grid=self.dataset.positions, - ) + values = xp.asarray(shear) + + if is_reduced: + convergence = xp.asarray( + lens_calc.convergence_2d_via_hessian_from( + grid=self.dataset.positions, xp=xp + ) + ) + if scale is not None: + values = (scale[:, None] * values) / ( + 1.0 - scale * convergence + )[:, None] + else: + values = values / (1.0 - convergence)[:, None] + else: + values = scale[:, None] * values + + if xp is np: + return ShearYX2DIrregular( + values=np.asarray(values), grid=self.dataset.positions + ) + return values @property - def residual_map(self) -> np.ndarray: + def residual_map(self): """``(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) + xp = self._xp + data = xp.asarray(np.asarray(self.dataset.shear_yx)) + return data - xp.asarray(self.model_shear) @property - def normalized_residual_map(self) -> np.ndarray: + def normalized_residual_map(self): """``(N, 2)`` residuals divided by the per-galaxy noise broadcast across both shear components.""" - noise = np.asarray(self.dataset.noise_map)[:, None] + xp = self._xp + noise = xp.asarray(np.asarray(self.dataset.noise_map))[:, None] return self.residual_map / noise @property - def chi_squared_map(self) -> np.ndarray: + def chi_squared_map(self): """``(N, 2)`` per-component chi-squared contributions.""" return self.normalized_residual_map**2 @property - def chi_squared(self) -> float: + def chi_squared(self): """Scalar chi-squared summed over all ``N x 2`` shear measurements.""" - return float(np.sum(self.chi_squared_map)) + xp = self._xp + chi_squared = xp.sum(self.chi_squared_map) + return float(chi_squared) if xp is np else chi_squared @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. + Always concrete (it depends only on the dataset). """ 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: + def log_likelihood(self): 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: + def figure_of_merit(self): """Quantity returned to non-linear searches; same as ``log_likelihood`` (no inversion / evidence).""" return self.log_likelihood diff --git a/autolens/weak/model/analysis.py b/autolens/weak/model/analysis.py index 874e7d0b8..bcc6f9f35 100644 --- a/autolens/weak/model/analysis.py +++ b/autolens/weak/model/analysis.py @@ -51,9 +51,9 @@ def __init__( contributes two independent shear measurements (gamma_1 and gamma_2), which `FitWeak` compares against the model shear field of the `Tracer`. - `use_jax` defaults to `False` because `FitWeak` is a NumPy-only fit (its `model_shear` is cached via - `functools.cached_property` and its statistics use `np.asarray`); JAX support requires pytree - registration of `FitWeak` and an `xp`-threaded fit path, which is deliberate future work. + `use_jax` defaults to `False`, the conservative choice for the newest Analysis class; pass + `use_jax=True` to run the `xp`-threaded fit path with `FitWeak` pytree registration (validated by + the `autolens_workspace_test` weak vmap-parity script). Parameters ---------- @@ -122,6 +122,9 @@ def fit_from(self, instance) -> FitWeak: ------- The fit of the lens model to the weak-lensing shear catalogue. """ + if self._use_jax: + self._register_fit_weak_pytrees() + tracer = self.tracer_via_instance_from( instance=instance, ) @@ -129,7 +132,27 @@ def fit_from(self, instance) -> FitWeak: return FitWeak( dataset=self.dataset, tracer=tracer, + xp=self._xp, + ) + + @staticmethod + def _register_fit_weak_pytrees() -> None: + """Register every type reachable from a ``FitWeak`` return value so + ``jax.jit(fit_from)`` can flatten its output. + + ``dataset`` and ``_xp`` are constants per analysis — ride as aux so JAX does + not recurse into them (the cached ``_redshift_scale_factors`` derives purely + from the dataset and plane redshifts, so it stays concrete). ``tracer`` is + dynamic per fit. + """ + from autoarray.abstract_ndarray import register_instance_pytree + from autolens.lens.tracer import Tracer + + register_instance_pytree( + FitWeak, + no_flatten=("dataset", "_xp", "_redshift_scale_factors"), ) + register_instance_pytree(Tracer, no_flatten=("cosmology",)) def save_attributes(self, paths: af.DirectoryPaths): """ diff --git a/test_autolens/weak/test_sigma_crit_scaling.py b/test_autolens/weak/test_sigma_crit_scaling.py new file mode 100644 index 000000000..d34b7c4f6 --- /dev/null +++ b/test_autolens/weak/test_sigma_crit_scaling.py @@ -0,0 +1,98 @@ +import numpy as np + +import autoarray as aa +import autolens as al + +from autolens.weak.fit import FitWeak + + +def _tracer(z_lens=0.5, z_source=1.0): + lens = al.Galaxy( + redshift=z_lens, + mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6), + ) + return al.Tracer(galaxies=[lens, al.Galaxy(redshift=z_source)]) + + +def _dataset(redshifts, is_reduced=False, positions=None): + positions = positions or [(1.0, 1.0), (1.0, 1.0), (1.0, 1.0)] + grid = aa.Grid2DIrregular(values=positions) + tracer = _tracer() + dataset = al.SimulatorShearYX(noise_sigma=0.0, seed=1).via_tracer_from( + tracer=tracer, grid=grid + ) + dataset.redshifts = aa.ArrayIrregular(values=list(redshifts)) + dataset.is_reduced = is_reduced + return dataset + + +def test__scale_factors__unity_at_ref_zero_at_lens_and_cosmology_between(): + dataset = _dataset(redshifts=[1.0, 0.5, 0.75]) + tracer = _tracer() + + fit = FitWeak(dataset=dataset, tracer=tracer) + factors = fit._redshift_scale_factors + + assert factors[0] == 1.0 # at the source plane + assert factors[1] == 0.0 # at the lens plane: unlensed + expected_mid = tracer.cosmology.scaling_factor_between_redshifts_from( + redshift_0=0.5, redshift_1=0.75, redshift_final=1.0 + ) + np.testing.assert_allclose(factors[2], expected_mid) + assert 0.0 < factors[2] < 1.0 + + +def test__model_shear__scales_per_galaxy(): + """Three galaxies at the SAME position, different redshifts: the model shear must be the + single-plane shear multiplied by each galaxy's beta ratio.""" + dataset = _dataset(redshifts=[1.0, 0.5, 0.75]) + tracer = _tracer() + + fit_scaled = FitWeak(dataset=dataset, tracer=tracer) + + dataset_plain = _dataset(redshifts=[1.0, 0.5, 0.75]) + dataset_plain.redshifts = None + fit_plain = FitWeak(dataset=dataset_plain, tracer=tracer) + + plain = np.asarray(fit_plain.model_shear) + scaled = np.asarray(fit_scaled.model_shear) + factors = fit_scaled._redshift_scale_factors + + np.testing.assert_allclose(scaled, factors[:, None] * plain, atol=1e-12) + np.testing.assert_allclose(scaled[1], 0.0, atol=1e-12) # z = z_lens galaxy + + +def test__no_redshifts__behaviour_unchanged(): + """A dataset without redshifts must reproduce the pre-scaling likelihood exactly.""" + grid = aa.Grid2DIrregular(values=[(0.7, 0.5), (1.0, 1.0), (-0.3, 0.6)]) + tracer = _tracer() + dataset = al.SimulatorShearYX(noise_sigma=0.3, seed=2).via_tracer_from( + tracer=tracer, grid=grid + ) + + fit = FitWeak(dataset=dataset, tracer=tracer) + assert fit._redshift_scale_factors is None + + # All galaxies AT the reference plane must also match the no-redshift fit exactly. + dataset.redshifts = aa.ArrayIrregular(values=[1.0, 1.0, 1.0]) + fit_ref = FitWeak(dataset=dataset, tracer=tracer) + np.testing.assert_allclose(fit_ref.log_likelihood, fit.log_likelihood, rtol=1e-12) + + +def test__reduced_and_scaled__g_uses_scaled_kappa(): + """For a reduced dataset with redshifts, g_i = s_i*gamma / (1 - s_i*kappa) — both the shear + and the convergence carry the efficiency factor.""" + from autogalaxy.operate.lens_calc import LensCalc + + dataset = _dataset(redshifts=[0.75, 1.0], is_reduced=True, positions=[(1.0, 1.0), (1.0, 1.0)]) + tracer = _tracer() + + fit = FitWeak(dataset=dataset, tracer=tracer) + + lens_calc = LensCalc.from_tracer(tracer) + gamma = np.asarray(lens_calc.shear_yx_2d_via_hessian_from(grid=dataset.positions)) + kappa = np.asarray(lens_calc.convergence_2d_via_hessian_from(grid=dataset.positions)) + s = fit._redshift_scale_factors + + expected = (s[:, None] * gamma) / (1.0 - s * kappa)[:, None] + np.testing.assert_allclose(np.asarray(fit.model_shear), expected, atol=1e-12)