Update build_evt.py to support latest pygama.build_evt()

legend-exp · Mar 25, 2024 · 1623699 · 1623699
1 parent f103828
commit 1623699
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 37 deletions.
diff --git a/scripts/build_evt.py b/scripts/build_evt.py
@@ -15,19 +15,6 @@
 sto = lh5.LH5Store()
 
 
-def replace_evt_with_key(dic, new_key):
-    for key, d in dic.items():
-        if isinstance(d, dict):
-            dic[key] = replace_evt_with_key(d, new_key)
-        elif isinstance(d, list):
-            dic[key] = [item.replace("evt", new_key) for item in d]
-        elif isinstance(d, str):
-            dic[key] = d.replace("evt", new_key)
-        else:
-            pass
-    return dic
-
-
 argparser = argparse.ArgumentParser()
 argparser.add_argument("--hit_file", help="hit file", type=str)
 argparser.add_argument("--dsp_file", help="dsp file", type=str)
@@ -45,8 +32,12 @@ def replace_evt_with_key(dic, new_key):
 argparser.add_argument("--output", help="output file", type=str)
 args = argparser.parse_args()
 
-pathlib.Path(os.path.dirname(args.log)).mkdir(parents=True, exist_ok=True)
-logging.basicConfig(level=logging.DEBUG, filename=args.log, filemode="w")
+if args.log is not None:
+    pathlib.Path(os.path.dirname(args.log)).mkdir(parents=True, exist_ok=True)
+    logging.basicConfig(level=logging.DEBUG, filename=args.log, filemode="w")
+else:
+    logging.basicConfig(level=logging.DEBUG)
+
 logging.getLogger("numba").setLevel(logging.INFO)
 logging.getLogger("parse").setLevel(logging.INFO)
 logging.getLogger("lgdo").setLevel(logging.INFO)
@@ -88,7 +79,8 @@ def replace_evt_with_key(dic, new_key):
                     else:
                         chans = []
                     _evt_config["channels"][field] = chans
-            evt_config[key] = replace_evt_with_key(_evt_config, f"evt/{key}")
+
+            evt_config[key] = _evt_config
 else:
     evt_config = {"all": Props.read_from(evt_config_file)}
     # block for snakemake to fill in channel lists
@@ -118,17 +110,16 @@ def replace_evt_with_key(dic, new_key):
 
 tables = {}
 for key, config in evt_config.items():
+    datainfo = {
+        "tcm": (args.tcm_file, "hardware_tcm_1", "ch{}"),
+        "dsp": (args.dsp_file, "dsp", "ch{}"),
+        "hit": (args.hit_file, "hit", "ch{}"),
+        "evt": (None, "evt"),
+    }
+
     tables[key] = build_evt(
-        f_tcm=args.tcm_file,
-        f_dsp=args.dsp_file,
-        f_hit=args.hit_file,
-        f_evt=None,
-        evt_config=config,
-        evt_group=f"evt/{key}" if key != "all" else "evt",
-        tcm_group="hardware_tcm_1",
-        dsp_group="dsp",
-        hit_group="hit",
-        tcm_id_table_pattern="ch{}",
+        datainfo,
+        config,
     )
 
 tbl = Table(col_dict=tables)

diff --git a/scripts/pars_hit_aoe.py b/scripts/pars_hit_aoe.py
@@ -221,7 +221,7 @@ def eres_func(x):
         pkl.dump(out_plot_dict, w, protocol=pkl.HIGHEST_PROTOCOL)
 
 pathlib.Path(os.path.dirname(args.hit_pars)).mkdir(parents=True, exist_ok=True)
-results_dict = dict(**ecal_dict["results"], aoe =  out_dict)
+results_dict = dict(**ecal_dict["results"], aoe=out_dict)
 with open(args.hit_pars, "w") as w:
     final_hit_dict = {
         "pars": {"operations": cal_dict},
@@ -231,8 +231,8 @@ def eres_func(x):
 
 pathlib.Path(os.path.dirname(args.aoe_results)).mkdir(parents=True, exist_ok=True)
 final_object_dict = dict(
-        **object_dict,
-        aoe=obj,
-    )
+    **object_dict,
+    aoe=obj,
+)
 with open(args.aoe_results, "wb") as w:
     pkl.dump(final_object_dict, w, protocol=pkl.HIGHEST_PROTOCOL)
diff --git a/scripts/pars_hit_ecal.py b/scripts/pars_hit_ecal.py
@@ -283,19 +283,19 @@ def energy_cal_th(
                     if plot in item:
                         param_dict.update({plot: item[plot]})
                 common_dict.update({key: param_dict})
-        plot_dict = {"ecal":plot_dict}
+        plot_dict = {"ecal": plot_dict}
         plot_dict["common"] = common_dict
 
         with open(args.plot_path, "wb") as f:
             pkl.dump(plot_dict, f, protocol=pkl.HIGHEST_PROTOCOL)
 
     # save output dictionary
-    output_dict = {"pars": out_dict, "results": {"ecal":result_dict}}
+    output_dict = {"pars": out_dict, "results": {"ecal": result_dict}}
     with open(args.save_path, "w") as fp:
         pathlib.Path(os.path.dirname(args.save_path)).mkdir(parents=True, exist_ok=True)
         json.dump(output_dict, fp, indent=4)
 
     # save calibration objects
     with open(args.results_path, "wb") as fp:
         pathlib.Path(os.path.dirname(args.results_path)).mkdir(parents=True, exist_ok=True)
-        pkl.dump({"ecal":ecal_object}, fp, protocol=pkl.HIGHEST_PROTOCOL)
+        pkl.dump({"ecal": ecal_object}, fp, protocol=pkl.HIGHEST_PROTOCOL)
diff --git a/scripts/pars_hit_lq.py b/scripts/pars_hit_lq.py
@@ -220,7 +220,7 @@ def eres_func(x):
         pkl.dump(out_plot_dict, w, protocol=pkl.HIGHEST_PROTOCOL)
 
 
-results_dict = dict(**eres_dict,lq =  out_dict)
+results_dict = dict(**eres_dict, lq=out_dict)
 pathlib.Path(os.path.dirname(args.hit_pars)).mkdir(parents=True, exist_ok=True)
 with open(args.hit_pars, "w") as w:
     final_hit_dict = {
@@ -231,8 +231,8 @@ def eres_func(x):
 
 pathlib.Path(os.path.dirname(args.lq_results)).mkdir(parents=True, exist_ok=True)
 final_object_dict = dict(
-        **object_dict,
-        lq=obj,
-    )
+    **object_dict,
+    lq=obj,
+)
 with open(args.lq_results, "wb") as w:
     pkl.dump(final_object_dict, w, protocol=pkl.HIGHEST_PROTOCOL)