From 7d54ccef0fc2fa8df07fa41423809a93d999cf3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20N=C3=B6the?= <maximilian.noethe@tu-dortmund.de>
Date: Mon, 10 Jan 2022 10:48:28 +0100
Subject: [PATCH] Use gammapy MapAxis and PSF class in PSF function

---
 pyirf/irf/psf.py            | 27 +++++++++++++++++++++------
 pyirf/irf/tests/test_psf.py | 22 ++++++++++++----------
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/pyirf/irf/psf.py b/pyirf/irf/psf.py
index eaa3c6114..5f90f7e29 100644
--- a/pyirf/irf/psf.py
+++ b/pyirf/irf/psf.py
@@ -1,10 +1,17 @@
 import numpy as np
 import astropy.units as u
+from gammapy.irf import PSF3D
+from gammapy.maps import MapAxis
 
 from ..utils import cone_solid_angle
 
 
-def psf_table(events, true_energy_bins, source_offset_bins, fov_offset_bins):
+def psf_table(
+    events,
+    true_energy_axis: MapAxis,
+    source_offset_axis: MapAxis,
+    fov_offset_axis: MapAxis
+):
     """
     Calculate the table based PSF (radially symmetrical bins around the true source)
     """
@@ -20,14 +27,22 @@ def psf_table(events, true_energy_bins, source_offset_bins, fov_offset_bins):
     hist, _ = np.histogramdd(
         array,
         [
-            true_energy_bins.to_value(u.TeV),
-            fov_offset_bins.to_value(u.deg),
-            source_offset_bins.to_value(u.deg),
+            true_energy_axis.edges.to_value(u.TeV),
+            fov_offset_axis.edges.to_value(u.deg),
+            source_offset_axis.edges.to_value(u.deg),
         ],
     )
 
-    psf = _normalize_psf(hist, source_offset_bins)
-    return psf
+    psf = _normalize_psf(hist, source_offset_axis.edges)
+
+    return PSF3D(
+        axes=[
+            true_energy_axis,
+            fov_offset_axis,
+            source_offset_axis,
+        ],
+        data=psf
+    )
 
 
 def _normalize_psf(hist, source_offset_bins):
diff --git a/pyirf/irf/tests/test_psf.py b/pyirf/irf/tests/test_psf.py
index f0384dda6..e7ddf826e 100644
--- a/pyirf/irf/tests/test_psf.py
+++ b/pyirf/irf/tests/test_psf.py
@@ -2,6 +2,8 @@
 from astropy.table import QTable
 import numpy as np
 
+from gammapy.maps import MapAxis
+
 
 def test_psf():
     from pyirf.irf import psf_table
@@ -25,25 +27,25 @@ def test_psf():
         }
     )
 
-    energy_bins = [0, 1.5, 3] * u.TeV
-    fov_bins = [0, 1] * u.deg
-    source_bins = np.linspace(0, 1, 201) * u.deg
+    energy_axis = MapAxis.from_edges([0, 1.5, 3], unit=u.TeV, name='energy_true')
+    fov_axis = MapAxis.from_edges([0, 1], unit=u.deg, name='offset')
+    source_axis = MapAxis.from_bounds(0, 1, 200, unit=u.deg, name='rad')
 
     # We return a table with one row as needed for gadf
-    psf = psf_table(events, energy_bins, source_bins, fov_bins)
+    psf = psf_table(events, energy_axis, source_axis, fov_axis)
 
     # 2 energy bins, 1 fov bin, 200 source distance bins
-    assert psf.shape == (2, 1, 200)
-    assert psf.unit == u.Unit("sr-1")
+    assert psf.quantity.shape == (2, 1, 200)
+    assert psf.quantity.unit == u.Unit("sr-1")
 
     # check that psf is normalized
-    bin_solid_angle = np.diff(cone_solid_angle(source_bins))
-    assert np.allclose(np.sum(psf * bin_solid_angle, axis=2), 1.0)
+    bin_solid_angle = np.diff(cone_solid_angle(source_axis.edges))
+    assert np.allclose(np.sum(psf.quantity * bin_solid_angle, axis=2), 1.0)
 
-    cumulated = np.cumsum(psf * bin_solid_angle, axis=2)
+    cumulated = np.cumsum(psf.quantity * bin_solid_angle, axis=2)
 
     # first energy and only fov bin
-    bin_centers = 0.5 * (source_bins[1:] + source_bins[:-1])
+    bin_centers = source_axis.center
     assert u.isclose(
         bin_centers[np.where(cumulated[0, 0, :] >= 0.68)[0][0]],
         TRUE_SIGMA_1 * u.deg,