Repository: gammapy/gammapy
Pull request #5387 (2024-07-22 21:16:49)

Introduce context manager to set energy range to mask fit
<!-- 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 introduces a context manager to set an energy range to the mask fit of a `Datasets` object. This is meant to replace the `Datasets.split_by_energy` method that `FluxPointsEstimator` is using. An advantage of this approach is to avoid building new evaluators for energy energy band. It is not clear it is sufficient to support multiprocessing though. 

Opinions @QRemy , @AtreyeeS , @adonath ?

Note that some code duplication can be further removed with the `create_energy_mask` method (namely the group table creation is also used in `Dataset.split_by_energy`)

**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? -->


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/gammapy/gammapy/pull/5387?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy) Report
Attention: Patch coverage is `90.00000%` with `3 lines` in your changes missing coverage. Please review.
> Project coverage is 94.47%. Comparing base [(`94fff07`)](https://app.codecov.io/gh/gammapy/gammapy/commit/94fff07ca7009c09414554d259b273b1b121e98b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy) to head [(`2b56ce9`)](https://app.codecov.io/gh/gammapy/gammapy/commit/2b56ce9260cb7288b91cd92eab29c70d02b9efc9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy).
> Report is 1 commits behind head on main.

| [Files](https://app.codecov.io/gh/gammapy/gammapy/pull/5387?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy) | Patch % | Lines |
|---|---|---|
| [gammapy/datasets/utils.py](https://app.codecov.io/gh/gammapy/gammapy/pull/5387?src=pr&el=tree&filepath=gammapy%2Fdatasets%2Futils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy#diff-Z2FtbWFweS9kYXRhc2V0cy91dGlscy5weQ==) | 90.00% | [3 Missing :warning: ](https://app.codecov.io/gh/gammapy/gammapy/pull/5387?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #5387      +/-   ##
==========================================
- Coverage   94.48%   94.47%   -0.01%     
==========================================
  Files         232      232              
  Lines       35982    36012      +30     
==========================================
+ Hits        33996    34023      +27     
- Misses       1986     1989       +3     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/gammapy/gammapy/pull/5387?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy).

adonath: Thanks @registerrier! I just wanted to add some context to this change: the reason why we decided to split the datasets was indeed performance. So evaluating a set of smaller dataset was faster than evaluating and masking the whole dataset. However it seems with all the caching of predicted counts and modifications of the `MapEvaluator` that we have introduced since this is not the case anymore and the model set-up became the bottleneck.

That being said, I think the context manager is a useful addition from the user side as well. 
MRegeard: Thanks @registerrier, it is a nice feature. However, I think it has to be exposed in tutorials, otherwise it will not be used a lot. Also adding an example in the docstring might help. 
I was also wondering why we prefer defining a class context manager over using contextlib ? 
bkhelifi: If a central "class context manager" is introduced, there are many features that could be added:
- logging level
- progress bar
- color-blind switch
- general multiprocessing options
- MD5 checks or writing
- maybe later some stuff associated to metadata
- maybe a switch of the FoValignement axes (#4667)
- AOB
And its representation into the HLI.

This obviously goes beyond this PR... It depends on what means 'context manager'.

Patch:
diff --git a/gammapy/datasets/utils.py b/gammapy/datasets/utils.py
index bb80e6347..d435b4d35 100644
--- a/gammapy/datasets/utils.py
+++ b/gammapy/datasets/utils.py
@@ -1,4 +1,5 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
+import logging
 import numpy as np
 from astropy.coordinates import SkyCoord
 from gammapy.data import Observation
@@ -14,6 +15,8 @@ __all__ = [
     "create_map_dataset_from_dl4",
 ]
 
+log = logging.getLogger(__name__)
+
 
 def apply_edisp(input_map, edisp):
     """Apply energy dispersion to map. Requires "energy_true" axis.
@@ -417,3 +420,58 @@ def create_global_dataset(
         per_decade=True,
     )
     return MapDataset.create(geom=geom, energy_axis_true=energy_axis_true, name=name)
+
+
+class set_and_restore_mask_fit:
+    """Context manager to set a `mask_fit` on dataset and restore the initial mask.
+
+    Parameters
+    ----------
+    datasets : `~gammapy.datasets.datasets`
+        the Datasets to apply the energy mask to.
+    energy_min : `~astropy.units.Quantity`, optional
+        minimum energy.
+    energy_min : `~astropy.units.Quantity`, optional
+        maximum energy.
+    round_to_edge: bool, optional
+        Whether to round `energy_min` and `energy_max` to the closest axis bin value.
+        See `~gammapy.maps.MapAxis.round`. Default is False.
+    """
+
+    def __init__(
+        self,
+        datasets,
+        mask_fit=None,
+        energy_min=None,
+        energy_max=None,
+        round_to_edge=False,
+    ):
+        self.energy_min = energy_min
+        self.energy_max = energy_max
+        self.round_to_edge = round_to_edge
+        self.mask_fit = mask_fit
+        self.datasets = datasets
+        self.mask_fits = [dataset.mask_fit for dataset in datasets]
+
+    def __enter__(self):
+        datasets = Datasets()
+        for dataset in self.datasets:
+            mask_fit = dataset._geom.energy_mask(
+                self.energy_min, self.energy_max, self.round_to_edge
+            )
+            if self.mask_fit is not None:
+                mask_fit *= self.mask_fit.interp_to_geom(
+                    mask_fit.geom, method="nearest"
+                )
+            dataset.mask_fit = mask_fit
+            if np.any(mask_fit.data):
+                datasets.append(dataset)
+            else:
+                log.info(
+                    f"Dataset {dataset.name} does not contribute in the energy range"
+                )
+        return datasets
+
+    def __exit__(self, type, value, traceback):
+        for dataset, mask_fit in zip(self.datasets, self.mask_fits):
+            dataset.mask_fit = mask_fit


Test patch:
diff --git a/gammapy/datasets/tests/test_utils.py b/gammapy/datasets/tests/test_utils.py
index 00932caf5..6714646e3 100644
--- a/gammapy/datasets/tests/test_utils.py
+++ b/gammapy/datasets/tests/test_utils.py
@@ -3,8 +3,13 @@ import pytest
 import numpy as np
 from numpy.testing import assert_allclose
 import astropy.units as u
-from gammapy.datasets import MapDataset
-from gammapy.datasets.utils import apply_edisp, split_dataset, create_global_dataset
+from gammapy.datasets import Datasets, MapDataset, SpectrumDatasetOnOff
+from gammapy.datasets.utils import (
+    apply_edisp,
+    set_and_restore_mask_fit,
+    split_dataset,
+    create_global_dataset,
+)
 from gammapy.irf import EDispKernel
 from gammapy.maps import Map, MapAxis
 from gammapy.modeling.models import (
@@ -161,3 +166,76 @@ def test_create_global_dataset():
     assert_allclose(global_dataset.counts.geom.axes[0].edges.max(), 0.2 * u.TeV)
     assert_allclose(global_dataset.exposure.geom.axes[0].edges.min(), 0.02 * u.TeV)
     assert_allclose(global_dataset.exposure.geom.axes[0].edges.max(), 0.4 * u.TeV)
+
+
+@requires_data()
+def test_set_and_restore_mask():
+    ds1 = SpectrumDatasetOnOff.read(
+        "$GAMMAPY_DATA/joint-crab/spectra/hess/pha_obs23523.fits"
+    )
+    ds2 = SpectrumDatasetOnOff.read(
+        "$GAMMAPY_DATA/joint-crab/spectra/hess/pha_obs23526.fits"
+    )
+
+    datasets = Datasets([ds1, ds2])
+    with set_and_restore_mask_fit(
+        datasets, None, 800 * u.GeV, 5 * u.TeV, round_to_edge=True
+    ) as masked_datasets:
+        range1 = masked_datasets[0].energy_range_fit
+        range2 = masked_datasets[1].energy_range_fit
+
+    assert_allclose(range1[0].quantity.to_value("TeV"), 0.7943282)
+    assert_allclose(range1[1].quantity.to_value("TeV"), 5.011872)
+    assert_allclose(range2[0].quantity.to_value("TeV"), 0.7943282)
+
+    range1 = datasets[0].energy_range_fit
+    range2 = datasets[1].energy_range_fit
+    assert_allclose(range1[0].quantity.to_value("TeV"), 0.01)
+    assert_allclose(range1[1].quantity.to_value("TeV"), 100.0)
+    assert_allclose(range2[0].quantity.to_value("TeV"), 0.01)
+
+    with set_and_restore_mask_fit(
+        datasets, None, 100 * u.TeV, 200 * u.TeV
+    ) as masked_datasets:
+        assert len(masked_datasets) == 0
+
+    mask = ds1.counts.geom.energy_mask(800 * u.GeV, 5 * u.TeV, round_to_edge=True)
+    datasets = Datasets([ds1, ds2])
+    with set_and_restore_mask_fit(datasets, mask) as masked_datasets:
+        range1 = masked_datasets[0].energy_range_fit
+        range2 = masked_datasets[1].energy_range_fit
+
+    assert_allclose(range1[0].quantity.to_value("TeV"), 0.7943282)
+    assert_allclose(range1[1].quantity.to_value("TeV"), 5.011872)
+    assert_allclose(range2[0].quantity.to_value("TeV"), 0.7943282)
+
+    range1 = datasets[0].energy_range_fit
+    range2 = datasets[1].energy_range_fit
+    assert_allclose(range1[0].quantity.to_value("TeV"), 0.01)
+    assert_allclose(range1[1].quantity.to_value("TeV"), 100.0)
+    assert_allclose(range2[0].quantity.to_value("TeV"), 0.01)
+
+
+@requires_data()
+def test_set_and_restore_mask_3d():
+    dataset = MapDataset.read("$GAMMAPY_DATA/fermi-3fhl-gc/fermi-3fhl-gc.fits.gz")
+
+    mask_fit = Map.from_geom(dataset.counts.geom, data=True)
+    mask_fit = mask_fit.binary_erode(2 * u.deg)
+
+    datasets = Datasets([dataset])
+    with set_and_restore_mask_fit(
+        datasets, mask_fit, 80 * u.GeV, 500 * u.GeV, round_to_edge=True
+    ) as masked_datasets:
+        range1 = masked_datasets[0].energy_range_fit
+
+    assert_allclose(range1[0].quantity[100, 220].to_value("GeV"), 84.47164)
+    assert_allclose(range1[1].quantity[100, 220].to_value("GeV"), 500.0)
+    assert np.isnan(range1[0].quantity[0, 0].to_value("GeV"))
+    assert np.isnan(range1[1].quantity[0, 0].to_value("GeV"))
+
+    range1 = datasets[0].energy_range_fit
+    assert_allclose(range1[0].quantity[100, 220].to_value("GeV"), 10.0)
+    assert_allclose(range1[1].quantity[100, 220].to_value("GeV"), 500.0)
+    assert_allclose(range1[0].quantity[0, 0].to_value("GeV"), 10.0)
+    assert_allclose(range1[1].quantity[0, 0].to_value("GeV"), 500.0)