Improve EDispMap and PSFMap stacking

gammapy/cube/edisp_map.py

@@ -81,8 +81,8 @@ class EDispMap(object):
         from astropy import units as u
         from astropy.coordinates import SkyCoord
         from gammapy.maps import Map, WcsGeom, MapAxis
-        from gammapy.irf import EnergyDispersion2D
-        from gammapy.cube import make_edisp_map, EDispMap
+        from gammapy.irf import EnergyDispersion2D, EffectiveAreaTable2D
+        from gammapy.cube import make_edisp_map, EDispMap, make_map_exposure_true_energy
 
         # Define energy axis. Note that the name is fixed.
         energy_axis = MapAxis.from_edges(np.logspace(-1., 1., 4), unit='TeV', name='energy')
@@ -100,9 +100,14 @@ class EDispMap(object):
         # Extract EnergyDispersion2D from CTA 1DC IRF
         filename = '$GAMMAPY_DATA/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits'
         edisp2D = EnergyDispersion2D.read(filename, hdu='ENERGY DISPERSION')
+        aeff2d = EffectiveAreaTable2D.read(filename, hdu='EFFECTIVE AREA')
+
+        # Create the exposure map
+        exposure_geom = geom.to_image().to_cube([energy_axis])
+        exposure_map = make_map_exposure_true_energy(pointing, "1 h", aeff2d, exposure_geom)
 
         # create the EDispMap for the specified pointing
-        edisp_map = make_edisp_map(edisp2D, pointing, geom, max_offset)
+        edisp_map = make_edisp_map(edisp2D, pointing, geom, max_offset, exposure_map)
 
         # Get an Energy Dispersion (1D) at any position in the image
         edisp = edisp_map.get_energy_dispersion(SkyCoord(2., 2.5, unit='deg'))
@@ -118,7 +123,7 @@ def __init__(self, edisp_map, exposure_map):
         if edisp_map.geom.axes[0].name.upper() != "MIGRA":
             raise ValueError("Incorrect migra axis position in input Map")
 
-        self._edisp_map = edisp_map
+        self.edisp_map = edisp_map
 
         if exposure_map is not None:
             # First adapt geometry, keep only energy axis
@@ -128,32 +133,7 @@ def __init__(self, edisp_map, exposure_map):
                     "EDispMap and exposure_map have inconsistent geometries"
                 )
 
-        self._exposure_map = exposure_map
-
-    @property
-    def edisp_map(self):
-        """the EDispMap itself (`~gammapy.maps.Map`)"""
-        return self._edisp_map
-
-    @property
-    def data(self):
-        """the EDispMap data"""
-        return self._edisp_map.data
-
-    @property
-    def quantity(self):
-        """the EDispMap data as a quantity"""
-        return self._edisp_map.quantity
-
-    @property
-    def exposure_map(self):
-        """the exposure map associated to the EDispMap."""
-        return self._exposure_map
-
-    @property
-    def geom(self):
-        """The EDispMap MapGeom object"""
-        return self._edisp_map.geom
+        self.exposure_map = exposure_map
 
     @classmethod
     def from_hdulist(
@@ -254,13 +234,13 @@ def get_energy_dispersion(self, position, e_reco, migra_step=5e-3):
             )
 
         # axes ordering fixed. Could be changed.
-        pix_ener = np.arange(self.geom.axes[1].nbin)
+        pix_ener = np.arange(self.edisp_map.geom.axes[1].nbin)
 
         # Define a vector of migration with mig_step step
-        mrec_min = self.geom.axes[0].edges[0]
-        mrec_max = self.geom.axes[0].edges[-1]
+        mrec_min = self.edisp_map.geom.axes[0].edges[0]
+        mrec_max = self.edisp_map.geom.axes[0].edges[-1]
         mig_array = np.arange(mrec_min, mrec_max, migra_step)
-        pix_migra = (mig_array - mrec_min) / mrec_max * self.geom.axes[0].nbin
+        pix_migra = (mig_array - mrec_min) / mrec_max * self.edisp_map.geom.axes[0].nbin
 
         # Convert position to pixels
         pix_lon, pix_lat = self.edisp_map.geom.to_image().coord_to_pix(position)
@@ -314,16 +294,17 @@ def get_energy_dispersion(self, position, e_reco, migra_step=5e-3):
             data=data,
         )
 
