diff --git a/pyproject.toml b/pyproject.toml index 83018bd..89e37f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -162,7 +162,7 @@ isort.required-imports = ["from __future__ import annotations"] [tool.pylint] -py-version = "3.9" +py-version = "3.11" ignore-paths = [".*/_version.py"] reports.output-format = "colorized" similarities.ignore-imports = "yes" diff --git a/src/legenddataflowscripts/par/geds/dsp/dplms.py b/src/legenddataflowscripts/par/geds/dsp/dplms.py index 085a6e8..f4947de 100644 --- a/src/legenddataflowscripts/par/geds/dsp/dplms.py +++ b/src/legenddataflowscripts/par/geds/dsp/dplms.py @@ -14,7 +14,7 @@ from pygama.pargen.dplms_ge_dict import dplms_ge_dict from pygama.pargen.dsp_optimize import run_one_dsp -from ....utils import build_log, convert_dict_np_to_float +from ....utils import build_log, convert_dict_np_to_float, require_config_keys def par_geds_dsp_dplms() -> None: @@ -107,6 +107,7 @@ def par_geds_dsp_dplms() -> None: dplms_dict = Props.read_from(args.config_file) db_dict = Props.read_from(args.database) + require_config_keys(dplms_dict, ["run_dplms"], f"dplms config ({args.config_file})") if dplms_dict["run_dplms"] is True: with Path(args.fft_raw_filelist).open() as f: diff --git a/src/legenddataflowscripts/par/geds/dsp/eopt.py b/src/legenddataflowscripts/par/geds/dsp/eopt.py index 9ca8d21..9c5c5bb 100644 --- a/src/legenddataflowscripts/par/geds/dsp/eopt.py +++ b/src/legenddataflowscripts/par/geds/dsp/eopt.py @@ -19,7 +19,7 @@ run_bayesian_optimisation, ) -from ....utils import build_log +from ....utils import build_log, require_config_keys warnings.filterwarnings(action="ignore", category=RuntimeWarning) try: @@ -107,7 +107,9 @@ def par_geds_dsp_eopt() -> None: argparser.add_argument( "--final-dsp-pars", help="final_dsp_pars", type=str, required=True ) - argparser.add_argument("--qbb-grid-path", help="qbb_grid_path", type=str) + argparser.add_argument( + "--qbb-grid-path", help="qbb_grid_path", type=str, required=True + ) argparser.add_argument("--plot-path", help="plot_path", type=str) argparser.add_argument( @@ -122,8 +124,10 @@ def par_geds_dsp_eopt() -> None: opt_dict = Props.read_from(args.config_file) db_dict = Props.read_from(args.decay_const) + require_config_keys(opt_dict, ["run_eopt"], f"eopt config ({args.config_file})") - if opt_dict.pop("run_eopt") is True: + run_eopt = opt_dict.pop("run_eopt") is True + if run_eopt: peaks_kev = np.array(opt_dict["peaks"]) kev_widths = [tuple(kev_width) for kev_width in opt_dict["kev_widths"]] @@ -431,20 +435,22 @@ def par_geds_dsp_eopt() -> None: else: plot_dict = {} - plot_dict["trap_optimisation"] = { - "kernel_space": bopt_trap.plot(init_samples=sample_x), - "acq_space": bopt_trap.plot_acq(init_samples=sample_x), - } - - plot_dict["cusp_optimisation"] = { - "kernel_space": bopt_cusp.plot(init_samples=sample_x), - "acq_space": bopt_cusp.plot_acq(init_samples=sample_x), - } - - plot_dict["zac_optimisation"] = { - "kernel_space": bopt_zac.plot(init_samples=sample_x), - "acq_space": bopt_zac.plot_acq(init_samples=sample_x), - } + # the optimiser plots only exist when the optimisation actually ran + if run_eopt: + plot_dict["trap_optimisation"] = { + "kernel_space": bopt_trap.plot(init_samples=sample_x), + "acq_space": bopt_trap.plot_acq(init_samples=sample_x), + } + + plot_dict["cusp_optimisation"] = { + "kernel_space": bopt_cusp.plot(init_samples=sample_x), + "acq_space": bopt_cusp.plot_acq(init_samples=sample_x), + } + + plot_dict["zac_optimisation"] = { + "kernel_space": bopt_zac.plot(init_samples=sample_x), + "acq_space": bopt_zac.plot_acq(init_samples=sample_x), + } Path(args.plot_path).parent.mkdir(parents=True, exist_ok=True) with Path(args.plot_path).open("wb") as w: diff --git a/src/legenddataflowscripts/par/geds/dsp/evtsel.py b/src/legenddataflowscripts/par/geds/dsp/evtsel.py index e4c9864..742c79b 100644 --- a/src/legenddataflowscripts/par/geds/dsp/evtsel.py +++ b/src/legenddataflowscripts/par/geds/dsp/evtsel.py @@ -17,7 +17,13 @@ from dspeed import build_dsp from pygama.pargen.data_cleaning import generate_cuts, get_keys -from ....utils import build_log, get_pulser_mask +from ....utils import ( + build_log, + check_pulser_mask, + expand_filelist, + get_pulser_mask, + require_config_keys, +) warnings.filterwarnings(action="ignore", category=RuntimeWarning) @@ -225,22 +231,20 @@ def par_geds_dsp_evtsel() -> None: peak_dict = Props.read_from(args.config_file) db_dict = Props.read_from(args.decay_const) + require_config_keys( + peak_dict, ["run_selection"], f"evtsel config ({args.config_file})" + ) Path(args.peak_file).parent.mkdir(parents=True, exist_ok=True) if peak_dict.pop("run_selection") is True: log.debug("Starting peak selection") + require_config_keys( + peak_dict, + ["peaks", "kev_widths", "cut_parameters", "n_events", "final_cut_field"], + f"evtsel config ({args.config_file})", + ) - if ( - isinstance(args.raw_filelist, list) - and args.raw_filelist[0].split(".")[-1] == "filelist" - ): - files = args.raw_filelist[0] - with Path(files).open() as f: - files = f.read().splitlines() - else: - files = args.raw_filelist - - raw_files = sorted(files) + raw_files = expand_filelist(args.raw_filelist, "--raw-filelist") peaks_kev = peak_dict["peaks"] kev_widths = peak_dict["kev_widths"] @@ -264,9 +268,13 @@ def par_geds_dsp_evtsel() -> None: ) if args.no_pulse is False: + if args.pulser_file is None: + msg = "either --pulser-file or --no-pulse is required" + raise ValueError(msg) mask = get_pulser_mask( args.pulser_file, ) + check_pulser_mask(mask, tb, args.raw_table_name) else: mask = np.full(len(tb), False) diff --git a/src/legenddataflowscripts/par/geds/dsp/nopt.py b/src/legenddataflowscripts/par/geds/dsp/nopt.py index f987d63..8ac0198 100644 --- a/src/legenddataflowscripts/par/geds/dsp/nopt.py +++ b/src/legenddataflowscripts/par/geds/dsp/nopt.py @@ -12,7 +12,7 @@ from dspeed import build_dsp from pygama.pargen.data_cleaning import generate_cuts, get_cut_indexes -from ....utils import build_log +from ....utils import build_log, require_config_keys def par_geds_dsp_nopt() -> None: @@ -91,6 +91,7 @@ def par_geds_dsp_nopt() -> None: opt_dict = Props.read_from(args.config_file) db_dict = Props.read_from(args.database) + require_config_keys(opt_dict, ["run_nopt"], f"nopt config ({args.config_file})") if opt_dict.pop("run_nopt") is True: with Path(args.raw_filelist).open() as f: @@ -162,4 +163,4 @@ def par_geds_dsp_nopt() -> None: pkl.dump(plot_dict, f, protocol=pkl.HIGHEST_PROTOCOL) Path(args.dsp_pars).parent.mkdir(parents=True, exist_ok=True) - Props.write_to(args.dsp_pars, dict(nopt_pars=out_dict, **db_dict)) + Props.write_to(args.dsp_pars, {**db_dict, "nopt_pars": out_dict}) diff --git a/src/legenddataflowscripts/par/geds/dsp/pz.py b/src/legenddataflowscripts/par/geds/dsp/pz.py index ad7472b..de1bcae 100644 --- a/src/legenddataflowscripts/par/geds/dsp/pz.py +++ b/src/legenddataflowscripts/par/geds/dsp/pz.py @@ -14,8 +14,11 @@ from ....utils import ( build_log, + check_pulser_mask, convert_dict_np_to_float, + expand_filelist, get_pulser_mask, + require_config_keys, ) @@ -99,31 +102,25 @@ def par_geds_dsp_pz() -> None: log = build_log(args.log_config, args.log) kwarg_dict = Props.read_from(args.config_file) + require_config_keys(kwarg_dict, ["run_tau"], f"pz config ({args.config_file})") if kwarg_dict["run_tau"] is True: + require_config_keys( + kwarg_dict, ["threshold"], f"pz config ({args.config_file})" + ) dsp_config = Props.read_from(args.processing_chain) kwarg_dict.pop("run_tau") + # prefer dedicated pz files when given; an empty pz filelist falls + # back to the raw files input_file = [] - 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() + if args.pz_files: + if len(args.pz_files) == 1 and Path(args.pz_files[0]).suffix == ".filelist": + with Path(args.pz_files[0]).open() as f: + input_file = [line for line in map(str.strip, f) if line] 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 + input_file = expand_filelist(args.raw_files, "--raw-files") msg = f"Reading Data for {args.raw_table_name} from:" log.debug(msg) @@ -139,7 +136,11 @@ def par_geds_dsp_pz() -> None: if args.no_pulse is False and ( args.pz_files is None or len(args.pz_files) == 0 ): + if args.pulser_file is None: + msg = "either --pulser-file or --no-pulse is required" + raise ValueError(msg) mask = get_pulser_mask(args.pulser_file) + check_pulser_mask(mask, data, args.raw_table_name) else: mask = np.full(len(data), False) diff --git a/src/legenddataflowscripts/par/geds/hit/aoe.py b/src/legenddataflowscripts/par/geds/hit/aoe.py index 001bc17..f7c77ca 100644 --- a/src/legenddataflowscripts/par/geds/hit/aoe.py +++ b/src/legenddataflowscripts/par/geds/hit/aoe.py @@ -17,9 +17,12 @@ from ....utils import ( build_log, + check_pulser_mask, convert_dict_np_to_float, + expand_filelist, fill_plot_dict, get_pulser_mask, + require_config_keys, ) log = logging.getLogger(__name__) @@ -338,6 +341,7 @@ def par_geds_hit_aoe() -> None: log = build_log(args.log_config, args.log) kwarg_dict = Props.read_from(args.config_file) + require_config_keys(kwarg_dict, ["run_aoe"], f"aoe config ({args.config_file})") ecal_dict = Props.read_from(args.ecal_file) cal_dict = ecal_dict["pars"] @@ -352,10 +356,14 @@ def par_geds_hit_aoe() -> None: else: out_plot_dict = {} - with Path(args.files[0]).open() as f: - files = sorted(f.read().splitlines()) + files = expand_filelist(args.files) if kwarg_dict["run_aoe"] is True: + require_config_keys( + kwarg_dict, + ["current_param", "energy_param", "cal_energy_param", "cut_field"], + f"aoe config ({args.config_file})", + ) params = [ kwarg_dict["current_param"], "tp_0_est", @@ -393,17 +401,15 @@ def par_geds_hit_aoe() -> None: else: mask = np.zeros(len(threshold_mask), dtype=bool) + check_pulser_mask(mask, threshold_mask, args.table_name) data["is_pulser"] = mask[threshold_mask] msg = f"{len(data.query('~is_pulser'))} non pulser events" log.info(msg) - with Path(args.files[0]).open() as f: - files = f.read().splitlines() - files = sorted(files) - data["run_timestamp"] = args.timestamp + override_dict = None if args.override_files: override_dict = Props.read_from(args.override_files) override_dict = override_dict.get(args.detector, None) diff --git a/src/legenddataflowscripts/par/geds/hit/ecal.py b/src/legenddataflowscripts/par/geds/hit/ecal.py index 94be187..e0c5bee 100644 --- a/src/legenddataflowscripts/par/geds/hit/ecal.py +++ b/src/legenddataflowscripts/par/geds/hit/ecal.py @@ -23,8 +23,11 @@ from ....utils import ( build_log, + check_pulser_mask, convert_dict_np_to_float, + expand_filelist, get_pulser_mask, + require_config_keys, ) mpl.use("agg") @@ -196,13 +199,13 @@ def plot_pulser_timemap( def get_median(x): - if len(x[~np.isnan(x)]) >= 10: + if len(x[~np.isnan(x)]) < 10: return np.nan return np.nanpercentile(x, 50) def get_err(x): - if len(x[~np.isnan(x)]) >= 10: + if len(x[~np.isnan(x)]) < 10: return np.nan return np.nanvar(x) / np.sqrt(len(x)) @@ -809,6 +812,8 @@ def par_geds_hit_ecal() -> None: log = build_log(args.log_config, args.log) + hit_dict = {} + in_results_dict = {} if args.in_hit_dict: hit_dict = Props.read_from(args.in_hit_dict) in_results_dict = hit_dict.get("results", {}) @@ -828,6 +833,18 @@ def par_geds_hit_ecal() -> None: hit_dict.update(database_dic["ctc_params"]) kwarg_dict = Props.read_from(args.config_file) + require_config_keys( + kwarg_dict, + [ + "plot_options", + "bl_plot_options", + "common_plots", + "energy_params", + "cut_param", + "threshold", + ], + f"ecal config ({args.config_file})", + ) # convert plot functions from strings to functions and split off baseline and common plots for field, item in kwarg_dict["plot_options"].items(): @@ -838,9 +855,7 @@ def par_geds_hit_ecal() -> None: bl_plots[field]["function"] = eval(item["function"]) common_plots = kwarg_dict.pop("common_plots") - with Path(args.files[0]).open() as f: - files = f.read().splitlines() - files = sorted(files) + files = expand_filelist(args.files) # load data in data, threshold_mask = load_data( @@ -865,6 +880,7 @@ def par_geds_hit_ecal() -> None: else: mask = np.zeros(len(threshold_mask), dtype=bool) + check_pulser_mask(mask, threshold_mask, args.table_name) data["is_pulser"] = mask[threshold_mask] pk_pars = [ @@ -920,9 +936,13 @@ def par_geds_hit_ecal() -> None: if kwarg_dict.get("guess_offset", False) is True: guess = np.array([0, guess]) - guess[0] = hit_dict.get( + # "guess_offset_param" is a fallback key into hit_dict, not a value + offset_cfg = hit_dict.get( "guess_offset_param", "is_valid_cuspEmax_classifier" - )["parameters"]["a"] + ) + if isinstance(offset_cfg, str): + offset_cfg = hit_dict[offset_cfg] + guess[0] = offset_cfg["parameters"]["a"] full_object_dict[cal_energy_param] = HPGeCalibration( energy_param, diff --git a/src/legenddataflowscripts/par/geds/hit/lq.py b/src/legenddataflowscripts/par/geds/hit/lq.py index f6f43d2..5aad304 100644 --- a/src/legenddataflowscripts/par/geds/hit/lq.py +++ b/src/legenddataflowscripts/par/geds/hit/lq.py @@ -19,9 +19,12 @@ from ....utils import ( build_log, + check_pulser_mask, convert_dict_np_to_float, + expand_filelist, fill_plot_dict, get_pulser_mask, + require_config_keys, ) log = logging.getLogger(__name__) @@ -187,7 +190,14 @@ def run_lq_calibration( if isinstance(configs, str | list): configs = Props.read_from(configs) + require_config_keys(configs, ["run_lq"], "lq calibration config") + if configs.pop("run_lq") is True: + require_config_keys( + configs, + ["energy_param", "cal_energy_param", "dt_param", "cut_field"], + "lq calibration config", + ) if "plot_options" in configs: for field, item in configs["plot_options"].items(): configs["plot_options"][field]["function"] = eval(item["function"]) @@ -248,9 +258,13 @@ def eres_func(x): lq_plot_dict = {} lq_obj = None + # mirror the aoe result shape: one (shared) lq results entry per timestamp + if not isinstance(out_dict, dict) or set(out_dict) != set(cal_dicts): + out_dict = dict.fromkeys(cal_dicts, out_dict) + out_result_dicts = {} for tstamp, result_dict in results_dicts.items(): - out_result_dicts[tstamp] = dict(**result_dict, lq=out_dict) + out_result_dicts[tstamp] = dict(**result_dict, lq=out_dict[tstamp]) out_object_dicts = {} for tstamp, object_dict in object_dicts.items(): @@ -358,6 +372,7 @@ def par_geds_hit_lq() -> None: build_log(args.log_config, args.log) kwarg_dict = Props.read_from(args.config_file) + require_config_keys(kwarg_dict, ["run_lq"], f"lq config ({args.config_file})") ecal_dict = Props.read_from(args.ecal_file) cal_dict = ecal_dict["pars"]["operations"] @@ -373,9 +388,12 @@ def par_geds_hit_lq() -> None: object_dict = pkl.load(o) if kwarg_dict["run_lq"] is True: - with Path(args.files[0]).open() as f: - files = f.read().splitlines() - files = sorted(files) + require_config_keys( + kwarg_dict, + ["energy_param", "cal_energy_param", "cut_field", "threshold"], + f"lq config ({args.config_file})", + ) + files = expand_filelist(args.files) params = [ "lq80", @@ -405,6 +423,7 @@ def par_geds_hit_lq() -> None: else: mask = np.zeros(len(threshold_mask), dtype=bool) + check_pulser_mask(mask, threshold_mask, args.table_name) data["is_pulser"] = mask[threshold_mask] msg = f"{len(data.query('~is_pulser'))} non pulser events" diff --git a/src/legenddataflowscripts/par/geds/hit/qc.py b/src/legenddataflowscripts/par/geds/hit/qc.py index fcac348..ee46f80 100644 --- a/src/legenddataflowscripts/par/geds/hit/qc.py +++ b/src/legenddataflowscripts/par/geds/hit/qc.py @@ -20,7 +20,9 @@ from ....utils import ( build_log, + check_pulser_mask, convert_dict_np_to_float, + expand_filelist, get_pulser_mask, ) @@ -110,6 +112,11 @@ def build_qc( ) if len(fft_files) > 0: + if kwarg_dict_fft is None: + msg = ( + "fft files were provided but the qc config has no 'fft_fields' section" + ) + raise ValueError(msg) fft_data = load_data( fft_files, table_name, @@ -161,12 +168,12 @@ def build_qc( for outname, info in cut_dict.items(): # convert to pandas eval exp = info["expression"] - for key in info.get("parameters", None): + for key in info.get("parameters", {}): exp = re.sub( f"(? None: else: overwrite = None + # an empty fft filelist is legitimate (no fft runs for this period) 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() + fft_files = [line for line in map(str.strip, f) if line] 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.cal_files + cal_files = expand_filelist(args.cal_files, "--cal-files") start = time.time() log.info("starting qc") diff --git a/src/legenddataflowscripts/tier/dsp.py b/src/legenddataflowscripts/tier/dsp.py index 1b34ce7..f22b154 100644 --- a/src/legenddataflowscripts/tier/dsp.py +++ b/src/legenddataflowscripts/tier/dsp.py @@ -107,7 +107,7 @@ def build_tier_dsp() -> None: argparser.add_argument( "--table-map", help="mapping from channel to table name", - required=False, + required=True, type=str, ) argparser.add_argument("--log", help="log file name") @@ -130,7 +130,7 @@ def build_tier_dsp() -> None: # set number of threads to use # set_num_threads(1) - table_map = json.loads(args.table_map) if args.table_map is not None else None + table_map = json.loads(args.table_map) df_configs = TextDB(args.configs, lazy=True) config_dict = df_configs.on(args.timestamp, system=args.datatype).snakemake_rules @@ -152,10 +152,9 @@ def build_tier_dsp() -> None: # now construct the dictionary of DSP configs for build_dsp() dsp_cfg_tbl_dict = {} for chan, file in chan_cfg_map.items(): - if chan in table_map: - input_tbl_name = table_map[chan] if table_map is not None else chan + "/raw" - else: + if chan not in table_map: continue + input_tbl_name = table_map[chan] # check if the raw tables are all existing if len(lh5.ls(args.input, input_tbl_name)) > 0: diff --git a/src/legenddataflowscripts/tier/hit.py b/src/legenddataflowscripts/tier/hit.py index add2a39..a686b65 100644 --- a/src/legenddataflowscripts/tier/hit.py +++ b/src/legenddataflowscripts/tier/hit.py @@ -51,10 +51,10 @@ def build_tier_hit() -> None: """ argparser = argparse.ArgumentParser() argparser.add_argument("--input") - argparser.add_argument("--pars-file", nargs="*") + argparser.add_argument("--pars-file", nargs="*", default=[]) argparser.add_argument("--configs", required=True) - argparser.add_argument("--table-map", required=False, type=str) + argparser.add_argument("--table-map", required=True, type=str) argparser.add_argument("--log") argparser.add_argument("--alias-table", help="Alias table", type=str, default=None) @@ -65,7 +65,7 @@ def build_tier_hit() -> None: argparser.add_argument("--output") args = argparser.parse_args() - table_map = json.loads(args.table_map) if args.table_map is not None else None + table_map = json.loads(args.table_map) df_config = ( TextDB(args.configs, lazy=True) @@ -98,10 +98,9 @@ def build_tier_hit() -> None: # get pars (to override hit config) hit_cfg = Props.add_to(hit_cfg, pars_dict.get(chan, {}).copy()) - if chan in table_map: - input_tbl_name = table_map[chan] if table_map is not None else chan + "/dsp" - else: + if chan not in table_map: continue + input_tbl_name = table_map[chan] # check if the raw tables are all existing if len(lh5.ls(args.input, input_tbl_name)) > 0: @@ -154,7 +153,7 @@ def build_tier_hit_single_channel() -> None: """ argparser = argparse.ArgumentParser() argparser.add_argument("--input") - argparser.add_argument("--pars-file", nargs="*") + argparser.add_argument("--pars-file", nargs="*", default=[]) argparser.add_argument("--configs", required=True) argparser.add_argument("--log") @@ -193,14 +192,11 @@ def build_tier_hit_single_channel() -> None: else chan_cfg_map ) - # now construct the dictionary of hit configs for build_hit() - channel_dict = {} + # now construct the hit config for build_hit() pars_dict = Props.read_from(args.pars_file) - pars_dict = ( - pars_dict[args.channel] - if args.channel is not None and args.channel in pars_dict - else pars_dict - ) + # mirror the multi-channel path: channel entries hold the pars under "pars" + if args.channel is not None and args.channel in pars_dict: + pars_dict = pars_dict[args.channel]["pars"] hit_cfg = Props.read_from(chan_cfg_map) hit_cfg = Props.add_to(hit_cfg, pars_dict.copy()) @@ -208,6 +204,6 @@ def build_tier_hit_single_channel() -> None: log.info("running build_hit()...") start = time.time() Path(args.output).parent.mkdir(parents=True, exist_ok=True) - build_hit(args.input, hit_config=channel_dict, outfile=args.output) + build_hit(args.input, hit_config=hit_cfg, outfile=args.output) msg = f"Hit built in {time.time() - start:.2f} seconds" log.info(msg) diff --git a/src/legenddataflowscripts/utils/__init__.py b/src/legenddataflowscripts/utils/__init__.py index 86e94a2..b36229c 100644 --- a/src/legenddataflowscripts/utils/__init__.py +++ b/src/legenddataflowscripts/utils/__init__.py @@ -1,17 +1,22 @@ from __future__ import annotations from .alias_table import alias_table -from .cfgtools import get_channel_config +from .cfgtools import get_channel_config, get_rule_config, require_config_keys from .convert_np import convert_dict_np_to_float +from .files import expand_filelist from .log import build_log from .plot_dict import fill_plot_dict -from .pulser_removal import get_pulser_mask +from .pulser_removal import check_pulser_mask, get_pulser_mask __all__ = [ "alias_table", "build_log", + "check_pulser_mask", "convert_dict_np_to_float", + "expand_filelist", "fill_plot_dict", "get_channel_config", "get_pulser_mask", + "get_rule_config", + "require_config_keys", ] diff --git a/src/legenddataflowscripts/utils/alias_table.py b/src/legenddataflowscripts/utils/alias_table.py index 25656e3..df53e2e 100644 --- a/src/legenddataflowscripts/utils/alias_table.py +++ b/src/legenddataflowscripts/utils/alias_table.py @@ -20,21 +20,28 @@ def convert_parents_to_structs(h5group): h5group : h5py.Group Leaf group whose parent hierarchy should be annotated. """ + child = h5group.name.split("/")[-1] if h5group.parent.name != "/" and len(h5group.parent.attrs) == 0: + h5group.parent.attrs.update({"datatype": "struct{" + child + "}"}) + elif len(h5group.parent.attrs) > 0: + datatype = h5group.parent.attrs["datatype"] + if not (datatype.startswith("struct{") and datatype.endswith("}")): + msg = ( + f"cannot register alias {child!r}: parent group " + f"{h5group.parent.name!r} has non-struct datatype {datatype!r}" + ) + raise ValueError(msg) + # membership must be tested against the parsed field list, not by + # substring (e.g. 'raw' is a substring of 'raw_blind') + fields = [ + field.strip() + for field in datatype[len("struct{") : -1].split(",") + if field.strip() + ] + if child in fields: + return h5group.parent.attrs.update( - {"datatype": "struct{" + h5group.name.split("/")[-1] + "}"} - ) - elif ( - len(h5group.parent.attrs) > 0 - and h5group.name.split("/")[-1] not in h5group.parent.attrs["datatype"] - ): - h5group.parent.attrs.update( - { - "datatype": h5group.parent.attrs["datatype"][:-1] - + "," - + h5group.name.split("/")[-1] - + "}" - } + {"datatype": "struct{" + ",".join([*fields, child]) + "}"} ) else: return diff --git a/src/legenddataflowscripts/utils/cfgtools.py b/src/legenddataflowscripts/utils/cfgtools.py index 8a7605d..e7d008a 100644 --- a/src/legenddataflowscripts/utils/cfgtools.py +++ b/src/legenddataflowscripts/utils/cfgtools.py @@ -1,10 +1,15 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Iterable, Mapping + +from dbetto import TextDB def get_channel_config( - mapping: Mapping, channel: str, default_key: str = "__default__" + mapping: Mapping, + channel: str, + default_key: str = "__default__", + name: str | None = None, ): """Return the configuration entry for *channel* with fallback to a default. @@ -23,6 +28,8 @@ def get_channel_config( default_key : str Fallback key used when *channel* is not present in *mapping*. Defaults to ``"__default__"``. + name : str, optional + Human-readable name of the mapping, used in the error message. Returns ------- @@ -35,4 +42,52 @@ def get_channel_config( KeyError If neither *channel* nor *default_key* is present in *mapping*. """ - return mapping.get(channel, mapping[default_key]) + if channel in mapping: + return mapping[channel] + try: + return mapping[default_key] + except KeyError: + msg = ( + f"channel {channel} has no entry in " + f"{name if name is not None else 'the channel config mapping'} " + f"and no {default_key!r} fallback entry is present" + ) + raise KeyError(msg) from None + + +def require_config_keys(config: Mapping, keys: Iterable[str], context: str) -> None: + """Raise ValueError listing all ``keys`` missing from ``config``. + + Parameters + ---------- + config : collections.abc.Mapping + Mapping to validate. + keys : iterable of str + Required key names. + context : str + Free-form string naming the config in the error message, e.g. + ``f"channel {channel} aoecal config ({config_file})"``. + """ + missing = [key for key in keys if key not in config] + if missing: + msg = f"{context} is missing required key(s) {missing}" + raise ValueError(msg) + + +def get_rule_config(configs_path, rule_name, timestamp, datatype): + """Resolve the dataflow config for one Snakemake rule. + + Wraps ``TextDB(...).on(...)["snakemake_rules"][rule_name]`` so that a + missing key names the rule, timestamp, datatype and config path instead + of raising a bare :class:`KeyError`. + """ + configs = TextDB(configs_path, lazy=True).on(timestamp, system=datatype) + try: + return configs["snakemake_rules"][rule_name] + except KeyError as err: + msg = ( + f"config resolved from {configs_path} for timestamp {timestamp} " + f"(system {datatype}) has no snakemake_rules.{rule_name} entry: " + f"missing key {err}" + ) + raise KeyError(msg) from None diff --git a/src/legenddataflowscripts/utils/files.py b/src/legenddataflowscripts/utils/files.py new file mode 100644 index 0000000..6f4b0d0 --- /dev/null +++ b/src/legenddataflowscripts/utils/files.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np + + +def expand_filelist(files, argname="--files"): + """Expand a CLI file argument that may be a single ``.filelist`` file. + + Parameters + ---------- + files : list of str or None + Value of an ``nargs="*"`` argparse argument: either the input files + themselves or a single ``.filelist`` file containing one path per + line. + argname : str + Name of the CLI argument, used in error messages. + + Returns + ------- + list of str + The sorted, de-duplicated input file paths. + """ + if not files: + msg = f"{argname}: no input files provided" + raise ValueError(msg) + if len(files) == 1 and Path(files[0]).suffix == ".filelist": + with Path(files[0]).open() as f: + expanded = [line.strip() for line in f if line.strip()] + if not expanded: + msg = f"{argname}: filelist {files[0]} is empty" + raise ValueError(msg) + files = expanded + return sorted(np.unique(files)) diff --git a/src/legenddataflowscripts/utils/log.py b/src/legenddataflowscripts/utils/log.py index c2de27c..7433857 100644 --- a/src/legenddataflowscripts/utils/log.py +++ b/src/legenddataflowscripts/utils/log.py @@ -2,7 +2,6 @@ import logging import sys -import traceback from logging.config import dictConfig from pathlib import Path @@ -75,7 +74,9 @@ def build_log( After the logger is created, :data:`sys.stderr` is redirected to it at :data:`logging.ERROR` level, and :data:`sys.excepthook` is overridden so - that unhandled exceptions are written to the same file handler. + that unhandled exceptions are logged as ``ERROR`` records (with the + traceback appended by the formatter). :class:`KeyboardInterrupt` is + passed through to the default excepthook. Parameters ---------- @@ -129,12 +130,7 @@ def build_log( dataflow.setdefault("class", "logging.FileHandler") dataflow.setdefault("level", "INFO") log_config.setdefault("version", 1) - if ( - "handlers" in log_config - and "dataflow" in log_config["handlers"] - and "root" not in log_config - and "loggers" not in log_config - ): + if "root" not in log_config and "loggers" not in log_config: dataflow_level = log_config["handlers"]["dataflow"].get("level", "INFO") log_config["root"] = { "level": dataflow_level, @@ -142,7 +138,7 @@ def build_log( } dictConfig(log_config) - log = logging.getLogger(config_dict["options"].get("logger", "prod")) + log = logging.getLogger(config_dict["options"].get("logger", fallback)) else: Path(log_file).parent.mkdir(parents=True, exist_ok=True) @@ -153,17 +149,13 @@ def build_log( # Redirect stderr to the logger (using the error level) sys.stderr = StreamToLogger(log, logging.ERROR) - # Extract the stream from the logger's file handler. - log_stream = None - for handler in log.handlers: - if hasattr(handler, "stream"): - log_stream = handler.stream - break - if log_stream is None: - log_stream = sys.stdout - + # log uncaught exceptions as ERROR records so the traceback goes through + # the formatter (and its severity token) into the configured handlers def excepthook(exc_type, exc_value, exc_traceback): - traceback.print_exception(exc_type, exc_value, exc_traceback, file=log_stream) + if issubclass(exc_type, KeyboardInterrupt): + sys.__excepthook__(exc_type, exc_value, exc_traceback) + return + log.error("uncaught exception", exc_info=(exc_type, exc_value, exc_traceback)) sys.excepthook = excepthook diff --git a/src/legenddataflowscripts/utils/pulser_removal.py b/src/legenddataflowscripts/utils/pulser_removal.py index 0dcfc73..1017c9c 100644 --- a/src/legenddataflowscripts/utils/pulser_removal.py +++ b/src/legenddataflowscripts/utils/pulser_removal.py @@ -26,7 +26,36 @@ def get_pulser_mask(pulser_file): mask = np.array([], dtype=bool) for file in pulser_file: pulser_dict = Props.read_from(file) + if "mask" not in pulser_dict: + msg = f"pulser file {file} does not contain a 'mask' key" + raise KeyError(msg) pulser_mask = np.array(pulser_dict["mask"]) mask = np.append(mask, pulser_mask) return mask + + +def check_pulser_mask(mask, threshold_mask, context) -> None: + """Validate that the pulser mask matches the loaded events. + + ``mask`` (from the pulser files) must have one entry per event read from + the input files, i.e. the same length as the ``threshold_mask`` + selection returned by ``load_data`` (only ``len()`` is used, so any + sized object is accepted). + + Parameters + ---------- + mask + Pulser mask array. + threshold_mask + Event selection mask (or any sized object with one entry per event). + context + Table/channel name used in the error message. + """ + if len(mask) != len(threshold_mask): + msg = ( + f"pulser mask length {len(mask)} != number of loaded events " + f"{len(threshold_mask)} for {context}; the pulser files and " + "input filelists likely cover different runs" + ) + raise ValueError(msg) diff --git a/src/legenddataflowscripts/workflow/execenv.py b/src/legenddataflowscripts/workflow/execenv.py index 8c9290e..2714349 100644 --- a/src/legenddataflowscripts/workflow/execenv.py +++ b/src/legenddataflowscripts/workflow/execenv.py @@ -173,6 +173,7 @@ def dataflow() -> None: "-v", "--verbose", help="increase verbosity", action="store_true" ) + parser.set_defaults(func=None) subparsers = parser.add_subparsers() parser_install = subparsers.add_parser( @@ -240,8 +241,22 @@ def dataflow() -> None: logger.setLevel(logging.DEBUG) logger.addHandler(handler) - if args.func: + if getattr(args, "func", None) is not None: args.func(args) + else: + parser.print_help(sys.stderr) + sys.exit(1) + + +def _select_execenv(config_dict, system, config_file): + try: + return config_dict["execenv"][system] + except KeyError: + msg = ( + f"system {system!r} not found under 'execenv' in {config_file}; " + f"available systems: {list(config_dict.get('execenv', {}))}" + ) + raise KeyError(msg) from None def install(args) -> None: @@ -270,7 +285,7 @@ def install(args) -> None: config_dict, var_values={"_": config_loc}, use_env=True, ignore_missing=False ) - config_dict["execenv"] = config_dict.execenv[args.system] + config_dict["execenv"] = _select_execenv(config_dict, args.system, args.config_file) # path to virtualenv location path_install = config_dict.paths.install @@ -373,7 +388,7 @@ def cmdexec(args) -> None: use_env=True, ignore_missing=False, ) - config_dict["execenv"] = config_dict["execenv"][args.system] + config_dict["execenv"] = _select_execenv(config_dict, args.system, args.config_file) exe_path = Path(config_dict.paths.install).resolve() / "bin" diff --git a/src/legenddataflowscripts/workflow/filedb.py b/src/legenddataflowscripts/workflow/filedb.py index 4e3c458..a0a0311 100644 --- a/src/legenddataflowscripts/workflow/filedb.py +++ b/src/legenddataflowscripts/workflow/filedb.py @@ -72,7 +72,11 @@ def build_filedb() -> None: log = logging.getLogger(__name__) if args.ignore_keys is not None: - ignore = Props.read_from(args.ignore_keys)["unprocessable"] + ignore_dict = Props.read_from(args.ignore_keys) + if "unprocessable" not in ignore_dict: + msg = f"ignore-keys file {args.ignore_keys} has no 'unprocessable' key" + raise KeyError(msg) + ignore = ignore_dict["unprocessable"] else: ignore = [] @@ -82,7 +86,6 @@ def build_filedb() -> None: except Exception as e: msg = f"error when building {args.output} from {args.scan_path}" raise RuntimeError(msg) from e - fdb.scan_files([args.scan_path]) fdb.scan_tables_columns(dir_files_conform=True) # augment dataframe with earliest timestamp found in file @@ -123,10 +126,9 @@ def build_filedb() -> None: if found and args.assume_nonsparse: break - if ( - (loc_timestamps == default).all() or not found - ) and row.raw_file not in ignore: - msg = "something went wrong! no valid first timestamp found. Likely: the file {row.raw_file} is empty" + # rows matching an ignore key were already dropped above + if (loc_timestamps == default).all() or not found: + msg = f"something went wrong! no valid first timestamp found. Likely: the file {row.raw_file} is empty" raise RuntimeError(msg) timestamps[i] = np.min(loc_timestamps) @@ -134,9 +136,7 @@ def build_filedb() -> None: msg = f"found {timestamps[i]}" log.info(msg) - if ( - timestamps[i] < 0 or timestamps[i] > 4102444800 - ) and row.raw_file not in ignore: + if timestamps[i] < 0 or timestamps[i] > 4102444800: msg = f"something went wrong! timestamp {timestamps[i]} does not make sense in {row.raw_file}" raise RuntimeError(msg) diff --git a/src/legenddataflowscripts/workflow/utils.py b/src/legenddataflowscripts/workflow/utils.py index 1793c44..eb34471 100644 --- a/src/legenddataflowscripts/workflow/utils.py +++ b/src/legenddataflowscripts/workflow/utils.py @@ -209,6 +209,6 @@ def as_ro(config, path): if isinstance(path, str): return re.sub(*sub_pattern, path) if isinstance(path, Path): - return Path(re.sub(*sub_pattern, path.name)) + return Path(re.sub(*sub_pattern, str(path))) return [as_ro(config, p) for p in path] diff --git a/tests/test_console_scripts.py b/tests/test_console_scripts.py new file mode 100644 index 0000000..b7c64f1 --- /dev/null +++ b/tests/test_console_scripts.py @@ -0,0 +1,36 @@ +"""Smoke tests for the console-script entry points in ``[project.scripts]``. + +Each entry point is imported and invoked with ``--help``. This catches broken +imports (e.g. moved third-party symbols), misconfigured entry points and +argparse setup errors without needing any input data. +""" + +from __future__ import annotations + +import importlib +import sys +import tomllib +from pathlib import Path + +import pytest + +with (Path(__file__).parent.parent / "pyproject.toml").open("rb") as f: + console_scripts = tomllib.load(f)["project"]["scripts"] + + +# the heavy scientific imports (pygama, matplotlib, ...) triggered by these +# modules may emit warnings; they are not what this smoke test is about +@pytest.mark.filterwarnings("ignore") +@pytest.mark.parametrize("script", sorted(console_scripts)) +def test_console_script_help(script, monkeypatch, capsys): + module_name, func_name = console_scripts[script].split(":") + + module = importlib.import_module(module_name) + func = getattr(module, func_name) + assert callable(func) + + monkeypatch.setattr(sys, "argv", [script, "--help"]) + with pytest.raises(SystemExit) as excinfo: + func() + assert excinfo.value.code == 0 + assert "usage" in capsys.readouterr().out.lower() diff --git a/tests/test_execenv.py b/tests/test_execenv.py index 9522d11..f7ee625 100644 --- a/tests/test_execenv.py +++ b/tests/test_execenv.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import sys import pytest from dbetto import AttrsDict @@ -100,3 +101,19 @@ def test_execenv_pyexe(config): "--image=legendexp/legend-base:latest " ".snakemake/software/bin/dio-boe " ) + + +def test_dataflow_no_subcommand(monkeypatch, capsys): + monkeypatch.setattr(sys, "argv", ["dataflow", "-v"]) + with pytest.raises(SystemExit) as excinfo: + execenv.dataflow() + assert excinfo.value.code == 1 + assert "usage" in capsys.readouterr().err.lower() + + +def test_select_execenv_missing_system(): + config = AttrsDict({"execenv": {"bare": {"cmd": "x"}, "lngs": {"cmd": "y"}}}) + assert execenv._select_execenv(config, "bare", "cfg.yaml") == {"cmd": "x"} + + with pytest.raises(KeyError, match=r"not found under 'execenv' in cfg\.yaml"): + execenv._select_execenv(config, "nersc", "cfg.yaml") diff --git a/tests/test_util.py b/tests/test_util.py index de4a7b1..fe89cc7 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -7,6 +7,7 @@ import yaml from legenddataflowscripts.workflow import ( + as_ro, subst_vars, subst_vars_impl, subst_vars_in_snakemake_config, @@ -193,3 +194,15 @@ def test_subst_vars_in_snakemake_config(mock_workflow, mock_os_environ): # noqa assert setup["execenv"]["env"] == {"VAR2": "val2"} assert setup["execenv"]["cmd"] == "apptainer exec" assert setup["execenv"]["arg"] == "prod/container.sif" + + +def test_as_ro(): + config = {"read_only_fs_sub_pattern": ["^/rw", "/ro"]} + + assert as_ro(config, "/rw/a/b") == "/ro/a/b" + # Path inputs must keep their directory components + assert as_ro(config, Path("/rw/a/b")) == Path("/ro/a/b") + assert as_ro(config, ["/rw/a", "/rw/b"]) == ["/ro/a", "/ro/b"] + + # passthrough when no substitution pattern is configured + assert as_ro({}, "/rw/a/b") == "/rw/a/b" diff --git a/tests/test_utils_helpers.py b/tests/test_utils_helpers.py new file mode 100644 index 0000000..109d1bc --- /dev/null +++ b/tests/test_utils_helpers.py @@ -0,0 +1,158 @@ +"""Tests for the shared helpers in ``legenddataflowscripts.utils``.""" + +from __future__ import annotations + +import logging +import sys + +import h5py +import numpy as np +import pytest +import yaml + +from legenddataflowscripts.utils import ( + alias_table, + build_log, + check_pulser_mask, + expand_filelist, + get_channel_config, + get_pulser_mask, + get_rule_config, + require_config_keys, +) + + +def test_get_channel_config(): + mapping = {"ch1000000": {"a": 1}, "__default__": {"a": 0}} + assert get_channel_config(mapping, "ch1000000") == {"a": 1} + assert get_channel_config(mapping, "V99999A") == {"a": 0} + + # channel present but no __default__ must not raise + assert get_channel_config({"ch1000000": {"a": 1}}, "ch1000000") == {"a": 1} + + with pytest.raises(KeyError, match="channel V99999A has no entry in qc_config"): + get_channel_config({"ch1000000": {}}, "V99999A", name="qc_config") + + +def test_require_config_keys(): + config = {"energy_param": "cuspEmax", "threshold": 900} + require_config_keys(config, ["energy_param", "threshold"], "test config") + + with pytest.raises( + ValueError, + match=r"test config is missing required key\(s\) \['final_cut_field'\]", + ): + require_config_keys(config, ["energy_param", "final_cut_field"], "test config") + + +def test_expand_filelist_plain_files(): + assert expand_filelist(["b.lh5", "a.lh5", "b.lh5"]) == ["a.lh5", "b.lh5"] + + +def test_expand_filelist_from_filelist(tmp_path): + filelist = tmp_path / "cal.filelist" + filelist.write_text("b.lh5\n\na.lh5\n \nb.lh5\n") + assert expand_filelist([str(filelist)]) == ["a.lh5", "b.lh5"] + + +def test_expand_filelist_empty_args(): + with pytest.raises(ValueError, match="--tcm-files: no input files"): + expand_filelist([], "--tcm-files") + with pytest.raises(ValueError, match="no input files"): + expand_filelist(None) + + +def test_expand_filelist_empty_filelist(tmp_path): + filelist = tmp_path / "cal.filelist" + filelist.write_text("\n \n") + with pytest.raises(ValueError, match="is empty"): + expand_filelist([str(filelist)]) + + +def test_check_pulser_mask(): + mask = np.array([True, False, True]) + check_pulser_mask(mask, np.array([True, True, False]), "ch1000000") + + with pytest.raises(ValueError, match="pulser mask length 3 != number of loaded"): + check_pulser_mask(mask, np.array([True, False]), "ch1000000") + + +def test_get_pulser_mask_missing_mask_key(tmp_path): + pulser_file = tmp_path / "pulser.yaml" + pulser_file.write_text(yaml.dump({"idxs": [1, 2]})) + with pytest.raises(KeyError, match="does not contain a 'mask' key"): + get_pulser_mask(str(pulser_file)) + + +@pytest.fixture +def configs_tree(tmp_path): + (tmp_path / "config.yaml").write_text( + yaml.dump({"snakemake_rules": {"tier_tcm": {"inputs": {"config": "c.yaml"}}}}) + ) + (tmp_path / "validity.yaml").write_text( + yaml.dump([{"valid_from": "20230101T000000Z", "apply": ["config.yaml"]}]) + ) + return tmp_path + + +def test_get_rule_config(configs_tree): + config = get_rule_config(configs_tree, "tier_tcm", "20230201T000000Z", "cal") + assert config["inputs"]["config"] == "c.yaml" + + with pytest.raises(KeyError, match=r"no snakemake_rules\.tier_dsp entry"): + get_rule_config(configs_tree, "tier_dsp", "20230201T000000Z", "cal") + + +def test_alias_table_membership(tmp_path): + lh5_file = tmp_path / "test.lh5" + with h5py.File(lh5_file, "w") as f: + grp = f.create_group("det/raw_blind") + grp.attrs["datatype"] = "table{a}" + grp.create_dataset("a", data=[1, 2, 3]) + f["det"].attrs["datatype"] = "struct{raw_blind}" + + # 'raw' is a substring of 'raw_blind' and must still be registered + alias_table(lh5_file, {"det/raw_blind": "det/raw"}) + + with h5py.File(lh5_file) as f: + fields = f["det"].attrs["datatype"] + assert fields == "struct{raw_blind,raw}" + assert list(f["det/raw/a"]) == [1, 2, 3] + + # re-registering an existing member must not duplicate it + alias_table(lh5_file, {"det/raw_blind": "det/raw2"}) + with h5py.File(lh5_file) as f: + assert f["det"].attrs["datatype"] == "struct{raw_blind,raw,raw2}" + + +def test_build_log_excepthook(tmp_path, monkeypatch): + monkeypatch.setattr(sys, "stderr", sys.stderr) + monkeypatch.setattr(sys, "excepthook", sys.excepthook) + + log_file = tmp_path / "job.log" + cfg = { + "version": 1, + "disable_existing_loggers": False, + "formatters": {"f": {"format": "%(levelname)s %(message)s"}}, + "handlers": {"dataflow": {"class": "logging.FileHandler", "formatter": "f"}}, + } + log = build_log(cfg, str(log_file)) + + sys.excepthook(ValueError, ValueError("boom"), None) + for handler in logging.getLogger().handlers: + handler.flush() + + text = log_file.read_text() + assert "ERROR" in text + assert "uncaught exception" in text + assert "ValueError: boom" in text + + # KeyboardInterrupt is delegated to the default excepthook and must not + # be logged as a second "uncaught exception" record (it may still reach + # the log via the stderr redirection) + sys.excepthook(KeyboardInterrupt, KeyboardInterrupt(), None) + for handler in logging.getLogger().handlers: + handler.flush() + assert log_file.read_text().count("uncaught exception") == 1 + + assert isinstance(log, logging.Logger)