Merge pull request #557 from legend-exp/dev

Allow `build_evt()` and `build_skm()` to just return the output table

.github/workflows/main.yml

@@ -13,7 +13,7 @@ concurrency:
   cancel-in-progress: true
 
 env:
-  TQDM_MININTERVAL: 10
+  TQDM_MININTERVAL: 100
 
 jobs:
 

docs/source/conf.py

@@ -58,7 +58,6 @@
     "scipy": ("https://docs.scipy.org/doc/scipy", None),
     "pandas": ("https://pandas.pydata.org/docs", None),
     "matplotlib": ("https://matplotlib.org/stable", None),
-    "iminuit": ("https://iminuit.readthedocs.io/en/stable", None),
     "h5py": ("https://docs.h5py.org/en/stable", None),
     "pint": ("https://pint.readthedocs.io/en/stable", None),
     "lgdo": ("https://legend-pydataobj.readthedocs.io/en/stable", None),

src/pygama/evt/build_evt.py

@@ -24,15 +24,15 @@ def build_evt(
     f_tcm: str,
     f_dsp: str,
     f_hit: str,
-    f_evt: str,
     evt_config: str | dict,
+    f_evt: str | None = None,
     wo_mode: str = "write_safe",
     evt_group: str = "evt",
     tcm_group: str = "hardware_tcm_1",
     dsp_group: str = "dsp",
     hit_group: str = "hit",
     tcm_id_table_pattern: str = "ch{}",
-) -> None:
+) -> None | Table:
     """Transform data from the `hit` and `dsp` levels which a channel sorted to a
     event sorted data format.
 
@@ -44,8 +44,6 @@ def build_evt(
         input LH5 file of the `dsp` level.
     f_hit
         input LH5 file of the `hit` level.
-    f_evt
-        name of the output file.
     evt_config
         name of configuration file or dictionary defining event fields. Channel
         lists can be defined by importing a metadata module.
@@ -108,6 +106,9 @@ def build_evt(
               }
             }
 
+    f_evt
+        name of the output file. If ``None``, return the output :class:`.Table`
+        instead of writing to disk.
     wo_mode
         writing mode.
     evt group
@@ -280,19 +281,25 @@ def build_evt(
     if "outputs" in tbl_cfg.keys():
         if len(tbl_cfg["outputs"]) < 1:
             log.warning("No output fields specified, no file will be written.")
+            return table
         else:
             clms_to_remove = [e for e in table.keys() if e not in tbl_cfg["outputs"]]
             for fld in clms_to_remove:
                 table.remove_field(fld, True)
-            store.write(
-                obj=table, name=f"/{evt_group}/", lh5_file=f_evt, wo_mode=wo_mode
-            )
+
+            if f_evt:
+                store.write(
+                    obj=table, name=f"/{evt_group}/", lh5_file=f_evt, wo_mode=wo_mode
+                )
+            else:
+                return table
     else:
         log.warning("No output fields specified, no file will be written.")
 
     key = re.search(r"\d{8}T\d{6}Z", f_hit).group(0)
     log.info(
-        f"Applied {len(tbl_cfg['operations'])} operations to key {key} and saved {len(tbl_cfg['outputs'])} evt fields across {len(chns)} channel groups"
+        f"Applied {len(tbl_cfg['operations'])} operations to key {key} and saved "
+        f"{len(tbl_cfg['outputs'])} evt fields across {len(chns)} channel groups"
     )
 
 

src/pygama/skm/build_skm.py

@@ -24,16 +24,16 @@ def build_skm(
     f_hit: str,
     f_dsp: str,
     f_tcm: str,
-    f_skm: str,
     skm_conf: dict | str,
-    wo_mode="w",
+    f_skm: str | None = None,
+    wo_mode: str = "w",
     skm_group: str = "skm",
     evt_group: str = "evt",
     tcm_group: str = "hardware_tcm_1",
     dsp_group: str = "dsp",
     hit_group: str = "hit",
     tcm_id_table_pattern: str = "ch{}",
-) -> None:
+) -> None | Table:
     """Builds a skimmed file from a (set) of `evt/hit/dsp` tier file(s).
 
     Parameters
@@ -46,8 +46,6 @@ def build_skm(
         path of `dsp` file.
     f_tcm
         path of `tcm` file.
-    f_skm
-        name of the `skm` output file.
     skm_conf
         name of configuration file or dictionary defining `skm` fields.
 
@@ -87,6 +85,9 @@ def build_skm(
                }
              }
            }
+    f_skm
+        name of the `skm` output file. If ``None``, return the output
+        class:`.Table` instead of writing to disk.
 
     wo_mode
         writing mode.
