Merge pull request #1846 from AtreyeeS/edisp_28_9

Allow different energy binning for true and reconstructed energy in Maps

gammapy/cube/fit.py

@@ -5,7 +5,7 @@
 import astropy.units as u
 from ..utils.fitting import Fit
 from ..stats import cash
-from ..maps import Map
+from ..maps import Map, MapAxis
 
 __all__ = ["MapFit", "MapEvaluator"]
 
@@ -118,6 +118,7 @@ def __init__(
 
     @lazyproperty
     def geom(self):
+        """This will give the energy axes in e_true"""
         return self.exposure.geom
 
     @lazyproperty
@@ -126,15 +127,15 @@ def geom_image(self):
 
     @lazyproperty
     def energy_center(self):
-        """Energy axis bin centers (`~astropy.units.Quantity`)"""
-        energy_axis = self.geom.axes[0]
+        """True energy axis bin centers (`~astropy.units.Quantity`)"""
+        energy_axis = self.geom.get_axis_by_name("energy")
         energy = energy_axis.center * energy_axis.unit
         return energy[:, np.newaxis, np.newaxis]
 
     @lazyproperty
     def energy_edges(self):
         """Energy axis bin edges (`~astropy.units.Quantity`)"""
-        energy_axis = self.geom.axes[0]
+        energy_axis = self.geom.get_axis_by_name("energy")
         energy = energy_axis.edges * energy_axis.unit
         return energy[:, np.newaxis, np.newaxis]
 
@@ -208,22 +209,46 @@ def apply_psf(self, npred):
         """Convolve npred cube with PSF"""
         return npred.convolve(self.psf)
 
-    def apply_edisp(self, data):
-        """Convolve map data with energy dispersion."""
-        data = np.rollaxis(data, 0, 3)
+    def apply_edisp(self, npred):
+        """Convolve map data with energy dispersion.
+
+        Parameters
+        ----------
+        npred : `~gammapy.maps.Map`
+            Predicted counts in true energy bins
+
+        Returns
+        ---------
+        npred_reco : `~gammapy.maps.Map`
+            Predicted counts in reco energy bins
+        """
+        loc = npred.geom.get_axis_index_by_name("energy")
+        data = np.rollaxis(npred.data, loc, len(npred.data))
         data = np.dot(data, self.edisp.pdf_matrix)
-        return np.rollaxis(data, 2, 0)
+        data = np.rollaxis(data, -1, loc)
+        e_reco_axis = MapAxis.from_edges(
+            self.edisp.e_reco.bins, unit=self.edisp.e_reco.unit
+        )
+        geom_ereco = self.exposure.geom.to_image().to_cube(axes=[e_reco_axis])
+        npred = Map.from_geom(geom_ereco, unit="")
+        npred.data = data
+        return npred
 
     def compute_npred(self):
-        """Evaluate model predicted counts.
+        """
+        Evaluate model predicted counts.
+
+        Returns
+        -------
+        npred.data : ~numpy.ndarray
+            array of the predicted counts in each bin (in reco energy)
         """
         flux = self.compute_flux()
         npred = self.apply_exposure(flux)
         if self.psf is not None:
             npred = self.apply_psf(npred)
-        # TODO: discuss and decide whether we need to make map objects in `apply_aeff` and `apply_psf`.
         if self.edisp is not None:
-            npred.data = self.apply_edisp(npred.data)
+            npred = self.apply_edisp(npred)
         if self.background:
             npred.data += self.background.data
         return npred.data

gammapy/cube/make.py

@@ -1,6 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 from __future__ import absolute_import, division, print_function, unicode_literals
 import logging
+import numpy as np
 from astropy.nddata.utils import NoOverlapError
 from astropy.coordinates import Angle
 from ..maps import Map, WcsGeom
@@ -19,21 +20,25 @@ class MapMaker(object):
     Parameters
     ----------
     geom : `~gammapy.maps.WcsGeom`
-        Reference image geometry
+        Reference image geometry in reco energy
     offset_max : `~astropy.coordinates.Angle`
         Maximum offset angle
+    geom_true : `~gammapy.maps.WcsGeom`
+        Reference image geometry in true energy, used for exposure maps and PSF.
+         If none, the same as geom is assumed
     exclusion_mask : `~gammapy.maps.Map`
         Exclusion mask
     """
 
-    def __init__(self, geom, offset_max, exclusion_mask=None):
+    def __init__(self, geom, offset_max, geom_true=None, exclusion_mask=None):
         if not isinstance(geom, WcsGeom):
             raise ValueError("MapMaker only works with WcsGeom")
 
         if geom.is_image:
             raise ValueError("MapMaker only works with geom with an energy axis")
 
         self.geom = geom
+        self.geom_true = geom_true if geom_true else geom
         self.offset_max = Angle(offset_max)
         self.maps = {}
 
@@ -63,8 +68,10 @@ def run(self, obs_list, selection=None):
 
         # Initialise zero-filled maps
         for name in selection:
-            unit = "m2 s" if name == "exposure" else ""
-            self.maps[name] = Map.from_geom(self.geom, unit=unit)
+            if name == "exposure":
+                self.maps[name] = Map.from_geom(self.geom_true, unit="m2 s")
+            else:
+                self.maps[name] = Map.from_geom(self.geom, unit="")
 
         for obs in obs_list:
             try:
@@ -82,6 +89,9 @@ def _process_obs(self, obs, selection):
         cutout_map = Map.from_geom(self.geom).cutout(
             position=obs.pointing_radec, width=2 * self.offset_max, mode="trim"
         )
+        cutout_map_etrue = Map.from_geom(self.geom_true).cutout(
+            position=obs.pointing_radec, width=2 * self.offset_max, mode="trim"
+        )
 
         log.info("Processing observation: OBS_ID = {}".format(obs.obs_id))
 
@@ -90,7 +100,13 @@ def _process_obs(self, obs, selection):
         offset = coords.skycoord.separation(obs.pointing_radec)
         fov_mask = offset >= self.offset_max
 
+        # Compute field of view mask on the cutout in true energy
+        coords_etrue = cutout_map_etrue.geom.get_coord()
+        offset_etrue = coords_etrue.skycoord.separation(obs.pointing_radec)
+        fov_mask_etrue = offset_etrue >= self.offset_max
+
         # Only if there is an exclusion mask, make a cutout
+        # Exclusion mask only on the background, so only in reco-energy
         exclusion_mask = self.maps.get("exclusion", None)
         if exclusion_mask is not None:
             exclusion_mask = exclusion_mask.cutout(
@@ -101,14 +117,19 @@ def _process_obs(self, obs, selection):
         maps_obs = MapMakerObs(
             obs=obs,
             geom=cutout_map.geom,
+            geom_true=cutout_map_etrue.geom,
             fov_mask=fov_mask,
+            fov_mask_etrue=fov_mask_etrue,
             exclusion_mask=exclusion_mask,
         ).run(selection)
 
         # Stack observation maps to total
         for name in selection:
             data = maps_obs[name].quantity.to(self.maps[name].unit).value
-            self.maps[name].fill_by_coord(coords, data)
+            if name == "exposure":
+                self.maps[name].fill_by_coord(coords_etrue, data)
+            else:
+                self.maps[name].fill_by_coord(coords, data)
 
     def make_images(self, spectrum=None):
         """Create 2D images by summing over the energy axis.
@@ -146,16 +167,29 @@ class MapMakerObs(object):
         Observation
     geom : `~gammapy.maps.WcsGeom`
         Reference image geometry
+    geom_true : `~gammapy.maps.WcsGeom`
+        Reference image geometry in true energy, used for exposure maps and PSF.
+        If none, the same as geom is assumed
     fov_mask : `~numpy.ndarray`
         Mask to select pixels in field of view
     exclusion_mask : `~gammapy.maps.Map`
         Exclusion mask (used by some background estimators)
     """
 
-    def __init__(self, obs, geom, fov_mask=None, exclusion_mask=None):
+    def __init__(
+        self,
+        obs,
+        geom,
+        geom_true=None,
+        fov_mask=None,
+        fov_mask_etrue=None,
+        exclusion_mask=None,
+    ):
         self.obs = obs
         self.geom = geom
+        self.geom_true = geom_true if geom_true else geom
         self.fov_mask = fov_mask
+        self.fov_mask_etrue = fov_mask_etrue
         self.exclusion_mask = exclusion_mask
         self.maps = {}
 
@@ -190,10 +224,10 @@ def _make_exposure(self):
             pointing=self.obs.pointing_radec,
             livetime=self.obs.observation_live_time_duration,
             aeff=self.obs.aeff,
-            geom=self.geom,
+            geom=self.geom_true,
         )
-        if self.fov_mask is not None:
-            exposure.data[..., self.fov_mask] = 0
+        if self.fov_mask_etrue is not None:
+            exposure.data[..., self.fov_mask_etrue] = 0
         self.maps["exposure"] = exposure
 
     def _make_background(self):

gammapy/cube/tests/test_fit.py

@@ -25,15 +25,23 @@ def geom():
 
 
 @pytest.fixture(scope="session")
-def exposure(geom):
+def geom_etrue():
+    axis = MapAxis.from_edges(np.logspace(-1., 1., 4), name="energy", unit=u.TeV)
+    return WcsGeom.create(
+        skydir=(0, 0), binsz=0.02, width=(2, 2), coordsys="GAL", axes=[axis]
+    )
+
+
+@pytest.fixture(scope="session")
+def exposure(geom_etrue):
     filename = "$GAMMAPY_EXTRA/datasets/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits"
     aeff = EffectiveAreaTable2D.read(filename, hdu="EFFECTIVE AREA")
 
     exposure_map = make_map_exposure_true_energy(
         pointing=SkyCoord(1, 0.5, unit="deg", frame="galactic"),
         livetime="1 hour",
         aeff=aeff,
-        geom=geom,
+        geom=geom_etrue,
     )
     return exposure_map
 
@@ -46,18 +54,19 @@ def background(geom):
 
 
 @pytest.fixture(scope="session")
-def edisp(geom):
-    e_true = geom.get_axis_by_name("energy").edges
-    return EnergyDispersion.from_diagonal_response(e_true=e_true)
+def edisp(geom, geom_etrue):
+    e_true = geom_etrue.get_axis_by_name("energy").edges
+    e_reco = geom.get_axis_by_name("energy").edges
+    return EnergyDispersion.from_diagonal_response(e_true=e_true, e_reco=e_reco)
 
 
 @pytest.fixture(scope="session")
-def psf(geom):
+def psf(geom_etrue):
     filename = "$GAMMAPY_EXTRA/datasets/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits"
     psf = EnergyDependentMultiGaussPSF.read(filename, hdu="POINT SPREAD FUNCTION")
 
     table_psf = psf.to_energy_dependent_table_psf(theta=0.5 * u.deg)
-    psf_kernel = PSFKernel.from_table_psf(table_psf, geom, max_radius=0.5 * u.deg)
+    psf_kernel = PSFKernel.from_table_psf(table_psf, geom_etrue, max_radius=0.5 * u.deg)
     return psf_kernel
 
 
@@ -85,7 +94,7 @@ def counts(sky_model, exposure, background, psf, edisp):
         model=sky_model, exposure=exposure, background=background, psf=psf, edisp=edisp
     )
     npred = evaluator.compute_npred()
