legend-dataflow-scripts 0.1.0__py3-none-any.whl

legenddataflowscripts/par/geds/dsp/pz.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import argparse
+import copy
+import pickle as pkl
+from pathlib import Path
+
+import numpy as np
+from dbetto import TextDB
+from dbetto.catalog import Props
+from lgdo import lh5
+from pygama.pargen.data_cleaning import get_cut_indexes
+from pygama.pargen.dsp_optimize import run_one_dsp
+from pygama.pargen.pz_correct import PZCorrect
+
+from ....utils import (
+    build_log,
+    convert_dict_np_to_float,
+    get_pulser_mask,
+)
+
+
+def par_geds_dsp_pz() -> None:
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--configs", help="configs path", type=str, required=True)
+    argparser.add_argument("--log", help="log file", type=str)
+    argparser.add_argument(
+        "-p", "--no-pulse", help="no pulser present", action="store_true"
+    )
+
+    argparser.add_argument("--datatype", help="Datatype", type=str, required=True)
+    argparser.add_argument("--timestamp", help="Timestamp", type=str, required=True)
+    argparser.add_argument("--channel", help="Channel", type=str, required=True)
+    argparser.add_argument(
+        "--raw-table-name", help="raw table name", type=str, required=True
+    )
+
+    argparser.add_argument("--plot-path", help="plot path", type=str, required=False)
+    argparser.add_argument("--output-file", help="output file", type=str, required=True)
+
+    argparser.add_argument(
+        "--pulser-file", help="pulser file", type=str, required=False
+    )
+
+    argparser.add_argument("--raw-files", help="input files", nargs="*", type=str)
+    argparser.add_argument("--pz-files", help="input files", nargs="*", type=str)
+    args = argparser.parse_args()
+
+    configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype)
+    config_dict = configs["snakemake_rules"]["pars_dsp_tau"]
+
+    log = build_log(config_dict, args.log)
+
+    channel_dict = config_dict["inputs"]["processing_chain"][args.channel]
+    kwarg_dict = config_dict["inputs"]["tau_config"][args.channel]
+
+    kwarg_dict = Props.read_from(kwarg_dict)
+
+    if kwarg_dict["run_tau"] is True:
+        dsp_config = Props.read_from(channel_dict)
+        kwarg_dict.pop("run_tau")
+        if args.pz_files is not None and len(args.pz_files) > 0:
+            if (
+                isinstance(args.pz_files, list)
+                and args.pz_files[0].split(".")[-1] == "filelist"
+            ):
+                input_file = args.pz_files[0]
+                with Path(input_file).open() as f:
+                    input_file = f.read().splitlines()
+            else:
+                input_file = args.pz_files
+        if len(input_file) == 0:
+            if (
+                isinstance(args.raw_files, list)
+                and args.raw_files[0].split(".")[-1] == "filelist"
+            ):
+                input_file = args.raw_files[0]
+                with Path(input_file).open() as f:
+                    input_file = f.read().splitlines()
+            else:
+                input_file = args.raw_files
+
+        msg = f"Reading Data for {args.raw_table_name} from:"
+        log.debug(msg)
+        log.debug(input_file)
+
+        data = lh5.read(
+            args.raw_table_name,
+            input_file,
+            field_mask=["daqenergy", "timestamp", "t_sat_lo"],
+        ).view_as("pd")
+        threshold = kwarg_dict.pop("threshold")
+
+        if args.no_pulse is False and (
+            args.pz_files is None or len(args.pz_files) == 0
+        ):
+            mask = get_pulser_mask(args.pulser_file)
+        else:
+            mask = np.full(len(data), False)
+
+        discharges = data["t_sat_lo"] > 0
+        discharge_timestamps = np.where(data["timestamp"][discharges])[0]
+        is_recovering = np.full(len(data), False, dtype=bool)
+        for tstamp in discharge_timestamps:
+            is_recovering = is_recovering | np.where(
+                (
+                    ((data["timestamp"] - tstamp) < 0.01)
+                    & ((data["timestamp"] - tstamp) > 0)
+                ),
+                True,
+                False,
+            )
+        cuts = np.where(
+            (data.daqenergy.to_numpy() > threshold) & (~mask) & (~is_recovering)
+        )[0]
+        msg = f"{len(cuts)} events passed threshold and pulser cuts"
+        log.debug(msg)
+        log.debug(cuts)
+        tb_data = lh5.read(
+            args.raw_table_name,
+            input_file,
+            idx=cuts,
+            n_rows=kwarg_dict["n_events"] * 2,
+        )
+
+        dsp_config_optimise_removed = copy.deepcopy(dsp_config)
+        if "tau1" in dsp_config["outputs"]:
+            dsp_config_optimise_removed["outputs"].remove("tau1")
+        if "tau2" in dsp_config["outputs"]:
+            dsp_config_optimise_removed["outputs"].remove("tau2")
+        if "frac" in dsp_config["outputs"]:
+            dsp_config_optimise_removed["outputs"].remove("frac")
+
+        tb_out = run_one_dsp(tb_data, dsp_config_optimise_removed)
+        log.debug("Processed Data")
+        cut_parameters = kwarg_dict.get("cut_parameters", None)
+        if cut_parameters is not None:
+            idxs = get_cut_indexes(tb_out, cut_parameters=cut_parameters)
+            log.debug("Applied cuts")
+            msg = f"{len(idxs)} events passed cuts"
+            log.debug(msg)
+            tb_data = lh5.read(
+                args.raw_table_name,
+                input_file,
+                idx=cuts[: 2 * kwarg_dict["n_events"]][idxs],
+                n_rows=kwarg_dict.pop("n_events"),
+            )
+
+        tau = PZCorrect(
+            dsp_config,
+            kwarg_dict["wf_field"],
+            debug_mode=kwarg_dict.get("debug_mode", False),
+        )
+        log.debug("Calculating pz constant")
+        if kwarg_dict["mode"] == "single":
+            tau.get_single_decay_constant(
+                tb_data, kwarg_dict.get("slope_param", "tail_slope")
+            )
+            msg = f"Found tau: {tau.output_dict['pz']['tau']}+- {tau.output_dict['pz']['tau_err']}"
+            log.debug(msg)
+        elif kwarg_dict["mode"] == "double":
+            tau.get_dpz_decay_constants(
+                tb_data,
+                kwarg_dict.get("percent_tau1_fit", 0.1),
+                kwarg_dict.get("percent_tau2_fit", 0.2),
+                kwarg_dict.get("offset_from_wf_max", 10),
+                kwarg_dict.get("superpulse_bl_idx", 25),
+                kwarg_dict.get("superpulse_window_width", 13),
+            )
+            log.debug("found dpz constants : ")
+            for entry in ["tau1", "tau2", "frac"]:
+                msg = f"{entry}:{tau.output_dict['pz'][entry]}+- {tau.output_dict['pz'][f'{entry}_err']}"
+                log.debug(msg)
+        else:
+            msg = f"Unknown mode: {kwarg_dict['mode']}, must be either single or double"
+            raise ValueError(msg)
+        tau.dsp_config = dsp_config_optimise_removed
+
+        if args.plot_path:
+            Path(args.plot_path).parent.mkdir(parents=True, exist_ok=True)
+
+            plot_dict = tau.plot_waveforms_after_correction(
+                tb_data,
+                kwarg_dict("wf_pz_field", "wf_pz"),
+                norm_param=kwarg_dict.get("norm_param", "pz_mean"),
+                xlim=[0, len(tb_data[kwarg_dict["wf_field"]]["values"].nda[0])],
+            )
+
+            zoomed = tau.plot_waveforms_after_correction(
+                tb_data,
+                kwarg_dict("wf_pz_field", "wf_pz"),
+                norm_param=kwarg_dict.get("norm_param", "pz_mean"),
+                xlim=[400, len(tb_data[kwarg_dict["wf_field"]]["values"].nda[0])],
+                ylim=[0.8, 1.1],
+            )
+
+            plot_dict.update({"waveforms_zoomed": zoomed["waveforms"]})
+
+            plot_dict.update(
+                tau.plot_slopes(
+                    tb_data, kwarg_dict.get("final_slope_param", "pz_slope")
+                )
+            )
+            plot_dict.update(
+                tau.plot_slopes(
+                    tb_data, kwarg_dict.get("final_slope_param", "pz_slope"), True
+                )
+            )
+
+            with Path(args.plot_path).open("wb") as f:
+                pkl.dump({"pz": plot_dict}, f, protocol=pkl.HIGHEST_PROTOCOL)
+        out_dict = convert_dict_np_to_float(tau.output_dict)
+    else:
+        out_dict = {}
+
+    Path(args.output_file).parent.mkdir(parents=True, exist_ok=True)
+    Props.write_to(args.output_file, out_dict)