@@ -229,9 +230,13 @@ def build_skm(
                     obj.attrs |= tbl_cfg["operations"][op]["lgdo_attrs"]
                 table.add_field(op, obj, True)
 
+    if not f_skm:
+        return table
+
     # last thing missing is writing it out
     if wo_mode not in ["w", "write_safe", "o", "overwrite", "a", "append"]:
         raise ValueError(f"wo_mode {wo_mode} not valid.")
+
     log.debug("saving skm file")
     if (wo_mode in ["w", "write_safe"]) and os.path.exists(f_skm):
         raise FileExistsError(f"Write_safe mode: {f_skm} exists.")

tests/evt/test_build_evt.py

@@ -18,18 +18,20 @@ def test_basics(lgnd_test_data, tmptestdir):
     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,
         evt_config=f"{config_dir}/basic-evt-config.json",
+        f_evt=outfile,
         wo_mode="o",
         evt_group="evt",
         hit_group="hit",
         dsp_group="dsp",
         tcm_group="hardware_tcm_1",
     )
+
     assert "statement" in store.read("/evt/multiplicity", outfile)[0].getattrs().keys()
     assert (
         store.read("/evt/multiplicity", outfile)[0].getattrs()["statement"]
@@ -74,8 +76,8 @@ def test_lar_module(lgnd_test_data, tmptestdir):
         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,
         evt_config=f"{config_dir}/module-test-evt-config.json",
+        f_evt=outfile,
         wo_mode="o",
         evt_group="evt",
         hit_group="hit",
@@ -103,8 +105,8 @@ def test_lar_t0_vov_module(lgnd_test_data, tmptestdir):
         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,
         evt_config=f"{config_dir}/module-test-t0-vov-evt-config.json",
+        f_evt=outfile,
         wo_mode="o",
         evt_group="evt",
         hit_group="hit",
@@ -136,8 +138,8 @@ def test_vov(lgnd_test_data, tmptestdir):
         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,
         evt_config=f"{config_dir}/vov-test-evt-config.json",
+        f_evt=outfile,
         wo_mode="o",
         evt_group="evt",
         hit_group="hit",
@@ -187,21 +189,21 @@ def test_graceful_crashing(lgnd_test_data, tmptestdir):
     f_config = f"{config_dir}/basic-evt-config.json"
 
     with pytest.raises(KeyError):
-        build_evt(f_dsp, f_tcm, f_hit, outfile, f_config)
+        build_evt(f_dsp, f_tcm, f_hit, f_config, outfile)
 
     with pytest.raises(KeyError):
-        build_evt(f_tcm, f_hit, f_dsp, outfile, f_config)
+        build_evt(f_tcm, f_hit, f_dsp, f_config, outfile)
 
     with pytest.raises(TypeError):
-        build_evt(f_tcm, f_dsp, f_hit, outfile, None)
+        build_evt(f_tcm, f_dsp, f_hit, None, outfile)
 
     conf = {"operations": {}}
     with pytest.raises(ValueError):
-        build_evt(f_tcm, f_dsp, f_hit, outfile, conf)
+        build_evt(f_tcm, f_dsp, f_hit, conf, outfile)
 
     conf = {"channels": {"geds_on": ["ch1084803", "ch1084804", "ch1121600"]}}
     with pytest.raises(ValueError):
-        build_evt(f_tcm, f_dsp, f_hit, outfile, conf)
+        build_evt(f_tcm, f_dsp, f_hit, conf, outfile)
 
     conf = {
         "channels": {"geds_on": ["ch1084803", "ch1084804", "ch1121600"]},
@@ -217,7 +219,7 @@ def test_graceful_crashing(lgnd_test_data, tmptestdir):
         },
     }
     with pytest.raises(ValueError):
-        build_evt(f_tcm, f_dsp, f_hit, outfile, conf)
+        build_evt(f_tcm, f_dsp, f_hit, conf, outfile)
 
 
 def test_query(lgnd_test_data, tmptestdir):
@@ -229,8 +231,8 @@ def test_query(lgnd_test_data, tmptestdir):
         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,
         evt_config=f"{config_dir}/query-test-evt-config.json",
+        f_evt=outfile,
         wo_mode="o",
         evt_group="evt",
         hit_group="hit",
@@ -277,7 +279,7 @@ def test_vector_sort(lgnd_test_data, tmptestdir):
             },
         },
     }
-    build_evt(f_tcm, f_dsp, f_hit, outfile, conf)
+    build_evt(f_tcm, f_dsp, f_hit, conf, outfile)
 
     assert os.path.exists(outfile)
     assert len(lh5.ls(outfile, "/evt/")) == 4
@@ -300,14 +302,14 @@ def test_tcm_id_table_pattern(lgnd_test_data, tmptestdir):
     f_config = f"{config_dir}/basic-evt-config.json"
 
     with pytest.raises(ValueError):
