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

legenddataflowscripts/tier/dsp.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import argparse
+import json
+import time
+import warnings
+from multiprocessing import Pool
+from pathlib import Path
+
+import numpy as np
+from dbetto import TextDB
+from dbetto.catalog import Props
+from dspeed import build_dsp
+from lgdo import lh5
+
+from ..utils import alias_table, build_log
+
+warnings.filterwarnings(action="ignore", category=RuntimeWarning)
+
+
+def _replace_list_with_array(dic):
+    for key, value in dic.items():
+        if isinstance(value, dict):
+            dic[key] = _replace_list_with_array(value)
+        elif isinstance(value, list):
+            dic[key] = np.array(value, dtype="float32")
+        else:
+            pass
+    return dic
+
+
+def build_dsp_wrapper(kwargs):
+    build_dsp(**kwargs)
+
+
+def build_tier_dsp() -> None:
+    # CLI config
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument(
+        "--configs", help="path to dataflow config files", required=True
+    )
+    argparser.add_argument(
+        "--table-map",
+        help="mapping from channel to table name",
+        required=False,
+        type=str,
+    )
+    argparser.add_argument("--log", help="log file name")
+    argparser.add_argument("--alias-table", help="Alias table", type=str, default=None)
+
+    argparser.add_argument("--n-processes", help="log file name", default=1, type=int)
+
+    argparser.add_argument("--datatype", help="datatype", required=True)
+    argparser.add_argument("--timestamp", help="timestamp", required=True)
+    argparser.add_argument("--tier", help="tier", required=True)
+
+    argparser.add_argument(
+        "--pars-file", help="database file for HPGes", nargs="*", default=[]
+    )
+    argparser.add_argument("--input", help="input file")
+
+    argparser.add_argument("--output", help="output file")
+    args = argparser.parse_args()
+
+    # 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
+
+    df_configs = TextDB(args.configs, lazy=True)
+    config_dict = df_configs.on(args.timestamp, system=args.datatype).snakemake_rules
+    config_dict = config_dict[f"tier_{args.tier}"]
+
+    log = build_log(config_dict, args.log, fallback=__name__)
+
+    settings_dict = config_dict.options.get("settings", {})
+    if isinstance(settings_dict, str):
+        settings_dict = Props.read_from(settings_dict)
+
+    chan_cfg_map = config_dict.inputs.processing_chain
+
+    # if the dictionary only contains one __default__ key, build the channel
+    # list from the (processable) channel map and assign the default config
+    if list(chan_cfg_map.keys()) == ["__default__"]:
+        chan_cfg_map = dict.fromkeys(table_map, chan_cfg_map.__default__)
+
+    # 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:
+            continue
+
+        # check if the raw tables are all existing
+        if len(lh5.ls(args.input, input_tbl_name)) > 0:
+            dsp_cfg_tbl_dict[input_tbl_name] = Props.read_from(file)
+        else:
+            msg = f"table {input_tbl_name} not found in {args.input} skipping"
+            log.info(msg)
+
+    if len(dsp_cfg_tbl_dict) == 0:
+        msg = f"could not find any of the requested channels in {args.input}"
+        raise RuntimeError(msg)
+
+    # par files
+    db_files = [
+        par_file
+        for par_file in args.pars_file
+        if Path(par_file).suffix in (".json", ".yaml", ".yml")
+    ]
+
+    database_dict = _replace_list_with_array(
+        Props.read_from(db_files, subst_pathvar=True)
+    )
+    database_dict = {
+        (table_map[chan].split("/")[0] if chan in table_map else chan): dic
+        for chan, dic in database_dict.items()
+    }
+    log.info("loaded database files")
+
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+
+    start = time.time()
+    if args.n_processes > 1:
+        # sort by table lengths, longest tables first
+        dsp_cfg_tbl_dict = dict(
+            sorted(
+                dsp_cfg_tbl_dict.items(),
+                key=lambda item: lh5.read_n_rows(item[0], args.input),
+                reverse=True,
+            )
+        )
+
+        chan_configs = [{} for _ in range(args.n_processes)]
+        for i, key in enumerate(dsp_cfg_tbl_dict):
+            chan_configs[i % args.n_processes][key] = dsp_cfg_tbl_dict[key]
+
+        dsp_files = [
+            f"{args.output}{i}" if i > 0 else args.output
+            for i in range(args.n_processes)
+        ]
+
+        # Process arguments for each worker
+        process_kwargs_list = []
+        for i, config in enumerate(chan_configs):
+            kwargs = {
+                "f_raw": args.input,
+                "f_dsp": dsp_files[i],
+                "chan_config": config,
+                "database": database_dict,
+                "write_mode": "r",
+                "buffer_len": settings_dict.get("buffer_len", 1000),
+                "block_width": settings_dict.get("block_width", 16),
+            }
+            process_kwargs_list.append(kwargs)
+
+        # Create a multiprocessing pool
+        with Pool(processes=args.n_processes) as pool:
+            # Use starmap to pass multiple arguments to the process function
+            pool.map(build_dsp_wrapper, process_kwargs_list)
+
+        # merge the DSPs
+        log.info("Merging DSPs")
+        for i, file in enumerate(dsp_files[1:]):
+            chans = chan_configs[i + 1]
+            for chan in chans:
+                tbl = lh5.read(chan.replace("raw", "dsp"), file)
+                lh5.write(tbl, chan.replace("raw", "dsp"), args.output, wo_mode="a")
+            Path(file).unlink()
+
+    else:
+        build_dsp(
+            args.input,
+            args.output,
+            database=database_dict,
+            chan_config=dsp_cfg_tbl_dict,
+            write_mode="r",
+            buffer_len=settings_dict.get("buffer_len", 1000),
+            block_width=settings_dict.get("block_width", 16),
+        )
+
+    msg = f"Finished building DSP in {time.time() - start:.2f} seconds"
+    log.info(msg)
+    if args.alias_table is not None:
+        log.info("Creating alias table")
+        alias_table(args.output, args.alias_table)
+
+
+def build_tier_dsp_single_channel() -> None:
+    # CLI config
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument(
+        "--configs", help="path to dataflow config files", required=True
+    )
+    argparser.add_argument(
+        "--channel",
+        help="channel to process",
+        required=False,
+        type=str,
+    )
+    argparser.add_argument("--log", help="log file name")
+
+    argparser.add_argument("--datatype", help="datatype", required=True)
+    argparser.add_argument("--timestamp", help="timestamp", required=True)
+    argparser.add_argument("--tier", help="tier", required=True)
+
+    argparser.add_argument(
+        "--pars-file", help="database file for HPGes", nargs="*", default=[]
+    )
+    argparser.add_argument("--input", help="input file")
+
+    argparser.add_argument("--output", help="output file")
+    args = argparser.parse_args()
+
+    df_configs = TextDB(args.configs, lazy=True)
+    config_dict = df_configs.on(args.timestamp, system=args.datatype).snakemake_rules
+    config_dict = config_dict[f"tier_{args.tier}"]
+    config_dict = (
+        config_dict[args.channel]
+        if args.channel is not None and args.channel in config_dict
+        else config_dict
+    )
+
+    log = build_log(config_dict, args.log, fallback=__name__)
+
+    settings_dict = config_dict.options.get("settings", {})
+    if isinstance(settings_dict, str):
+        settings_dict = Props.read_from(settings_dict)
+
+    proc_chain = config_dict.inputs.processing_chain
+
+    # par files
+    db_files = [
+        par_file
+        for par_file in args.pars_file
+        if Path(par_file).suffix in (".json", ".yaml", ".yml")
+    ]
+
+    database_dict = _replace_list_with_array(
+        Props.read_from(db_files, subst_pathvar=True)
+    )
+    database_dict = (
+        database_dict[args.channel]
+        if args.channel is not None and args.channel in database_dict
+        else database_dict
+    )
+
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+
+    start = time.time()
+
+    build_dsp(
+        args.input,
+        args.output,
+        proc_chain,
+        database=database_dict,
+        write_mode="r",
+        buffer_len=settings_dict.get("buffer_len", 1000),
+        block_width=settings_dict.get("block_width", 16),
+    )
+    msg = f"Finished building DSP in {time.time() - start:.2f} seconds"
+    log.info(msg)

