Repository: gammapy/gammapy
Pull request #5365 (2024-07-05 23:10:59)

Backport PR #5297 on branch v1.2.x (Precompute the PSF Kernel if the PSFMap has only one bin)
Backport PR #5297: Precompute the PSF Kernel if the PSFMap has only one bin


Patch:
diff --git a/gammapy/datasets/evaluator.py b/gammapy/datasets/evaluator.py
index 3631b6a69..f88981ec3 100644
--- a/gammapy/datasets/evaluator.py
+++ b/gammapy/datasets/evaluator.py
@@ -7,7 +7,7 @@ from astropy.coordinates import angular_separation
 from astropy.utils import lazyproperty
 from regions import CircleSkyRegion
 import matplotlib.pyplot as plt
-from gammapy.irf import EDispKernel
+from gammapy.irf import EDispKernel, PSFKernel
 from gammapy.maps import HpxNDMap, Map, RegionNDMap, WcsNDMap
 from gammapy.modeling.models import PointSpatialModel, TemplateNPredModel
 from .utils import apply_edisp
@@ -205,9 +205,14 @@ class MapEvaluator:
             self.edisp = edisp.get_edisp_kernel(
                 position=self.position, energy_axis=energy_axis
             )
+            del self._edisp_diagonal
 
         # lookup psf
-        if psf and self.model.spatial_model:
+        if (
+            psf
+            and self.model.spatial_model
+            and not (isinstance(self.psf, PSFKernel) and psf.has_single_spatial_bin)
+        ):
             energy_name = psf.energy_name
             if energy_name == "energy":
                 geom_psf = geom
@@ -219,9 +224,6 @@ class MapEvaluator:
                 kwargs = {energy_name: energy_values, "rad": geom.region.radius}
                 self.psf_containment = psf.containment(**kwargs)
             else:
