Merge pull request #149 from Jammy2211/feature/noise_scaling_snr

Feature/noise scaling snr

autoarray/dataset/imaging/dataset.py

@@ -340,11 +340,27 @@ def apply_mask(self, mask: Mask2D) -> "Imaging":
 
         return dataset
 
-    def apply_noise_scaling(self, mask: Mask2D, noise_value: float = 1e8) -> "Imaging":
+    def apply_noise_scaling(
+        self,
+        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.
@@ -358,16 +374,36 @@ def apply_noise_scaling(self, mask: Mask2D, noise_value: float = 1e8) -> "Imagin
         ----------
         mask
             The 2D mask that is applied to the image and noise-map, to scale the noise-map values to large values.
+        noise_value
+            The value that the noise-map values are set to in the masked region where noise scaling is applied.
+        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.
         """
-        data = np.where(np.invert(mask), 0.0, self.data.native)
+
+        if should_zero_data:
+            data = np.where(np.invert(mask), 0.0, self.data.native)
+        else:
+            data = self.data.native
+
         data = Array2D.no_mask(
             values=data,
             shape_native=self.data.shape_native,
             pixel_scales=self.data.pixel_scales,
         )
 
-        noise_map = self.noise_map.native
-        noise_map[mask == False] = noise_value
+        if signal_to_noise_value is None:
+            noise_map = self.noise_map.native
+            noise_map[mask == False] = noise_value
+        else:
+            noise_map = np.where(
+                mask == False,
+                np.abs(data.native) / signal_to_noise_value,
+                self.noise_map.native,
+            )
+
         noise_map = Array2D.no_mask(
             values=noise_map,
             shape_native=self.data.shape_native,

autoarray/inversion/pixelization/mappers/abstract.py

@@ -394,12 +394,16 @@ def mapped_to_source_from(self, array: Array2D) -> np.ndarray:
         )
 
     def extent_from(
-        self, values: np.ndarray = None, zoom_to_brightest: bool = True
+        self,
+        values: np.ndarray = None,
+        zoom_to_brightest: bool = True,
+        zoom_percent: Optional[float] = None,
     ) -> Tuple[float, float, float, float]:
         if zoom_to_brightest and values is not None:
-            zoom_percent = conf.instance["visualize"]["general"]["zoom"][
-                "inversion_percent"
-            ]
+            if zoom_percent is None:
+                zoom_percent = conf.instance["visualize"]["general"]["zoom"][
+                    "inversion_percent"
+                ]
 
             fractional_value = np.max(values) * zoom_percent
             fractional_bool = values > fractional_value

test_autoarray/dataset/imaging/test_dataset.py

@@ -1,3 +1,4 @@
+import copy
 import os
 from os import path
 
@@ -158,6 +159,21 @@ def test__apply_noise_scaling(imaging_7x7, mask_2d_7x7):
     assert masked_imaging_7x7.noise_map.native[4, 4] == 1e5
 
 
+def test__apply_noise_scaling__use_signal_to_noise_value(imaging_7x7, mask_2d_7x7):
+    imaging_7x7 = copy.copy(imaging_7x7)
+
+    imaging_7x7.data[24] = 2.0
+
+    masked_imaging_7x7 = imaging_7x7.apply_noise_scaling(
+        mask=mask_2d_7x7, signal_to_noise_value=0.1, should_zero_data=False
+    )
+
+    assert masked_imaging_7x7.data.native[3, 4] == 1.0
+    assert masked_imaging_7x7.noise_map.native[3, 4] == 10.0
+    assert masked_imaging_7x7.data.native[3, 3] == 2.0
+    assert masked_imaging_7x7.noise_map.native[3, 3] == 20.0
+
+
 def test__apply_mask__noise_covariance_matrix():
     image = aa.Array2D.ones(shape_native=(2, 2), pixel_scales=(1.0, 1.0))