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

legenddataflowscripts/par/geds/hit/lq.py

@@ -0,0 +1,213 @@
+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.math.distributions import gaussian
+from pygama.pargen.AoE_cal import *  # noqa: F403
+from pygama.pargen.lq_cal import *  # noqa: F403
+from pygama.pargen.lq_cal import LQCal
+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(lq_class):
+    return {
+        "cal_energy_param": lq_class.cal_energy_param,
+        "DEP_means": lq_class.timecorr_df.to_dict("index"),
+        "rt_correction": lq_class.dt_fit_pars,
+        "cut_fit_pars": lq_class.cut_fit_pars.to_dict(),
+        "cut_value": lq_class.cut_val,
+        "sfs": lq_class.low_side_sf.to_dict("index"),
+    }
+
+
+def fill_plot_dict(lq_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"](lq_class, data, **item["options"])
+            else:
+                plot_dict[key] = item["function"](lq_class, data)
+    else:
+        plot_dict = {}
+    return plot_dict
+
+
+def par_geds_hit_lq() -> 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("--lq-results", help="lq_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_lqcal"]
+
+    log = build_log(config_dict, args.log)
+
+    channel_dict = config_dict["inputs"]["lqcal_config"][args.channel]
+    kwarg_dict = Props.read_from(channel_dict)
+
+    ecal_dict = Props.read_from(args.ecal_file)
+    cal_dict = ecal_dict["pars"]["operations"]
+    eres_dict = ecal_dict["results"]["ecal"]
+
+    with Path(args.eres_file).open("rb") as o:
+        object_dict = pkl.load(o)
+
+    if kwarg_dict["run_lq"] is True:
+        kwarg_dict.pop("run_lq")
+
+        cdf = eval(kwarg_dict.pop("cdf")) if "cdf" in kwarg_dict else gaussian
+
+        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 = [
+            "lq80",
+            "dt_eff",
+            kwarg_dict["energy_param"],
+            kwarg_dict["cal_energy_param"],
+            kwarg_dict["cut_field"],
+        ]
+
+        # 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]
+
+        lq = LQCal(
+            cal_dict,
+            kwarg_dict["cal_energy_param"],
+            kwarg_dict["dt_param"],
+            eres_func,
+            cdf,
+            selection_string=f"{kwarg_dict.pop('cut_field')}&(~is_pulser)",
+            debug_mode=args.debug | kwarg_dict.get("debug_mode", False),
+        )
+
+        data["LQ_Ecorr"] = np.divide(data["lq80"], data[kwarg_dict["energy_param"]])
+
+        lq.update_cal_dicts(
+            {
+                "LQ_Ecorr": {
+                    "expression": f"lq80/{kwarg_dict['energy_param']}",
+                    "parameters": {},
+                }
+            }
+        )
+
+        lq.calibrate(data, "LQ_Ecorr")
+        log.info("Calibrated LQ")
+
+        out_dict = get_results_dict(lq)
+        plot_dict = fill_plot_dict(lq, data, kwarg_dict.get("plot_options", None))
+
+        # need to change eres func as can't pickle lambdas
+        try:
+            lq.eres_func = eres_dict[kwarg_dict["cal_energy_param"]][
+                "eres_linear"
+            ].copy()
+        except KeyError:
+            lq.eres_func = {}
+    else:
+        out_dict = {}
+        plot_dict = {}
+        lq = 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({"lq": plot_dict})
+        else:
+            out_plot_dict = {"lq": 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)
+
+    final_hit_dict = convert_dict_np_to_float(
+        {
+            "pars": {"operations": cal_dict},
+            "results": dict(**ecal_dict["results"], lq=out_dict),
+        }
+    )
+    Path(args.hit_pars).parent.mkdir(parents=True, exist_ok=True)
+    Props.write_to(args.hit_pars, final_hit_dict)
+
+    final_object_dict = dict(
+        **object_dict,
+        lq=lq,
+    )
+    Path(args.lq_results).parent.mkdir(parents=True, exist_ok=True)
+    with Path(args.lq_results).open("wb") as w:
+        pkl.dump(final_object_dict, w, protocol=pkl.HIGHEST_PROTOCOL)

legenddataflowscripts/par/geds/hit/qc.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+import argparse
+import json
+import pickle as pkl
+import re
+import warnings
+from pathlib import Path
+
+import numpy as np
+from dbetto import TextDB
+from dbetto.catalog import Props
+from lgdo.lh5 import ls
+from pygama.pargen.data_cleaning import (
+    generate_cut_classifiers,
+    get_keys,
+)
+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 par_geds_hit_qc() -> None:
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--cal-files", help="cal_files", nargs="*", type=str)
+    argparser.add_argument("--fft-files", help="fft_files", nargs="*", type=str)
+
+    argparser.add_argument(
+        "--tcm-filelist", help="tcm_filelist", type=str, required=False
+    )
+    argparser.add_argument(
+        "--pulser-file", help="pulser_file", type=str, required=False
+    )
+    argparser.add_argument(
+        "--overwrite-files",
+        help="overwrite_files",
+        type=str,
+        required=False,
+        nargs="*",
+    )
+
+    argparser.add_argument("--configs", help="config", 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("--tier", help="tier", type=str, default="hit")
+
+    argparser.add_argument("--plot-path", help="plot_path", type=str, required=False)
+    argparser.add_argument("--save-path", help="save_path", type=str)
+    args = argparser.parse_args()
+
+    configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype)
+    if args.tier == "hit":
+        config_dict = configs["snakemake_rules"]["pars_hit_qc"]
+    elif args.tier == "pht":
+        config_dict = configs["snakemake_rules"]["pars_pht_qc"]
+    else:
+        msg = f"tier {args.tier} not recognized"
+        raise ValueError(msg)
+
+    log = build_log(config_dict, args.log)
+
+    # get metadata dictionary
+    channel_dict = config_dict["inputs"]["qc_config"][args.channel]
+    kwarg_dict = Props.read_from(channel_dict)
+
+    if args.overwrite_files:
+        overwrite = Props.read_from(args.overwrite_files)
+        if args.channel in overwrite:
+            overwrite = overwrite[args.channel]["pars"]["operations"]
+        else:
+            overwrite = None
+    else:
+        overwrite = None
+
+    if len(args.fft_files) == 1 and Path(args.fft_files[0]).suffix == ".filelist":
+        with Path(args.fft_files[0]).open() as f:
+            fft_files = f.read().splitlines()
+    else:
+        fft_files = args.fft_files
+
+    if len(args.cal_files) == 1 and Path(args.cal_files[0]).suffix == ".filelist":
+        with Path(args.cal_files[0]).open() as f:
+            cal_files = f.read().splitlines()
+    else:
+        cal_files = args.fft_files
+
+    search_name = (
+        args.table_name if args.table_name[-1] == "/" else args.table_name + "/"
+    )
+
+    kwarg_dict_fft = kwarg_dict["fft_fields"]
+    kwarg_dict_cal = kwarg_dict["cal_fields"]
+
+    cut_fields = get_keys(
+        [key.replace(search_name, "") for key in ls(cal_files[0], search_name)],
+        kwarg_dict_cal["cut_parameters"],
+    )
+    cut_fields += get_keys(
+        [key.replace(search_name, "") for key in ls(cal_files[0], search_name)],
+        kwarg_dict_fft["cut_parameters"],
+    )
+
+    if "initial_cal_cuts" in kwarg_dict:
+        init_cal = kwarg_dict["initial_cal_cuts"]
+        cut_fields += get_keys(
+            [key.replace(search_name, "") for key in ls(cal_files[0], search_name)],
+            init_cal["cut_parameters"],
+        )
+
+    if len(fft_files) > 0:
+        fft_data = load_data(
+            fft_files,
+            args.table_name,
+            {},
+            [*cut_fields, "t_sat_lo", "timestamp", "trapTmax"],
+        )
+
+        discharges = fft_data["t_sat_lo"] > 0
+        discharge_timestamps = np.where(fft_data["timestamp"][discharges])[0]
+        is_recovering = np.full(len(fft_data), False, dtype=bool)
+        for tstamp in discharge_timestamps:
+            is_recovering = is_recovering | np.where(
+                (
+                    ((fft_data["timestamp"] - tstamp) < 0.01)
+                    & ((fft_data["timestamp"] - tstamp) > 0)
+                ),
+                True,
+                False,
+            )
+        fft_data["is_recovering"] = is_recovering
+
+        hit_dict_fft = {}
+        plot_dict_fft = {}
+        cut_data = fft_data.query("is_recovering==0")
+        msg = f"cut_data shape: {len(cut_data)}"
+        log.debug(msg)
+        for name, cut in kwarg_dict_fft["cut_parameters"].items():
+            cut_dict, cut_plots = generate_cut_classifiers(
+                cut_data,
+                {name: cut},
+                kwarg_dict.get("rounding", 4),
+                display=1 if args.plot_path else 0,
+            )
+            hit_dict_fft.update(cut_dict)
+            plot_dict_fft.update(cut_plots)
+
+            msg = f"{name} calculated cut_dict is: {json.dumps(convert_dict_np_to_float(cut_dict), indent=2)}"
+            log.debug(msg)
+
+            ct_mask = np.full(len(fft_data), True, dtype=bool)
+            for outname, info in cut_dict.items():
+                # convert to pandas eval
+                exp = info["expression"]
+                for key in info.get("parameters", None):
+                    exp = re.sub(
+                        f"(?<![a-zA-Z0-9]){key}(?![a-zA-Z0-9])", f"@{key}", exp
+                    )
+                fft_data[outname] = fft_data.eval(
+                    exp, local_dict=info.get("parameters", None)
+                )
+                if "_classifier" not in outname:
+                    ct_mask = ct_mask & fft_data[outname]
+            cut_data = fft_data[ct_mask]
+
+        log.debug("fft cuts applied")
+        msg = f"cut_dict is: {json.dumps(convert_dict_np_to_float(hit_dict_fft), indent=2)}"
+        log.debug(msg)
+
+    else:
+        hit_dict_fft = {}
+        plot_dict_fft = {}
+
+    if overwrite is not None:
+        for name in kwarg_dict_fft["cut_parameters"]:
+            for cut_name, cut_dict in overwrite.items():
+                if name in cut_name:
+                    hit_dict_fft.update({cut_name: cut_dict})
+
+    # load data in
+    data, threshold_mask = load_data(
+        cal_files,
+        args.table_name,
+        {},
+        [*cut_fields, "timestamp", "trapTmax", "t_sat_lo"],
+        threshold=kwarg_dict_cal.get("threshold", 0),
+        return_selection_mask=True,
+        cal_energy_param="trapTmax",
+    )
+
+    mask = get_pulser_mask(
+        pulser_file=args.pulser_file,
+    )
+
+    data["is_pulser"] = mask[threshold_mask]
+
+    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,
+        )
+    data["is_recovering"] = is_recovering
+
+    rng = np.random.default_rng()
+    mask = np.full(len(data.query("~is_pulser & ~is_recovering")), False, dtype=bool)
+    mask[
+        rng.choice(len(data.query("~is_pulser & ~is_recovering")), 4000, replace=False)
+    ] = True
+
+    if "initial_cal_cuts" in kwarg_dict:
+        init_cal = kwarg_dict["initial_cal_cuts"]
+        hit_dict_init_cal, plot_dict_init_cal = generate_cut_classifiers(
+            data.query("~is_pulser & ~is_recovering")[mask],
+            init_cal["cut_parameters"],
+            init_cal.get("rounding", 4),
+            display=1 if args.plot_path else 0,
+        )
+        ct_mask = np.full(len(data), True, dtype=bool)
+        for outname, info in hit_dict_init_cal.items():
+            # convert to pandas eval
+            exp = info["expression"]
+            for key in info.get("parameters", None):
+                exp = re.sub(f"(?<![a-zA-Z0-9]){key}(?![a-zA-Z0-9])", f"@{key}", exp)
+            data[outname] = data.eval(exp, local_dict=info.get("parameters", None))
+            if "classifier" not in outname:
+                ct_mask = ct_mask & data[outname]
+
+        mask = mask[ct_mask[(~data["is_pulser"] & ~data["is_recovering"]).to_numpy()]]
+        data = data[ct_mask]
+        log.debug("initial cal cuts applied")
+        msg = f"cut_dict is: {json.dumps(convert_dict_np_to_float(hit_dict_init_cal), indent=2)}"
+        log.debug(msg)
+
+    else:
+        hit_dict_init_cal = {}
+        plot_dict_init_cal = {}
+
+    if len(data.query("~is_pulser & ~is_recovering")) < 500:
+        log.info("Less than 500 pulser events")
+        cal_data = data.query("~is_pulser & ~is_recovering")
+    else:
+        cal_data = data.query("~is_pulser & ~is_recovering")[mask]
+
+    hit_dict_cal, plot_dict_cal = generate_cut_classifiers(
+        cal_data,
+        kwarg_dict_cal["cut_parameters"],
+        kwarg_dict.get("rounding", 4),
+        display=1 if args.plot_path else 0,
+    )
+
+    if overwrite is not None:
+        for name in kwarg_dict_cal["cut_parameters"]:
+            for cut_name, cut_dict in overwrite.items():
+                if name in cut_name:
+                    hit_dict_cal.update({cut_name: cut_dict})
+
+    hit_dict = {**hit_dict_fft, **hit_dict_init_cal, **hit_dict_cal}
+    plot_dict = {**plot_dict_fft, **plot_dict_init_cal, **plot_dict_cal}
+
+    hit_dict = convert_dict_np_to_float(hit_dict)
+
+    for outname, info in hit_dict.items():
+        # convert to pandas eval
+        exp = info["expression"]
+        for key in info.get("parameters", None):
+            exp = re.sub(f"(?<![a-zA-Z0-9]){key}(?![a-zA-Z0-9])", f"@{key}", exp)
+        if outname not in fft_data:
+            fft_data[outname] = fft_data.eval(
+                exp, local_dict=info.get("parameters", None)
+            )
+        if outname not in data:
+            data[outname] = data.eval(exp, local_dict=info.get("parameters", None))
+
+    qc_results = {}
+    for entry in hit_dict:
+        if "classifier" not in entry:
+            sf_cal = len(data.query(f"{entry}& ~is_pulser & ~is_recovering")) / len(
+                data.query("~is_pulser & ~is_recovering")
+            )
+            sf_cal_err = 100 * np.sqrt(
+                ((sf_cal) * (1 - sf_cal))
+                / len(data.query("~is_pulser & ~is_recovering"))
+            )
+            sf_fft = len(fft_data.query(f"{entry} & ~is_recovering")) / len(
+                fft_data.query("~is_recovering")
+            )
+            sf_fft_err = 100 * np.sqrt(
+                ((sf_fft) * (1 - sf_fft)) / len(fft_data.query("~is_recovering"))
+            )
+            sf_cal *= 100
+            sf_fft *= 100
+            msg = f"{entry} cut applied: {sf_cal:.2f}% of events passed the cut for cal data, {sf_fft:.2f}% for fft data"
+            log.info(msg)
+            qc_results[entry] = {
+                "sf_cal": sf_cal,
+                "sf_cal_err": sf_cal_err,
+                "sf_fft": sf_fft,
+                "sf_fft_err": sf_fft_err,
+            }
+    qc_results = convert_dict_np_to_float(qc_results)
+
+    Path(args.save_path).parent.mkdir(parents=True, exist_ok=True)
+    Props.write_to(
+        args.save_path, {"operations": hit_dict, "results": {"qc": qc_results}}
+    )
+
+    if args.plot_path:
+        Path(args.plot_path).parent.mkdir(parents=True, exist_ok=True)
+        with Path(args.plot_path).open("wb") as f:
+            pkl.dump({"qc": plot_dict}, f, protocol=pkl.HIGHEST_PROTOCOL)