more pre-commit formatting fixes, fixed snakefile imports and init file

Author: ggmarshall

Snakefile

@@ -11,9 +11,18 @@ This includes:
 """
 
 import pathlib, os, json, sys
-import scripts as ds
+import scripts.util as ds
 from scripts.util.pars_loading import pars_catalog
 from scripts.util.patterns import *
+from scripts.util.utils import (
+    subst_vars_in_snakemake_config,
+    config_path,
+    chan_map_path,
+    pars_path,
+    filelist_path,
+    log_path,
+    runcmd,
+)
 from datetime import datetime
 from collections import OrderedDict
 

scripts/pars_hit_aoe.py

@@ -6,8 +6,9 @@
 import pickle as pkl
 
 import numpy as np
-import pygama.pargen.AoE_cal as aoe
 from legendmeta import LegendMetadata
+from pygama.pargen.AoE_cal import *  # noqa: F403
+from pygama.pargen.AoE_cal import aoe_calibration, pol1, sigma_fit, standard_aoe
 
 argparser = argparse.ArgumentParser()
 argparser.add_argument("files", help="files", nargs="*", type=str)
@@ -53,13 +54,11 @@
 if kwarg_dict["run_aoe"] is True:
     kwarg_dict.pop("run_aoe")
 
-    pdf = eval(kwarg_dict.pop("pdf")) if "pdf" in kwarg_dict else aoe.standard_aoe
+    pdf = eval(kwarg_dict.pop("pdf")) if "pdf" in kwarg_dict else standard_aoe
 
-    sigma_func = (
-        eval(kwarg_dict.pop("sigma_func")) if "sigma_func" in kwarg_dict else aoe.sigma_fit
-    )
+    sigma_func = eval(kwarg_dict.pop("sigma_func")) if "sigma_func" in kwarg_dict else sigma_fit
 
-    mean_func = eval(kwarg_dict.pop("mean_func")) if "mean_func" in kwarg_dict else aoe.pol1
+    mean_func = eval(kwarg_dict.pop("mean_func")) if "mean_func" in kwarg_dict else pol1
 
     if "plot_options" in kwarg_dict:
         for field, item in kwarg_dict["plot_options"].items():
@@ -80,7 +79,7 @@ def eres_func(x):
         def eres_func(x):
             return x * np.nan
 
-    cal_dict, out_dict, plot_dict, obj = aoe.aoe_calibration(
+    cal_dict, out_dict, plot_dict, obj = aoe_calibration(
         files,
         lh5_path=f"{args.channel}/dsp",
         cal_dicts=cal_dict,

scripts/pars_hit_ecal.py

@@ -10,10 +10,11 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pygama.math.histogram as pgh
-import pygama.pargen.ecal_th as ect
 from legendmeta import LegendMetadata
 from legendmeta.catalog import Props
 from matplotlib.colors import LogNorm
+from pygama.pargen.ecal_th import *  # noqa: F403
+from pygama.pargen.ecal_th import energy_cal_th
 from scipy.stats import binned_statistic
 
 log = logging.getLogger(__name__)
@@ -164,7 +165,7 @@ def baseline_tracking_plots(files, lh5_path, plot_options=None):
     common_plots = kwarg_dict.pop("common_plots")
 
     if args.plot_path:
-        out_dict, result_dict, plot_dict, ecal_object = ect.energy_cal_th(
+        out_dict, result_dict, plot_dict, ecal_object = energy_cal_th(
             sorted(args.files),
             lh5_path=f"{args.channel}/dsp",
             hit_dict=hit_dict,
@@ -194,7 +195,7 @@ def baseline_tracking_plots(files, lh5_path, plot_options=None):
         with open(args.plot_path, "wb") as f:
             pkl.dump(plot_dict, f, protocol=pkl.HIGHEST_PROTOCOL)
     else:
-        out_dict, result_dict, _ = ect.energy_cal_th(
+        out_dict, result_dict, _ = energy_cal_th(
             sorted(args.files),
             lh5_path=f"{args.channel}/dsp",
             hit_dict=hit_dict,

scripts/pars_pht.py

@@ -6,9 +6,11 @@
 import pickle as pkl
 
 import numpy as np
-import pygama.pargen.AoE_cal as aoe
-import pygama.pargen.ecal_th as ect
 from legendmeta import LegendMetadata
+from pygama.pargen.AoE_cal import *  # noqa: F403
+from pygama.pargen.AoE_cal import aoe_calibration, pol1, sigma_fit, standard_aoe
+from pygama.pargen.ecal_th import *  # noqa: F403
+from pygama.pargen.ecal_th import partition_energy_cal_th
 from util.FileKey import ChannelProcKey, ProcessingFileKey
 
 argparser = argparse.ArgumentParser()
@@ -143,23 +145,23 @@ def run_splitter(files):
     final_dict[timestamp] = sorted(filelist)
 
 # run energy supercal
-ecal_results, ecal_plots, ecal_obj = ect.partition_energy_cal_th(
+ecal_results, ecal_plots, ecal_obj = partition_energy_cal_th(
     final_dict, lh5_path=f"{args.channel}/dsp", hit_dict=cal_dict, **ecal_options
 )
 
 # run aoe cal
 if aoe_options.pop("run_aoe") is True:
-    pdf = eval(aoe_options.pop("pdf")) if "pdf" in aoe_options else aoe.standard_aoe
+    pdf = eval(aoe_options.pop("pdf")) if "pdf" in aoe_options else standard_aoe
 
     if "mean_func" in aoe_options:
         mean_func = eval(aoe_options.pop("mean_func"))
     else:
-        mean_func = aoe.pol1
+        mean_func = pol1
 
     if "sigma_func" in aoe_options:
         sigma_func = eval(aoe_options.pop("sigma_func"))
     else:
-        sigma_func = aoe.sigma_fit
+        sigma_func = sigma_fit
 
     try:
         eres = ecal_results[aoe_options["cal_energy_param"]]["eres_linear"].copy()
@@ -183,7 +185,7 @@ def eres_func(x):
             def eres_func(x):
                 return x * np.nan
 
-    cal_dict, out_dict, plot_dict, aoe_obj = aoe.aoe_calibration(
+    cal_dict, out_dict, plot_dict, aoe_obj = aoe_calibration(
         final_dict,
         lh5_path=f"{args.channel}/dsp",
         cal_dicts=cal_dict,

scripts/util/__init__.py

@@ -1,6 +1,29 @@
-from .CalibCatalog import Props, PropsStream, CalibCatalog
+from .CalibCatalog import CalibCatalog, Props, PropsStream
 from .create_pars_keylist import pars_key_resolve
 from .dataset_cal import dataset_file
-from .FileKey import FileKey, ProcessingFileKey, ChannelProcKey
+from .FileKey import ChannelProcKey, FileKey, ProcessingFileKey
 from .pars_loading import pars_catalog
-from .utils import unix_time, subst_vars, subst_vars_in_snakemake_config, subst_vars_impl, runcmd
+from .utils import (
+    runcmd,
+    subst_vars,
+    subst_vars_impl,
+    subst_vars_in_snakemake_config,
+    unix_time,
+)
+
+__all__ = [
+    "Props",
+    "PropsStream",
+    "CalibCatalog",
+    "pars_key_resolve",
+    "dataset_file",
+    "FileKey",
+    "ProcessingFileKey",
+    "ChannelProcKey",
+    "pars_catalog",
+    "unix_time",
+    "runcmd",
+    "subst_vars_impl",
+    "subst_vars",
+    "subst_vars_in_snakemake_config",
+]

tests/dummy_cycle/generated/par/dsp/validity.jsonl

@@ -1,3 +1,3 @@
 {"valid_from": "20230101T123456Z", "category": "all", "apply": ["cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json"]}
 {"valid_from": "20230110T123456Z", "category": "all", "apply": ["lar/p00/r000/l200-p00-r000-lar-20230110T123456Z-par_dsp.json", "cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json"]}
-{"valid_from": "20230202T004321Z", "category": "all", "apply": ["cal/p00/r001/l200-p00-r001-cal-20230202T004321Z-par_dsp.json","lar/p00/r000/l200-p00-r000-lar-20230110T123456Z-par_dsp.json"]}
+{"valid_from": "20230202T004321Z", "category": "all", "apply": ["cal/p00/r001/l200-p00-r001-cal-20230202T004321Z-par_dsp.json","lar/p00/r000/l200-p00-r000-lar-20230110T123456Z-par_dsp.json"]}

tests/test_util.py

@@ -12,12 +12,10 @@
 )
 from scripts.util.patterns import get_pattern_tier_daq, get_pattern_tier_dsp
 from scripts.util.utils import (
-    tier_path,
     par_dsp_path,
-    par_hit_path,
     par_overwrite_path,
-    pars_path,
     tier_dsp_path,
+    tier_path,
 )
 
 testprod = Path(__file__).parent / "dummy_cycle"
@@ -59,7 +57,9 @@ def test_filekey():
 
 def test_create_pars_keylist():
     key1 = FileKey("l200", "p00", "r000", "cal", "20230101T123456Z")
-    assert pars_key_resolve.from_filekey(key1, {"cal": ["par_dsp"]}).valid_from == "20230101T123456Z"
+    assert (
+        pars_key_resolve.from_filekey(key1, {"cal": ["par_dsp"]}).valid_from == "20230101T123456Z"
+    )
     key2 = FileKey("l200", "p00", "r000", "cal", "20230102T123456Z")
     assert pars_key_resolve.match_keys(key1, key2) == key1
     key3 = FileKey("l200", "p00", "r000", "cal", "20230101T000000Z")
@@ -69,9 +69,7 @@ def test_create_pars_keylist():
     pkey2 = pars_key_resolve.from_filekey(
         FileKey("l200", "p00", "r000", "lar", "20230102T123456Z"), {"lar": ["par_dsp"]}
     )
-    assert pkey2.apply == [
-        "lar/p00/r000/l200-p00-r000-lar-20230102T123456Z-par_dsp.json"
-    ]
+    assert pkey2.apply == ["lar/p00/r000/l200-p00-r000-lar-20230102T123456Z-par_dsp.json"]
     pars_key_resolve.match_entries(pkey1, pkey2)
     assert set(pkey2.apply) == {
         "cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json",
@@ -97,7 +95,11 @@ def test_create_pars_keylist():
 
     pkeylist = pars_key_resolve.generate_par_keylist(keylist)
     assert pkeylist == keylist
-    assert set(pars_key_resolve.match_all_entries(pkeylist, {"cal": ["par_dsp"], "lar": ["par_dsp"]})[1].apply) == {
+    assert set(
+        pars_key_resolve.match_all_entries(pkeylist, {"cal": ["par_dsp"], "lar": ["par_dsp"]})[
+            1
+        ].apply
+    ) == {
         "cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json",
         "lar/p00/r000/l200-p00-r000-lar-20230110T123456Z-par_dsp.json",
     }
@@ -107,9 +109,7 @@ def test_pars_loading():
     pars_files = CalibCatalog.get_calib_files(
         os.path.join(par_dsp_path(setup), "validity.jsonl"), "20230101T123456Z"
     )
-    assert pars_files == [
-        "cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json"
-    ]
+    assert pars_files == ["cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json"]
 
     par_override_files = CalibCatalog.get_calib_files(
         os.path.join(par_overwrite_path(setup), "dsp", "validity.jsonl"), "20230101T123456Z"
@@ -119,10 +119,8 @@ def test_pars_loading():
         pars_files, par_override_files
     )
 
-    assert pars_files == [
-        "cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json"
-    ]
-    
+    assert pars_files == ["cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.json"]
+
     assert set(pars_catalog.get_par_file(setup, "20230101T123456Z", "dsp")) == {
         os.path.join(
             par_dsp_path(setup),