-    def stack(self, other):
+    def stack(self, other, copy=True):
         """Stack EdispMap with another one.
 
         The current EdispMap is unchanged and a new one is created and returned.
-        For the moment, this works only if the EDispMap to be stacked contain compatible exposure maps.
 
         Parameters
         ----------
         other : `~gammapy.cube.EDispMap`
             the edispmap to be stacked with this one.
+        copy : bool
+            if set to True returns a new EdispMap
 
         Returns
         -------
@@ -333,13 +314,38 @@ def stack(self, other):
         if self.exposure_map is None or other.exposure_map is None:
             raise ValueError("Missing exposure map for EdispMap.stack")
 
-        total_exposure = self.exposure_map + other.exposure_map
-        exposure = self.exposure_map.quantity[:, np.newaxis, :, :]
-        stacked_edisp_quantity = self.quantity * exposure
-        other_exposure = other.exposure_map.quantity[:, np.newaxis, :, :]
-        stacked_edisp_quantity += other.quantity * other_exposure
-        stacked_edisp_quantity /= total_exposure.quantity[:, np.newaxis, :, :]
-        stacked_edisp = Map.from_geom(
-            self.geom, data=stacked_edisp_quantity.to("").value, unit=""
+        # Reproject other exposure
+        exposure_coord = self.exposure_map.geom.get_coord()
+        reproj_exposure = Map.from_geom(
+            self.exposure_map.geom, unit=self.exposure_map.unit
         )
-        return EDispMap(stacked_edisp, total_exposure)
+        reproj_exposure.fill_by_coord(
+            exposure_coord, other.exposure_map.get_by_coord(exposure_coord)
+        )
+
+        # Reproject other psfmap using same geom
+        edispmap_coord = self.edisp_map.geom.get_coord()
+        reproj_edispmap = Map.from_geom(self.edisp_map.geom, unit=self.edisp_map.unit)
+        reproj_edispmap.fill_by_coord(
+            edispmap_coord, other.edisp_map.get_by_coord(edispmap_coord)
+        )
+
+        exposure = self.exposure_map.quantity[:, np.newaxis, :, :]
+        stacked_edisp_quantity = self.edisp_map.quantity * exposure
+
+        other_exposure = reproj_exposure.quantity[:, np.newaxis, :, :]
+        stacked_edisp_quantity += reproj_edispmap.quantity * other_exposure
+
+        total_exposure = exposure + other_exposure
+        stacked_edisp_quantity /= total_exposure
+
+        reproj_edispmap.quantity = stacked_edisp_quantity
+        # We need to remove the extra axis in the total exposure
+        reproj_exposure.quantity = total_exposure[:, 0, :, :]
+
+        if copy:
+            return EDispMap(reproj_edispmap, reproj_exposure)
+        else:
+            self.edisp_map = reproj_edispmap
+            self.exposure_map = reproj_exposure
+            return self

gammapy/cube/psf_map.py

@@ -79,8 +79,8 @@ class PSFMap:
         from astropy import units as u
         from astropy.coordinates import SkyCoord
         from gammapy.maps import Map, WcsGeom, MapAxis
-        from gammapy.irf import EnergyDependentMultiGaussPSF
-        from gammapy.cube import make_psf_map, PSFMap
+        from gammapy.irf import EnergyDependentMultiGaussPSF, EffectiveAreaTable2D
+        from gammapy.cube import make_psf_map, PSFMap, make_map_exposure_true_energy
 
         # Define energy axis. Note that the name is fixed.
         energy_axis = MapAxis.from_edges(np.logspace(-1., 1., 4), unit='TeV', name='energy')
@@ -99,9 +99,14 @@ class PSFMap:
         filename = '$GAMMAPY_DATA/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits'
         psf = EnergyDependentMultiGaussPSF.read(filename, hdu='POINT SPREAD FUNCTION')
         psf3d = psf.to_psf3d(rads)
+        aeff2d = EffectiveAreaTable2D.read(filename, hdu='EFFECTIVE AREA')
+
+        # Create the exposure map
+        exposure_geom = geom.to_image().to_cube([energy_axis])
+        exposure_map = make_map_exposure_true_energy(pointing, "1 h", aeff2d, exposure_geom)
 
         # create the PSFMap for the specified pointing
-        psf_map = make_psf_map(psf3d, pointing, geom, max_offset)
+        psf_map = make_psf_map(psf3d, pointing, geom, max_offset, exposure_map)
 
         # Get an EnergyDependentTablePSF at any position in the image
         psf_table = psf_map.get_energy_dependent_table_psf(SkyCoord(2., 2.5, unit='deg'))
@@ -117,40 +122,15 @@ def __init__(self, psf_map, exposure_map=None):
         if psf_map.geom.axes[0].name.upper() != "THETA":
             raise ValueError("Incorrect theta axis position in input Map")
 
-        self._psf_map = psf_map
+        self.psf_map = psf_map
 
         if exposure_map is not None:
             # First adapt geometry, keep only energy axis
             expected_geom = psf_map.geom.to_image().to_cube([psf_map.geom.axes[1]])
             if exposure_map.geom != expected_geom:
                 raise ValueError("PSFMap and exposure_map have inconsistent geometries")
 
-        self._exposure_map = exposure_map
-
-    @property
-    def psf_map(self):
-        """the PSFMap itself (`~gammapy.maps.Map`)"""
-        return self._psf_map
-
-    @property
-    def data(self):
-        """the PSFMap data"""
-        return self._psf_map.data
-
-    @property
-    def quantity(self):
-        """the PSFMap data as a quantity"""
-        return self._psf_map.quantity
-
-    @property
-    def exposure_map(self):
-        """the exposure map associated to the PSFMap."""
-        return self._exposure_map
-
-    @property
-    def geom(self):
-        """The PSFMap MapGeom object"""
-        return self._psf_map.geom
+        self.exposure_map = exposure_map
 
     @classmethod
     def from_hdulist(
@@ -246,8 +226,8 @@ def get_energy_dependent_table_psf(self, position):
             )
 
         # axes ordering fixed. Could be changed.
-        pix_ener = np.arange(self.geom.axes[1].nbin)
-        pix_rad = np.arange(self.geom.axes[0].nbin)
+        pix_ener = np.arange(self.psf_map.geom.axes[1].nbin)
+        pix_rad = np.arange(self.psf_map.geom.axes[0].nbin)
 
         # Convert position to pixels
         pix_lon, pix_lat = self.psf_map.geom.to_image().coord_to_pix(position)
@@ -305,8 +285,8 @@ def containment_radius_map(self, energy, fraction=0.68):
         containment_radius_map : `~gammapy.maps.Map`
             Containment radius map
         """
