Repository: gammapy/gammapy (python)

Issue:
Modify `MapAxis.downsample`
Modify `MapAxis.downsample` to have an option strict=False option which when set to True will simply put the remainder in the last bin
 https://github.com/gammapy/gammapy/pull/5007#discussion_r1455938626_

Pull request description:
Add strict option to `MapAxis.downsample`
This PR resolves #5047

Adding a strict option to the downsample, such that when it is true, it places the remainder in the last bin

Patch:
diff --git a/gammapy/maps/axes.py b/gammapy/maps/axes.py
index 3a0ee0c4f..e359d94ba 100644
--- a/gammapy/maps/axes.py
+++ b/gammapy/maps/axes.py
@@ -1110,10 +1110,10 @@ class MapAxis:
             nodes = self.pix_to_coord(pix_new)
             return self.from_nodes(nodes, name=self.name, interp=self.interp)
 
-    def downsample(self, factor):
+    def downsample(self, factor, strict=True):
         """Downsample map axis by a given factor.
 
-        When downsampling each n-th (given by the factor) bin is selected from
+        When downsampling, each n-th (given by the factor) bin is selected from
         the axis while preserving the axis limits. For node type "edges" this
         requires nbin to be dividable by the factor, for node type "center" this
         requires nbin - 1 to be dividable by the factor.
@@ -1122,6 +1122,11 @@ class MapAxis:
         ----------
         factor : int
             Downsampling factor.
+        strict : bool
+            Whether the number of bins is strictly divisible by the factor.
+            If True, ``nbin`` must be divisible by the ``factor``.
+            If False, the reminder bins are put into the last bin of the new axis.
+            Default is True.
 
         Returns
         -------
@@ -1129,25 +1134,26 @@ class MapAxis:
             Downsampled map axis.
         """
         if self.node_type == "edges":
-            nbin = self.nbin / factor
-
-            if np.mod(nbin, 1) > 0:
-                raise ValueError(
-                    f"Number of {self.name} bins is not divisible by {factor}"
-                )
-
             edges = self.edges[::factor]
+            if edges[-1] != self.edges[-1]:
+                if strict is True:
+                    raise ValueError(
+                        f"Number of {self.name} bins ({self.nbin}) is not divisible by {factor}"
+                    )
+                else:
+                    edges = np.append(edges, self.edges[-1])
             return self.from_edges(edges, name=self.name, interp=self.interp)
-        else:
-            nbin = (self.nbin - 1) / factor
 
-            if np.mod(nbin, 1) > 0:
-                raise ValueError(
-                    f"Number of {self.name} bins - 1 is not divisible by {factor}"
-                )
-
-            nodes = self.center[::factor]
-            return self.from_nodes(nodes, name=self.name, interp=self.interp)
+        elif self.node_type == "center":
+            centers = self.center[::factor]
+            if centers[-1] != self.center[-1]:
+                if strict is True:
+                    raise ValueError(
+                        f"Number of {self.name} bins - 1 ({self.nbin-1}) is not divisible by {factor}"
+                    )
+                else:
+                    centers = np.append(centers, self.center[-1])
+            return self.from_nodes(centers, name=self.name, interp=self.interp)
 
     def to_header(self, format="ogip", idx=0):
         """Create FITS header.


Test patch:
diff --git a/gammapy/maps/tests/test_axes.py b/gammapy/maps/tests/test_axes.py
index bd086e4e1..223879786 100644
--- a/gammapy/maps/tests/test_axes.py
+++ b/gammapy/maps/tests/test_axes.py
@@ -125,6 +125,8 @@ def test_downsample():
         node_type="edges",
         interp="lin",
     )
+
+    # Node_type edge, divisible, strict=True
     axis_down = axis.downsample(2)
 
     assert_allclose(axis_down.nbin, 0.5 * axis.nbin)
@@ -132,6 +134,75 @@ def test_downsample():
     assert_allclose(axis_down.edges[-1], axis.edges[-1])
     assert axis_down.node_type == axis.node_type
 
+    # Node_type edge, divisible, strict=False
+    axis_down1 = axis.downsample(2, strict=False)
+    assert axis_down == axis_down1
+
+    axis = MapAxis(
+        nodes=[0, 1, 2, 3, 4, 5, 6, 7],
+        unit="TeV",
+        name="energy",
+        node_type="edges",
+        interp="lin",
+    )
+    # Node_type edge, not divisible, strict=True
+    with pytest.raises(ValueError) as exc_info:
+        axis.downsample(2, strict=True)
+    assert str(exc_info.value) == "Number of energy bins (7) is not divisible by 2"
+
+    # Node_type edge, not divisible, strict=False
+    axis_down = axis.downsample(2, strict=False)
+    assert_allclose(axis_down.nbin, np.ceil(0.5 * axis.nbin))
+    assert_allclose(axis_down.edges[0], axis.edges[0])
+    assert_allclose(axis_down.edges[-1], axis.edges[-1])
+    assert axis_down.node_type == axis.node_type
+
+    axis = MapAxis(
+        nodes=[
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+        ],
+        unit="TeV",
+        name="energy",
+        node_type="center",
+        interp="lin",
+    )
+
+    # Node_type center, not divisible, strict=True
+    with pytest.raises(ValueError) as exc_info:
+        axis.downsample(2, strict=True)
+    assert str(exc_info.value) == "Number of energy bins - 1 (5) is not divisible by 2"
+    # Node_type center, not divisible, strict=False
+    axis_down = axis.downsample(2, strict=False)
+    assert_allclose(axis_down.nbin, 4)
+    assert_allclose(axis_down.center, [0.0, 2.0, 4.0, 5] * u.TeV)
+    assert axis_down.node_type == "center"
+
+    axis = MapAxis(
+        nodes=[
+            0,
+            1,
+            2,
+            3,
+            4,
+        ],
+        unit="TeV",
+        name="energy",
+        node_type="center",
+        interp="lin",
+    )
+    # Node_type center, divisible, strict=True
+    axis_down = axis.downsample(2, strict=True)
+    assert_allclose(axis_down.nbin, 3)
+    assert_allclose(axis_down.center, [0.0, 2.0, 4.0] * u.TeV)
+    # Node_type center, divisible, strict=False
+    axis_down1 = axis.downsample(2, strict=False)
+    assert axis_down == axis_down1
+
 
 def test_upsample_non_regular():
     axis = MapAxis.from_edges([0, 1, 3, 7], name="test", interp="lin")