diff --git a/autolens/point/fit/positions/image/pair_all.py b/autolens/point/fit/positions/image/pair_all.py index dbb651db3..9a7295e98 100644 --- a/autolens/point/fit/positions/image/pair_all.py +++ b/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) diff --git a/autolens/point/fit/positions/source/separations.py b/autolens/point/fit/positions/source/separations.py index 15bfa3e65..cdf0a620b 100644 --- a/autolens/point/fit/positions/source/separations.py +++ b/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: diff --git a/autolens/point/fit/solved.py b/autolens/point/fit/solved.py index 14814c3fe..e17f98cb9 100644 --- a/autolens/point/fit/solved.py +++ b/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) 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 4017f5b79..a6bfbdab8 100644 --- a/test_autolens/point/fit/positions/image/test_pair_all.py +++ b/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.""" diff --git a/test_autolens/point/fit/positions/source/test_separations.py b/test_autolens/point/fit/positions/source/test_separations.py index 4f61283be..fc1509269 100644 --- a/test_autolens/point/fit/positions/source/test_separations.py +++ b/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)