legenddataflowscripts/tier/hit.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+from dbetto.catalog import Props
+from legendmeta import TextDB
+from lgdo import lh5
+from pygama.hit.build_hit import build_hit
+
+from ..utils import alias_table, build_log
+
+
+def build_tier_hit() -> None:
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--input")
+    argparser.add_argument("--pars-file", nargs="*")
+
+    argparser.add_argument("--configs", required=True)
+    argparser.add_argument("--table-map", required=False, type=str)
+    argparser.add_argument("--log")
+    argparser.add_argument("--alias-table", help="Alias table", type=str, default=None)
+
+    argparser.add_argument("--datatype", required=True)
+    argparser.add_argument("--timestamp", required=True)
+    argparser.add_argument("--tier", required=True)
+
+    argparser.add_argument("--output")
+    args = argparser.parse_args()
+
+    table_map = json.loads(args.table_map) if args.table_map is not None else None
+
+    df_config = (
+        TextDB(args.configs, lazy=True)
+        .on(args.timestamp, system=args.datatype)
+        .snakemake_rules[f"tier_{args.tier}"]
+    )
+    log = build_log(df_config, args.log, fallback=__name__)
+    log.info("initializing")
+
+    settings_dict = df_config.options.get("settings", {})
+
+    if isinstance(settings_dict, str):
+        settings_dict = Props.read_from(settings_dict)
+
+    # mapping channel -> hit config file
+    chan_cfg_map = df_config.inputs.hit_config
+
+    log.info("building the build_hit config")
+    # if the mapping only contains one __default__ key, build the channel
+    # list from the (processable) channel map and assign the default config
+    if list(chan_cfg_map.keys()) == ["__default__"]:
+        chan_cfg_map = dict.fromkeys(table_map, chan_cfg_map.__default__)
+
+    # now construct the dictionary of hit configs for build_hit()
+    channel_dict = {}
+    pars_dict = {ch: chd["pars"] for ch, chd in Props.read_from(args.pars_file).items()}
+    for chan, file in chan_cfg_map.items():
+        hit_cfg = Props.read_from(file)
+
+        # get pars (to override hit config)
+        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:
+            continue
+
+        # check if the raw tables are all existing
+        if len(lh5.ls(args.input, input_tbl_name)) > 0:
+            channel_dict[input_tbl_name] = hit_cfg
+        else:
+            msg = f"table {input_tbl_name} not found in {args.input} skipping"
+            log.warning(msg)
+
+    log.info("running build_hit()...")
+    start = time.time()
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    build_hit(args.input, lh5_tables_config=channel_dict, outfile=args.output)
+    msg = f"Hit built in {time.time() - start:.2f} seconds"
+    log.info(msg)
+    if args.alias_table is not None:
+        log.info("Creating alias table")
+        alias_table(args.output, args.alias_table)
+
+
+def build_tier_hit_single_channel() -> None:
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--input")
+    argparser.add_argument("--pars-file", nargs="*")
+
+    argparser.add_argument("--configs", required=True)
+    argparser.add_argument("--log")
+
+    argparser.add_argument(
+        "--channel",
+        help="channel to process",
+        required=False,
+        type=str,
+    )
+    argparser.add_argument("--datatype", required=True)
+    argparser.add_argument("--timestamp", required=True)
+    argparser.add_argument("--tier", required=True)
+
+    argparser.add_argument("--output")
+    args = argparser.parse_args()
+
+    df_config = (
+        TextDB(args.configs, lazy=True)
+        .on(args.timestamp, system=args.datatype)
+        .snakemake_rules[f"tier_{args.tier}"]
+    )
+    log = build_log(df_config, args.log, fallback=__name__)
+    log.info("initializing")
+
+    settings_dict = df_config.options.get("settings", {})
+
+    if isinstance(settings_dict, str):
+        settings_dict = Props.read_from(settings_dict)
+
+    # mapping channel -> hit config file
+    chan_cfg_map = df_config.inputs.hit_config
+    chan_cfg_map = (
+        chan_cfg_map[args.channel]
+        if args.channel is not None and args.channel in chan_cfg_map
+        else chan_cfg_map
+    )
+
+    # now construct the dictionary of hit configs for build_hit()
+    channel_dict = {}
+    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
+    )
+
+    hit_cfg = Props.read_from(chan_cfg_map)
+    Props.add_to(hit_cfg, pars_dict.copy())
+
+    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)
+    msg = f"Hit built in {time.time() - start:.2f} seconds"
+    log.info(msg)