legenddataflowscripts/par/geds/dsp/svm.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from dbetto.catalog import Props
+
+
+def par_geds_dsp_svm() -> None:
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--log", help="log file", type=str)
+    argparser.add_argument(
+        "--output-file", help="output par file", type=str, required=True
+    )
+    argparser.add_argument(
+        "--input-file", help="input par file", type=str, required=True
+    )
+    argparser.add_argument("--svm-file", help="svm file", required=True)
+    args = argparser.parse_args()
+
+    par_data = Props.read_from(args.input_file)
+
+    file = f"'$_/{Path(args.svm_file).name}'"
+
+    par_data["svm"] = {"model_file": file}
+
+    Path(args.output_file).parent.mkdir(parents=True, exist_ok=True)
+    Props.write_to(args.output_file, par_data)

legenddataflowscripts/par/geds/dsp/svm_build.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import argparse
+import pickle as pkl
+from pathlib import Path
+
+from dbetto import TextDB
+from dbetto.catalog import Props
+from lgdo import lh5
+from sklearn.svm import SVC
+
+from ....utils import build_log
+
+
+def par_geds_dsp_svm_build() -> None:
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--log", help="log file", type=str)
+    argparser.add_argument("--configs", help="config file", type=str)
+
+    argparser.add_argument("--datatype", help="Datatype", type=str, required=True)
+    argparser.add_argument("--timestamp", help="Timestamp", type=str, required=True)
+
+    argparser.add_argument(
+        "--output-file", help="output SVM file", type=str, required=True
+    )
+    argparser.add_argument(
+        "--train-data", help="input data file", nargs="*", default=None
+    )
+    argparser.add_argument(
+        "--train-hyperpars", help="input hyperparameter file", nargs="*", default=None
+    )
+    args = argparser.parse_args()
+
+    configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype)
+    config_dict = configs["snakemake_rules"]["pars_dsp_build_svm"]
+
+    log = build_log(config_dict, args.log)
+
+    if args.train_data is not None and len(args.train_data) > 0:
+        # Load files
+        tb = lh5.read("ml_train/dsp", args.train_data)
+        log.debug("loaded data")
+
+        hyperpars = Props.read_from(args.train_hyperpars)
+
+        # Define training inputs
+        dwts_norm = tb["dwt_norm"].nda
+        labels = tb["dc_label"].nda
+
+        log.debug("training model")
+        # Initialize and train SVM
+        svm = SVC(
+            random_state=int(hyperpars["random_state"]),
+            kernel=hyperpars["kernel"],
+            decision_function_shape=hyperpars["decision_function_shape"],
+            class_weight=hyperpars["class_weight"],
+            C=float(hyperpars["C"]),
+            gamma=float(hyperpars["gamma"]),
+            cache_size=1000,
+        )
+
+        svm.fit(dwts_norm, labels)
+        log.debug("trained model")
+    else:
+        svm = None
+
+    # Save trained model with pickle
+    with Path(args.output_file).open("wb") as svm_file:
+        pkl.dump(svm, svm_file, protocol=pkl.HIGHEST_PROTOCOL)

