Implement ReflectedRegionsBackgroundMaker

gammapy/spectrum/reflected.py

@@ -7,8 +7,14 @@
 from gammapy.maps import Map, WcsGeom, WcsNDMap
 from gammapy.maps.geom import frame_to_coordsys
 from .background_estimate import BackgroundEstimate
+from .dataset import SpectrumDatasetOnOff
+from .core import CountsSpectrum
 
-__all__ = ["ReflectedRegionsFinder", "ReflectedRegionsBackgroundEstimator"]
+__all__ = [
+    "ReflectedRegionsFinder",
+    "ReflectedRegionsBackgroundEstimator",
+    "ReflectedRegionsBackgroundMaker",
+]
 
 log = logging.getLogger(__name__)
 
@@ -44,7 +50,7 @@ class ReflectedRegionsFinder:
     exclusion_mask : `~gammapy.maps.WcsNDMap`, optional
         Exclusion mask
     binsz : `~astropy.coordinates.Angle`
-        Bin size of the reference map used for region finding. Default : 0.01 deg
+        Bin size of the reference map used for region finding.
 
     Examples
     --------
@@ -132,9 +138,9 @@ def make_reference_map(region, center, binsz="0.01 deg", min_width="0.3 deg"):
         center : `~astropy.coordinates.SkyCoord`
             Rotation point
         binsz : `~astropy.coordinates.Angle`
-            Reference map bin size. Default : 0.01 deg
+            Reference map bin size.
         min_width : `~astropy.coordinates.Angle`
-            Minimal map width. Default : 0.3 deg
+            Minimal map width.
 
         Returns
         -------
@@ -265,7 +271,7 @@ class ReflectedRegionsBackgroundEstimator:
     observations : `~gammapy.data.Observations`
         Observations to process
     binsz : `~astropy.coordinates.Angle`
-        Optional, bin size of the maps used to compute the regions, Default '0.01 deg'
+        Optional, bin size of the maps used to compute the regions.
     kwargs : dict
         Forwarded to `~gammapy.spectrum.ReflectedRegionsFinder`
     """
@@ -345,9 +351,9 @@ def plot(self, fig=None, ax=None, cmap=None, idx=None, add_legend=False):
         cmap : `~matplotlib.colors.ListedColormap`, optional
             Color map to use
         idx : int, optional
-            Observations to include in the plot, default: all
+            Observations to include in the plot.
         add_legend : boolean, optional
-            Enable/disable legend in the plot, default: False
+            Enable/disable legend in the plot.
         """
         import matplotlib.pyplot as plt
 
@@ -414,3 +420,126 @@ def plot(self, fig=None, ax=None, cmap=None, idx=None, add_legend=False):
             ax.legend(handles=handles)
 
         return fig, ax, cbar
+
+
+class ReflectedRegionsBackgroundMaker:
+    """Reflected regions background maker.
+
+    Parameters
+    ----------
+    region: `~regions.SkyRegion`
+        On region to compute off regions for.
+    angle_increment : `~astropy.coordinates.Angle`, optional
+        Rotation angle applied when a region falls in an excluded region.
+    min_distance : `~astropy.coordinates.Angle`, optional
+        Minimal distance between two consecutive reflected regions
+    min_distance_input : `~astropy.coordinates.Angle`, optional
+        Minimal distance from input region
+    max_region_number : int, optional
+        Maximum number of regions to use
+    exclusion_mask : `~gammapy.maps.WcsNDMap`, optional
+        Exclusion mask
+    binsz : `~astropy.coordinates.Angle`
+        Bin size of the reference map used for region finding.
+    """
+
+    def __init__(
+        self,
+        region,
+        angle_increment="0.1 rad",
+        min_distance="0 rad",
+        min_distance_input="0.1 rad",
+        max_region_number=10000,
+        exclusion_mask=None,
+        binsz="0.01 deg",
+    ):
+        self.region = region
+        self.binsz = binsz
+        self.exclusion_mask = exclusion_mask
+        self.angle_increment = Angle(angle_increment)
+        self.min_distance = Angle(min_distance)
+        self.min_distance_input = Angle(min_distance_input)
+        self.max_region_number = max_region_number
+
+    def _get_finder(self, observation):
+        return ReflectedRegionsFinder(
+            binsz=self.binsz,
+            exclusion_mask=self.exclusion_mask,
+            center=observation.pointing_radec,
+            region=self.region,
+            min_distance=self.min_distance,
+            min_distance_input=self.min_distance_input,
+            max_region_number=self.max_region_number,
+            angle_increment=self.angle_increment,
+        )
+
+    def make_counts_off(self, dataset, observation):
+        """Make off counts.
+
+        Parameters
+        ----------
+        dataset : `SpectrumDataset`
+            Spectrum dataset.
+        observation : `DatastoreObservation`
+            Data store observation.
+
+
+        Returns
+        -------
+        counts_off : `CountsSpectrum`
+            Off counts.
+        """
+        finder = self._get_finder(observation)
+
+        finder.run()
+        regions = finder.reflected_regions
+
+        region_union = regions[0]
+
+        for region in regions[1:]:
+            region_union = region_union.union(region)
+
+        wcs = finder.reference_map.geom.wcs
+        events_off = observation.events.select_region(region_union, wcs)
+
+        edges = dataset.counts.energy.edges
+        counts_off = CountsSpectrum(
+            energy_hi=edges[1:],
+            energy_lo=edges[:-1],
+        )
+        counts_off.fill_events(events_off)
+
+        # TODO: temporarily attach the region info to the CountsSpectrum.
+        #  but CountsSpectrum should be modified to hold this info anyway.
+        counts_off.regions = regions
+        return counts_off
+
+    def run(self, dataset, observation):
+        """Run reflected regions background maker
+
+        Parameters
+        ----------
+        dataset : `SpectrumDataset`
+            Spectrum dataset.
+        observation : `DatastoreObservation`
+            Data store observation.
+
+        Returns
+        -------
+        dataset_on_off : `SpectrumDatasetOnOff`
+            On off dataset.
+        """
+        counts_off = self.make_counts_off(dataset, observation)
+
+        return SpectrumDatasetOnOff(
+            counts=dataset.counts,
+            counts_off=counts_off,
+            gti=dataset.gti,
+            name=dataset.name,
+            livetime=dataset.livetime,
+            edisp=dataset.edisp,
+            aeff=dataset.aeff,
+            acceptance_off=len(counts_off.regions),
+            mask_safe=dataset.mask_safe,
+            mask_fit=dataset.mask_fit,
+        )

