perf: speed and memory optimisation of the par-geds dsp and hit scripts

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")
 

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,
             },
         ]
 

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,

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
         )