Use gammapy classes directly for energy dispersion

pyirf/irf/effective_area.py

@@ -47,6 +47,8 @@ def effective_area_per_energy(
         The overall statistics of the simulated events
     true_energy_axis: gammapy.maps.MapAxis
         The bin edges in which to calculate effective area.
+    fov_offset_bins: astropy.units.Quantity[angle]
+        The field of view radial offset axis. This must only have a single bin.
     """
     area = np.pi * simulation_info.max_impact ** 2
 
@@ -87,10 +89,10 @@ def effective_area_per_energy_and_fov(
         - `true_source_fov_offset`
     simulation_info: pyirf.simulations.SimulatedEventsInfo
         The overall statistics of the simulated events
-    true_energy_bins: astropy.units.Quantity[energy]
-        The true energy bin edges in which to calculate effective area.
-    fov_offset_bins: astropy.units.Quantity[angle]
-        The field of view radial bin edges in which to calculate effective area.
+    true_energy_axis: MapAxis[energy]
+        The true energy axis in which to calculate effective area.
+    fov_offset_axis: MapAxis[angle]
+        The field of view radial offset axis in which to calculate effective area.
     """
     area = np.pi * simulation_info.max_impact ** 2
 

pyirf/irf/energy_dispersion.py

@@ -1,12 +1,13 @@
 import warnings
 import numpy as np
 import astropy.units as u
-from ..binning import resample_histogram1d
+
+from gammapy.irf import EnergyDispersion2D
+from gammapy.maps import MapAxis
 
 
 __all__ = [
     "energy_dispersion",
-    "energy_dispersion_to_migration"
 ]
 
 
@@ -28,7 +29,10 @@ def _normalize_hist(hist):
 
 
 def energy_dispersion(
-    selected_events, true_energy_bins, fov_offset_bins, migration_bins,
+    selected_events,
+    true_energy_axis: MapAxis,
+    fov_offset_axis: MapAxis,
+    migration_axis: MapAxis,
 ):
     """
     Calculate energy dispersion for the given DL2 event list.
@@ -41,12 +45,12 @@ def energy_dispersion(
     selected_events: astropy.table.QTable
         Table of the DL2 events.
         Required columns: ``reco_energy``, ``true_energy``, ``true_source_fov_offset``.
-    true_energy_bins: astropy.units.Quantity[energy]
+    true_energy_aix: MapAxis[energy]
         Bin edges in true energy
-    migration_bins: astropy.units.Quantity[energy]
+    migration_axis: MapAxis[dimensionless]
         Bin edges in relative deviation, recommended range: [0.2, 5]
-    fov_offset_bins: astropy.units.Quantity[angle]
-        Bin edges in the field of view offset.
+    fov_offset_axis: MapAxis[angle]
+        Field of view offset axis.
         For Point-Like IRFs, only giving a single bin is appropriate.
 
     Returns
@@ -68,91 +72,19 @@ def energy_dispersion(
             ]
         ),
         bins=[
-            true_energy_bins.to_value(u.TeV),
-            migration_bins,
-            fov_offset_bins.to_value(u.deg),
+            true_energy_axis.edges.to_value(u.TeV),
+            migration_axis.edges.to_value(u.one),
+            fov_offset_axis.edges.to_value(u.deg),
         ],
     )
 
-    n_events_per_energy = energy_dispersion.sum(axis=1)
     energy_dispersion = _normalize_hist(energy_dispersion)
 
-    return energy_dispersion
-
-
-def energy_dispersion_to_migration(
-    dispersion_matrix,
-    disp_true_energy_edges,
-    disp_migration_edges,
-    new_true_energy_edges,
-    new_reco_energy_edges,
-):
-    """
-    Construct a energy migration matrix from a energy
-    dispersion matrix.
-    Depending on the new energy ranges, the sum over the first axis
-    can be smaller than 1.
-    The new true energy bins need to be a subset of the old range,
-    extrapolation is not supported.
-    New reconstruction bins outside of the old migration range are filled with
-    zeros.
-
-    Parameters
-    ----------
-    dispersion_matrix: numpy.ndarray
-        Energy dispersion_matrix
-    disp_true_energy_edges: astropy.units.Quantity[energy]
-        True energy edges matching the first dimension of the dispersion matrix
-    disp_migration_edges: numpy.ndarray
-        Migration edges matching the second dimension of the dispersion matrix
-    new_true_energy_edges: astropy.units.Quantity[energy]
-        True energy edges matching the first dimension of the output
-    new_reco_energy_edges: astropy.units.Quantity[energy]
-        Reco energy edges matching the second dimension of the output
-
-    Returns:
-    --------
-    migration_matrix: numpy.ndarray
-        Three-dimensional energy migration matrix. The third dimension
-        equals the fov offset dimension of the energy dispersion matrix.
-    """
-
-    migration_matrix = np.zeros((
-        len(new_true_energy_edges) - 1,
-        len(new_reco_energy_edges) - 1,
-        dispersion_matrix.shape[2],
-    ))
-
-    true_energy_interpolation = resample_histogram1d(
-        dispersion_matrix,
-        disp_true_energy_edges,
-        new_true_energy_edges,
-        axis=0,
+    return EnergyDispersion2D(
+        axes=[
+            true_energy_axis,
+            migration_axis,
+            fov_offset_axis,
+        ],
+        data=energy_dispersion
     )
-
-    norm = np.sum(true_energy_interpolation, axis=1, keepdims=True)
-    norm[norm == 0] = 1
-    true_energy_interpolation /= norm
-
-    for idx, e_true in enumerate(
-        (new_true_energy_edges[1:] + new_true_energy_edges[:-1]) / 2
-    ):
-
-        # get migration for the new true energy bin
-        e_true_dispersion = true_energy_interpolation[idx]
-
-        with warnings.catch_warnings():
-            # silence inf/inf division warning
-            warnings.filterwarnings('ignore', 'invalid value encountered in true_divide')
-            interpolation_edges = new_reco_energy_edges / e_true
-
-        y = resample_histogram1d(
-            e_true_dispersion,
-            disp_migration_edges,
-            interpolation_edges,
-            axis=0,
-        )
-
-        migration_matrix[idx, :, :] = y
-
-    return migration_matrix

pyirf/irf/tests/test_energy_dispersion.py

@@ -1,5 +1,6 @@
 import astropy.units as u
 import numpy as np
+from gammapy.maps import MapAxis
 from astropy.table import QTable
 
 
@@ -41,19 +42,20 @@ def test_energy_dispersion():
         }
     )
 
-    true_energy_bins = np.array([0.1, 1.0, 10.0, 100]) * u.TeV
-    fov_offset_bins = np.array([0, 1, 2]) * u.deg
-    migration_bins = np.linspace(0, 2, 1001)
+    true_energy_axis = MapAxis.from_edges([0.1, 1.0, 10.0, 100], unit=u.TeV, name='energy_true')
+    fov_offset_axis = MapAxis.from_edges([0, 1, 2], unit=u.deg, name='offset')
+    migration_axis = MapAxis.from_bounds(0.2, 5, 1000, interp='log', name='migra')
 
-    result = energy_dispersion(
-        selected_events, true_energy_bins, fov_offset_bins, migration_bins
+    edisp2d = energy_dispersion(
+        selected_events, true_energy_axis, fov_offset_axis, migration_axis
     )
 
-    assert result.shape == (3, 1000, 2)
-    assert np.isclose(result.sum(), 6.0)
+    edisp = edisp2d.quantity
+    assert edisp.shape == (3, 1000, 2)
+    assert np.isclose(edisp.sum(), 6.0)
 
-    cumulative_sum = np.cumsum(result, axis=1)
-    bin_centers = 0.5 * (migration_bins[1:] + migration_bins[:-1])
+    cumulative_sum = np.cumsum(edisp, axis=1)
+    bin_centers = migration_axis.center
     assert np.isclose(
         TRUE_SIGMA_1,
         (
@@ -81,125 +83,3 @@ def test_energy_dispersion():
         / 2,
         rtol=0.1,
     )
-
-
-def test_energy_dispersion_to_migration():
-    from pyirf.irf import energy_dispersion
-    from pyirf.irf.energy_dispersion import energy_dispersion_to_migration
-
-    np.random.seed(0)
-    N = 10000
-    true_energy_bins = 10**np.arange(np.log10(0.2), np.log10(200), 1/10) * u.TeV
-
-    fov_offset_bins = np.array([0, 1, 2]) * u.deg
-    migration_bins = np.linspace(0, 2, 101)
-
-    true_energy = np.random.uniform(
-        true_energy_bins[0].value,
-        true_energy_bins[-1].value,
-        size=N
-    ) * u.TeV
-    reco_energy = true_energy * np.random.uniform(0.5, 1.5, size=N)
-
-    selected_events = QTable(
-        {
-            "reco_energy": reco_energy,
-            "true_energy": true_energy,
-            "true_source_fov_offset": np.concatenate(
-                [
-                    np.full(N // 2, 0.2),
-                    np.full(N // 2, 1.5),
-                ]
-            )
-            * u.deg,
-        }
-    )
-
-    dispersion_matrix = energy_dispersion(
-        selected_events, true_energy_bins, fov_offset_bins, migration_bins
-    )
-
-    # migration matrix selecting a limited energy band with different binsizes
-    new_true_energy_bins = 10**np.arange(np.log10(2), np.log10(20), 1/5) * u.TeV
-    new_reco_energy_bins = 10**np.arange(np.log10(2), np.log10(20), 1/5) * u.TeV
-    migration_matrix = energy_dispersion_to_migration(
-        dispersion_matrix,
-        true_energy_bins,
-        migration_bins,
-        new_true_energy_bins,
-        new_reco_energy_bins,
-
-    )
-
-    # test dimension
-    assert migration_matrix.shape[0] == len(new_true_energy_bins) - 1
-    assert migration_matrix.shape[1] == len(new_reco_energy_bins) - 1
-    assert migration_matrix.shape[2] == dispersion_matrix.shape[2]
-
-    # test that all migrations are included for central energies
-    assert np.isclose(migration_matrix.sum(axis=1).max(), 1, rtol=0.01)
-
-    # test that migrations dont always sum to 1 (since some energies are
-    # not included in the matrix)
-    assert migration_matrix.sum() < (len(new_true_energy_bins) - 1) * (len(fov_offset_bins) - 1)
-    assert np.all(np.isfinite(migration_matrix))
-
-
-def test_overflow_bins_migration_matrix():
-    from pyirf.irf import energy_dispersion
-    from pyirf.irf.energy_dispersion import energy_dispersion_to_migration
-    from pyirf.binning import add_overflow_bins
-
-    np.random.seed(0)
-    N = 10000
-
-    # add under/overflow bins
-    bins = 10 ** np.arange(
-        np.log10(0.2),
-        np.log10(200),
-        1 / 10,
-    ) * u.TeV
-    true_energy_bins = add_overflow_bins(bins, positive=False)
-
-    fov_offset_bins = np.array([0, 1, 2]) * u.deg
-    migration_bins = np.linspace(0, 2, 101)
-
-    true_energy = np.random.uniform(
-        true_energy_bins[1].value,
-        true_energy_bins[-2].value,
-        size=N
-    ) * u.TeV
-    reco_energy = true_energy * np.random.uniform(0.5, 1.5, size=N)
-
-    selected_events = QTable(
-        {
-            "reco_energy": reco_energy,
-            "true_energy": true_energy,
-            "true_source_fov_offset": np.concatenate(
-                [
-                    np.full(N // 2, 0.2),
-                    np.full(N // 2, 1.5),
-                ]
-            )
-            * u.deg,
-        }
-    )
-
-    dispersion_matrix = energy_dispersion(
-        selected_events, true_energy_bins, fov_offset_bins, migration_bins
-    )
-
-    migration_matrix = energy_dispersion_to_migration(
-        dispersion_matrix,
-        true_energy_bins,
-        migration_bins,
-        true_energy_bins,
-        true_energy_bins,
-    )
-
-    # migration into underflow bin is not empty
-    assert np.sum(migration_matrix[:, 0, :]) > 0
-    # migration from underflow bin is empty
-    assert np.sum(migration_matrix[0, :, :]) == 0
-
-    assert np.all(np.isfinite(migration_matrix))