diff --git a/src/pygama/evt/build_evt.py b/src/pygama/evt/build_evt.py index 8febc5f55..f606e3774 100644 --- a/src/pygama/evt/build_evt.py +++ b/src/pygama/evt/build_evt.py @@ -100,15 +100,7 @@ def evaluate_expression( exprl = re.findall(r"[a-zA-Z_$][\w$]*", expr) var_ph = {} if os.path.exists(f_evt): - var_ph = store.load_nda( - f_evt, - [ - e.split("/")[-1] - for e in store.ls(f_evt, group) - if e.split("/")[-1] in exprl - ], - group, - ) + var_ph = load_vars_to_nda(f_evt, group, exprl) if para: var_ph = var_ph | para @@ -292,29 +284,41 @@ def find_parameters( hit_group: str, ) -> dict: # find fields in either dsp, hit - var = store.load_nda( - f_hit, - [ - e.split("/")[-1] - for e in store.ls(f_hit, ch + hit_group) - if e.split("/")[-1] in exprl - ], - ch + hit_group, - idx_ch, - ) - dsp_dic = store.load_nda( - f_dsp, - [ - e.split("/")[-1] - for e in store.ls(f_dsp, ch + dsp_group) - if e.split("/")[-1] in exprl - ], - ch + dsp_group, - idx_ch, - ) + var = load_vars_to_nda(f_hit, ch + hit_group, exprl) + dsp_dic = load_vars_to_nda(f_dsp, ch + dsp_group, exprl) + return dsp_dic | var +def load_vars_to_nda(f_evt: str, group: str, exprl: list) -> dict: + lstore = store.LH5Store() + flds = [ + e.split("/")[-1] for e in store.ls(f_evt, group) if e.split("/")[-1] in exprl + ] + var = {e: lstore.read_object(group + e, f_evt)[0] for e in flds} + + # to make any operations to VoVs we have to blow it up to a table (future change to more intelligant way) + arr_keys = [] + for key, value in var.items(): + if isinstance(value, VectorOfVectors): + var[key] = value.to_aoesa().nda + elif isinstance(value, Array): + var[key] = value.nda + if var[key].ndim > 2: + raise ValueError("Dim > 2 not supported") + if var[key].ndim == 1: + arr_keys.append(key) + else: + raise ValueError(f"{type(value)} not supported") + + # now we also need to set dimensions if we have an expression + # consisting of a mix of VoV and Arrays + if len(arr_keys) > 0 and not set(arr_keys) == set(var.keys()): + for key in arr_keys: + var[key] = var[key][:, None] + return var + + def evaluate_to_first( idx: np.ndarray, ids: np.ndarray, @@ -907,27 +911,7 @@ def build_evt( exprl = re.findall(r"[a-zA-Z_$][\w$]*", v["expression"]) var = {} if os.path.exists(f_evt): - flds = [ - e.split("/")[-1] - for e in store.ls(f_evt, group) - if e.split("/")[-1] in exprl - ] - var = {e: lstore.read_object(group + e, f_evt)[0] for e in flds} - - # to make any operations to VoVs we have to blow it up to a table (future change to more intelligant way) - arr_keys = [] - for key, value in var.items(): - if isinstance(value, VectorOfVectors): - var[key] = value.to_aoesa().nda - elif isinstance(value, Array): - var[key] = value.nda - arr_keys.append(key) - - # now we also need to set dimensions if we have an expression - # consisting of a mix of VoV and Arrays - if len(arr_keys) > 0 and not set(arr_keys) == set(var.keys()): - for key in arr_keys: - var[key] = var[key][:, None] + var = load_vars_to_nda(f_evt, group, exprl) if "parameters" in v.keys(): var = var | v["parameters"] diff --git a/src/pygama/evt/modules/spm.py b/src/pygama/evt/modules/spm.py index b43bf134d..7bd530531 100644 --- a/src/pygama/evt/modules/spm.py +++ b/src/pygama/evt/modules/spm.py @@ -13,11 +13,24 @@ import lgdo.lh5_store as store import numpy as np +from lgdo import Array, VectorOfVectors # get LAr energy per event over all channels def get_energy(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): - trig = np.where(np.isnan(trgr), tdefault, trgr) + trig = trgr + if isinstance(trgr, VectorOfVectors): + trig = trig.to_aoesa().nda + elif isinstance(trgr, Array): + trig = trig.nda + if isinstance(trig, np.ndarray) and trig.ndim == 2: + trig = np.where(np.isnan(trig).all(axis=1)[:, None], tdefault, trig) + trig = np.nanmin(trig, axis=1) + + elif isinstance(trig, np.ndarray) and trig.ndim == 1: + trig = np.where(np.isnan(trig), tdefault, trig) + else: + raise ValueError(f"Can't deal with t0 of type {type(trgr)}") tmi = trig - tmin tma = trig + tmax sum = np.zeros(len(trig)) @@ -46,7 +59,19 @@ def get_energy(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): # get LAr majority per event over all channels def get_majority(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): - trig = np.where(np.isnan(trgr), tdefault, trgr) + trig = trgr + if isinstance(trgr, VectorOfVectors): + trig = trig.to_aoesa().nda + elif isinstance(trgr, Array): + trig = trig.nda + if isinstance(trig, np.ndarray) and trig.ndim == 2: + trig = np.where(np.isnan(trig).all(axis=1)[:, None], tdefault, trig) + trig = np.nanmin(trig, axis=1) + + elif isinstance(trig, np.ndarray) and trig.ndim == 1: + trig = np.where(np.isnan(trig), tdefault, trig) + else: + raise ValueError(f"Can't deal with t0 of type {type(trgr)}") tmi = trig - tmin tma = trig + tmax maj = np.zeros(len(trig)) @@ -76,7 +101,19 @@ def get_majority(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): # get LAr energy per event over all channels def get_energy_dplms(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): - trig = np.where(np.isnan(trgr), tdefault, trgr) + trig = trgr + if isinstance(trgr, VectorOfVectors): + trig = trig.to_aoesa().nda + elif isinstance(trgr, Array): + trig = trig.nda + if isinstance(trig, np.ndarray) and trig.ndim == 2: + trig = np.where(np.isnan(trig).all(axis=1)[:, None], tdefault, trig) + trig = np.nanmin(trig, axis=1) + + elif isinstance(trig, np.ndarray) and trig.ndim == 1: + trig = np.where(np.isnan(trig), tdefault, trig) + else: + raise ValueError(f"Can't deal with t0 of type {type(trgr)}") tmi = trig - tmin tma = trig + tmax sum = np.zeros(len(trig)) @@ -105,7 +142,19 @@ def get_energy_dplms(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): # get LAr majority per event over all channels def get_majority_dplms(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): - trig = np.where(np.isnan(trgr), tdefault, trgr) + trig = trgr + if isinstance(trgr, VectorOfVectors): + trig = trig.to_aoesa().nda + elif isinstance(trgr, Array): + trig = trig.nda + if isinstance(trig, np.ndarray) and trig.ndim == 2: + trig = np.where(np.isnan(trig).all(axis=1)[:, None], tdefault, trig) + trig = np.nanmin(trig, axis=1) + + elif isinstance(trig, np.ndarray) and trig.ndim == 1: + trig = np.where(np.isnan(trig), tdefault, trig) + else: + raise ValueError(f"Can't deal with t0 of type {type(trgr)}") tmi = trig - tmin tma = trig + tmax maj = np.zeros(len(trig)) @@ -146,7 +195,20 @@ def get_etc(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax, swin, tra pes = np.zeros([len(chs), peshape[0], peshape[1]]) times = np.zeros([len(chs), peshape[0], peshape[1]]) - tge = np.where(np.isnan(trgr), tdefault, trgr) + tge = trgr + if isinstance(trgr, VectorOfVectors): + tge = tge.to_aoesa().nda + elif isinstance(trgr, Array): + tge = tge.nda + if isinstance(tge, np.ndarray) and tge.ndim == 2: + tge = np.where(np.isnan(tge).all(axis=1)[:, None], tdefault, tge) + tge = np.nanmin(tge, axis=1) + + elif isinstance(tge, np.ndarray) and tge.ndim == 1: + tge = np.where(np.isnan(tge), tdefault, tge) + else: + raise ValueError(f"Can't deal with t0 of type {type(trgr)}") + tmi = tge - tmin tma = tge + tmax @@ -213,7 +275,20 @@ def get_time_shift(f_hit, f_dsp, f_tcm, chs, lim, trgr, tdefault, tmin, tmax): peshape = (predf["energy_in_pe"]).shape times = np.zeros([len(chs), peshape[0], peshape[1]]) - tge = np.where(np.isnan(trgr), tdefault, trgr) + tge = trgr + if isinstance(trgr, VectorOfVectors): + tge = tge.to_aoesa().nda + elif isinstance(trgr, Array): + tge = tge.nda + if isinstance(tge, np.ndarray) and tge.ndim == 2: + tge = np.where(np.isnan(tge).all(axis=1)[:, None], tdefault, tge) + tge = np.nanmin(tge, axis=1) + + elif isinstance(tge, np.ndarray) and tge.ndim == 1: + tge = np.where(np.isnan(tge), tdefault, tge) + else: + raise ValueError(f"Can't deal with t0 of type {type(trgr)}") + tmi = tge - tmin tma = tge + tmax diff --git a/tests/evt/configs/module-test-t0-vov-evt-config.json b/tests/evt/configs/module-test-t0-vov-evt-config.json new file mode 100644 index 000000000..da162733c --- /dev/null +++ b/tests/evt/configs/module-test-t0-vov-evt-config.json @@ -0,0 +1,53 @@ +{ + "channels": { + "spms_on": ["ch1057600", "ch1059201", "ch1062405"], + "geds_on": ["ch1084803", "ch1084804", "ch1121600"] + }, + "operations": { + "energy": { + "channels": "geds_on", + "mode": "vov>25", + "get_ch": true, + "expression": "cuspEmax_ctc_cal" + }, + "t0": { + "channels": ["geds_on"], + "mode": "energy_id", + "expression": "tp_0_est", + "initial": 0.0 + }, + "lar_energy": { + "channels": "spms_on", + "mode": "func", + "expression": ".modules.spm.get_energy(0.5,t0,48000,1000,5000)" + }, + "lar_multiplicity": { + "channels": "spms_on", + "mode": "func", + "expression": ".modules.spm.get_majority(0.5,t0,48000,1000,5000)" + }, + "is_lar_rejected": { + "expression": "(lar_energy >4) | (lar_multiplicity > 4) " + }, + "lar_classifier": { + "channels": "spms_on", + "mode": "func", + "expression": ".modules.spm.get_etc(0.5,t0,48000,100,6000,80,1)" + }, + "lar_energy_dplms": { + "channels": "spms_on", + "mode": "func", + "expression": ".modules.spm.get_energy_dplms(0.5,t0,48000,1000,5000)" + }, + "lar_multiplicity_dplms": { + "channels": "spms_on", + "mode": "func", + "expression": ".modules.spm.get_majority_dplms(0.5,t0,48000,1000,5000)" + }, + "lar_time_shift": { + "channels": "spms_on", + "mode": "func", + "expression": ".modules.spm.get_time_shift(0.5,t0,48000,1000,5000)" + } + } +} diff --git a/tests/evt/test_build_evt.py b/tests/evt/test_build_evt.py index a08848934..128833e5b 100644 --- a/tests/evt/test_build_evt.py +++ b/tests/evt/test_build_evt.py @@ -79,6 +79,33 @@ def test_lar_module(lgnd_test_data, tmptestdir): assert ((nda["lar_time_shift"] + nda["t0"]) >= 0).all() +def test_lar_t0_vov_module(lgnd_test_data, tmptestdir): + outfile = f"{tmptestdir}/l200-p03-r001-phy-20230322T160139Z-tier_evt.lh5" + tcm_path = "lh5/prod-ref-l200/generated/tier/tcm/phy/p03/r001/l200-p03-r001-phy-20230322T160139Z-tier_tcm.lh5" + if os.path.exists(outfile): + os.remove(outfile) + build_evt( + f_tcm=lgnd_test_data.get_path(tcm_path), + f_dsp=lgnd_test_data.get_path(tcm_path.replace("tcm", "dsp")), + f_hit=lgnd_test_data.get_path(tcm_path.replace("tcm", "hit")), + f_evt=outfile, + meta_path=None, + evt_config=f"{config_dir}/module-test-t0-vov-evt-config.json", + wo_mode="o", + group="/evt/", + ) + + assert os.path.exists(outfile) + assert len(ls(outfile, "/evt/")) == 10 + nda = load_nda( + outfile, + ["lar_multiplicity", "lar_multiplicity_dplms", "lar_time_shift"], + "/evt/", + ) + assert np.max(nda["lar_multiplicity"]) <= 3 + assert np.max(nda["lar_multiplicity_dplms"]) <= 3 + + def test_vov(lgnd_test_data, tmptestdir): outfile = f"{tmptestdir}/l200-p03-r001-phy-20230322T160139Z-tier_evt.lh5" tcm_path = "lh5/prod-ref-l200/generated/tier/tcm/phy/p03/r001/l200-p03-r001-phy-20230322T160139Z-tier_tcm.lh5"