legenddataflowscripts/utils/__init__.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from .alias_table import alias_table
+from .cfgtools import get_channel_config
+from .convert_np import convert_dict_np_to_float
+from .log import build_log
+from .pulser_removal import get_pulser_mask
+
+__all__ = [
+    "alias_table",
+    "build_log",
+    "convert_dict_np_to_float",
+    "get_channel_config",
+    "get_pulser_mask",
+]

legenddataflowscripts/utils/alias_table.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import json
+
+import h5py
+
+
+def alias_table(file, mapping):
+    """
+    Create an alias table for the given file and mapping.
+
+    Args:
+        file (str): Path to the input file.
+        mapping (dict): Mapping of current table name and alias table name.
+
+    Returns:
+        dict: A dictionary containing the alias table.
+    """
+    if isinstance(mapping, str):
+        mapping = json.loads(mapping)
+    with h5py.File(file, "a") as f:
+        for raw_id, alias in mapping.items():
+            if raw_id in f:
+                if isinstance(alias, list | tuple):
+                    for a in alias:
+                        f[a] = f[raw_id]
+                else:
+                    f[alias] = f[raw_id]

legenddataflowscripts/utils/cfgtools.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+
+def get_channel_config(
+    mapping: Mapping, channel: str, default_key: str = "__default__"
+):
+    """Get channel key from mapping with default.
+
+    Returns the value at key `channel`, if existing, otherwise return value at
+    `default_key`.
+    """
+    return mapping.get(channel, mapping[default_key])

