fix: Keep masked fit utility gradients finite

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)

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])