-    return WcsNDMap(exposure.geom, npred)
+    return WcsNDMap(background.geom, npred)
 
 
 @requires_dependency("scipy")
@@ -113,18 +122,18 @@ def test_cube_fit(sky_model, counts, exposure, psf, background, mask, edisp):
     assert result.success
     assert "minuit" in repr(result)
 
-    stat_expected = 3840.0605649268496
+    stat_expected = 5417.350078
     assert_allclose(result.total_stat, stat_expected, rtol=1e-2)
 
     pars = result.model.parameters
     assert_allclose(pars["lon_0"].value, 0.2, rtol=1e-2)
-    assert_allclose(pars.error("lon_0"), 0.005895, rtol=1e-2)
+    assert_allclose(pars.error("lon_0"), 0.004177, rtol=1e-2)
 
     assert_allclose(pars["index"].value, 3, rtol=1e-2)
-    assert_allclose(pars.error("index"), 0.05614, rtol=1e-2)
+    assert_allclose(pars.error("index"), 0.033947, rtol=1e-2)
 
     assert_allclose(pars["amplitude"].value, 1e-11, rtol=1e-2)
-    assert_allclose(pars.error("amplitude"), 3.936e-13, rtol=1e-2)
+    assert_allclose(pars.error("amplitude"), 4.03049e-13, rtol=1e-2)
 
     assert result.model.spectral_model.parameters.covariance is not None
     assert result.model.spatial_model.parameters.covariance is not None

