Merge pull request #591 from PyAutoLabs/feature/weak-sigma-crit-jax

feat: FitWeak per-galaxy sigma_crit scaling + JAX support (weak series step 10)

Author: Jammy2211

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

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,14 +122,37 @@ 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,
         )
 
         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):
         """

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)