fix of edge cases

legend-exp · May 6, 2024 · 40eec51 · 40eec51
1 parent e992543
commit 40eec51
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git a/src/pygama/evt/modules/geds.py b/src/pygama/evt/modules/geds.py
@@ -64,6 +64,8 @@ def apply_xtalk_correction(
     ----------
     datainfo, tcm, table_names
         positional arguments automatically supplied by :func:`.build_evt`.
+    mode
+        string which can be either energy to return corrected energy or tcm_id
     uncalibrated_energy_name
         expression for the pulse parameter to be gathered, can be a combination of different fields.
     calibrated_energy_name
@@ -107,6 +109,7 @@ def apply_xtalk_correction(
     )
     energies_corr = ak.from_regular(energies_corr)
     multiplicity_mask = ak.from_regular(multiplicity_mask)
+    tcm_id_array=ak.from_regular(tcm_id_array)
 
     if mode == "energy":
         return types.VectorOfVectors(energies_corr[multiplicity_mask])
@@ -142,12 +145,18 @@ def apply_xtalk_correction_and_calibrate(
     ----------
     datainfo, tcm, table_names
         positional arguments automatically supplied by :func:`.build_evt`.
+    mode
+        string which can be either energy to return corrected energy or tcm_id
     uncalibrated_energy_name
         expression for the pulse parameter to be gathered, can be a combination of different fields.
     calibrated_energy_name
         name of the pulse parameter for calibrated energy to be gathered, optionally prefixed by tier
         name (e.g. ``hit.cusp_Emax``). If no tier is specified, it defaults
         to ``hit``.
+    par_files
+        path to the generated par files used to recalibrate the data
+    multiplicity_logic:
+        string containing the logic used to define the multiplicity
     threshold
         threshold used for xtalk correction, hits below this energy will not
         be used to correct the other hits.
@@ -159,6 +168,8 @@ def apply_xtalk_correction_and_calibrate(
         name of the lh5 object containing the positive polarity xtalk matrix
     xtalk_matrix_rawids
         name of the lh5 object containing the name of the rawids
+    out_param
+        name of the energy variable to use for recalibration (default None)
     """
 
     xtalk_matrix_rawids = lh5.read_as(xtalk_rawid_name, xtalk_matrix_filename, "np")
@@ -197,8 +208,11 @@ def apply_xtalk_correction_and_calibrate(
         xtalk_matrix_rawids,
     )
 
+
     calibrated_corr = ak.from_regular(calibrated_corr)
     multiplicity_mask = ak.from_regular(multiplicity_mask)
+    tcm_id_array = ak.from_regular(tcm_id_array)
+
 
     if mode == "energy":
         return types.VectorOfVectors(calibrated_corr[multiplicity_mask])

diff --git a/src/pygama/evt/modules/xtalk.py b/src/pygama/evt/modules/xtalk.py
@@ -88,16 +88,14 @@ def build_energy_array(
         idx_events = ak.to_numpy(tcm.idx[tcm.id == channel])
 
         for name, file, group, column in tier_params:
-            keys = ls(file)
             try:
                 # read the energy data
-                if f"ch{channel}" in keys:
-                    data = lh5.read(
+                data = lh5.read(
                         f"ch{channel}/{group}/{column}", file, idx=idx_events
                     )
                 tbl.add_column(name, data)
-            except KeyError:
-                tbl.add_column(name, np.full_like(idx_events, np.nan))
+            except (lh5.exceptions.LH5DecodeError,KeyError):
+                tbl.add_column(name, types.Array(np.full_like(idx_events, np.nan)))
 
         res = tbl.eval(observable)
         energies_out[idx_events, idx_chan] = res.nda
@@ -154,17 +152,14 @@ def filter_hits(
         idx_events = ak.to_numpy(tcm.idx[tcm.id == channel])
 
         for name, file, group, column in tier_params:
-            keys = ls(file)
-
             try:
                 # read the energy data
-                if f"ch{channel}" in keys:
-                    data = lh5.read(
-                        f"ch{channel}/{group}/{column}", file, idx=idx_events
-                    )
+                data = lh5.read(
+                        f"ch{channel}/{group}/{column}", file, idx=idx_events)
+
                 tbl.add_column(name, data)
-            except KeyError:
-                tbl.add_column(name, np.full_like(idx_events, np.nan))
+            except (lh5.exceptions.LH5DecodeError,KeyError):
+                tbl.add_column(name, types.Array(np.full_like(idx_events, np.nan)))
 
         # add the corrected energy to the table
         tbl.add_column(