Repository: gammapy/gammapy
Pull request #5297 (2024-06-11 07:56:23)

Precompute the PSF Kernel if the PSFMap has only one bin
Improve PSF  kernel computation performance
- precompute psf kernel on the map dataset if the psf map has a region geom
- add a multiscale option on get_psf_kernel, if true it returns a list of kernels with energy dependent max_radius (having instead a single irregular map with different spatial shape for each energy bins is broken and it would take a lot more work to fix it than this patch) 

This should drastically speed-up Fermi-LAT analyses.
For now I only tested it on the dataset that David Green sent to @AtreyeeS
Could you also test it for your analysis @facero ?


adonath: Thanks @QRemy for PR! A couple of general comments first:

- In general I am skeptical about these very specialized performance improvements. They add a lot of code, with high complexity that sooner or later will break and needs to be refactored or removed again for maintenance.
- The list of PSF kernels does not really offer a performance improvement for the actual convolution. The performance of the fft based convolution barely depends on the kernel size. However I would indeed expect that the computation of the kernel is faster. 
- It is almost always better to solve an issue with an abstraction...

So my concrete proposal here would be to work on a minimal version of the `PSFKernelMap` instead (see also #3689). It can remain very simple, so just support a region geom and a single PSF kernel. This would move the computation of the kernel to the data reduction. For the fitting step the kernel is just extracted using indexing. This avoids the re-computation for different model components in the same dataset. To speed up the data reduction you can still use the "multi scale" approach, but instead of returning a list I would just insert into a kernel map and zero pad for the small kernels. For Fermi analysis there can also be a simple way to just create a kernel map from a normal PSF map in the gtpsf format. 
QRemy: making the upsampling factor energy dependent helps a lot, the kernel computation is no longer the bottleneck
<img width="1461" alt="Screenshot 2024-06-12 at 23 28 03" src="https://github.com/gammapy/gammapy/assets/25845654/0ec71743-756f-4332-9cb3-fc2125a45877">
 
facero: Hey,
I tested the computation time with your latest PR of this morning and it's an amazing speed improvment !
I have more than 30 sources in my  Fermi ROI and previously the first Npred call was taking several minutes.
Now it's about a few sec !
Then the fit take a similar time which is expected as I fit only spectral parameters.
But this would have a huge impact if I were to fit the spatial parameters as well.

QRemy: Thanks @facero for the feedback.
The computation time for the flux points could be further improved. 
For now the slicing of the dataset in energy range  of each flux points impose to reset the models and clear the psf convolution cache. Using making instead of slicing to select the energy range would prevent that
https://github.com/gammapy/gammapy/blob/ddf72e06a15cf995363c86d1acb9dce4eabdbb91/gammapy/estimators/points/sed.py#L201-L206
This will be a separate PR.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/gammapy/gammapy/pull/5297?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 `97.05882%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 94.46%. Comparing base [(`cfe0659`)](https://app.codecov.io/gh/gammapy/gammapy/commit/cfe0659b3f407dde6054588fe60d8ade424a38e1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy) to head [(`02fbf2d`)](https://app.codecov.io/gh/gammapy/gammapy/commit/02fbf2d24c35a4376e863b9b56e1da1ddcc405cd?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy).
> Report is 3 commits behind head on main.

| [Files](https://app.codecov.io/gh/gammapy/gammapy/pull/5297?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/map.py](https://app.codecov.io/gh/gammapy/gammapy/pull/5297?src=pr&el=tree&filepath=gammapy%2Fdatasets%2Fmap.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gammapy#diff-Z2FtbWFweS9kYXRhc2V0cy9tYXAucHk=) | 90.00% | [1 Missing :warning: ](https://app.codecov.io/gh/gammapy/gammapy/pull/5297?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    #5297      +/-   ##
==========================================
- Coverage   94.46%   94.46%   -0.01%     
==========================================
  Files         232      232              
  Lines       35650    35678      +28     
==========================================
+ Hits        33678    33704      +26     
- Misses       1972     1974       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/gammapy/gammapy/pull/5297?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).

registerrier: How much backporting is needed here? Performance improvement wouldn't require a backport a priori.
QRemy: > How much backporting is needed here? Performance improvement wouldn't require a backport a priori.

At least for 1.2 it should go in without much conflicts.

Patch:
diff --git a/gammapy/datasets/evaluator.py b/gammapy/datasets/evaluator.py
index 7961da006..c246517e6 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
@@ -204,9 +204,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
             geom_psf = geom if energy_name == "energy" else exposure.geom
 
@@ -215,9 +220,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 8fc890730..14761daa7 100644
--- a/gammapy/datasets/map.py
+++ b/gammapy/datasets/map.py
@@ -420,7 +420,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)):
@@ -439,6 +438,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
@@ -539,15 +549,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 41e8f3168..4482fccf7 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 4a0c75495..3115d5133 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 62acc434c..bde72a9bd 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,
 )
@@ -1944,6 +1945,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)
@@ -1957,6 +1960,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]
     )