diff --git a/src/legenddataflowscripts/par/geds/dsp/dplms.py b/src/legenddataflowscripts/par/geds/dsp/dplms.py index b8f3a5c..409733a 100644 --- a/src/legenddataflowscripts/par/geds/dsp/dplms.py +++ b/src/legenddataflowscripts/par/geds/dsp/dplms.py @@ -21,6 +21,8 @@ expand_filelist, prepare_output_paths, require_config_keys, + require_peaks_present, + take_table_rows, ) @@ -62,7 +64,13 @@ def par_geds_dsp_dplms() -> None: Processing chain configuration file(s). ``--config-file`` : list of str DPLMS configuration file(s). Must contain ``run_dplms`` (bool), - ``n_baselines`` (int), and ``peaks_kev`` (list). + ``n_baselines`` (int), and ``peaks_kev`` (list). The optional key + ``use_log_pdf`` (bool, default false) is forwarded to + :func:`pygama.pargen.dplms_ge_dict.dplms_ge_dict` and builds the + grid-search FOM fits from the model's log-density (iminuit + ``log=True``) — substantially faster, results differ at + machine-precision level; requires a pygama version with + ``use_log_pdf`` support. ``--channel`` : str Channel identifier; used as the HDF5 group name in *lh5-path*. ``--raw-table-name`` : str @@ -142,17 +150,20 @@ def par_geds_dsp_dplms() -> None: log.info("\nRunning event selection") peaks_kev = np.array(dplms_dict["peaks_kev"]) - # kev_widths = [tuple(kev_width) for kev_width in dplms_dict["kev_widths"]] peaks_rounded = [int(peak) for peak in peaks_kev] peaks = lh5.read_as(f"{args.raw_table_name}/peak", args.peak_file, library="np") + require_peaks_present( + np.unique(peaks), + peaks_rounded, + f"peak file {args.peak_file} requested by dplms config ({args.config_file})", + ) ids = np.isin(peaks, peaks_rounded) - peaks = peaks[ids] - # idx_list = [np.where(peaks == peak)[0] for peak in peaks_rounded] raw_cal = lh5.read(args.raw_table_name, args.peak_file, idx=ids) msg = f"Time to run event selection {(time.time() - t1):.2f} s, total events {len(raw_cal)}" log.info(msg) + del peaks, ids if isinstance(dsp_config, str | list): dsp_config = Props.read_from(dsp_config) @@ -168,13 +179,12 @@ def par_geds_dsp_dplms() -> None: for cut in cut_dict: idxs = dsp_fft[cut].nda & idxs - selected_eidxs = eidxs[: len(dsp_fft)] - raw_fft = lh5.read( - args.raw_table_name, - fft_files, - n_rows=dplms_dict["n_baselines"], - idx=selected_eidxs[idxs], - ) + # idxs is a boolean mask over the rows already loaded in raw_fft, so + # subset in memory instead of re-reading the FFT files from disk + # (row-identical: the first read was idx=eidxs capped at n_baselines, + # so the re-applied n_rows cap was a no-op) + raw_fft = take_table_rows(raw_fft, idxs) + del dsp_fft, energies, eidxs, cut_dict, idxs log.debug("Applied Cuts") diff --git a/src/legenddataflowscripts/par/geds/dsp/eopt.py b/src/legenddataflowscripts/par/geds/dsp/eopt.py index bcaf77d..50a63b3 100644 --- a/src/legenddataflowscripts/par/geds/dsp/eopt.py +++ b/src/legenddataflowscripts/par/geds/dsp/eopt.py @@ -24,6 +24,7 @@ check_input_files, prepare_output_paths, require_config_keys, + require_peaks_present, ) warnings.filterwarnings(action="ignore", category=RuntimeWarning) @@ -69,7 +70,11 @@ def par_geds_dsp_eopt() -> None: Energy optimisation configuration file(s). Must contain ``run_eopt`` (bool), ``peaks`` (list of keV values), ``kev_widths``, ``fom``, ``fom_field``, ``fom_err_field``, ``initial_samples``, ``acq_func``, - ``batch_size``, and ``n_iter``. + ``batch_size``, and ``n_iter``. Optional: ``use_log_pdf`` (bool, + default false) builds the FOM's unbinned fits from the model's + log-density (iminuit ``log=True``) — substantially faster, results + differ at machine-precision level; requires a pygama version with + ``use_log_pdf`` support. ``--log-config`` : str, optional Logging configuration file. ``--raw-table-name`` : str @@ -144,10 +149,9 @@ def par_geds_dsp_eopt() -> None: kwarg_dicts_cusp = [] kwarg_dicts_trap = [] kwarg_dicts_zac = [] - for peak in peaks_kev: - peak_idx = np.where(peaks_kev == peak)[0][0] - kev_width = kev_widths[peak_idx] - + # strict: a length mismatch between the peaks and kev_widths config + # lists is a config error and must not silently drop peaks + for peak, kev_width in zip(peaks_kev, kev_widths, strict=True): kwarg_dicts_cusp.append( { "parameter": "cuspEmax", @@ -178,6 +182,11 @@ def par_geds_dsp_eopt() -> None: peaks_rounded = [int(peak) for peak in peaks_kev] peaks = lh5.read_as(f"{args.raw_table_name}/peak", args.peak_file, library="np") + require_peaks_present( + np.unique(peaks), + peaks_rounded, + f"peak file {args.peak_file} requested by eopt config ({args.config_file})", + ) ids = np.isin(peaks, peaks_rounded) peaks = peaks[ids] idx_list = [np.where(peaks == peak)[0] for peak in peaks_rounded] @@ -205,29 +214,42 @@ def par_geds_dsp_eopt() -> None: db_dict["etrap"] = {"flat": flat_val} tb_data.add_column("dt_eff", init_data["dt_eff"]) + # add_column stores the Array object itself, so dt_eff survives freeing + # the rest of the init DSP output (otherwise held for the whole + # optimisation) + del init_data, full_dt dsp_config["processors"].pop("dt_eff") dsp_config["outputs"] = ["zacEmax", "cuspEmax", "trapEmax", "dt_eff"] + # opt-in: build the unbinned NLL fits from the model's log-density + # (iminuit log=True mode) — much faster on large samples, results + # differ at machine-precision level. Requires pygama with + # use_log_pdf support; older versions silently ignore the key. + use_log_pdf = opt_dict.get("use_log_pdf", False) + kwarg_dict = [ { "peak_dicts": kwarg_dicts_cusp, "ctc_param": "dt_eff", "idx_list": idx_list, "peaks_kev": peaks_kev, + "use_log_pdf": use_log_pdf, }, { "peak_dicts": kwarg_dicts_zac, "ctc_param": "dt_eff", "idx_list": idx_list, "peaks_kev": peaks_kev, + "use_log_pdf": use_log_pdf, }, { "peak_dicts": kwarg_dicts_trap, "ctc_param": "dt_eff", "idx_list": idx_list, "peaks_kev": peaks_kev, + "use_log_pdf": use_log_pdf, }, ] diff --git a/src/legenddataflowscripts/par/geds/dsp/evtsel.py b/src/legenddataflowscripts/par/geds/dsp/evtsel.py index 9d781b6..41b11a0 100644 --- a/src/legenddataflowscripts/par/geds/dsp/evtsel.py +++ b/src/legenddataflowscripts/par/geds/dsp/evtsel.py @@ -22,6 +22,7 @@ check_pulser_mask, expand_filelist, get_channel_config, + get_is_recovering_mask, get_pulser_mask, require_config_keys, ) @@ -126,10 +127,56 @@ def get_out_data( "trapTmax_cal": lgdo.Array( dsp_data["trapTmax"].nda[final_mask] * ecal_pars ), - "peak": lgdo.Array(np.full(len(np.where(final_mask)[0]), int(peak))), + "peak": lgdo.Array(np.full(int(np.count_nonzero(final_mask)), int(peak))), } ) - return out_tbl, len(np.where(final_mask)[0]) + return out_tbl, int(np.count_nonzero(final_mask)) + + +def build_peak_dicts(peaks_kev, kev_widths, masks, log=None): + """Build the per-peak accumulator state for the file x peak read loop. + + Peaks for which *masks* holds no candidate event are dropped: with an + empty index array nothing is ever read into the peak's buffer, and the + end-of-file flush would then hand ``None`` to + :func:`dspeed.build_dsp`. This happens for channels whose DAQ trigger + threshold sits above a configured gamma line. + + Parameters + ---------- + peaks_kev : sequence of float + Nominal gamma-line energies in keV. + kev_widths : sequence of (float, float) + Lower/upper selection widths in keV, one pair per peak. + masks : dict + Mapping of peak energy to the array of candidate event indices. + log : logging.Logger, optional + Logger used to warn about skipped peaks. + + Returns + ------- + dict + Mapping of peak energy to its accumulator state, with peaks that have + no candidate events omitted. + """ + pk_dicts = {} + for peak, kev_width in zip(peaks_kev, kev_widths, strict=False): + if len(masks[peak]) == 0: + if log is not None: + msg = ( + f"no events in the rough energy range for {peak}, " + "skipping this peak" + ) + log.warning(msg) + continue + pk_dicts[peak] = { + "idxs": (masks[peak],), + "n_rows_read": 0, + "obj_buf_start": 0, + "obj_buf": None, + "kev_width": kev_width, + } + return pk_dicts def par_geds_dsp_evtsel() -> None: @@ -281,16 +328,9 @@ def par_geds_dsp_evtsel() -> None: discharges = tb["t_sat_lo"].nda > 0 discharge_timestamps = np.where(tb["timestamp"].nda[discharges])[0] - is_recovering = np.full(len(tb), False, dtype=bool) - for tstamp in discharge_timestamps: - is_recovering = is_recovering | np.where( - ( - ((tb["timestamp"].nda - tstamp) < 0.01) - & ((tb["timestamp"].nda - tstamp) > 0) - ), - True, - False, - ) + is_recovering = get_is_recovering_mask( + tb["timestamp"].nda, discharge_timestamps + ) if args.raw_cal_curve: raw_dict = get_channel_config( @@ -342,6 +382,10 @@ def par_geds_dsp_evtsel() -> None: input_data = lh5.read( f"{lh5_path}", raw_files, n_rows=10000, idx=np.where(~mask)[0] ) + # the all-events scalar table and its derived masks are no longer + # needed; free them before the per-peak waveform buffers accumulate + del tb, rough_energy, e_mask, mask, is_recovering, discharges + del discharge_timestamps if isinstance(dsp_config, str): dsp_config = Props.read_from(dsp_config) @@ -360,23 +404,19 @@ def par_geds_dsp_evtsel() -> None: log.debug(msg) else: cut_dict = None + # free the cut-derivation waveform buffer and its DSP output — they + # would otherwise stay resident alongside every per-peak obj_buf for + # the whole file x peak loop below + del input_data, tb_data - pk_dicts = {} - for peak, kev_width in zip(peaks_kev, kev_widths, strict=False): - pk_dicts[peak] = { - "idxs": (masks[peak],), - "n_rows_read": 0, - "obj_buf_start": 0, - "obj_buf": None, - "kev_width": kev_width, - } + pk_dicts = build_peak_dicts(peaks_kev, kev_widths, masks, log=log) for file in raw_files: log.debug(Path(file).name) + # idx is a long continuous array + n_rows_i = sto.read_n_rows(lh5_path, file) for peak, peak_dict in pk_dicts.items(): if peak_dict["idxs"] is not None: - # idx is a long continuous array - n_rows_i = sto.read_n_rows(lh5_path, file) # find the length of the subset of idx that contains indices # that are less than n_rows_i n_rows_to_read_i = bisect_left(peak_dict["idxs"][0], n_rows_i) @@ -401,7 +441,13 @@ def par_geds_dsp_evtsel() -> None: log.debug(msg) peak_dict["obj_buf_start"] += n_rows_read_i if peak_dict["n_rows_read"] >= 10000 or file == raw_files[-1]: - if "e_lower_lim" not in peak_dict: + if ( + peak_dict["obj_buf"] is None + or len(peak_dict["obj_buf"]) == 0 + ): + # nothing accumulated for this peak yet, nothing to flush + pass + elif "e_lower_lim" not in peak_dict: tb_out = build_dsp( raw_in=peak_dict["obj_buf"], dsp_config=dsp_config, @@ -499,10 +545,7 @@ def par_geds_dsp_evtsel() -> None: peak_dict["n_events"] = n_wfs msg = f"found {peak_dict['n_events']} events for {peak}" log.debug(msg) - elif ( - peak_dict["obj_buf"] is not None - and len(peak_dict["obj_buf"]) > 0 - ): + else: tb_out = build_dsp( raw_in=peak_dict["obj_buf"], dsp_config=dsp_config, diff --git a/src/legenddataflowscripts/par/geds/dsp/nopt.py b/src/legenddataflowscripts/par/geds/dsp/nopt.py index cdafcf1..e90df3e 100644 --- a/src/legenddataflowscripts/par/geds/dsp/nopt.py +++ b/src/legenddataflowscripts/par/geds/dsp/nopt.py @@ -110,6 +110,7 @@ def par_geds_dsp_nopt() -> None: f"{args.raw_table_name}/daqenergy", raw_files, library="np" ) idxs = np.where(energies <= 10)[0] + del energies tb_data = lh5.read( args.raw_table_name, raw_files, n_rows=opt_dict["n_events"], idx=idxs ) diff --git a/src/legenddataflowscripts/par/geds/dsp/pz.py b/src/legenddataflowscripts/par/geds/dsp/pz.py index 8759185..3baeea2 100644 --- a/src/legenddataflowscripts/par/geds/dsp/pz.py +++ b/src/legenddataflowscripts/par/geds/dsp/pz.py @@ -17,9 +17,11 @@ check_pulser_mask, convert_dict_np_to_float, expand_filelist, + get_is_recovering_mask, get_pulser_mask, prepare_output_paths, require_config_keys, + take_table_rows, ) @@ -134,7 +136,10 @@ def par_geds_dsp_pz() -> None: args.raw_table_name, input_file, field_mask=["daqenergy", "timestamp", "t_sat_lo"], - ).view_as("pd") + ) + daqenergy = data["daqenergy"].nda + timestamps = data["timestamp"].nda + t_sat_lo = data["t_sat_lo"].nda threshold = kwarg_dict.pop("threshold") if args.no_pulse is False and ( @@ -148,24 +153,15 @@ def par_geds_dsp_pz() -> None: 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] + discharges = t_sat_lo > 0 + discharge_timestamps = np.where(timestamps[discharges])[0] + is_recovering = get_is_recovering_mask(timestamps, discharge_timestamps) + cuts = np.where((daqenergy > threshold) & (~mask) & (~is_recovering))[0] msg = f"{len(cuts)} events passed threshold and pulser cuts" log.debug(msg) log.debug(cuts) + del data, daqenergy, timestamps, t_sat_lo, mask, is_recovering, discharges + del discharge_timestamps tb_data = lh5.read( args.raw_table_name, input_file, @@ -189,12 +185,17 @@ def par_geds_dsp_pz() -> None: 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"), - ) + # tb_data already holds exactly the rows cuts[:2*n_events], so + # subset it in memory instead of re-reading the raw files + # (row-identical: the read idx was cuts[:2*n_events][idxs] with + # the first n_events kept, both in ascending row order) + del tb_out + sel = np.asarray(idxs) + # get_cut_indexes returns a boolean mask; also accept an index + # array in case that ever changes + if sel.dtype == np.bool_: + sel = np.flatnonzero(sel) + tb_data = take_table_rows(tb_data, sel[: kwarg_dict.pop("n_events")]) tau = PZCorrect( dsp_config, diff --git a/src/legenddataflowscripts/par/geds/dsp/svm_build.py b/src/legenddataflowscripts/par/geds/dsp/svm_build.py index ddca364..bb5e1c0 100644 --- a/src/legenddataflowscripts/par/geds/dsp/svm_build.py +++ b/src/legenddataflowscripts/par/geds/dsp/svm_build.py @@ -70,8 +70,10 @@ def par_geds_dsp_svm_build() -> None: raise ValueError(msg) check_input_files(args.train_data, "--train-data") - # Load files - tb = lh5.read("ml_train/dsp", args.train_data) + # Load files (only the two columns used below) + tb = lh5.read( + "ml_train/dsp", args.train_data, field_mask=["dwt_norm", "dc_label"] + ) log.debug("loaded data") hyperpars = Props.read_from(args.train_hyperpars) diff --git a/src/legenddataflowscripts/par/geds/hit/aoe.py b/src/legenddataflowscripts/par/geds/hit/aoe.py index b0ca8f2..716c25c 100644 --- a/src/legenddataflowscripts/par/geds/hit/aoe.py +++ b/src/legenddataflowscripts/par/geds/hit/aoe.py @@ -2,6 +2,7 @@ import argparse import copy +import inspect import logging import pickle as pkl import re @@ -110,7 +111,12 @@ def run_aoe_calibration( ``energy_param``; optional keys are ``dt_param``, ``dt_corr``, ``dep_correct``, ``dt_cut``, ``pdf``, ``mean_func``, ``sigma_func``, ``high_cut_val``, ``suffix`` (appended to the output parameter names - so entries don't collide), and ``plot_options``. + so entries don't collide), and ``plot_options``. Optional top-level + key ``use_log_pdf`` (bool, default false): build the unbinned A/E + and survival-fraction fits from the models' log-densities (iminuit + ``log=True``) — substantially faster, results differ at + machine-precision level; requires a pygama version with + ``use_log_pdf`` support (silently ignored otherwise). debug_mode : bool Activates additional diagnostic output in :class:`~pygama.pargen.AoE_cal.CalAoE`. Defaults to ``False``. @@ -169,6 +175,15 @@ def eres_func(x): def eres_func(x): return x * np.nan + # opt-in: build the unbinned A/E and survival-fraction fits from the + # models' log-densities (iminuit log=True mode) — faster, results + # differ at machine-precision level. Needs a pygama version with + # use_log_pdf support; fall back silently on older versions. + use_log_pdf = config.get("use_log_pdf", False) + if use_log_pdf and "use_log_pdf" not in inspect.signature(CalAoE).parameters: + log.warning("installed pygama does not support use_log_pdf, ignoring") + use_log_pdf = False + for name, param_config in config["params"].items(): require_config_keys( param_config, @@ -219,6 +234,7 @@ def eres_func(x): high_cut_val=param_config.get("high_cut_val", 3), compt_bands_width=config.get("debug_mode", 20), debug_mode=debug_mode | config.get("debug_mode", False), + **({"use_log_pdf": True} if use_log_pdf else {}), ) aoe.update_cal_dicts( { diff --git a/src/legenddataflowscripts/par/geds/hit/ecal.py b/src/legenddataflowscripts/par/geds/hit/ecal.py index 74c9fbc..6b51d71 100644 --- a/src/legenddataflowscripts/par/geds/hit/ecal.py +++ b/src/legenddataflowscripts/par/geds/hit/ecal.py @@ -2,6 +2,7 @@ import argparse import copy +import inspect import pickle as pkl import warnings from datetime import datetime @@ -758,7 +759,11 @@ def par_geds_hit_ecal() -> None: ``--log-config`` : str, optional Logging configuration file. ``--config-file`` : list of str - Energy calibration configuration file(s). + Energy calibration configuration file(s). Optional key + ``use_log_pdf`` (bool, default false): build the unbinned peak fits + from the model's log-density (iminuit ``log=True``) — substantially + faster, results differ at machine-precision level; requires a pygama + version with ``use_log_pdf`` support (silently ignored otherwise). ``--det-status`` : str Detector status (``"on"`` or other); affects peak-finding tolerances. Defaults to ``"on"``. @@ -919,6 +924,20 @@ def par_geds_hit_ecal() -> None: selection_string = f"~is_pulser&{kwarg_dict['cut_param']}" + # opt-in: build the unbinned peak fits from the model's log-density + # (iminuit log=True mode) — much faster on large samples, results differ + # at machine-precision level. Needs a pygama version with use_log_pdf + # support; fall back silently on older versions. + use_log_pdf = kwarg_dict.get("use_log_pdf", False) + if ( + use_log_pdf + and "use_log_pdf" + not in inspect.signature(HPGeCalibration.hpge_fit_energy_peaks).parameters + ): + log.warning("installed pygama does not support use_log_pdf, ignoring") + use_log_pdf = False + fit_kwargs = {"use_log_pdf": True} if use_log_pdf else {} + results_dict = {} plot_dict = {} full_object_dict = {} @@ -996,6 +1015,7 @@ def par_geds_hit_ecal() -> None: allowed_p_val=kwarg_dict.get("p_val", 0), update_cal_pars=bool(args.det_status == "on"), bin_width_kev=0.5, + **fit_kwargs, ) full_object_dict[cal_energy_param].hpge_fit_energy_peaks( e_uncal, @@ -1006,6 +1026,7 @@ def par_geds_hit_ecal() -> None: allowed_p_val=kwarg_dict.get("p_val", 0), update_cal_pars=kwarg_dict.get("use_all_peaks", False), bin_width_kev=0.5, + **fit_kwargs, ) full_object_dict[cal_energy_param].get_energy_res_curve( diff --git a/src/legenddataflowscripts/par/geds/hit/lq.py b/src/legenddataflowscripts/par/geds/hit/lq.py index 3c8bd7e..96fdb06 100644 --- a/src/legenddataflowscripts/par/geds/hit/lq.py +++ b/src/legenddataflowscripts/par/geds/hit/lq.py @@ -2,6 +2,7 @@ import argparse import copy +import inspect import logging import pickle as pkl import time @@ -68,6 +69,7 @@ def lq_calibration( selection_string: str = "", plot_options: dict | None = None, debug_mode: bool = False, + use_log_pdf: bool = False, ): """Calibrate the LQ (late-charge) pulse-shape discriminant. @@ -102,6 +104,11 @@ def lq_calibration( passed to :func:`~legenddataflowscripts.utils.fill_plot_dict`. debug_mode : bool Activates additional diagnostic output. Defaults to ``False``. + use_log_pdf : bool + Build the survival-fraction unbinned NLL fits from the models' + log-densities (iminuit ``log=True`` mode) — faster, results differ + at machine-precision level. Silently ignored when the installed + pygama does not support it. Returns ------- @@ -115,6 +122,10 @@ def lq_calibration( Calibrated LQ object. """ + if use_log_pdf and "use_log_pdf" not in inspect.signature(LQCal).parameters: + log.warning("installed pygama does not support use_log_pdf, ignoring") + use_log_pdf = False + lq = LQCal( cal_dicts, cal_energy_param, @@ -123,6 +134,7 @@ def lq_calibration( cdf, selection_string, debug_mode=debug_mode, + **({"use_log_pdf": True} if use_log_pdf else {}), ) data["LQ_Ecorr"] = np.divide(data["lq80"], data[energy_param]) @@ -173,7 +185,11 @@ def run_lq_calibration( configs : dict or str or list LQ calibration configuration. Must contain ``run_lq`` (bool), ``energy_param``, ``cal_energy_param``, ``cut_field``, and - ``dt_param``. + ``dt_param``. Optional key ``use_log_pdf`` (bool, default false): + build the survival-fraction unbinned fits from the models' + log-densities (iminuit ``log=True``) — substantially faster, results + differ at machine-precision level; requires a pygama version with + ``use_log_pdf`` support (silently ignored otherwise). debug_mode : bool Activates additional diagnostic output. Defaults to ``False``. @@ -244,6 +260,7 @@ def eres_func(x): selection_string=f"{configs.pop('cut_field')}&(~is_pulser)", plot_options=configs.get("plot_options", None), debug_mode=debug_mode | configs.get("debug_mode", False), + use_log_pdf=configs.get("use_log_pdf", False), ) msg = f"lq calibration took {time.time() - start:.2f} seconds" log.info(msg) diff --git a/src/legenddataflowscripts/par/geds/hit/qc.py b/src/legenddataflowscripts/par/geds/hit/qc.py index 6a0b4d1..3e9a138 100644 --- a/src/legenddataflowscripts/par/geds/hit/qc.py +++ b/src/legenddataflowscripts/par/geds/hit/qc.py @@ -23,6 +23,7 @@ check_pulser_mask, convert_dict_np_to_float, expand_filelist, + get_is_recovering_mask, get_pulser_mask, prepare_output_paths, ) @@ -128,19 +129,12 @@ def build_qc( msg = f"{len(fft_data)} events" log.info(msg) - 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 + discharges = fft_data["t_sat_lo"].to_numpy() > 0 + timestamps = fft_data["timestamp"].to_numpy() + discharge_timestamps = np.where(timestamps[discharges])[0] + fft_data["is_recovering"] = get_is_recovering_mask( + timestamps, discharge_timestamps + ) msg = f"{len(fft_data.query('is_recovering'))} discharge recovery events" log.info(msg) @@ -219,24 +213,18 @@ def build_qc( msg = f"{len(data.query('~is_pulser'))} non pulser events" log.info(msg) - 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 + discharges = data["t_sat_lo"].to_numpy() > 0 + timestamps = data["timestamp"].to_numpy() + discharge_timestamps = np.where(timestamps[discharges])[0] + data["is_recovering"] = get_is_recovering_mask(timestamps, discharge_timestamps) msg = f"{len(data.query('is_recovering'))} discharge recovery events" log.info(msg) - rng = np.random.default_rng() + # fixed seed: the qc cut values must be reproducible between runs of the + # same data (an unseeded generator made the 4000-event cut-derivation + # sample — and thus the derived cut boundaries — vary run to run) + rng = np.random.default_rng(1) n_clean = len(data.query("~is_pulser & ~is_recovering")) mask = np.full(n_clean, False, dtype=bool) mask[rng.choice(n_clean, min(4000, n_clean), replace=False)] = True diff --git a/src/legenddataflowscripts/utils/__init__.py b/src/legenddataflowscripts/utils/__init__.py index 016dcba..bf47fd7 100644 --- a/src/legenddataflowscripts/utils/__init__.py +++ b/src/legenddataflowscripts/utils/__init__.py @@ -1,14 +1,21 @@ from __future__ import annotations from .alias_table import alias_table -from .cfgtools import get_channel_config, get_rule_config, require_config_keys +from .cfgtools import ( + get_channel_config, + get_rule_config, + require_config_keys, + require_peaks_present, +) from .convert_np import convert_dict_np_to_float +from .discharge import get_is_recovering_mask from .files import ( check_input_files, expand_filelist, parse_json_arg, prepare_output_paths, ) +from .lgdo_utils import take_table_rows from .log import build_log from .plot_dict import fill_plot_dict from .pulser_removal import check_pulser_mask, get_pulser_mask @@ -22,9 +29,12 @@ "expand_filelist", "fill_plot_dict", "get_channel_config", + "get_is_recovering_mask", "get_pulser_mask", "get_rule_config", "parse_json_arg", "prepare_output_paths", "require_config_keys", + "require_peaks_present", + "take_table_rows", ] diff --git a/src/legenddataflowscripts/utils/cfgtools.py b/src/legenddataflowscripts/utils/cfgtools.py index bae80d0..043789b 100644 --- a/src/legenddataflowscripts/utils/cfgtools.py +++ b/src/legenddataflowscripts/utils/cfgtools.py @@ -75,6 +75,33 @@ def require_config_keys(config: Mapping, keys: Iterable[str], context: str) -> N raise ValueError(msg) +def require_peaks_present( + available: Iterable, required: Iterable, context: str +) -> None: + """Raise ValueError listing all ``required`` peaks absent from ``available``. + + The peak files written by ``par-geds-dsp-evtsel`` tag every row with the + nominal gamma-line energy in a ``peak`` column, so a peak for which no + events were selected is simply absent rather than empty. Consumers filter + on that column and would otherwise operate silently on zero rows. + + Parameters + ---------- + available : iterable + Peak labels actually present in the peak file. + required : iterable + Peak labels the consumer's configuration asks for. + context : str + Free-form string naming the peak file and requesting config in the + error message. + """ + present = set(available) + missing = sorted({peak for peak in required if peak not in present}) + if missing: + msg = f"{context} is missing required peak(s) {missing}" + raise ValueError(msg) + + def get_rule_config(configs_path, rule_name, timestamp, datatype): """Resolve the dataflow config for one Snakemake rule. diff --git a/src/legenddataflowscripts/utils/discharge.py b/src/legenddataflowscripts/utils/discharge.py new file mode 100644 index 0000000..b137a91 --- /dev/null +++ b/src/legenddataflowscripts/utils/discharge.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import numpy as np + +__all__ = ["get_is_recovering_mask"] + + +def get_is_recovering_mask( + timestamps: np.ndarray, + discharge_timestamps: np.ndarray, + window: float = 0.01, +) -> np.ndarray: + """Mark events that follow a discharge within a recovery window. + + An event at time ``t`` is flagged when any discharge time ``d`` satisfies + ``0 < t - d < window``. Results are bit-identical to the original + ``O(n_discharges * n_events)`` loop: a sorted search collects a slightly + widened candidate superset of (event, discharge) pairs, and the original + floating-point conditions are then evaluated exactly on those pairs, in + ``O((n + m) log m + n_candidates)``. + + Parameters + ---------- + timestamps + Event timestamps (any order). + discharge_timestamps + Timestamps of discharge events (any order). + window + Length of the recovery window in the same units as the timestamps. + + Returns + ------- + is_recovering + Boolean array, one entry per event in *timestamps*. + """ + ts = np.asarray(timestamps) + ds = np.sort(np.asarray(discharge_timestamps)) + out = np.zeros(ts.shape, dtype=bool) + if ds.size == 0 or ts.size == 0: + return out + + # candidate discharges for event t: d < t (exact: t - d > 0 iff d < t) + # and d >= t - margin, with margin slightly wider than the window so no + # pair with fl(t - d) < window can be missed by the rounding of t - margin + margin = window * 1.001 + lo = np.searchsorted(ds, ts - margin, side="left") + hi = np.searchsorted(ds, ts, side="left") + + counts = hi - lo + cand = counts > 0 + if not cand.any(): + return out + + ev_rows = np.nonzero(cand)[0] + ln = counts[ev_rows] + total = int(ln.sum()) + # flat indices of each candidate pair's discharge in ds + seg_starts = np.cumsum(ln) - ln + offsets = np.arange(total) - np.repeat(seg_starts, ln) + d_idx = np.repeat(lo[ev_rows], ln) + offsets + + # evaluate the original per-pair conditions with the original arithmetic + diff = ts[np.repeat(ev_rows, ln)] - ds[d_idx] + hit = (diff < window) & (diff > 0) + + out[ev_rows] = np.logical_or.reduceat(hit, seg_starts) + return out diff --git a/src/legenddataflowscripts/utils/lgdo_utils.py b/src/legenddataflowscripts/utils/lgdo_utils.py new file mode 100644 index 0000000..29dc97d --- /dev/null +++ b/src/legenddataflowscripts/utils/lgdo_utils.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from lgdo import Array, Table, WaveformTable + +__all__ = ["take_table_rows"] + + +def take_table_rows(tbl: Table, idx) -> Table: + """Build a new :class:`lgdo.Table` from a row subset of *tbl*. + + Row-identical to re-reading the same rows from disk, so it can replace a + second ``lh5.read`` of data that is already in memory. + + Parameters + ---------- + tbl + Input table. Supported column types are :class:`lgdo.WaveformTable` + (with equal-sized ``values``) and :class:`lgdo.Array` (including + ``ArrayOfEqualSizedArrays``). + idx + Integer index array or boolean mask selecting the rows to keep. + + Returns + ------- + subset + New table holding copies of the selected rows; ``units`` attributes + are preserved. + + Raises + ------ + NotImplementedError + If a column type is not supported (e.g. ``VectorOfVectors``). + """ + col_dict = {} + for name in tbl: + col = tbl[name] + if isinstance(col, WaveformTable): + values = col.values + if not hasattr(values, "nda"): + msg = ( + f"take_table_rows does not support waveform values of type " + f"{type(values).__name__} (column {name!r})" + ) + raise NotImplementedError(msg) + col_dict[name] = WaveformTable( + t0=col.t0.nda[idx], + t0_units=col.t0.attrs.get("units"), + dt=col.dt.nda[idx], + dt_units=col.dt.attrs.get("units"), + values=values.nda[idx], + values_units=values.attrs.get("units"), + ) + elif isinstance(col, Array): + attrs = {k: v for k, v in col.attrs.items() if k != "datatype"} + col_dict[name] = type(col)(nda=col.nda[idx], attrs=attrs) + else: + msg = ( + f"take_table_rows does not support column type " + f"{type(col).__name__} (column {name!r})" + ) + raise NotImplementedError(msg) + return Table(col_dict=col_dict) diff --git a/src/legenddataflowscripts/utils/pulser_removal.py b/src/legenddataflowscripts/utils/pulser_removal.py index 1017c9c..0df43f1 100644 --- a/src/legenddataflowscripts/utils/pulser_removal.py +++ b/src/legenddataflowscripts/utils/pulser_removal.py @@ -23,16 +23,15 @@ def get_pulser_mask(pulser_file): """ if not isinstance(pulser_file, list): pulser_file = [pulser_file] - mask = np.array([], dtype=bool) + masks = [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) + masks.append(np.array(pulser_dict["mask"])) - return mask + return np.concatenate(masks) def check_pulser_mask(mask, threshold_mask, context) -> None: diff --git a/tests/test_discharge.py b/tests/test_discharge.py new file mode 100644 index 0000000..f356ec4 --- /dev/null +++ b/tests/test_discharge.py @@ -0,0 +1,60 @@ +"""Tests for legenddataflowscripts.utils.discharge.get_is_recovering_mask.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from legenddataflowscripts.utils import get_is_recovering_mask + + +def _reference_loop(timestamps, discharge_timestamps, window=0.01): + """Verbatim structure of the original per-discharge loop.""" + is_recovering = np.full(len(timestamps), False, dtype=bool) + for tstamp in discharge_timestamps: + is_recovering = is_recovering | np.where( + (((timestamps - tstamp) < window) & ((timestamps - tstamp) > 0)), + True, + False, + ) + return is_recovering + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_matches_loop_on_random_data(seed): + rng = np.random.default_rng(seed) + ts = np.sort(rng.uniform(1.7e9, 1.7e9 + 20, 5000)) + discharges = rng.choice(ts, 40, replace=False) + ref = _reference_loop(ts, discharges) + fast = get_is_recovering_mask(ts, discharges) + assert np.array_equal(fast, ref) + # discharges must actually flag something for the test to be meaningful + assert np.count_nonzero(ref) > 0 + + +def test_edge_cases(): + ts = np.array([10.0, 10.005, 10.01, 10.02]) + # no discharges + assert np.array_equal( + get_is_recovering_mask(ts, np.array([])), np.zeros(4, dtype=bool) + ) + # event exactly at the discharge time is NOT recovering (t - d == 0) + d = np.array([10.0]) + assert np.array_equal(get_is_recovering_mask(ts, d), _reference_loop(ts, d)) + # duplicate and unsorted discharges + d = np.array([10.02, 10.0, 10.0]) + assert np.array_equal(get_is_recovering_mask(ts, d), _reference_loop(ts, d)) + + +def test_production_call_pattern_stays_all_false(): + """The call sites currently pass ``np.where(ts[discharges])[0]`` (indices, + not timestamps); with epoch-scale event timestamps the mask is all-False. + The vectorized form must reproduce that exactly.""" + rng = np.random.default_rng(3) + ts = np.sort(rng.uniform(1.7e9, 1.7e9 + 3600, 2000)) + discharges = rng.random(2000) < 0.01 + discharge_timestamps = np.where(ts[discharges])[0] + ref = _reference_loop(ts, discharge_timestamps) + fast = get_is_recovering_mask(ts, discharge_timestamps) + assert np.array_equal(fast, ref) + assert not fast.any() diff --git a/tests/test_lgdo_utils.py b/tests/test_lgdo_utils.py new file mode 100644 index 0000000..d5e2b00 --- /dev/null +++ b/tests/test_lgdo_utils.py @@ -0,0 +1,66 @@ +"""Tests for legenddataflowscripts.utils.lgdo_utils.take_table_rows.""" + +from __future__ import annotations + +import numpy as np +import pytest +from lgdo import Array, ArrayOfEqualSizedArrays, Table, VectorOfVectors, WaveformTable + +from legenddataflowscripts.utils import take_table_rows + +RNG = np.random.default_rng(11) +N = 20 + + +def _make_table(): + wf = WaveformTable( + t0=np.zeros(N), + t0_units="ns", + dt=np.full(N, 16.0), + dt_units="ns", + values=RNG.normal(0, 1, (N, 32)).astype("float32"), + values_units="ADC", + ) + return Table( + col_dict={ + "waveform": wf, + "daqenergy": Array(RNG.integers(0, 5000, N).astype("uint32")), + "matrix": ArrayOfEqualSizedArrays(nda=RNG.normal(0, 1, (N, 4))), + } + ) + + +@pytest.mark.parametrize( + "idx", + [ + np.array([0, 3, 7, 19]), + RNG.random(N) < 0.4, + ], + ids=["int-index", "bool-mask"], +) +def test_take_table_rows_matches_fancy_indexing(idx): + tbl = _make_table() + sub = take_table_rows(tbl, idx) + + n_sel = int(np.count_nonzero(idx)) if idx.dtype == bool else len(idx) + assert len(sub) == n_sel + assert isinstance(sub["waveform"], WaveformTable) + assert np.array_equal(sub["waveform"].values.nda, tbl["waveform"].values.nda[idx]) + assert np.array_equal(sub["waveform"].dt.nda, tbl["waveform"].dt.nda[idx]) + assert np.array_equal(sub["waveform"].t0.nda, tbl["waveform"].t0.nda[idx]) + assert sub["waveform"].dt.attrs["units"] == "ns" + assert sub["waveform"].values.attrs["units"] == "ADC" + assert np.array_equal(sub["daqenergy"].nda, tbl["daqenergy"].nda[idx]) + assert sub["daqenergy"].nda.dtype == np.uint32 + assert isinstance(sub["matrix"], ArrayOfEqualSizedArrays) + assert np.array_equal(sub["matrix"].nda, tbl["matrix"].nda[idx]) + + +def test_take_table_rows_rejects_unsupported_columns(): + tbl = Table( + col_dict={ + "vov": VectorOfVectors([[1, 2], [3], [4, 5, 6]]), + } + ) + with pytest.raises(NotImplementedError, match="vov"): + take_table_rows(tbl, np.array([0, 1])) diff --git a/tests/test_peak_selection.py b/tests/test_peak_selection.py new file mode 100644 index 0000000..f24e26e --- /dev/null +++ b/tests/test_peak_selection.py @@ -0,0 +1,83 @@ +"""Tests for the peak-selection helpers used by the dsp par scripts. + +Covers ``build_peak_dicts`` (which drops peaks with no candidate events, so a +channel whose DAQ threshold sits above a configured gamma line does not hand +``None`` to ``build_dsp``) and ``require_peaks_present`` (which makes the peak +file consumers fail loudly instead of filtering down to zero rows). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from legenddataflowscripts.par.geds.dsp.evtsel import build_peak_dicts +from legenddataflowscripts.utils import require_peaks_present + +PEAKS = [583.191, 727.33, 860.564, 1620.5, 2614.553] +KEV_WIDTHS = [(25, 30), (25, 35), (25, 40), (15, 55), (70, 70)] + + +def _masks(counts): + return { + peak: np.arange(count, dtype=int) + for peak, count in zip(PEAKS, counts, strict=True) + } + + +def test_all_peaks_populated(): + counts = [33123, 10172, 6538, 2519, 12999] + pk_dicts = build_peak_dicts(PEAKS, KEV_WIDTHS, _masks(counts)) + + assert list(pk_dicts) == PEAKS + for peak, kev_width, count in zip(PEAKS, KEV_WIDTHS, counts, strict=True): + entry = pk_dicts[peak] + assert entry["kev_width"] == kev_width + assert len(entry["idxs"][0]) == count + assert entry["n_rows_read"] == 0 + assert entry["obj_buf_start"] == 0 + assert entry["obj_buf"] is None + + +def test_empty_peaks_are_skipped(): + # V01404A in p18: the DAQ threshold sits above the three lowest lines + counts = [0, 0, 0, 7315, 12112] + pk_dicts = build_peak_dicts(PEAKS, KEV_WIDTHS, _masks(counts)) + + assert list(pk_dicts) == [1620.5, 2614.553] + assert len(pk_dicts[1620.5]["idxs"][0]) == 7315 + assert pk_dicts[1620.5]["kev_width"] == (15, 55) + + +def test_no_peaks_at_all(): + assert build_peak_dicts(PEAKS, KEV_WIDTHS, _masks([0] * 5)) == {} + + +def test_skipped_peaks_are_logged(): + class _Recorder: + def __init__(self): + self.messages = [] + + def warning(self, msg): + self.messages.append(msg) + + log = _Recorder() + build_peak_dicts(PEAKS, KEV_WIDTHS, _masks([0, 0, 0, 7315, 12112]), log=log) + + assert len(log.messages) == 3 + assert all("skipping this peak" in msg for msg in log.messages) + assert "583.191" in log.messages[0] + + +def test_require_peaks_present_passes(): + require_peaks_present(np.array([1620, 2614]), [2614], "peak file") + + +def test_require_peaks_present_lists_all_missing(): + with pytest.raises(ValueError, match=r"missing required peak\(s\) \[583, 727\]"): + require_peaks_present(np.array([1620, 2614]), [583, 727, 2614], "peak file") + + +def test_require_peaks_present_names_the_context(): + with pytest.raises(ValueError, match=r"my-peaks\.lh5"): + require_peaks_present(np.array([2614]), [583], "peak file my-peaks.lh5")