legenddataflowscripts/par/geds/hit/aoe.py

@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import argparse
+import pickle as pkl
+import warnings
+from pathlib import Path
+
+import numpy as np
+from dbetto import TextDB
+from dbetto.catalog import Props
+from pygama.pargen.AoE_cal import *  # noqa: F403
+from pygama.pargen.AoE_cal import CalAoE, Pol1, SigmaFit, aoe_peak
+from pygama.pargen.utils import load_data
+
+from ....utils import (
+    build_log,
+    convert_dict_np_to_float,
+    get_pulser_mask,
+)
+
+warnings.filterwarnings(action="ignore", category=RuntimeWarning)
+
+
+def get_results_dict(aoe_class):
+    return {
+        "cal_energy_param": aoe_class.cal_energy_param,
+        "dt_param": aoe_class.dt_param,
+        "rt_correction": aoe_class.dt_corr,
+        "1000-1300keV": aoe_class.timecorr_df.to_dict("index"),
+        "correction_fit_results": aoe_class.energy_corr_res_dict,
+        "low_cut": aoe_class.low_cut_val,
+        "high_cut": aoe_class.high_cut_val,
+        "low_side_sfs": aoe_class.low_side_sfs.to_dict("index"),
+        "2_side_sfs": aoe_class.two_side_sfs.to_dict("index"),
+    }
+
+
+def fill_plot_dict(aoe_class, data, plot_options, plot_dict=None):
+    if plot_dict is not None:
+        for key, item in plot_options.items():
+            if item["options"] is not None:
+                plot_dict[key] = item["function"](aoe_class, data, **item["options"])
+            else:
+                plot_dict[key] = item["function"](aoe_class, data)
+    else:
+        plot_dict = {}
+    return plot_dict
+
+
+def par_geds_hit_aoe() -> None:
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("files", help="files", nargs="*", type=str)
+    argparser.add_argument(
+        "--pulser-file", help="pulser_file", type=str, required=False
+    )
+    argparser.add_argument(
+        "--tcm-filelist", help="tcm_filelist", type=str, required=False
+    )
+
+    argparser.add_argument("--ecal-file", help="ecal_file", type=str, required=True)
+    argparser.add_argument("--eres-file", help="eres_file", type=str, required=True)
+    argparser.add_argument("--inplots", help="in_plot_path", type=str, required=False)
+
+    argparser.add_argument("--configs", help="configs", type=str, required=True)
+    argparser.add_argument("--log", help="log_file", type=str)
+
+    argparser.add_argument("--datatype", help="Datatype", type=str, required=True)
+    argparser.add_argument("--timestamp", help="Timestamp", type=str, required=True)
+    argparser.add_argument("--channel", help="Channel", type=str, required=True)
+    argparser.add_argument("--table-name", help="table name", type=str, required=True)
+
+    argparser.add_argument("--plot-file", help="plot_file", type=str, required=False)
+    argparser.add_argument("--hit-pars", help="hit_pars", type=str)
+    argparser.add_argument("--aoe-results", help="aoe_results", type=str)
+
+    argparser.add_argument("-d", "--debug", help="debug_mode", action="store_true")
+    args = argparser.parse_args()
+
+    configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype)
+    config_dict = configs["snakemake_rules"]["pars_hit_aoecal"]
+
+    log = build_log(config_dict, args.log)
+
+    channel_dict = config_dict["inputs"]["aoecal_config"][args.channel]
+    kwarg_dict = Props.read_from(channel_dict)
+
+    ecal_dict = Props.read_from(args.ecal_file)
+    cal_dict = ecal_dict["pars"]
+    eres_dict = ecal_dict["results"]["ecal"]
+
+    with Path(args.eres_file).open("rb") as o:
+        object_dict = pkl.load(o)
+
+    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_peak
+
+        sigma_func = (
+            eval(kwarg_dict.pop("sigma_func"))
+            if "sigma_func" in kwarg_dict
+            else SigmaFit
+        )
+
+        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():
+                kwarg_dict["plot_options"][field]["function"] = eval(item["function"])
+
+        with Path(args.files[0]).open() as f:
+            files = f.read().splitlines()
+        files = sorted(files)
+
+        try:
+            eres = eres_dict[kwarg_dict["cal_energy_param"]]["eres_linear"].copy()
+
+            def eres_func(x):
+                return eval(eres["expression"], dict(x=x, **eres["parameters"]))
+
+        except KeyError:
+
+            def eres_func(x):
+                return x * np.nan
+
+        params = [
+            kwarg_dict["current_param"],
+            "tp_0_est",
+            "tp_99",
+            kwarg_dict["energy_param"],
+            kwarg_dict["cal_energy_param"],
+            kwarg_dict["cut_field"],
+            "timestamp",
+        ]
+
+        if "dt_param" in kwarg_dict:
+            params += kwarg_dict["dt_param"]
+        else:
+            params.append("dt_eff")
+
+        if "dt_cut" in kwarg_dict and kwarg_dict["dt_cut"] is not None:
+            cal_dict.update(kwarg_dict["dt_cut"]["cut"])
+            params.append(kwarg_dict["dt_cut"]["out_param"])
+
+        # load data in
+        data, threshold_mask = load_data(
+            files,
+            args.table_name,
+            cal_dict,
+            params=params,
+            threshold=kwarg_dict.pop("threshold"),
+            return_selection_mask=True,
+        )
+
+        mask = get_pulser_mask(
+            pulser_file=args.pulser_file,
+        )
+
+        data["is_pulser"] = mask[threshold_mask]
+
+        data["AoE_Uncorr"] = (
+            data[kwarg_dict["current_param"]] / data[kwarg_dict["energy_param"]]
+        )
+        aoe = CalAoE(
+            cal_dicts=cal_dict,
+            cal_energy_param=kwarg_dict["cal_energy_param"],
+            eres_func=eres_func,
+            pdf=pdf,
+            mean_func=mean_func,
+            sigma_func=sigma_func,
+            selection_string=f"{kwarg_dict.pop('cut_field')}&(~is_pulser)",
+            dt_corr=kwarg_dict.get("dt_corr", False),
+            dep_correct=kwarg_dict.get("dep_correct", False),
+            dt_cut=kwarg_dict.get("dt_cut", None),
+            dt_param=kwarg_dict.get("dt_param", 3),
+            high_cut_val=kwarg_dict.get("high_cut_val", 3),
+            compt_bands_width=kwarg_dict.get("debug_mode", 20),
+            debug_mode=args.debug | kwarg_dict.get("debug_mode", False),
+        )
+        aoe.update_cal_dicts(
+            {
+                "AoE_Uncorr": {
+                    "expression": f"{kwarg_dict['current_param']}/{kwarg_dict['energy_param']}",
+                    "parameters": {},
+                }
+            }
+        )
+        aoe.calibrate(data, "AoE_Uncorr")
+
+        log.info("Calibrated A/E")
+        out_dict = get_results_dict(aoe)
+        plot_dict = fill_plot_dict(aoe, data, kwarg_dict.get("plot_options", None))
+
+        aoe.pdf = aoe.pdf.name
+
+        # need to change eres func as can't pickle lambdas
+        try:
+            aoe.eres_func = eres_dict[kwarg_dict["cal_energy_param"]][
+                "eres_linear"
+            ].copy()
+        except KeyError:
+            aoe.eres_func = {}
+    else:
+        out_dict = {}
+        plot_dict = {}
+        aoe = None
+
+    if args.plot_file:
+        common_dict = plot_dict.pop("common") if "common" in list(plot_dict) else None
+        if args.inplots:
+            with Path(args.inplots).open("rb") as r:
+                out_plot_dict = pkl.load(r)
+            out_plot_dict.update({"aoe": plot_dict})
+        else:
+            out_plot_dict = {"aoe": plot_dict}
+
+        if "common" in list(out_plot_dict) and common_dict is not None:
+            out_plot_dict["common"].update(common_dict)
+        elif common_dict is not None:
+            out_plot_dict["common"] = common_dict
+
+        Path(args.plot_file).parent.mkdir(parents=True, exist_ok=True)
+        with Path(args.plot_file).open("wb") as w:
+            pkl.dump(out_plot_dict, w, protocol=pkl.HIGHEST_PROTOCOL)
+
+    Path(args.hit_pars).parent.mkdir(parents=True, exist_ok=True)
+    results_dict = dict(**ecal_dict["results"], aoe=out_dict)
+    final_hit_dict = {
+        "pars": {"operations": cal_dict},
+        "results": results_dict,
+    }
+
+    final_hit_dict = convert_dict_np_to_float(final_hit_dict)
+
+    Props.write_to(args.hit_pars, final_hit_dict)
+
+    Path(args.aoe_results).parent.mkdir(parents=True, exist_ok=True)
+    final_object_dict = dict(
+        **object_dict,
+        aoe=aoe,
+    )
+    with Path(args.aoe_results).open("wb") as w:
+        pkl.dump(final_object_dict, w, protocol=pkl.HIGHEST_PROTOCOL)