-        build_evt(f_tcm, f_dsp, f_hit, outfile, f_config, tcm_id_table_pattern="ch{{}}")
+        build_evt(f_tcm, f_dsp, f_hit, f_config, outfile, tcm_id_table_pattern="ch{{}}")
     with pytest.raises(ValueError):
-        build_evt(f_tcm, f_dsp, f_hit, outfile, f_config, tcm_id_table_pattern="ch{}{}")
+        build_evt(f_tcm, f_dsp, f_hit, f_config, outfile, tcm_id_table_pattern="ch{}{}")
     with pytest.raises(NotImplementedError):
         build_evt(
-            f_tcm, f_dsp, f_hit, outfile, f_config, tcm_id_table_pattern="ch{tcm_id}"
+            f_tcm, f_dsp, f_hit, f_config, outfile, tcm_id_table_pattern="ch{tcm_id}"
         )
     with pytest.raises(ValueError):
         build_evt(
-            f_tcm, f_dsp, f_hit, outfile, f_config, tcm_id_table_pattern="apple{}banana"
+            f_tcm, f_dsp, f_hit, f_config, outfile, tcm_id_table_pattern="apple{}banana"
         )

tests/skm/configs/basic-skm-config.json

@@ -13,7 +13,7 @@
     },
     "energy": {
       "forward_field": "hit.cuspEmax_ctc_cal",
-      "missing_value": "np.nan",
+      "missing_value": 0.0,
       "tcm_idx": "evt.energy_idx"
     },
     "energy_id": {

tests/skm/test_build_skm.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 
 import awkward as ak
-import numpy as np
+import lgdo
 from lgdo.lh5 import LH5Store
 
 from pygama.evt import build_evt
@@ -18,12 +18,13 @@ def test_basics(lgnd_test_data, tmptestdir):
     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,
         evt_config=f"{evt_config_dir}/vov-test-evt-config.json",
+        f_evt=outfile,
         wo_mode="o",
         evt_group="evt",
         hit_group="hit",
@@ -33,18 +34,33 @@ def test_basics(lgnd_test_data, tmptestdir):
 
     skm_conf = f"{config_dir}/basic-skm-config.json"
     skm_out = f"{tmptestdir}/l200-p03-r001-phy-20230322T160139Z-tier_skm.lh5"
+
+    result = build_skm(
+        outfile,
+        lgnd_test_data.get_path(tcm_path.replace("tcm", "hit")),
+        lgnd_test_data.get_path(tcm_path.replace("tcm", "dsp")),
+        lgnd_test_data.get_path(tcm_path),
+        skm_conf,
+    )
+
+    assert isinstance(result, lgdo.Table)
+
     build_skm(
         outfile,
         lgnd_test_data.get_path(tcm_path.replace("tcm", "hit")),
         lgnd_test_data.get_path(tcm_path.replace("tcm", "dsp")),
         lgnd_test_data.get_path(tcm_path),
-        skm_out,
         skm_conf,
+        skm_out,
         wo_mode="o",
     )
 
     assert os.path.exists(skm_out)
-    df = store.read("/skm/", skm_out)[0].view_as("pd")
+    obj, _ = store.read("/skm/", skm_out)
+
+    assert obj == result
+
+    df = obj.view_as("pd")
     assert "timestamp" in df.keys()
     assert "energy_0" in df.keys()
     assert "energy_1" in df.keys()
@@ -56,9 +72,7 @@ def test_basics(lgnd_test_data, tmptestdir):
     assert "energy_sum" in df.keys()
     assert (df.multiplicity.to_numpy() <= 3).all()
     assert (
-        np.nan_to_num(df.energy_0.to_numpy())
-        + np.nan_to_num(df.energy_1.to_numpy())
-        + np.nan_to_num(df.energy_2.to_numpy())
+        df.energy_0.to_numpy() + df.energy_1.to_numpy() + df.energy_2.to_numpy()
         == df.energy_sum.to_numpy()
     ).all()
 
@@ -81,12 +95,13 @@ def test_attribute_passing(lgnd_test_data, tmptestdir):
     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,
         evt_config=f"{evt_config_dir}/vov-test-evt-config.json",
+        f_evt=outfile,
         wo_mode="o",
         evt_group="evt",
         hit_group="hit",
@@ -103,8 +118,8 @@ def test_attribute_passing(lgnd_test_data, tmptestdir):
         lgnd_test_data.get_path(tcm_path.replace("tcm", "hit")),
         lgnd_test_data.get_path(tcm_path.replace("tcm", "dsp")),
         lgnd_test_data.get_path(tcm_path),
-        skm_out,
         skm_conf,
+        f_skm=skm_out,
         wo_mode="o",
     )