From e5e05217af976242df9e47117cf1e77c0835304b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 2 Oct 2026 11:08:59 +0100 Subject: [PATCH] fix: keep masked fit utility gradients finite --- autoarray/fit/fit_util.py | 15 +++++++--- test_autoarray/fit/test_fit_util.py | 44 +++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/autoarray/fit/fit_util.py b/autoarray/fit/fit_util.py index f8d7cc9da..133ba7018 100644 --- a/autoarray/fit/fit_util.py +++ b/autoarray/fit/fit_util.py @@ -248,7 +248,9 @@ def chi_squared_map_with_mask_from( mask The mask applied to the residual-map, where `False` entries are included in the calculation. """ - return xp.where(mask == 0, xp.square(residual_map / noise_map), 0) + included = mask == 0 + safe_noise_map = xp.where(included, noise_map, 1.0) + return xp.where(included, xp.square(residual_map / safe_noise_map), 0) def chi_squared_with_mask_from( @@ -449,7 +451,9 @@ def residual_flux_fraction_map_from( data The data of the dataset. """ - return xp.where(data != 0, residual_map / data, 0) + included = data != 0 + safe_data = xp.where(included, data, 1.0) + return xp.where(included, residual_map / safe_data, 0) def residual_flux_fraction_map_with_mask_from( @@ -460,7 +464,8 @@ def residual_flux_fraction_map_with_mask_from( Residual_Flux_Fraction = Residuals / Data = (Data - Model)/Data - The residual flux fraction map values in masked pixels are returned as zero. + The residual flux fraction map values in masked pixels or pixels with zero data + are returned as zero, matching the unmasked helper. Parameters ---------- @@ -471,4 +476,6 @@ def residual_flux_fraction_map_with_mask_from( mask The mask applied to the residual-map, where `False` entries are included in the calculation. """ - return xp.where(mask == 0, residual_map / data, 0) + included = (mask == 0) & (data != 0) + safe_data = xp.where(included, data, 1.0) + return xp.where(included, residual_map / safe_data, 0) diff --git a/test_autoarray/fit/test_fit_util.py b/test_autoarray/fit/test_fit_util.py index a4d1a7bfb..c9875215e 100644 --- a/test_autoarray/fit/test_fit_util.py +++ b/test_autoarray/fit/test_fit_util.py @@ -549,3 +549,47 @@ def test__residual_flux_fraction_map_with_mask_from__different_model__correct_un ) assert (residual_flux_fraction_map == np.array([0.0, 0.1, 0.2, 0.0])).all() + + +@pytest.mark.parametrize( + "mask", [np.array([False, False, True, True]), np.ones(4, dtype=bool)] +) +def test__chi_squared_masked_zero_noise__safe_excluded_divisions(mask): + residual = np.array([2.0, -3.0, 4.0, 5.0]) + noise = np.array([1.0, 2.0, 0.0, 0.0]) + with np.errstate(divide="raise", invalid="raise"): + result = aa.util.fit.chi_squared_map_with_mask_from( + residual_map=residual, noise_map=noise, mask=mask + ) + np.testing.assert_allclose(result, np.where(mask, 0.0, [4.0, 2.25, 0.0, 0.0])) + + +@pytest.mark.parametrize("masked", [False, True]) +@pytest.mark.parametrize( + "mask", [np.array([False, False, True, True]), np.ones(4, dtype=bool)] +) +def test__residual_fraction_zero_data__safe_divisions(masked, mask): + residual = np.array([2.0, -3.0, 4.0, 5.0]) + data = np.array([1.0, 0.0, -2.0, 0.0]) + expected = np.array([2.0, 0.0, -2.0, 0.0]) + with np.errstate(divide="raise", invalid="raise"): + if masked: + result = aa.util.fit.residual_flux_fraction_map_with_mask_from( + residual_map=residual, data=data, mask=mask + ) + expected = np.where(mask, 0.0, expected) + else: + result = aa.util.fit.residual_flux_fraction_map_from( + residual_map=residual, data=data + ) + np.testing.assert_allclose(result, expected) + + +def test__chi_squared_included_zero_noise__remains_undefined(): + with np.errstate(divide="ignore"): + result = aa.util.fit.chi_squared_map_with_mask_from( + residual_map=np.array([2.0]), + noise_map=np.array([0.0]), + mask=np.array([False]), + ) + assert np.isinf(result[0])