add event skimming function

src/pygama/evt/__init__.py

@@ -2,8 +2,8 @@
 Utilities for grouping hit data into events.
 """
 
-from .build_evt import build_evt
+from .build_evt import build_evt, skim_evt
 from .build_tcm import build_tcm
 from .tcm import generate_tcm_cols
 
-__all__ = ["build_tcm", "generate_tcm_cols", "build_evt"]
+__all__ = ["build_tcm", "generate_tcm_cols", "build_evt", "skim_evt"]

src/pygama/evt/build_evt.py

@@ -1012,3 +1012,93 @@ def build_evt(
                 )
 
     log.info("Done")
+
+
+def skim_evt(
+    f_evt: str,
+    expression: str,
+    params: dict = None,
+    f_out: str = None,
+    wo_mode="n",
+    evt_group="/evt/",
+) -> None:
+    """
+    Skimms events from a evt file which are fullfling the expression, discards all other events.
+
+    Parameters
+    ----------
+    f_evt
+        input LH5 file of the evt level
+    expression
+        skimming expression. Can contain variabels from event file or from the params dictionary.
+    f_out
+        output LH5 file. Can be None if wo_mode is set to overwrite f_evt.
+    wo_mode
+        Write mode: "o"/"overwrite" overwrites f_evt. "n"/"new" writes to a new file specified in f_out.
+    evt_group
+        lh5 root group of the evt file
+    """
+
+    if wo_mode not in ["o", "overwrite", "n", "new"]:
+        raise ValueError(
+            wo_mode
+            + " is a invalid writing mode. Valid options are: 'o', 'overwrite','n','new'"
+        )
+    lstore = store.LH5Store()
+    fields = store.ls(f_evt, evt_group)
+    nrows = lstore.read_n_rows(fields[0], f_evt)
+    # load fields in expression
+    exprl = re.findall(r"[a-zA-Z_$][\w$]*", expression)
+    var = {}
+
+    flds = [
+        e.split("/")[-1]
+        for e in store.ls(f_evt, evt_group)
+        if e.split("/")[-1] in exprl
+    ]
+    var = {e: lstore.read_object(evt_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]
+
+    if params is not None:
+        var = var | params
+    res = eval(expression, var)
+
+    if res.shape != (nrows,):
+        raise ValueError(
+            f"The expression must result to 1D with length = event number. Current shape is {res.shape}"
+        )
+
+    res = res.astype(bool)
+    idx_list = np.arange(nrows, dtype=int)[res]
+
+    of = f_out
+    if wo_mode in ["o", "overwrite"]:
+        of = f_evt
+    of_tmp = of.replace(of.split("/")[-1], ".tmp_" + of.split("/")[-1])
+
+    for fld in fields:
+        ob, _ = lstore.read_object(fld, f_evt, idx=idx_list)
+        lstore.write_object(
+            obj=ob,
+            name=fld,
+            lh5_file=of_tmp,
+            wo_mode="o",
+        )
+
+    if os.path.exists(of):
+        os.remove(of)
+    os.rename(of_tmp, of)

tests/evt/test_build_evt.py

@@ -6,7 +6,7 @@
 import pytest
 from lgdo import Array, VectorOfVectors, load_nda, ls
 
-from pygama.evt import build_evt
+from pygama.evt import build_evt, skim_evt
 
 config_dir = Path(__file__).parent / "configs"
 
@@ -192,3 +192,28 @@ def test_graceful_crashing(lgnd_test_data, tmptestdir):
     }
     with pytest.raises(ValueError):
         build_evt(f_tcm, f_dsp, f_hit, outfile, conf, meta_path)
+
+
+def test_skimming(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)
+    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"))
+    meta_path = None
+    f_config = f"{config_dir}/vov-test-evt-config.json"
+    build_evt(f_tcm, f_dsp, f_hit, outfile, f_config, meta_path)
+
+    lstore = store.LH5Store()
+    ac = lstore.read_object("/evt/multiplicity", outfile)[0].nda
+    ac = len(ac[ac == 3])
+
+    outfile_skm = f"{tmptestdir}/l200-p03-r001-phy-20230322T160139Z-tier_skm.lh5"
+
+    skim_evt(outfile, "multiplicity == 3", None, outfile_skm, "n")
+    assert ac == len(lstore.read_object("/evt/energy", outfile_skm)[0].to_aoesa().nda)
+
+    skim_evt(outfile, "multiplicity == 3", None, None, "o")
+    assert ac == len(lstore.read_object("/evt/energy", outfile)[0].to_aoesa().nda)