-        coords = self.geom.to_image().get_coord().skycoord.flatten()
-        m = Map.from_geom(self.geom.to_image(), unit="deg")
+        coords = self.psf_map.geom.to_image().get_coord().skycoord.flatten()
+        m = Map.from_geom(self.psf_map.geom.to_image(), unit="deg")
 
         for coord in coords:
             psf_table = self.get_energy_dependent_table_psf(coord)
@@ -315,16 +295,17 @@ def containment_radius_map(self, energy, fraction=0.68):
 
         return m
 
-    def stack(self, other):
+    def stack(self, other, copy=True):
         """Stack PSFMap with another one.
 
-        The current PSFMap is unchanged and a new one is created and returned.
-        For the moment, this works only if the PSFMap to be stacked contain compatible exposure maps.
+        The other PSFMap is projected on the current PSFMap geometry.
 
         Parameters
         ----------
         other : `~gammapy.cube.PSFMap`
             the psfmap to be stacked with this one.
+        copy : bool
+            if set to True returns a new PSFMap
 
         Returns
         -------
@@ -334,13 +315,38 @@ def stack(self, other):
         if self.exposure_map is None or other.exposure_map is None:
             raise ValueError("Missing exposure map for PSFMap.stack")
 
-        total_exposure = self.exposure_map + other.exposure_map
-        exposure = self.exposure_map.quantity[:, np.newaxis, :, :]
-        stacked_psf_quantity = self.quantity * exposure
-        other_exposure = other.exposure_map.quantity[:, np.newaxis, :, :]
-        stacked_psf_quantity += other.quantity * other_exposure
-        stacked_psf_quantity /= total_exposure.quantity[:, np.newaxis, :, :]
-        stacked_psf = Map.from_geom(
-            self.geom, data=stacked_psf_quantity.to("1/sr").value, unit="1/sr"
+        # Reproject other exposure
+        exposure_coord = self.exposure_map.geom.get_coord()
+        reproj_exposure = Map.from_geom(
+            self.exposure_map.geom, unit=self.exposure_map.unit
         )
-        return PSFMap(stacked_psf, total_exposure)
+        reproj_exposure.fill_by_coord(
+            exposure_coord, other.exposure_map.get_by_coord(exposure_coord)
+        )
+
+        # Reproject other psfmap using same geom
+        psfmap_coord = self.psf_map.geom.get_coord()
+        reproj_psfmap = Map.from_geom(self.psf_map.geom, unit=self.psf_map.unit)
+        reproj_psfmap.fill_by_coord(
+            psfmap_coord, other.psf_map.get_by_coord(psfmap_coord)
+        )
+
+        exposure = self.exposure_map.quantity[:, np.newaxis, :, :]
+        stacked_psf_quantity = self.psf_map.quantity * exposure
+
+        other_exposure = reproj_exposure.quantity[:, np.newaxis, :, :]
+        stacked_psf_quantity += reproj_psfmap.quantity * other_exposure
+
+        total_exposure = exposure + other_exposure
+        stacked_psf_quantity /= total_exposure
+
+        reproj_psfmap.quantity = stacked_psf_quantity
+        # We need to remove the extra axis in the total exposure
+        reproj_exposure.quantity = total_exposure[:, 0, :, :]
+
+        if copy:
+            return PSFMap(reproj_psfmap, reproj_exposure)
+        else:
+            self.psf_map = reproj_psfmap
+            self.exposure_map = reproj_exposure
+            return self

gammapy/cube/tests/test_edisp_map.py

@@ -45,12 +45,12 @@ def make_edisp_map_test():
     edisp2d = EnergyDispersion2D.from_gauss(etrue, migra, 0.0, 0.2, offsets)
 
     geom = WcsGeom.create(
-        skydir=pointing, binsz=1., width=5., axes=[migra_axis, energy_axis]
+        skydir=pointing, binsz=1.0, width=5.0, axes=[migra_axis, energy_axis]
     )
 
     aeff2d = fake_aeff2d()
     exposure_geom = WcsGeom.create(
-        skydir=pointing, binsz=1., width=5., axes=[energy_axis]
+        skydir=pointing, binsz=1.0, width=5.0, axes=[energy_axis]
     )
 
     exposure_map = make_map_exposure_true_energy(pointing, "1 h", aeff2d, exposure_geom)
@@ -74,7 +74,7 @@ def test_make_edisp_map():
     assert edmap.edisp_map.geom.axes[0] == migra_axis
     assert edmap.edisp_map.geom.axes[1] == energy_axis
     assert edmap.edisp_map.unit == Unit("")
-    assert edmap.data.shape == (4, 50, 5, 5)
+    assert edmap.edisp_map.data.shape == (4, 50, 5, 5)
 
 
 def test_edisp_map_to_from_hdulist():
@@ -89,7 +89,7 @@ def test_edisp_map_to_from_hdulist():
         hdulist, edisp_hdu="EDISP", edisp_hdubands="BANDSEDISP"
     )
     assert_allclose(edmap.edisp_map.data, new_edmap.edisp_map.data)
-    assert new_edmap.geom == edmap.geom
+    assert new_edmap.edisp_map.geom == edmap.edisp_map.geom
     assert new_edmap.exposure_map.geom == edmap.exposure_map.geom
 
 
@@ -121,6 +121,9 @@ def test_edisp_map_stacking():
     edmap2 = make_edisp_map_test()
     edmap2.exposure_map.quantity *= 2
 
-    edmap_stack = edmap1.stack(edmap2)
-    assert_allclose(edmap_stack.data, edmap1.data)
+    edmap_stack = edmap1.stack(edmap2, True)
+    assert_allclose(edmap_stack.edisp_map.data, edmap1.edisp_map.data)
     assert_allclose(edmap_stack.exposure_map.data, edmap1.exposure_map.data * 3)
+
+    edmap1.stack(edmap2, False)
+    assert_allclose(edmap1.edisp_map.data, edmap_stack.edisp_map.data)