feat: stabilize PairAll mixture via log-sum-exp; add FitPositionsSource weighting option

Phase A of the point-source defaults campaign (#678):

- FitPositionsImagePairAll.all_permutations_log_likelihoods now reduces via a
  max-shifted log-sum-exp: identical where the literal log(sum(exp(...))) was
  finite, finite (not -inf) at >~38 sigma mismatch, restoring gradient flow
  across the approach to the basin. Zero-model-position and inf-padded solver
  rows behave exactly as before.
- FitPositionsSource gains a `weighting` class attribute: "magnification"
  (default) is unchanged behaviour; "jacobian" opts into the per-image
  precision-tensor chi-squared with the observed-plane normalization matching
  FitPositionsSourceSolved, reusing precision_tensor_components_from.

Author: Jammy2211

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

@@ -106,26 +106,46 @@ def all_permutations_log_likelihoods(self) -> np.ndarray:
         P(data_0 | model_1) * P(data_1 | model_1)
 
         This is every way in which the coordinates generated by the model can explain the observed coordinates.
+
+        The reduction over model positions is a max-shifted log-sum-exp rather than a literal
+        `log(sum(exp(...)))`: exponentiating first underflows to 0 once the best model/observed pairing is
+        ~38 sigma or worse, turning the log likelihood into `-inf` and killing gradient flow across the
+        exact region gradient searches must traverse to find the basin. The shifted form is mathematically
+        identical wherever the literal form is finite, and stays finite (the max term contributes exactly 0
+        after the shift) at arbitrarily large mismatch.
         """
 
         model_data = self.model_data.array
 
+        def log_sum_exp(log_ps):
+            # `initial` covers the zero-model-positions case (an empty `log_ps`), where a bare `max`
+            # raises: the -inf sentinel is clamped to 0 below, giving `log(sum of nothing) = -inf`,
+            # exactly the literal form's result, which `chi_squared`'s `has_image` fallback replaces.
+            max_log_p = self._xp.max(log_ps, initial=-np.inf)
+            # With no finite model position every log_p is -inf, and shifting by a -inf max would
+            # produce NaN (`-inf - -inf`) inside `exp` — including under `jax.grad`, where a NaN in
+            # the branch `chi_squared`'s `xp.where` discards still poisons the gradient. Clamp the
+            # shift to 0 so this case reduces to the literal form's `log(0) = -inf`, which the
+            # `has_image` fallback in `chi_squared` then replaces.
+            max_log_p = self._xp.where(
+                self._xp.isfinite(max_log_p), max_log_p, 0.0
+            )
+            return max_log_p + self._xp.log(
+                self._xp.sum(self._xp.exp(log_ps - max_log_p))
+            )
+
         return self._xp.array(
             [
-                self._xp.log(
-                    self._xp.sum(
-                        self._xp.array(
-                            [
-                                self._xp.exp(
-                                    self.log_p(
-                                        data_position,
-                                        model_position,
-                                        sigma,
-                                    )
-                                )
-                                for model_position in model_data
-                            ]
-                        )
+                log_sum_exp(
+                    self._xp.array(
+                        [
+                            self.log_p(
+                                data_position,
+                                model_position,
+                                sigma,
+                            )
+                            for model_position in model_data
+                        ]
                     )
                 )
                 for data_position, sigma in zip(self.data, self.noise_map)

autolens/point/fit/positions/source/separations.py

@@ -29,6 +29,13 @@
 
 
 class FitPositionsSource(AbstractFitPositions):
+    #: How each back-traced position's residual from the source-plane centre is weighted:
+    #: `"magnification"` — the traditional scalar `µᵢ²/σᵢ²` weighting with the magnified-noise
+    #: normalization (the long-standing behaviour of this class, and the Lenstool convention);
+    #: `"jacobian"` — the per-image precision tensor `Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹` with the observed-plane
+    #: normalization, matching `FitPositionsSourceSolved` but with the centre a free parameter.
+    weighting = "magnification"
+
     def __init__(
         self,
         name: str,
@@ -66,6 +73,13 @@ def __init__(
 
         7) Sum the chi-squared values to compute the overall log likelihood of the fit.
 
+        Steps 4-6 describe the default `weighting = "magnification"` scalar convention. Setting the
+        `weighting` class attribute to `"jacobian"` instead weights each vector residual `β̂ᵢ − c` (with `c`
+        the profile's free `centre`) by the per-image precision tensor `Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹` (see
+        `autolens.point.fit.solved.precision_tensor_components_from`), with the observed-plane noise
+        normalization matching `FitPositionsSourceSolved` — the same tensor likelihood as that class, but
+        with the centre sampled as a free parameter rather than solved and marginalized.
+
         Point source fitting uses name pairing, whereby the `name` of the `Point` object is paired to the name of the
         point source dataset to ensure that point source datasets are fitted to the correct point source.
 
@@ -136,32 +150,68 @@ def residual_map(self) -> aa.ArrayIrregular:
             coordinate=self.source_plane_coordinate
         )
 
+    @property
+    def residual_vectors(self) -> np.ndarray:
+        """
+        The (n_positions, 2) array of vector residuals `β̂ᵢ − c`: the back-traced source-plane positions
+        minus the source-plane centre `c` (here the profile's free `centre`; `FitPositionsSourceSolved`
+        overrides this to use the solved `β*` via `_beta_hat`, tolerating plain-ndarray test inputs).
+        """
+        beta_hat = self.model_data.array
+        centre_y, centre_x = self.source_plane_coordinate
+        centre = self._xp.array([centre_y, centre_x])
+        return beta_hat - centre
+
     @property
     def chi_squared_map(self) -> float:
         """
-        Returns the chi-squared of the point-source source-plane fit, which is the sum of the squared residuals
-        multiplied by the magnifications squared, divided by the noise-map values squared.
+        Returns the chi-squared of the point-source source-plane fit.
+
+        For `weighting = "magnification"` this is the squared residuals multiplied by the magnifications
+        squared, divided by the noise-map values squared. For `weighting = "jacobian"` it is the per-image
+        quadratic form `(β̂ᵢ−c)ᵀ Wᵢ (β̂ᵢ−c)` with the precision tensor `Wᵢ = Aᵢ⁻ᵀΘᵢAᵢ⁻¹`.
         """
+        if self.weighting == "magnification":
+            return self.residual_map**2.0 / (
+                self.magnifications_at_positions.array**-2.0
+                * self.noise_map.array**2.0
+            )
 
-        return self.residual_map**2.0 / (
-            self.magnifications_at_positions.array**-2.0 * self.noise_map.array**2.0
-        )
+        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 noise_normalization(self) -> float:
         """
-        Returns the normalization of the noise-map, which is the sum of the noise-map values squared.
+        Returns the noise normalization of the fit's Gaussian likelihood.
+
+        For `weighting = "magnification"` this is the long-standing magnified-noise source-plane-data
+        convention `Σᵢ log(2π µᵢ⁻²σᵢ²)`. For `weighting = "jacobian"` it is the observed-plane
+        (model-independent) convention `Σᵢ log((2π)² σᵢ⁴)` matching `FitPositionsSourceSolved` (see that
+        class's docstring for why a model-dependent normalization would spuriously favour
+        high-magnification models).
         """
-        return self._xp.sum(
-            self._xp.log(
-                2
-                * np.pi
-                * (
-                    self.magnifications_at_positions.array**-2.0
-                    * self.noise_map.array**2.0
+        if self.weighting == "magnification":
+            return self._xp.sum(
+                self._xp.log(
+                    2
+                    * np.pi
+                    * (
+                        self.magnifications_at_positions.array**-2.0
+                        * self.noise_map.array**2.0
+                    )
                 )
             )
-        )
+
+        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 log_likelihood(self) -> float:

autolens/point/fit/solved.py

@@ -113,9 +113,9 @@ def precision_tensor_components_from(fit, weighting: str) -> Tuple:
 
     if weighting != "jacobian":
         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)."
+            f"Unsupported weighting '{weighting}' for the source-plane position residuals. "
+            f"Valid options are 'jacobian' (tensor weighting, the default of the *Solved fit classes) or "
+            f"'magnification' (scalar isotropic weighting, the default of `FitPositionsSource`)."
         )
 
     lens_calc = _lens_calc_for(fit)

test_autolens/point/fit/positions/image/test_pair_all.py

@@ -215,6 +215,73 @@ def test__no_model_positions__finite_no_image_floor_matching_siblings(data, nois
     )
 
 
+def test__extreme_mismatch__log_sum_exp_stays_finite(data, noise_map):
+    """
+    Regression: the literal `log(sum(exp(log_p)))` underflows to `log(0) = -inf` once the best
+    model/observed pairing is ~38 sigma or worse (`exp` underflows below the smallest float64),
+    strangling gradient flow across the exact region gradient searches traverse to find the basin.
+    The max-shifted log-sum-exp must stay finite at arbitrarily large mismatch and equal the
+    directly-computed shifted reduction.
+    """
+    model_positions = al.Grid2DIrregular([(40.0, 40.0), (50.0, 50.0)])
+
+    fit = al.FitPositionsImagePairAll(
+        name="point_0",
+        data=data,
+        noise_map=noise_map,
+        tracer=tracer,
+        solver=al.mock.MockPointSolver(model_positions),
+    )
+
+    log_likelihoods = fit.all_permutations_log_likelihoods()
+
+    assert np.all(np.isfinite(log_likelihoods))
+    assert np.isfinite(fit.chi_squared)
+
+    # ~56 sigma worst pairing: every log_p is far below the ~-745 underflow threshold of exp().
+    for data_position, sigma, log_likelihood in zip(data, noise_map, log_likelihoods):
+        log_ps = np.array(
+            [
+                fit.log_p(data_position, model_position, sigma)
+                for model_position in model_positions.array
+            ]
+        )
+        assert np.all(log_ps < -745.0)
+        expected = log_ps.max() + np.log(np.sum(np.exp(log_ps - log_ps.max())))
+        assert log_likelihood == pytest.approx(expected, rel=1.0e-12)
+
+
+def test__moderate_mismatch__log_sum_exp_matches_literal_form(data, noise_map):
+    """Where the literal `log(sum(exp(...)))` is finite, the shifted form must equal it."""
+
+    model_positions = al.Grid2DIrregular([(-1.0749, -1.1), (1.19117, 1.175)])
+
+    fit = al.FitPositionsImagePairAll(
+        name="point_0",
+        data=data,
+        noise_map=noise_map,
+        tracer=tracer,
+        solver=al.mock.MockPointSolver(model_positions),
+    )
+
+    for data_position, sigma, log_likelihood in zip(
+        data, noise_map, fit.all_permutations_log_likelihoods()
+    ):
+        literal = np.log(
+            np.sum(
+                np.exp(
+                    np.array(
+                        [
+                            fit.log_p(data_position, model_position, sigma)
+                            for model_position in model_positions.array
+                        ]
+                    )
+                )
+            )
+        )
+        assert log_likelihood == pytest.approx(literal, rel=1.0e-14)
+
+
 def test__model_positions_present__chi_squared_unchanged(data, noise_map):
     """The no-image branch must not perturb the ordinary path."""
 

test_autolens/point/fit/positions/source/test_separations.py

@@ -80,6 +80,115 @@ def test__fit_positions_source__multi_plane_tracer__model_data_traces_to_correct
     assert (fit_1.model_data == traced_grids[2]).all()
 
 
+class FitPositionsSourceJacobian(al.FitPositionsSource):
+    weighting = "jacobian"
+
+
+def test__fit_positions_source__default_weighting_is_magnification():
+    assert al.FitPositionsSource.weighting == "magnification"
+
+
+def test__fit_positions_source__jacobian_weighting__matches_solved_at_solved_centre():
+    """
+    The free-centre tensor fit evaluated with its `centre` fixed at the solved centre `β*` must
+    reproduce `FitPositionsSourceSolved`'s chi-squared and noise normalization exactly — the two
+    likelihoods differ only by the solved class's analytic-marginalization term.
+    """
+    galaxy_mass = al.Galaxy(
+        redshift=0.5, mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.1)
+    )
+    positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0), (1.0, 0.0)])
+    noise_map = al.ArrayIrregular([0.5, 1.0, 0.8])
+
+    tracer_solved = al.Tracer(
+        galaxies=[
+            galaxy_mass,
+            al.Galaxy(redshift=1.0, point_0=al.ps.PointSolved()),
+        ]
+    )
+    fit_solved = al.FitPositionsSourceSolved(
+        name="point_0",
+        data=positions,
+        noise_map=noise_map,
+        tracer=tracer_solved,
+        solver=None,
+    )
+
+    beta_star = fit_solved.source_plane_coordinate
+
+    tracer_free = al.Tracer(
+        galaxies=[
+            galaxy_mass,
+            al.Galaxy(redshift=1.0, point_0=al.ps.Point(centre=beta_star)),
+        ]
+    )
+    fit_free = FitPositionsSourceJacobian(
+        name="point_0",
+        data=positions,
+        noise_map=noise_map,
+        tracer=tracer_free,
+        solver=None,
+    )
+
+    assert fit_free.chi_squared_map.in_list == pytest.approx(
+        fit_solved.chi_squared_map.in_list, rel=1.0e-8
+    )
+    assert fit_free.chi_squared == pytest.approx(fit_solved.chi_squared, rel=1.0e-8)
+    assert fit_free.noise_normalization == pytest.approx(
+        fit_solved.noise_normalization, rel=1.0e-8
+    )
+    assert fit_free.log_likelihood == pytest.approx(
+        fit_solved.log_likelihood - fit_solved.marginalization_term, rel=1.0e-8
+    )
+
+
+def test__fit_positions_source__jacobian_weighting__observed_plane_noise_normalization():
+    galaxy_mass = al.Galaxy(
+        redshift=0.5, mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=0.1)
+    )
+    positions = al.Grid2DIrregular([(0.0, 1.0), (0.0, 2.0)])
+    noise_map = al.ArrayIrregular([0.5, 1.0])
+
+    fit = FitPositionsSourceJacobian(
+        name="point_0",
+        data=positions,
+        noise_map=noise_map,
+        tracer=al.Tracer(
+            galaxies=[
+                galaxy_mass,
+                al.Galaxy(redshift=1.0, point_0=al.ps.Point(centre=(0.0, 0.0))),
+            ]
+        ),
+        solver=None,
+    )
+
+    sigma_sq = noise_map.array**2.0
+    assert fit.noise_normalization == pytest.approx(
+        np.sum(np.log((2.0 * np.pi) ** 2.0 * sigma_sq**2.0)), rel=1.0e-12
+    )
+
+
+def test__fit_positions_source__unknown_weighting_raises():
+    class FitPositionsSourceTypo(al.FitPositionsSource):
+        weighting = "magnificaton"
+
+    fit = FitPositionsSourceTypo(
+        name="point_0",
+        data=al.Grid2DIrregular([(0.0, 1.0)]),
+        noise_map=al.ArrayIrregular([0.5]),
+        tracer=al.Tracer(
+            galaxies=[
+                al.Galaxy(redshift=0.5),
+                al.Galaxy(redshift=1.0, point_0=al.ps.Point(centre=(0.0, 0.0))),
+            ]
+        ),
+        solver=None,
+    )
+
+    with pytest.raises(al.exc.PointProfileMismatchException):
+        fit.chi_squared_map
+
+
 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)