allow VoV t0 in spm module

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"]

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
 

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)"
+    }
+  }
+}

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"