gammapy/cube/tests/test_make.py

@@ -38,6 +38,7 @@ def geom(ebounds):
         {
             # Default, normal test case
             "geom": geom(ebounds=[0.1, 1, 10]),
+            "geom_true": None,
             "counts": 34366,
             "exposure": 3.99815e+11,
             "exposure_image": 7.921993e+10,
@@ -46,24 +47,35 @@ def geom(ebounds):
         {
             # Test single energy bin
             "geom": geom(ebounds=[0.1, 10]),
+            "geom_true": None,
             "counts": 34366,
             "exposure": 1.16866e+11,
             "exposure_image": 1.16866e+11,
             "background": 1988492.8,
         },
         {
-            # Test single energy bin
+            # Test single energy bin with exclusion mask
             "geom": geom(ebounds=[0.1, 10]),
+            "geom_true": None,
             "exclusion_mask": Map.from_geom(geom(ebounds=[0.1, 10])),
             "counts": 34366,
             "exposure": 1.16866e+11,
             "exposure_image": 1.16866e+11,
             "background": 1988492.8,
         },
+        {
+            # Test for different e_true and e_reco bins
+            "geom": geom(ebounds=[0.1, 1, 10]),
+            "geom_true": geom(ebounds=[0.1, 0.5, 2.5, 10.0]),
+            "counts": 34366,
+            "exposure": 5.971096e+11,
+            "exposure_image": 6.492968e+10,
+            "background": 187528.89,
+        },
     ],
 )
 def test_map_maker(pars, obs_list):
-    maker = MapMaker(geom=pars["geom"], offset_max="2 deg")
+    maker = MapMaker(geom=pars["geom"], geom_true=pars["geom_true"], offset_max="2 deg")
 
     maps = maker.run(obs_list)