add option should_zero_data

autoarray/dataset/imaging/dataset.py

@@ -345,11 +345,22 @@ def apply_noise_scaling(
         mask: Mask2D,
         noise_value: float = 1e8,
         signal_to_noise_value: Optional[float] = None,
+        should_zero_data : bool = True
     ) -> "Imaging":
         """
-        Apply a mask to the imaging dataset using noise scaling, whereby the mask increases noise-map values to be
-        extremely large such that they are never included in the likelihood calculation, but it does
-        not remove the image data values, which are set to zero.
+        Apply a mask to the imaging dataset using noise scaling, whereby the maskmay zero the data and increase
+        noise-map values to change how they enter the likelihood calculation.
+
+        Given this data region is masked, it is likely thr data itself should not be included and therefore
+        the masked data values are set to zero. This can be disabled by setting `should_zero_data=False`.
+
+        Two forms of scaling are supported depending on whether the `signal_to_noise_value` is input:
+
+        - `noise_value`: The noise-map values in the masked region are set to this value, typically a very large value,
+        such that they are never included in the likelihood calculation.
+
+        - `signal_to_noise_value`: The noise-map values in the masked region are set to values such that they give
+        this signal-to-noise ratio. This overwrites the `noise_value` parameter.
 
         For certain modeling tasks, the mask defines regions of the data that are used to calculate the likelihood.
         For example, all data points in a mask may be used to create a pixel-grid, which is used in the likelihood.
@@ -368,20 +379,25 @@ def apply_noise_scaling(
         signal_to_noise_value
             The noise-map values are instead set to values such that they give this signal-to-noise_maps ratio.
             This overwrites the noise_value parameter.
+        should_zero_data
+            If True, the data values in the masked region are set to zero.
         """
 
-        if signal_to_noise_value is None:
+        if should_zero_data:
             data = np.where(np.invert(mask), 0.0, self.data.native)
-            data = Array2D.no_mask(
-                values=data,
-                shape_native=self.data.shape_native,
-                pixel_scales=self.data.pixel_scales,
-            )
+        else:
+            data = self.data.native
+
+        data = Array2D.no_mask(
+            values=data,
+            shape_native=self.data.shape_native,
+            pixel_scales=self.data.pixel_scales,
+        )
+
+        if signal_to_noise_value is None:
             noise_map = self.noise_map.native
             noise_map[mask == False] = noise_value
-
         else:
-            data = self.data.native
             noise_map = np.where(
                 mask == False,
                 np.abs(data.native) / signal_to_noise_value,

test_autoarray/dataset/imaging/test_dataset.py

@@ -165,7 +165,7 @@ def test__apply_noise_scaling__use_signal_to_noise_value(imaging_7x7, mask_2d_7x
     imaging_7x7.data[24] = 2.0
 
     masked_imaging_7x7 = imaging_7x7.apply_noise_scaling(
-        mask=mask_2d_7x7, signal_to_noise_value=0.1
+        mask=mask_2d_7x7, signal_to_noise_value=0.1, should_zero_data=False
     )
 
     assert masked_imaging_7x7.data.native[3, 4] == 1.0