Merge pull request #273 from cta-observatory/HealthyPear-fix_energy_d…

…ispersion_to_migration

Update energy_dispersion_to_migration to account for fix in energy dispersion normalization

docs/changes/273.bugfix.rst

@@ -0,0 +1,4 @@
+Fix ``pyirf.irfs.energy_dispersion.energy_dispersion_to_migration``.
+This function was not adapted to the now correct normalization of the
+energy dispersion matrix, resulting in wrong results on the now correct
+matrices.

pyirf/irf/energy_dispersion.py

@@ -6,7 +6,8 @@
 
 __all__ = [
     "energy_dispersion",
-    "energy_dispersion_to_migration"
+    "energy_migration_matrix",
+    "energy_dispersion_to_migration",
 ]
 
 
@@ -29,7 +30,10 @@ def _normalize_hist(hist, migration_bins):
 
 
 def energy_dispersion(
-    selected_events, true_energy_bins, fov_offset_bins, migration_bins,
+    selected_events,
+    true_energy_bins,
+    fov_offset_bins,
+    migration_bins,
 ):
     """
     Calculate energy dispersion for the given DL2 event list.
@@ -80,6 +84,56 @@ def energy_dispersion(
     return energy_dispersion
 
 
+@u.quantity_input(true_energy_bins=u.TeV, reco_energy_bins=u.TeV, fov_offset_bins=u.deg)
+def energy_migration_matrix(
+    events, true_energy_bins, reco_energy_bins, fov_offset_bins
+):
+    """Compute the energy migration matrix directly from the events.
+
+    Parameters
+    ----------
+    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`
+        Bin edges in true energy.
+    reco_energy_bins : `~astropy.units.Quantity`
+        Bin edges in reconstructed energy.
+
+    Returns
+    -------
+    matrix : array-like
+        Migration matrix as probabilities along the reconstructed energy axis.
+        energy axis with shape
+        (n_true_energy_bins, n_reco_energy_bins, n_fov_offset_bins)
+        containing energies in TeV.
+    """
+
+    hist, _ = np.histogramdd(
+        np.column_stack(
+            [
+                events["true_energy"].to_value(u.TeV),
+                events["reco_energy"].to_value(u.TeV),
+                events["true_source_fov_offset"].to_value(u.deg),
+            ]
+        ),
+        bins=[
+            true_energy_bins.to_value(u.TeV),
+            reco_energy_bins.to_value(u.TeV),
+            fov_offset_bins.to_value(u.deg),
+        ],
+    )
+
+    with np.errstate(invalid="ignore"):
+        hist /= hist.sum(axis=1)[:, np.newaxis, :]
+        # the nans come from the fact that the sum along the reconstructed energy axis
+        # might sometimes be 0 when there are no events in that given true energy bin
+        # and fov offset bin
+        hist[np.isnan(hist)] = 0
+
+    return hist
+
+
 def energy_dispersion_to_migration(
     dispersion_matrix,
     disp_true_energy_edges,
@@ -88,8 +142,8 @@ def energy_dispersion_to_migration(
     new_reco_energy_edges,
 ):
     """
-    Construct a energy migration matrix from a energy
-    dispersion matrix.
+    Construct a energy migration matrix from an 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,
@@ -110,21 +164,25 @@ def energy_dispersion_to_migration(
     new_reco_energy_edges: astropy.units.Quantity[energy]
         Reco energy edges matching the second dimension of the output
 
-    Returns:
-    --------
+    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],
+        )
+    )
 
-    migration_matrix = np.zeros((
-        len(new_true_energy_edges) - 1,
-        len(new_reco_energy_edges) - 1,
-        dispersion_matrix.shape[2],
-    ))
+    migra_width = np.diff(disp_migration_edges)
+    probability = dispersion_matrix * migra_width[np.newaxis, :, np.newaxis]
 
     true_energy_interpolation = resample_histogram1d(
-        dispersion_matrix,
+        probability,
         disp_true_energy_edges,
         new_true_energy_edges,
         axis=0,
@@ -137,13 +195,14 @@ def energy_dispersion_to_migration(
     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')
+            warnings.filterwarnings(
+                "ignore", "invalid value encountered in true_divide"
+            )
             interpolation_edges = new_reco_energy_edges / e_true
 
         y = resample_histogram1d(

pyirf/irf/tests/test_energy_dispersion.py

@@ -64,17 +64,29 @@ def ppf(cdf, bins, value):
 
     assert np.isclose(
         TRUE_SIGMA_1,
-        0.5 * (ppf(cdf[0, :, 0], migration_bins, 0.84) - ppf(cdf[0, :, 0], migration_bins, 0.16)),
+        0.5
+        * (
+            ppf(cdf[0, :, 0], migration_bins, 0.84)
+            - ppf(cdf[0, :, 0], migration_bins, 0.16)
+        ),
         rtol=0.1,
     )
     assert np.isclose(
         TRUE_SIGMA_2,
-        0.5 * (ppf(cdf[1, :, 0], migration_bins, 0.84) - ppf(cdf[1, :, 0], migration_bins, 0.16)),
+        0.5
+        * (
+            ppf(cdf[1, :, 0], migration_bins, 0.84)
+            - ppf(cdf[1, :, 0], migration_bins, 0.16)
+        ),
         rtol=0.1,
     )
     assert np.isclose(
         TRUE_SIGMA_3,
-        0.5 * (ppf(cdf[2, :, 0], migration_bins, 0.84) - ppf(cdf[2, :, 0], migration_bins, 0.16)),
+        0.5
+        * (
+            ppf(cdf[2, :, 0], migration_bins, 0.84)
+            - ppf(cdf[2, :, 0], migration_bins, 0.16)
+        ),
         rtol=0.1,
     )
 
@@ -85,16 +97,15 @@ def test_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
+    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
+    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(
@@ -116,15 +127,14 @@ def test_energy_dispersion_to_migration():
     )
 
     # 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
+    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
@@ -137,10 +147,53 @@ def test_energy_dispersion_to_migration():
 
     # 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 migration_matrix.sum() < (len(new_true_energy_bins) - 1) * (
+        len(fov_offset_bins) - 1
+    )
     assert np.all(np.isfinite(migration_matrix))
 
 
+def test_energy_migration_matrix_from_events():
+    from pyirf.irf.energy_dispersion import energy_migration_matrix
+
+    np.random.seed(0)
+    N = 10000
+    true_energy_bins = 10 ** np.arange(np.log10(0.2), np.log10(200), 1 / 10) * u.TeV
+    reco_energy_bins = 10 ** np.arange(np.log10(2), np.log10(20), 1 / 5) * u.TeV
+    fov_offset_bins = np.array([0, 1, 2]) * u.deg
+
+    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)
+
+    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,
+        }
+    )
+
+    matrix = energy_migration_matrix(
+        events, true_energy_bins, reco_energy_bins, fov_offset_bins
+    )
+
+    assert matrix.shape == (
+        len(true_energy_bins) - 1,
+        len(reco_energy_bins) - 1,
+        len(fov_offset_bins) - 1,
+    )
+    assert np.allclose(matrix.sum(axis=1).max(), 1, rtol=0.1)
+
+
 def test_overflow_bins_migration_matrix():
     from pyirf.irf import energy_dispersion
     from pyirf.irf.energy_dispersion import energy_dispersion_to_migration
@@ -150,21 +203,24 @@ def test_overflow_bins_migration_matrix():
     N = 10000
 
     # add under/overflow bins
-    bins = 10 ** np.arange(
-        np.log10(0.2),
-        np.log10(200),
-        1 / 10,
-    ) * u.TeV
+    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
+    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(