Repository: gammapy/gammapy
Pull request #5699 (2025-02-11 12:20:56)

Add `mask` option to `plot_distribution()`
<!-- These are hidden commments. You can leave them or delete them. -->

<!-- Don't be shy! Github issues or pull requests are welcome any time. -->
<!-- But if you first want to chat, use the Gammapy Slack. -->
<!-- See https://gammapy.org/contact.html -->

**Description**
<!-- What is the motivation for this pull request? -->
<!-- Briefly: what changes are done here? -->
<!-- If this is to address a Github issue, mention its number and use the specific keywords to create a link -->
<!-- See docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue -->
<!-- Example: This PR is to resolve #42 -->

This pull request solves issue #5698 . An optional mask parameter is added to `plot_distribution()`, it takes a 2D mask to mask regions in the map. If the mask geometry or data type is wrong, an error is thrown.

A mask was also included in the corresponding test function.

**Dear reviewer**
<!-- Let the reviewer and Gammapy team know what you want: -->
<!-- * Is this ready for review? Or is it work in progress and you want some feedback? -->
<!-- * Do you want to go through review here? Or if someone should just finish this up and merge it? -->
<!-- Do you have any specific questions, e.g. about API or implementation? -->
<!-- Do you include a test executing the new code you're adding (to make sure it runs)? -->
<!-- Do you include some documentation? Is it needed? -->

This PR is ready for review. Please let me know what you think!



Patch:
diff --git a/gammapy/visualization/utils.py b/gammapy/visualization/utils.py
index 464bc2d18..a6fc051f3 100644
--- a/gammapy/visualization/utils.py
+++ b/gammapy/visualization/utils.py
@@ -249,6 +249,7 @@ def plot_theta_squared_table(table):
 
 def plot_distribution(
     wcs_map,
+    mask=None,
     ax=None,
     ncols=3,
     func=None,
@@ -264,6 +265,8 @@ def plot_distribution(
     ----------
     wcs_map : `~gammapy.maps.WcsNDMap`
         A map that contains data to be plotted.
+    mask : `~gammapy.maps.WcsNDMap`, optional
+        2D mask defining the input data region.
     ax : `~matplotlib.axes.Axes` or list of `~matplotlib.axes.Axes`
         Axis object to plot on. If a list of Axis is provided it has to be the same length as the length of _map.data.
     ncols : int
@@ -321,8 +324,18 @@ def plot_distribution(
     kwargs_hist.setdefault("density", True)
     kwargs_fit.setdefault("full_output", True)
 
-    cutout, mask = wcs_map.cutout_and_mask_region()
-    idx_x, idx_y = np.where(mask)
+    cutout, cutout_mask = wcs_map.cutout_and_mask_region()
+
+    if mask:
+        if mask.geom != wcs_map.geom.to_image():
+            raise ValueError("Map and mask spatial geometry must agree!")
+
+        if not mask.is_mask:
+            raise ValueError(f"Mask map must be of boolean type, got {mask.data.dtype} instead.")
+
+        cutout_mask.data = np.logical_and(cutout_mask.data, mask.data)
+
+    idx_x, idx_y = np.where(cutout_mask)
 
     data = cutout.data[..., idx_x, idx_y]
 
@@ -350,7 +363,6 @@ def plot_distribution(
     result_list = []
 
     for idx in range(cells_in_grid):
-
         axe = axes.flat[idx]
         if idx > len(data) - 1:
             axe.set_visible(False)


Test patch:
diff --git a/gammapy/visualization/tests/test_utils.py b/gammapy/visualization/tests/test_utils.py
index ef13dcbb6..af858f712 100644
--- a/gammapy/visualization/tests/test_utils.py
+++ b/gammapy/visualization/tests/test_utils.py
@@ -114,6 +114,8 @@ def test_plot_distribution():
 
     map_ = WcsNDMap.create(npix=(100, 100), axes=[energy_axis])
     map_.data = array_2d
+    mask_ = WcsNDMap.create(npix=(100, 100), dtype="bool")
+    mask_.data[:50, :] = True
 
     energy_axis_10 = MapAxis.from_energy_bounds(1 * u.TeV, 10 * u.TeV, 10)
     map_empty = WcsNDMap.create(npix=(100, 100), axes=[energy_axis_10])
@@ -123,7 +125,7 @@ def test_plot_distribution():
 
     with mpl_plot_check():
         axes, res = plot_distribution(
-            wcs_map=map_, func=fit_func, kwargs_hist={"bins": 40}
+            wcs_map=map_, mask=mask_, func=fit_func, kwargs_hist={"bins": 40}
         )
 
         assert axes.shape == (1,)
@@ -140,8 +142,23 @@ def test_plot_distribution():
         assert axes.shape == (4, 3)
 
         axes, res = plot_distribution(
-            wcs_map=map_, func="norm", kwargs_hist={"bins": 40}
+            wcs_map=map_, mask=mask_, func="norm", kwargs_hist={"bins": 40}
         )
 
         fig, ax = plt.subplots(nrows=1, ncols=1)
         plot_distribution(map_empty, ax=ax)
+
+        # test mask exceptions
+        mask_2 = WcsNDMap.create(npix=(200, 200), dtype="bool")
+        mask_2.data[:50, :] = True
+        with pytest.raises(ValueError):
+            axes, res = plot_distribution(
+                wcs_map=map_, mask=mask_2, func="norm", kwargs_hist={"bins": 40}
+            )
+
+        mask_3 = WcsNDMap.create(npix=(100, 100), dtype="float")
+        mask_3.data = np.ones_like(mask_3.data)
+        with pytest.raises(ValueError):
+            axes, res = plot_distribution(
+                wcs_map=map_, mask=mask_3, func="norm", kwargs_hist={"bins": 40}
+            )