legenddataflowscripts/utils/convert_np.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import numpy as np
+
+
+def convert_dict_np_to_float(dic: dict) -> dict:
+    """
+    Convert all numpy floats in a dictionary to Python floats.
+
+    Parameters
+    ----------
+    dic : dict
+        The dictionary to convert.
+
+    Returns
+    -------
+    dict
+        The dictionary with all numpy floats converted to Python floats.
+    """
+    for key, value in dic.items():
+        if isinstance(value, Mapping):
+            convert_dict_np_to_float(value)
+        elif isinstance(value, np.float32 | np.float64):
+            dic[key] = float(value)
+        elif isinstance(dic[key], Sequence):
+            dic[key] = [
+                float(x) if isinstance(x, np.float32 | np.float64) else x for x in value
+            ]
+    return dic

legenddataflowscripts/utils/log.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import logging
+import sys
+import traceback
+from logging.config import dictConfig
+from pathlib import Path
+
+from dbetto import Props
+
+
+class StreamToLogger:
+    """File-like stream object that redirects writes to a logger instance."""
+
+    def __init__(self, logger, log_level=logging.ERROR):
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ""
+
+    def write(self, buf):
+        for line in buf.rstrip().splitlines():
+            self.logger.log(self.log_level, line.rstrip())
+
+    def flush(self):
+        pass
+
+
+def build_log(
+    config_dict: dict, log_file: str | None = None, fallback: str = "prod"
+) -> logging.Logger:
+    """Build a logger from a configuration dictionary.
+
+    If a log file is provided, the logger will write to that file.
+
+    Parameters
+    ----------
+    config_dict
+        A dictionary containing the logging configuration.
+    log_file
+        The path to the log file.
+    """
+    if "logging" in config_dict["options"]:
+        log_config = config_dict["options"]["logging"]
+        log_config = Props.read_from(log_config)
+
+        if log_file is not None:
+            Path(log_file).parent.mkdir(parents=True, exist_ok=True)
+            log_config["handlers"]["dataflow"]["filename"] = log_file
+
+        dictConfig(log_config)
+        log = logging.getLogger(config_dict["options"].get("logger", "prod"))
+
+    else:
+        if log_file is not None:
+            Path(log_file).parent.mkdir(parents=True, exist_ok=True)
+            logging.basicConfig(level=logging.INFO, filename=log_file, filemode="w")
+
+        log = logging.getLogger(fallback)
+
+    # 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
+
+    def excepthook(exc_type, exc_value, exc_traceback):
+        traceback.print_exception(exc_type, exc_value, exc_traceback, file=log_stream)
+
+    sys.excepthook = excepthook
+
+    return log

legenddataflowscripts/utils/pulser_removal.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+import numpy as np
+from dbetto.catalog import Props
+
+
+def get_pulser_mask(pulser_file):
+    if not isinstance(pulser_file, list):
+        pulser_file = [pulser_file]
+    mask = np.array([], dtype=bool)
+    for file in pulser_file:
+        pulser_dict = Props.read_from(file)
+        pulser_mask = np.array(pulser_dict["mask"])
+        mask = np.append(mask, pulser_mask)
+
+    return mask

legenddataflowscripts/workflow/__init__.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from .execenv import execenv_prefix, execenv_pyexe
+from .utils import (
+    as_ro,
+    set_last_rule_name,
+    subst_vars,
+    subst_vars_impl,
+    subst_vars_in_snakemake_config,
+)
+
+__all__ = [
+    "as_ro",
+    "execenv_prefix",
+    "execenv_pyexe",
+    "set_last_rule_name",
+    "subst_vars",
+    "subst_vars_impl",
+    "subst_vars_in_snakemake_config",
+]