gammapy/spectrum/tests/test_reflected.py

@@ -1,5 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 import pytest
+import numpy as np
+from numpy.testing import assert_allclose
 import astropy.units as u
 from astropy.coordinates import Angle, SkyCoord
 from regions import (
@@ -10,12 +12,12 @@
 )
 from gammapy.data import DataStore
 from gammapy.maps import WcsGeom, WcsNDMap
-from gammapy.spectrum import ReflectedRegionsBackgroundEstimator, ReflectedRegionsFinder
+from gammapy.spectrum import ReflectedRegionsBackgroundMaker, ReflectedRegionsFinder
+from gammapy.spectrum.make import SpectrumDatasetMaker
 from gammapy.utils.testing import (
     assert_quantity_allclose,
     mpl_plot_check,
     requires_data,
-    requires_dependency,
 )
 
 
@@ -46,17 +48,18 @@ def observations():
     return datastore.get_observations(obs_ids)
 
 
-@pytest.fixture(scope="session")
-def bkg_estimator(observations, exclusion_mask, on_region):
-    """Example background estimator for testing."""
-    maker = ReflectedRegionsBackgroundEstimator(
-        observations=observations,
-        on_region=on_region,
-        exclusion_mask=exclusion_mask,
-        min_distance_input="0.2 deg",
+@pytest.fixture()
+def spectrum_dataset_maker(on_region):
+    e_reco = np.logspace(0, 2, 5) * u.TeV
+    e_true = np.logspace(-0.5, 2, 11) * u.TeV
+    return SpectrumDatasetMaker(region=on_region, e_reco=e_reco, e_true=e_true)
+
+
+@pytest.fixture()
+def reflected_bkg_maker(on_region, exclusion_mask):
+    return ReflectedRegionsBackgroundMaker(
+        region=on_region, exclusion_mask=exclusion_mask
     )
-    maker.run()
-    return maker
 
 
 region_finder_param = [
@@ -161,16 +164,16 @@ def bad_on_region(exclusion_mask, on_region):
 
 
 @requires_data()
-class TestReflectedRegionBackgroundEstimator:
-    def test_basic(self, bkg_estimator):
-        assert "ReflectedRegionsBackgroundEstimator" in str(bkg_estimator)
-
-    def test_run(self, bkg_estimator):
-        assert len(bkg_estimator.result[1].off_region) == 11
-        assert "Reflected" in str(bkg_estimator.result[1])
-
-    @requires_dependency("matplotlib")
-    def test_plot(self, bkg_estimator):
-        with mpl_plot_check():
-            bkg_estimator.plot(idx=1, add_legend=True)
-            bkg_estimator.plot(idx=[0, 1])
+def test_reflected_bkg_maker(spectrum_dataset_maker, reflected_bkg_maker, observations):
+    datasets = []
+
+    for obs in observations:
+        dataset = spectrum_dataset_maker.run(obs, selection=["counts"])
+        dataset_on_off = reflected_bkg_maker.run(dataset, obs)
+        datasets.append(dataset_on_off)
+
+    assert_allclose(datasets[0].counts_off.data.sum(), 76)
+    assert_allclose(datasets[1].counts_off.data.sum(), 60)
+
+    assert_allclose(len(datasets[0].counts_off.regions), 11)
+    assert_allclose(len(datasets[1].counts_off.regions), 11)