diff --git a/src/km3io/tools.py b/src/km3io/tools.py index c2d8a829b387c0672631cd6f7c04c1e4749a152d..30e32c7c2e0bdc5591229ecf344ce04a4ac241ae 100644 --- a/src/km3io/tools.py +++ b/src/km3io/tools.py @@ -5,6 +5,7 @@ import numpy as np import awkward as ak import uproot3 +import km3io.definitions from km3io.definitions import reconstruction as krec from km3io.definitions import trigger as ktrg from km3io.definitions import fitparameters as kfit @@ -538,6 +539,24 @@ def is_3dmuon(trigger_mask): return is_bit_set(trigger_mask, ktrg.JTRIGGER3DMUON) +def get_w2list_idx(f): + """ + Get the correct w2list_idx for the given file, or None if there is none. + + Parameters + ---------- + f : km3io.OfflineReader + The file. + + """ + w2s_idx = { + "genhen": km3io.definitions.w2list_genhen_idx, + "gseagen": km3io.definitions.w2list_gseagen_idx, + } + sim_program = f.header.simul.program.lower() + return w2s_idx.get(sim_program) + + def is_nanobeacon(trigger_mask): """Returns True if the trigger mask contains the nano-beacon flag.