Update energy_dispersion_to_migration

pyirf/irf/energy_dispersion.py

@@ -88,8 +88,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,49 +110,55 @@ 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],
+        )
+    )
 
     true_energy_interpolation = resample_histogram1d(
         dispersion_matrix,
         disp_true_energy_edges,
         new_true_energy_edges,
         axis=0,
+        normalize=True
     )
 
-    norm = np.sum(true_energy_interpolation, axis=1, keepdims=True)
-    norm[norm == 0] = 1
-    true_energy_interpolation /= norm
+    true_energy_interpolation /= np.diff(disp_migration_edges)[np.newaxis, :, np.newaxis]
 
     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(
             e_true_dispersion,
             disp_migration_edges,
             interpolation_edges,
             axis=0,
+            normalize=True
         )
 
         migration_matrix[idx, :, :] = y
 
+    with np.errstate(invalid="ignore"):
+        migration_matrix = migration_matrix / migration_matrix.sum(axis=2).sum(axis=1)[:, np.newaxis, np.newaxis]
+    migration_matrix[np.isnan(migration_matrix)] = 0
+
     return migration_matrix