-                if geom_psf.is_region or geom_psf.is_hpx:
-                    geom_psf = geom_psf.to_wcs_geom()
-
                 self.psf = psf.get_psf_kernel(
                     position=self.position,
                     geom=geom_psf,
diff --git a/gammapy/datasets/map.py b/gammapy/datasets/map.py
index 9660eb5da..1ad87d19d 100644
--- a/gammapy/datasets/map.py
+++ b/gammapy/datasets/map.py
@@ -410,7 +410,6 @@ class MapDataset(Dataset):
             raise ValueError(
                 f"'psf' must be a 'PSFMap' or `HDULocation` object, got {type(psf)}"
             )
-
         self.psf = psf
 
         if edisp and not isinstance(edisp, (EDispMap, EDispKernelMap, HDULocation)):
@@ -429,6 +428,17 @@ class MapDataset(Dataset):
         else:
             self._meta = meta
 
+    @property
+    def _psf_kernel(self):
+        """Precompute PSFkernel if there is only one spatial bin in the PSFmap"""
+        if self.psf and self.psf.has_single_spatial_bin:
+            if self.psf.energy_name == "energy_true":
+                map_ref = self.exposure
+            else:
+                map_ref = self.counts
+            if map_ref and not map_ref.geom.is_region:
+                return self.psf.get_psf_kernel(map_ref.geom)
+
     @property
     def meta(self):
         return self._meta
@@ -529,15 +539,16 @@ class MapDataset(Dataset):
     def models(self, models):
         """Models setter."""
         self._evaluators = {}
-
         if models is not None:
             models = DatasetModels(models)
             models = models.select(datasets_names=self.name)
-
+            if models:
+                psf = self._psf_kernel
             for model in models:
                 if not isinstance(model, FoVBackgroundModel):
                     evaluator = MapEvaluator(
                         model=model,
+                        psf=psf,
                         evaluation_mode=EVALUATION_MODE,
                         gti=self.gti,
                         use_cache=USE_NPRED_CACHE,
diff --git a/gammapy/irf/core.py b/gammapy/irf/core.py
index 42dda97d5..8097863c8 100644
--- a/gammapy/irf/core.py
+++ b/gammapy/irf/core.py
@@ -662,6 +662,10 @@ class IRFMap:
     def required_axes(self):
         pass
 
+    @lazyproperty
+    def has_single_spatial_bin(self):
+        return self._irf_map.geom.to_image().data_shape == (1, 1)
+
     # TODO: add mask safe to IRFMap as a regular attribute and don't derive it from the data
     @property
     def mask_safe_image(self):
diff --git a/gammapy/irf/edisp/map.py b/gammapy/irf/edisp/map.py
index 9f4b4e90f..4f3e30225 100644
--- a/gammapy/irf/edisp/map.py
+++ b/gammapy/irf/edisp/map.py
@@ -88,6 +88,7 @@ class EDispMap(IRFMap):
 
     @edisp_map.setter
     def edisp_map(self, value):
+        del self.has_single_spatial_bin
         self._irf_map = value
 
     def normalize(self):
diff --git a/gammapy/irf/psf/map.py b/gammapy/irf/psf/map.py
index 900035f4c..88eaeec1f 100644
--- a/gammapy/irf/psf/map.py
+++ b/gammapy/irf/psf/map.py
@@ -109,6 +109,7 @@ class PSFMap(IRFMap):
 
     @psf_map.setter
     def psf_map(self, value):
+        del self.has_single_spatial_bin
         self._irf_map = value
 
     def normalize(self):


Test patch:
diff --git a/gammapy/datasets/tests/test_map.py b/gammapy/datasets/tests/test_map.py
index 0cdbccf71..38f232e33 100644
--- a/gammapy/datasets/tests/test_map.py
+++ b/gammapy/datasets/tests/test_map.py
@@ -21,6 +21,7 @@ from gammapy.irf import (
     EffectiveAreaTable2D,
     EnergyDependentMultiGaussPSF,
     EnergyDispersion2D,
+    PSFKernel,
     PSFMap,
     RecoPSFMap,
 )
@@ -1909,6 +1910,8 @@ def test_dataset_mixed_geom(tmpdir):
     dataset = MapDataset.from_geoms(
         geom=geom, geom_exposure=geom_exposure, geom_psf=geom_psf, geom_edisp=geom_edisp
     )
+    assert isinstance(dataset.psf, PSFMap)
+    assert isinstance(dataset._psf_kernel, PSFKernel)
 
     filename = tmpdir / "test.fits"
     dataset.write(filename)
@@ -1922,6 +1925,33 @@ def test_dataset_mixed_geom(tmpdir):
     assert isinstance(dataset.psf.psf_map.geom.region, CircleSkyRegion)
     assert isinstance(dataset.edisp.edisp_map.geom.region, CircleSkyRegion)
 
+    assert isinstance(dataset.psf, PSFMap)
+    assert dataset.psf.has_single_spatial_bin
+    assert isinstance(dataset._psf_kernel, PSFKernel)
+
+    geom_psf = WcsGeom.create(npix=1, axes=[rad_axis, energy_axis_true])
+    dataset = MapDataset.from_geoms(
+        geom=geom, geom_exposure=geom_exposure, geom_psf=geom_psf, geom_edisp=geom_edisp
+    )
+
+    assert isinstance(dataset.psf, PSFMap)
+    assert dataset.psf.has_single_spatial_bin
+    assert isinstance(dataset._psf_kernel, PSFKernel)
+    psf_1bin = dataset.psf.copy()
+
+    geom_psf = WcsGeom.create(npix=3, axes=[rad_axis, energy_axis_true])
+    dataset = MapDataset.from_geoms(
+        geom=geom, geom_exposure=geom_exposure, geom_psf=geom_psf, geom_edisp=geom_edisp
+    )
+
+    assert isinstance(dataset.psf, PSFMap)
+    assert not dataset.psf.has_single_spatial_bin
+    assert dataset._psf_kernel is None
+
+    dataset.psf.psf_map = psf_1bin.psf_map
+    assert dataset.psf.has_single_spatial_bin
+    assert isinstance(dataset._psf_kernel, PSFKernel)
+
     geom_psf_reco = RegionGeom.create(
         "icrs;circle(0, 0, 0.2)", axes=[rad_axis, energy_axis]
     )