diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index bb5d58a4..5f150f5f 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -76,6 +76,23 @@ jobs: - name: Run data production tests run: ./tests/runprod/run-all.sh + build-docs: + name: Build documentation + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + - uses: actions/setup-python@v6 + with: + python-version: "3.12" + - name: Install docs dependencies + run: | + python -m pip install --upgrade pip wheel setuptools + python -m pip install --upgrade .[docs] + - name: Build docs + run: make -C docs + test-coverage: name: Calculate and upload test coverage runs-on: ubuntu-latest diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..48ec0b98 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,97 @@ +# LEGEND Dataflow — Agent / Developer Reference + +`legend-dataflow` orchestrates the LEGEND L200 data processing pipeline — raw +digitiser output to physics-ready events — with +[Snakemake](https://snakemake.readthedocs.io/). Full docs: +. AI-assisted contributions must follow +`AI_POLICY.md` (human review + disclosure). + +Layout: `workflow/Snakefile` is the main workflow (tcm → … → skm); +`workflow/Snakefile-build-raw` covers DAQ → raw. Both include rule files from +`workflow/rules/`. The Python package lives at `workflow/src/legenddataflow/` +(note the non-standard `package-dir` in `pyproject.toml`): `methods/` is the +library used by rules (file keys, patterns, pars catalogs, cal grouping); +`scripts/` holds the executables rules invoke (`flow/` plumbing, `tier/` +builders, `par/geds|spms/` calibration). + +## Split Across Two Packages + +Many console scripts invoked by rules (`build-tier-dsp`, `build-tier-hit`, +`build-tier-pht`, …) and the `dataflow`/`dataprod` CLI live in the external +`legend-dataflow-scripts` package (imported as `legenddataflowscripts`), pinned in +`pyproject.toml`. This repo's own entry points are listed under `[project.scripts]`. +If a command's implementation isn't here, check that package before concluding it's +missing. + +## Processing Tiers + +`raw` (DAQ→LH5, blinding) → `tcm` (time coincidence map) → `dsp` (signal +processing) → `hit` (energy calibration, PSD) → `evt` (event building, +cross-talk) → `skm` (physics skim). Partition-level variants, averaged over +calibration runs: `psp` (dsp), `pht` (hit), `ann`/`pan` (ANN cuts, coax only), +`pet` (evt). Rule files are named after their tier. Per tier, calibration +parameters are derived from `cal` runs, then applied to `phy` data. + +## Key Conventions + +- **File keys:** `{experiment}-{period}-{run}-{datatype}-{timestamp}`, e.g. + `l200-p03-r001-phy-20230401T000000Z`; parsed/generated by `methods/FileKey.py`. +- **Targets:** `[all|valid|]-{experiment}-{period}-{run}-{datatype}-{tier}.gen`; + `*` wildcards and `_`-separated multi-selectors work for most components. +- **Patterns/paths:** all filename patterns live in `methods/patterns.py`, path + resolvers in `methods/paths.py`; new tiers need entries in both plus + `dataflow-config.yaml`. +- **Parameter catalogs:** per-tier `validity.yaml` files map parameter sets to time + ranges; resolved by `ParsKeyResolve` and `ParsCatalog` in `methods/`. +- **Two rule→script styles:** `script:` directives (receive the `snakemake` object) + vs `shell:` + `execenv_pyexe(config, "")`, which runs inside the + configured execenv. Don't convert between them casually. +- **Config/environment:** the `PRODENV` env var is required. In + `dataflow-config.yaml`, `$_` expands to the config file's directory and other + `$VARS` come from the environment; inputs live under `inputs/`, outputs under + `generated/`. `execenv:` defines named environments (bare/lngs/sator/nersc); + non-bare ones run in containers. + +## Commands + +```bash +uv venv --python 3.12 && source .venv/bin/activate +uv pip install -e ".[dev]" # add ".[docs]" to build docs + +python -m pytest # unit tests +./tests/runprod/run-all.sh # integration tests; run from repo root, + # needs LEGEND_METADATA set +pre-commit run --all-files # lint/format (once: pre-commit install) +make -C docs # Sphinx docs (regenerates apidoc; -W: warnings are errors) + +# run the workflow +snakemake --profile workflow/profiles/default all-l200-p03-r001-phy-skm.gen +``` + +## Style + +- **ruff** (lint + format), **snakefmt** for Snakefiles (`channel_merge.smk` + excluded), **shellcheck**, **prettier** (YAML/MD/JSON/TOML) — all via pre-commit; + pre-commit.ci pushes autofixes to PRs. +- `from __future__ import annotations` is required in all Python files. +- **NumPy-style docstrings** (Sphinx Napoleon is configured NumPy-only; apidoc + publishes undocumented modules as empty stubs). +- Versioning is **setuptools_scm** from git tags — never edit `_version.py`. +- Commit messages loosely follow conventional-commit prefixes (`fix(evt): …`). + +## Testing + +- pytest promotes warnings to errors and uses strict markers/config; + `xfail_strict = true`. +- Unit tests cover only `methods/` (file keys, pars catalogs, cal grouping) and + load the `tests/dummy_cycle/` fixture config at module import time. +- `tests/runprod/*.sh` are DAG smoke tests using empty `touch`ed dummy files — by + design: real processing is ~1 TB per run, far beyond CI capacity. Never propose + processing real data in GitHub Actions; test pure logic with unit tests instead. + +## Extending the Pipeline + +Step-by-step guides for adding processors, calibration scripts, tiers, and +execution environments are in `docs/source/developer.rst` (published in the +Developer Guide on ReadTheDocs) — follow those rather than reverse-engineering +existing rules. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 120000 index 00000000..47dc3e3d --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 00000000..39c4ba4b --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,59 @@ +# Contributing + +Contributions are welcome via pull requests against `main`. + +## Development setup + +```bash +git clone https://github.com/legend-exp/legend-dataflow.git +cd legend-dataflow +uv venv --python 3.12 +source .venv/bin/activate +uv pip install -e ".[dev]" +pre-commit install +``` + +## Code style + +Linting and formatting are enforced by [pre-commit](https://pre-commit.com/): +ruff (lint + format) for Python, snakefmt for Snakemake files, shellcheck for +shell scripts, and prettier for YAML/Markdown/JSON/TOML. pre-commit.ci runs on +every pull request and pushes autofix commits, but running the hooks locally +gives faster feedback: + +```bash +pre-commit run --all-files +``` + +Docstrings use the NumPy style — the API reference is auto-generated from them +with Sphinx, so undocumented code shows up as empty stubs in the published docs. + +Commit messages loosely follow conventional-commit prefixes, e.g. +`fix(evt): ...`, `docs: ...`, `chore: ...`. + +## Tests + +```bash +python -m pytest # unit tests +./tests/runprod/run-all.sh # integration tests (from the repo root; + # requires LEGEND_METADATA to be set) +``` + +The integration tests validate the Snakemake DAG wiring with dummy files; real +data processing (~1 TB per run) is far beyond CI capacity, so new processing +logic should be covered by unit tests where practical. + +## Documentation + +```bash +uv pip install -e ".[docs]" +make -C docs # builds HTML into docs/build (warnings are errors) +``` + +See the [Developer Guide](https://legend-dataflow.readthedocs.io) for +step-by-step instructions on extending the pipeline (new processors, +calibration scripts, tiers, and execution environments). + +## AI-assisted contributions + +Permitted with human review and disclosure — see [AI_POLICY.md](AI_POLICY.md). diff --git a/README.md b/README.md index 4ee2a330..57e281f3 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,9 @@ Full documentation is available at the project's ReadTheDocs page: - **Developer Guide** – how to extend and contribute to the pipeline - **API Reference** – Python API documentation +See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup and contribution +guidelines. + ## Related projects - [legend-pydataobj](https://legend-pydataobj.readthedocs.io) – LEGEND data objects (LH5 format) diff --git a/dataflow-config.yaml b/dataflow-config.yaml index f28efbf4..3ecee60c 100644 --- a/dataflow-config.yaml +++ b/dataflow-config.yaml @@ -1,18 +1,35 @@ +# Site configuration for the LEGEND dataflow. Key settings are documented in +# docs/source/user_manual.rst (Configuration section). + +# Tag or branch of legend-metadata to check out automatically; +# comment out to manage the metadata checkout manually. legend_metadata_version: v1.4.1 +# If false, abort when calibration parameter generation fails; +# if true, continue with default or overridden parameters. allow_none_par: false +# Generate pygama FileDB databases after a successful production run. build_file_dbs: true +# Scan log files for errors/warnings after each run. check_log_files: true +# Parallel processing within a single Snakemake job, its mode, and the +# process cap. multiprocess: false -mutliprocess_mode: max_usage +multiprocess_mode: max_usage max_processes: 1 +# Input tier for pht building: dsp or psp. pht_intier: psp +# $_ expands to the directory containing this file; other $VARS +# (e.g. $PRODENV) are taken from the environment. paths: + # Staging area for incoming DAQ files (Snakefile-build-raw). sandbox_path: $_/sandbox + # Blinded raw output (physics data with signal region removed). tier_raw_blind: $_/generated/tier/raw-blind workflow: $_/workflow + # Inputs: metadata, processing configs, overrides, hardware maps. metadata: $_/inputs config: $_/inputs/dataprod/config par_overwrite: $_/inputs/dataprod/overrides @@ -20,6 +37,7 @@ paths: detector_status: $_/inputs/datasets detector_db: $_/inputs/hardware/detectors + # Data tier outputs. tier: $_/generated/tier tier_daq: $_/generated/tier/daq tier_raw: $_/generated/tier/raw @@ -34,6 +52,7 @@ paths: tier_pet: $_/generated/tier/pet tier_skm: $_/generated/tier/skm + # Calibration parameter outputs (validity catalogs + par files). par: $_/generated/par par_raw: $_/generated/par/raw par_tcm: $_/generated/par/tcm @@ -44,18 +63,23 @@ paths: par_pht: $_/generated/par/pht par_pet: $_/generated/par/pet + # Monitoring plots and logs. plt: $_/generated/plt log: $_/generated/log + # Scratch space, cleaned up on workflow success. tmp_plt: $_/generated/tmp/plt tmp_log: $_/generated/tmp/log tmp_benchmark: $_/generated/tmp/benchmark tmp_filelists: $_/generated/tmp/filelists tmp_par: $_/generated/tmp/par + # Python source checkout and venv created by `dataflow install`. src: $_/software/python/src install: $_/.snakemake/legend-dataflow/venv +# LH5 table name inside output files, per tier +# ({ch} = channel rawid, {grp} = detector group). table_format: raw: ch{ch:07d}/raw dsp: ch{ch:07d}/dsp @@ -67,6 +91,9 @@ table_format: skm: "{grp}/skm" tcm: hardware_tcm_1 +# Named execution environments, selected with `--config system=` +# (profiles set this). Non-bare environments wrap commands in a container +# via `cmd` + `arg`; `env` sets variables for every job. execenv: bare: env: diff --git a/docs/source/pipeline.rst b/docs/source/pipeline.rst index 25e8043b..227f26a2 100644 --- a/docs/source/pipeline.rst +++ b/docs/source/pipeline.rst @@ -76,12 +76,13 @@ Converts raw DAQ files to LH5 format: - **ORCA format** (``build_raw_orca``) – decodes data from ORCA. - **FCIO format** (``build_raw_fcio``) – decodes data from FLASHCAM -used for some special data taking. + used for some special data taking. - **Blinding** (``build_raw_blind``) – removes events with any hit within 25keV of Qbb, it -uses the daqenergy estimator to do this. The blinding is checked in each calibration using -``rules/blinding_check.smk`` to verify the parameters are still valid. If any channels fail the -processing will stop until a new calibration is put in. The calibrations are all stored in -``legend-dataflow-overrides`` under ``raw``. + uses the daqenergy estimator to do this. The blinding is checked in each calibration using + ``rules/blinding_check.smk`` to verify the parameters are still valid. If any channels fail the + processing will stop until a new calibration is put in. The calibrations are all stored in + ``legend-dataflow-overrides`` under ``raw``. + The helper rules in ``rules/blinding_calibration.smk`` can be used to derive all the daqenergy calibration coefficients. diff --git a/docs/source/user_manual.rst b/docs/source/user_manual.rst index cb8d1bde..bd460e0e 100644 --- a/docs/source/user_manual.rst +++ b/docs/source/user_manual.rst @@ -80,15 +80,21 @@ Key settings +-------------------------------+--------------------------------------------------------------+ | ``multiprocess`` | Enable parallel processing within a single Snakemake job. | +-------------------------------+--------------------------------------------------------------+ -| ``pht_intier `` | Which tier to use as input for pht can be ``dsp`` or ``psp`` | +| ``multiprocess_mode`` | Strategy used to share work when ``multiprocess`` is on | +| | (default ``max_usage``). | ++-------------------------------+--------------------------------------------------------------+ +| ``max_processes`` | Maximum number of processes when multiprocessing. | ++-------------------------------+--------------------------------------------------------------+ +| ``pht_intier`` | Which tier to use as input for pht can be ``dsp`` or ``psp`` | +-------------------------------+--------------------------------------------------------------+ Paths ----- -All path values support the ``$_`` placeholder, which is substituted with the value of -the ``$PRODENV`` environment variable at runtime. This allows a single config file to -be used across different machines by setting ``PRODENV`` appropriately. +All path values support the ``$_`` placeholder, which is substituted with the directory +containing the config file (the production cycle root). Other ``$VAR`` placeholders are +taken from the environment, e.g. ``$PRODENV``. This allows a single config file to be +used across different machines. The key path categories are: @@ -182,7 +188,7 @@ The target format is: where: - ``all`` / ``valid`` – process all data, or only data selected for analysis, -any keyword in ``runlists.yaml`` in ``legend-datasets`` is a possible option. + any keyword in ``runlists.yaml`` in ``legend-datasets`` is a possible option. - ``experiment`` – experiment name (e.g. ``l200``) - ``period`` – data-taking period (e.g. ``p03``) - ``run`` – run number (e.g. ``r001``) @@ -268,6 +274,6 @@ Two files in the ``detector_status`` directory control which data are included: - ``ignored_daq_cycles.yaml`` – lists DAQ cycles to exclude from processing entirely - ``run_override.yaml`` – allows overriding the validity window for specific runs, -e.g. to apply a previous valid calibration to a subsequent run + e.g. to apply a previous valid calibration to a subsequent run These files are part of the ``legend-datasets`` repository. diff --git a/pyproject.toml b/pyproject.toml index 5c227097..ee0bed42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,7 @@ dynamic = ["version"] dependencies = [ "colorlog", - "dbetto==1.3.7", + "dbetto>=1.4", "pygama==2.4.1", "dspeed==2.1.4", "pylegendmeta==1.4.4", @@ -95,7 +95,6 @@ docs = [ ] [project.scripts] -create-chankeylist = "legenddataflow.scripts.create_chankeylist:create_chankeylist" merge-channels = "legenddataflow.scripts.flow.merge_channels:merge_channels" merge-in-channels = "legenddataflow.scripts.flow.merge_in_channel:merge_in_channel" build-tier-evt = "legenddataflow.scripts.tier.evt:build_tier_evt" @@ -107,6 +106,7 @@ build-tier-tcm = "legenddataflow.scripts.tier.tcm:build_tier_tcm" par-geds-pht-aoe = "legenddataflow.scripts.par.geds.pht.aoe:par_geds_pht_aoe" par-geds-pht-ecal-part = "legenddataflow.scripts.par.geds.pht.ecal_part:par_geds_pht_ecal_part" par-geds-pht-fast = "legenddataflow.scripts.par.geds.pht.fast:par_geds_pht_fast" +par-geds-pht-lq = "legenddataflow.scripts.par.geds.pht.lq:par_geds_pht_lq" par-geds-pht-qc-phy = "legenddataflow.scripts.par.geds.pht.qc_phy:par_geds_pht_qc_phy" par-geds-pht-qc = "legenddataflow.scripts.par.geds.pht.qc:par_geds_pht_qc" par-geds-psp-average = "legenddataflow.scripts.par.geds.psp.average:par_geds_psp_average" diff --git a/tests/test_build_chanlist.py b/tests/test_build_chanlist.py new file mode 100644 index 00000000..754a2d89 --- /dev/null +++ b/tests/test_build_chanlist.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import pytest +from dbetto import AttrsDict + +from legenddataflow.scripts.flow.build_chanlist import ( + get_chanlist, + get_par_chanlist, + get_plt_chanlist, +) + + +class FakeValidity: + """Stand-in for a pre-compiled metadata catalog with a valid_for method.""" + + def __init__(self, mapping): + self.mapping = AttrsDict(mapping) + + def valid_for(self, timestamp, category=None): # noqa: ARG002 + return self.mapping + + +CHANNELMAP = FakeValidity( + { + "DET2": {"name": "DET2", "system": "geds"}, + "DET1": {"name": "DET1", "system": "geds"}, + "DET3": {"name": "DET3", "system": "geds"}, + "SIPM1": {"name": "SIPM1", "system": "spms"}, + } +) + +DET_STATUS = FakeValidity( + { + "DET1": {"processable": True}, + "DET2": {"processable": True}, + "DET3": {"processable": False}, + "SIPM1": {"processable": True}, + } +) + +TIMESTAMP = "20230101T123456Z" + + +@pytest.fixture +def setup(tmp_path): + return {"paths": {"tmp_par": str(tmp_path / "tmp/par")}} + + +def test_get_chanlist(): + # sorted, restricted to the system, non-processable channels dropped + assert get_chanlist(TIMESTAMP, "cal", DET_STATUS, CHANNELMAP, "geds") == [ + "DET1", + "DET2", + ] + assert get_chanlist(TIMESTAMP, "cal", DET_STATUS, CHANNELMAP, "spms") == ["SIPM1"] + + +def test_get_chanlist_missing_status(): + det_status = FakeValidity({"DET1": {"processable": True}}) + with pytest.raises(RuntimeError, match="not found in the status map"): + get_chanlist(TIMESTAMP, "cal", det_status, CHANNELMAP, "geds") + + +def test_get_par_chanlist(setup): + files = get_par_chanlist( + setup, + f"all-l200-p03-r000-cal-{TIMESTAMP}-channels", + "dsp", + DET_STATUS, + CHANNELMAP, + "geds", + name="eopt", + ) + assert len(files) == 2 + assert files[0].endswith(f"l200-p03-r000-cal-{TIMESTAMP}-DET1-par_dsp_eopt.yaml") + assert "DET2" in files[1] + + +def test_get_plt_chanlist(tmp_path): + setup = {"paths": {"tmp_plt": str(tmp_path / "tmp/plt")}} + files = get_plt_chanlist( + setup, + f"all-l200-p03-r000-cal-{TIMESTAMP}-channels", + "dsp", + DET_STATUS, + CHANNELMAP, + "geds", + ) + assert len(files) == 2 + assert all(f.endswith(".pkl") for f in files) + assert files[0].endswith(f"l200-p03-r000-cal-{TIMESTAMP}-DET1-plt_dsp.pkl") diff --git a/tests/test_build_filelist.py b/tests/test_build_filelist.py new file mode 100644 index 00000000..824d2ac6 --- /dev/null +++ b/tests/test_build_filelist.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import yaml + +from legenddataflow.methods import FileKey, patterns +from legenddataflow.scripts.flow.build_filelist import ( + build_filelist, + concat_phy_filenames, + expand_runs, + get_analysis_runs, + get_filelist, + get_keys, + get_pattern, +) + + +@pytest.fixture +def config(tmp_path): + paths = {"tier": str(tmp_path / "generated/tier")} + for tier in ("daq", "raw", "dsp", "evt", "pet"): + paths[f"tier_{tier}"] = str(tmp_path / f"generated/tier/{tier}") + paths["tier_raw_blind"] = str(tmp_path / "generated/tier/raw-blind") + paths["sandbox_path"] = str(tmp_path / "sandbox") + return {"paths": paths} + + +def test_expand_runs(): + assert expand_runs({"phy": {"p01": "r001..r003"}}) == { + "phy": {"p01": ["r001", "r002", "r003"]} + } + assert expand_runs({"phy": {"p01": ["r000", "r002..r004"]}}) == { + "phy": {"p01": ["r000", "r002", "r003", "r004"]} + } + assert expand_runs({"cal": {"p02": "all"}}) == {"cal": {"p02": "all"}} + + +def test_get_analysis_runs_defaults(): + assert get_analysis_runs() == ({}, []) + + +def test_get_analysis_runs_yaml(tmp_path): + runs_file = tmp_path / "runlists.yaml" + runs_file.write_text( + yaml.dump({"sel": {"phy": {"p03": "r000..r001"}}, "other": {}}) + ) + analysis_runs, ignore_keys = get_analysis_runs( + analysis_runs_file=runs_file, file_selection="sel" + ) + assert analysis_runs == {"phy": {"p03": ["r000", "r001"]}} + assert ignore_keys == [] + + with pytest.raises(ValueError, match="unknown file selection"): + get_analysis_runs( + analysis_runs_file=runs_file, file_selection="not_a_selection" + ) + + +def test_get_analysis_runs_ignore_keylist(tmp_path): + keylist = tmp_path / "ignore.keylist" + keylist.write_text( + "l200-p00-r000-cal-20230101T123456Z # some comment\n" + "l200-p00-r001-cal-20230102T123456Z\n" + ) + _, ignore_keys = get_analysis_runs(keylist) + assert ignore_keys == [ + "l200-p00-r000-cal-20230101T123456Z", + "l200-p00-r001-cal-20230102T123456Z", + ] + + with pytest.raises(ValueError, match="ignore_keys file not found"): + get_analysis_runs(tmp_path / "missing.yaml") + + bad = tmp_path / "ignore.txt" + bad.write_text("key") + with pytest.raises(ValueError, match="unsupported ignore_keys file extension"): + get_analysis_runs(bad) + + +def test_get_keys(): + keys = get_keys("-l200-p00-r000-cal-20230101T123456Z") + assert len(keys) == 1 + assert keys[0].name == "l200-p00-r000-cal-20230101T123456Z" + + # `_`-separated multi-selectors expand to all combinations + keys = get_keys("-l200-p00-r000_r001-cal_phy") + assert sorted(k.name for k in keys) == [ + "l200-p00-r000-cal-*", + "l200-p00-r000-phy-*", + "l200-p00-r001-cal-*", + "l200-p00-r001-phy-*", + ] + + +def test_get_pattern(config): + assert get_pattern(config, "dsp") == patterns.get_pattern_tier( + config, "dsp", check_in_cycle=False + ) + # tiers that search a different tier's files + assert get_pattern(config, "blind") == patterns.get_pattern_tier( + config, "raw", check_in_cycle=False + ) + assert get_pattern(config, "skm") == patterns.get_pattern_tier( + config, "pet", check_in_cycle=False + ) + assert get_pattern(config, "evt_concat") == patterns.get_pattern_tier( + config, "evt", check_in_cycle=False + ) + + +def test_concat_phy_filenames(config): + pattern = patterns.get_pattern_tier(config, "evt", check_in_cycle=False) + fnames = [ + str(pattern) + .replace("{experiment}", "l200") + .replace("{period}", "p03") + .replace("{run}", run) + .replace("{datatype}", "phy") + .replace("{timestamp}", ts) + for run, ts in [ + ("r000", "20230101T000000Z"), + ("r000", "20230101T010000Z"), + ("r001", "20230102T000000Z"), + ] + ] + out = concat_phy_filenames(config, fnames, "evt_concat") + # one output per run, with no timestamp component + assert len(out) == 2 + assert out[0].endswith("l200-p03-r000-phy-tier_evt.lh5") + assert out[1].endswith("l200-p03-r001-phy-tier_evt.lh5") + + +def _touch_raw_files(config, keys): + raw_pattern = patterns.get_pattern_tier(config, "raw", check_in_cycle=False) + for key in keys: + f = Path(key.get_path_from_filekey(raw_pattern)[0]) + f.parent.mkdir(parents=True, exist_ok=True) + f.touch() + + +def test_build_filelist(config): + keys = [ + FileKey("l200", "p03", "r000", "cal", "20230101T000000Z"), + FileKey("l200", "p03", "r000", "phy", "20230101T010000Z"), + FileKey("l200", "p03", "r001", "phy", "20230102T000000Z"), + ] + _touch_raw_files(config, keys) + search_pattern = patterns.get_pattern_tier(config, "raw", check_in_cycle=False) + wildcard_keys = get_keys("-l200-p03-*-*") + + out = build_filelist(config, wildcard_keys, search_pattern, "dsp") + assert len(out) == 3 + assert all("tier_dsp.lh5" in f for f in out) + + # ignore_keys drops the matching key + out = build_filelist( + config, + wildcard_keys, + search_pattern, + "dsp", + ignore_keys={"unprocessable": ["l200-p03-r000-cal-20230101T000000Z"]}, + ) + assert len(out) == 2 + assert all("phy" in f for f in out) + + # a flat list (as produced from .keylist files) works too + out = build_filelist( + config, + wildcard_keys, + search_pattern, + "dsp", + ignore_keys=["l200-p03-r000-cal-20230101T000000Z"], + ) + assert len(out) == 2 + + # analysis_runs restricts to the selected runs + out = build_filelist( + config, + wildcard_keys, + search_pattern, + "dsp", + analysis_runs={"phy": {"p03": ["r001"]}}, + ) + assert len(out) == 1 + assert "r001" in out[0] + + +def test_build_filelist_concat(config): + keys = [ + FileKey("l200", "p03", "r000", "phy", "20230101T000000Z"), + FileKey("l200", "p03", "r000", "phy", "20230101T010000Z"), + ] + evt_pattern = patterns.get_pattern_tier(config, "evt", check_in_cycle=False) + for key in keys: + f = Path(key.get_path_from_filekey(evt_pattern)[0]) + f.parent.mkdir(parents=True, exist_ok=True) + f.touch() + + out = build_filelist(config, get_keys("-l200-p03-*-*"), evt_pattern, "evt_concat") + # both timestamps of the run collapse into one concatenated target + assert len(out) == 1 + assert out[0].endswith("l200-p03-r000-phy-tier_evt.lh5") + + +def test_get_filelist_wildcards(config): + keys = [FileKey("l200", "p03", "r000", "cal", "20230101T000000Z")] + _touch_raw_files(config, keys) + search_pattern = patterns.get_pattern_tier(config, "raw", check_in_cycle=False) + + wildcards = SimpleNamespace(label="all-l200-p03-r000-cal", tier="dsp") + out = get_filelist(wildcards, config, search_pattern) + assert len(out) == 1 + assert out[0].endswith("l200-p03-r000-cal-20230101T000000Z-tier_dsp.lh5") diff --git a/tests/test_cal_grouping.py b/tests/test_cal_grouping.py index 52c51861..b9547f99 100644 --- a/tests/test_cal_grouping.py +++ b/tests/test_cal_grouping.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import logging from pathlib import Path import pytest @@ -134,3 +135,54 @@ def test_get_wildcard_constraints(input_file_json): cal_grouping = CalGrouping(setup, input_file_json) constraints = cal_grouping.get_wildcard_constraints("calgroup001a", "default") assert isinstance(constraints, str) + + +def test_get_dataset_unknown_partition(input_file_json): + cal_grouping = CalGrouping(setup, input_file_json) + with pytest.raises(KeyError, match="not defined for channel"): + cal_grouping.get_dataset("no_such_partition", "default") + + +def test_get_dataset_missing_default(tmp_path): + file = tmp_path / "input.json" + file.write_text(json.dumps({"V01234A": {"part1": {"p01": ["r001"]}}})) + cal_grouping = CalGrouping(setup, file) + with pytest.raises(KeyError, match="no 'default' section"): + cal_grouping.get_dataset("part1", "V01234A") + + +def test_expand_runs_malformed_range_yields_empty(tmp_path): + # pylegendmeta's expand_runs returns [] for malformed ranges (no error raised) + file = tmp_path / "input.json" + file.write_text(json.dumps({"default": {"part": {"p01": "001..r005"}}})) + cal_grouping = CalGrouping(setup, file) + assert cal_grouping.datasets["default"]["part"]["p01"] == [] + + file.write_text(json.dumps({"default": {"part": {"p01": ["r001..x005"]}}})) + cal_grouping = CalGrouping(setup, file) + assert cal_grouping.datasets["default"]["part"]["p01"] == [] + + +def test_expand_runs_all_preserved(tmp_path): + # "all" must not be expanded to [] + file = tmp_path / "input.json" + file.write_text(json.dumps({"default": {"part": {"p01": "all"}}})) + cal_grouping = CalGrouping(setup, file) + assert cal_grouping.datasets["default"]["part"]["p01"] == "all" + + +def test_empty_partition_warns(input_file_json, caplog): + cal_grouping = CalGrouping(setup, input_file_json) + empty_catalog = Catalog.get([{"valid_from": "20210101T000000Z", "apply": []}]) + + with caplog.at_level(logging.WARNING, logger="legenddataflow.methods.cal_grouping"): + timestamp = cal_grouping.get_timestamp( + empty_catalog, "calgroup001a", "default", "dsp" + ) + log_file = cal_grouping.get_log_file( + empty_catalog, "calgroup001a", "default", "dsp", "timestamp", "par_dsp" + ) + + assert timestamp == "20000101T000000Z" + assert log_file == "log.log" + assert caplog.text.count("resolved to no par files") == 2 diff --git a/tests/test_complete_run.py b/tests/test_complete_run.py new file mode 100644 index 00000000..3145f4be --- /dev/null +++ b/tests/test_complete_run.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import runpy +from types import SimpleNamespace + +import pytest +import snakemake.script + +import legenddataflow.scripts.flow.complete_run as complete_run_module +from legenddataflow.scripts.flow.complete_run import ( + build_file_db_config, + check_log_files, + fformat, + find_gen_runs, +) + + +@pytest.fixture +def setup(tmp_path): + paths = {"tier": str(tmp_path / "generated/tier")} + for tier in ("raw", "dsp", "evt"): + paths[f"tier_{tier}"] = str(tmp_path / f"generated/tier/{tier}") + return { + "paths": paths, + "table_format": {"raw": "ch{ch:07d}/raw", "evt": "{grp}/evt"}, + } + + +def _make_logs(log_path): + sub = log_path / "20230101T000000Z" / "build_dsp" + sub.mkdir(parents=True) + (sub / "bad.log").write_text( + "INFO: starting\nERROR: something broke\nWARNING: be careful\n" + ) + (sub / "clean.log").write_text("INFO: all fine\n") + return sub + + +def test_check_log_files_with_warning_file(tmp_path): + log_path = tmp_path / "log" + _make_logs(log_path) + summary = tmp_path / "summary" / "summary.log" + warning = tmp_path / "summary" / "warning.log" + + check_log_files(log_path, summary, "all-l200.gen", warning_file=warning) + + summary_text = summary.read_text() + assert "generated at" in summary_text + assert "with errors" in summary_text + assert "bad.log : ERROR: something broke" in summary_text + warning_text = warning.read_text() + assert "with warnings" in warning_text + assert "bad.log : WARNING: be careful" in warning_text + # log files are consumed and empty directories cleaned up + assert not log_path.exists() + + +def test_check_log_files_clean(tmp_path): + log_path = tmp_path / "log" + sub = log_path / "build_dsp" + sub.mkdir(parents=True) + (sub / "clean.log").write_text("INFO: all fine\n") + summary = tmp_path / "summary.log" + warning = tmp_path / "warning.log" + + check_log_files(log_path, summary, "all-l200.gen", warning_file=warning) + + assert "with no errors" in summary.read_text() + assert "with no warnings" in warning.read_text() + assert not log_path.exists() + + +def test_check_log_files_no_warning_file(tmp_path): + log_path = tmp_path / "log" + _make_logs(log_path) + summary = tmp_path / "summary.log" + + check_log_files(log_path, summary, "all-l200.gen") + + summary_text = summary.read_text() + assert "with errors" in summary_text + assert "ERROR: something broke" in summary_text + assert not log_path.exists() + + +def test_find_gen_runs(tmp_path): + # parent path contains a hyphen on purpose: run parsing for concat tiers + # must use the filename, not the full path + gen = tmp_path / "prod-blind" / "generated" / "tier" + (gen / "dsp" / "phy" / "p03" / "r000").mkdir(parents=True) + (gen / "evt" / "phy").mkdir(parents=True) + (gen / "evt" / "phy" / "l200-p03-r001-phy-tier_evt.lh5").touch() + + assert find_gen_runs(gen) == {"phy/p03/r000", "phy/p03/r001"} + + +def test_fformat(setup): + out = fformat(setup, "raw") + # tier path prefix is stripped, leaving the relative file pattern + assert out.startswith("/{datatype}/{period}/{run}/") + assert out.endswith("-tier_raw.lh5") + assert setup["paths"]["tier_raw"] not in out + + +def test_build_file_db_config_without_prodenv(setup, tmp_path, monkeypatch): + monkeypatch.delenv("PRODENV", raising=False) + + cfg = build_file_db_config(setup, tmp_path / "filedb") + + assert cfg["data_dir"] == "/" + assert cfg["tier_dirs"] == { + "raw": setup["paths"]["tier_raw"], + "evt": setup["paths"]["tier_evt"], + } + assert cfg["table_format"] == setup["table_format"] + assert cfg["file_format"]["raw"].startswith("/{datatype}") + + +def test_build_file_db_config_with_prodenv(setup, tmp_path, monkeypatch): + monkeypatch.setenv("PRODENV", str(tmp_path)) + + cfg = build_file_db_config(setup, tmp_path / "filedb") + + assert cfg["data_dir"] == "$PRODENV" + # tier dirs are relative to $PRODENV + assert cfg["tier_dirs"]["raw"] == "/generated/tier/raw" + assert cfg["tier_dirs"]["evt"] == "/generated/tier/evt" + + +def test_snakemake_glue(tmp_path, monkeypatch): + # emulate a snakemake script: execution by injecting the snakemake object; + # FileDB building is disabled so only the log check and touch run + log_path = tmp_path / "log" + _make_logs(log_path) + setup = { + "paths": {"tmp_log": str(log_path)}, + "build_file_dbs": False, + "check_log_files": True, + } + fake = SimpleNamespace( + params=SimpleNamespace( + setup=setup, filedb_path=str(tmp_path / "filedb"), ignore_keys_file=None + ), + wildcards=SimpleNamespace(tier="dsp", label="all-l200-p03-r000-phy-dsp"), + threads=1, + output=SimpleNamespace( + summary_log=str(tmp_path / "summary.log"), + warning_log=str(tmp_path / "warning.log"), + gen_output=str(tmp_path / "all-l200-p03-r000-phy-dsp.gen"), + ), + ) + monkeypatch.setattr(snakemake.script, "snakemake", fake, raising=False) + + runpy.run_path(complete_run_module.__file__) + + assert "with errors" in (tmp_path / "summary.log").read_text() + assert "with warnings" in (tmp_path / "warning.log").read_text() + assert (tmp_path / "all-l200-p03-r000-phy-dsp.gen").exists() + assert not log_path.exists() + + +def test_check_log_files_traceback(tmp_path): + log_path = tmp_path / "log" + sub = log_path / "build_dsp" + sub.mkdir(parents=True) + (sub / "crashed.log").write_text( + "INFO: starting\n" + "Traceback (most recent call last):\n" + ' File "script.py", line 10, in main\n' + " threshold = kwarg_dict['threshold']\n" + " ~~~~~~~~~~^^^^^^^^^^^^^\n" + "KeyError: 'threshold'\n" + ) + summary = tmp_path / "summary.log" + warning = tmp_path / "warning.log" + + check_log_files(log_path, summary, "all-l200.gen", warning_file=warning) + + summary_text = summary.read_text() + assert "with errors" in summary_text + assert "crashed.log : KeyError: 'threshold'" in summary_text + # the traceback frames themselves are not reported + assert "script.py" not in summary_text + assert "with no warnings" in warning.read_text() + + +def test_check_log_files_chained_traceback(tmp_path): + log_path = tmp_path / "log" + sub = log_path / "build_dsp" + sub.mkdir(parents=True) + (sub / "crashed.log").write_text( + "Traceback (most recent call last):\n" + ' File "script.py", line 3, in inner\n' + "ValueError: inner failure\n" + "\n" + "During handling of the above exception, another exception occurred:\n" + "\n" + "Traceback (most recent call last):\n" + ' File "script.py", line 10, in main\n' + "RuntimeError: outer failure\n" + ) + summary = tmp_path / "summary.log" + + check_log_files(log_path, summary, "all-l200.gen") + + summary_text = summary.read_text() + assert "crashed.log : ValueError: inner failure" in summary_text + assert "crashed.log : RuntimeError: outer failure" in summary_text + + +def test_check_log_files_error_and_traceback(tmp_path): + log_path = tmp_path / "log" + sub = log_path / "build_dsp" + sub.mkdir(parents=True) + (sub / "mixed.log").write_text( + "ERROR: something logged\n" + "Traceback (most recent call last):\n" + ' File "script.py", line 10, in main\n' + "OSError: disk full\n" + ) + summary = tmp_path / "summary.log" + + check_log_files(log_path, summary, "all-l200.gen") + + summary_text = summary.read_text() + assert "mixed.log : ERROR: something logged" in summary_text + assert "mixed.log : OSError: disk full" in summary_text diff --git a/tests/test_console_scripts.py b/tests/test_console_scripts.py new file mode 100644 index 00000000..b7c64f16 --- /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_create_pars.py b/tests/test_create_pars.py index 055b752c..f50ca546 100644 --- a/tests/test_create_pars.py +++ b/tests/test_create_pars.py @@ -2,6 +2,7 @@ from pathlib import Path +import pytest import yaml from dbetto import time from legenddataflowscripts.workflow import subst_vars @@ -11,6 +12,7 @@ ParsKeyResolve, patterns, ) +from legenddataflow.methods.pars_loading import ParsCatalog testprod = Path(__file__).parent / "dummy_cycle" @@ -73,3 +75,28 @@ def test_create_pars_keylist(): "cal/p00/r000/l200-p00-r000-cal-20230101T123456Z-par_dsp.yaml", "lar/p00/r000/l200-p00-r000-lar-20230110T123456Z-par_dsp.yaml", } + + +def test_apply_run_override_missing_all(tmp_path): + overwrite_validity = tmp_path / "validity.yaml" + overwrite_validity.write_text( + yaml.dump( + [ + { + "valid_from": "20230101T000000Z", + "apply": ["a.yaml"], + "category": "other", + } + ] + ) + ) + hit_catalog = ParsCatalog({"all": []}) + with pytest.raises(ValueError, match="has no 'all' system entry"): + ParsKeyResolve.apply_run_override(hit_catalog, {}, overwrite_validity) + + overwrite_validity.write_text( + yaml.dump([{"valid_from": "20230101T000000Z", "apply": ["a.yaml"]}]) + ) + hit_catalog = ParsCatalog({"other": []}) + with pytest.raises(ValueError, match="par catalog passed to apply_run_override"): + ParsKeyResolve.apply_run_override(hit_catalog, {}, overwrite_validity) diff --git a/tests/test_filekey.py b/tests/test_filekey.py index f4b326d7..f1061793 100644 --- a/tests/test_filekey.py +++ b/tests/test_filekey.py @@ -2,10 +2,11 @@ from pathlib import Path +import pytest import yaml from legenddataflowscripts.workflow import subst_vars -from legenddataflow.methods import FileKey, paths, patterns +from legenddataflow.methods import ChannelProcKey, FileKey, paths, patterns testprod = Path(__file__).parent / "dummy_cycle" @@ -38,3 +39,33 @@ def test_filekey(): ).name == key.name ) + + +def test_get_filekey_from_pattern_no_match(): + with pytest.raises(ValueError, match="does not match"): + FileKey.get_filekey_from_pattern( + "stray-file.txt", patterns.get_pattern_tier(setup, "dsp") + ) + + +def test_parse_keypart_invalid(): + with pytest.raises(ValueError, match="cannot be parsed as a FileKey keypart"): + FileKey.parse_keypart("garbage") + # keypart without the leading dash + with pytest.raises(ValueError, match="cannot be parsed"): + FileKey.parse_keypart("l200-p00-r000-cal") + # ChannelProcKey keyparts must start with 'all' + with pytest.raises(ValueError, match="ChannelProcKey keypart"): + ChannelProcKey.parse_keypart("-l200-p00-r000-cal") + + +def test_get_path_from_filekey_dict_kwargs(): + key = FileKey("l200", "p00", "r000", "cal", "20230101T123456Z") + # dict-valued kwargs are resolved via the key's field values + assert key.get_path_from_filekey( + "{experiment}-{run}-{extra}", extra={"cal": "x"} + ) == ["l200-r000-x"] + # non-intersecting dict kwargs are dropped instead of crashing + assert key.get_path_from_filekey("{experiment}-{run}", extra={"nomatch": "x"}) == [ + "l200-r000" + ] diff --git a/tests/test_flow_utils.py b/tests/test_flow_utils.py new file mode 100644 index 00000000..453d0ffe --- /dev/null +++ b/tests/test_flow_utils.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import pytest + +from legenddataflow.scripts.flow.utils import replace_path +from legenddataflow.scripts.par.geds.pht.util import ( + save_dict_to_files, + split_files_by_run, +) + + +def test_replace_path_nested(): + d = { + "file": "/old/dir/file.lh5", + "nested": {"list": ["/old/dir/file.lh5", 42]}, + "untouched": 1, + } + out = replace_path(d, "/old/dir/file.lh5", "/new/dir/merged.lh5") + assert out["file"] == "$_/merged.lh5" + assert out["nested"]["list"] == ["$_/merged.lh5", 42] + assert out["untouched"] == 1 + + +def test_replace_path_by_name(): + # strings containing only the basename of the old path are also replaced + out = replace_path({"file": "file.lh5"}, "/old/dir/file.lh5", "/new/merged.lh5") + assert out == {"file": "merged.lh5"} + + +def test_split_files_by_run_missing_members(tmp_path): + good = tmp_path / "l200-p03-r000-cal-20230101T000000Z-tier_dsp.lh5" + good.touch() + filelist = tmp_path / "cal.filelist" + filelist.write_text(f"{good}\n{tmp_path}/missing-file.lh5\n") + + with pytest.raises(FileNotFoundError, match="--input-files: input file"): + split_files_by_run([str(filelist)]) + + filelist.write_text(f"{good}\n") + final_dict, files = split_files_by_run([str(filelist)]) + assert files == [str(good)] + assert list(final_dict.values()) == [[str(good)]] + + +def test_save_dict_to_files_creates_parents(tmp_path): + out = ( + tmp_path + / "new" + / "dir" + / "l200-p03-r000-cal-20230101T000000Z-ch1000000-par_dsp.yaml" + ) + save_dict_to_files([str(out)], {"20230101T000000Z": {"a": 1}}) + assert out.is_file() diff --git a/tests/test_merge_channels.py b/tests/test_merge_channels.py new file mode 100644 index 00000000..2ac066bd --- /dev/null +++ b/tests/test_merge_channels.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import pickle as pkl +import shelve + +import lh5 +import pytest +import yaml +from lgdo.types import Array, Struct + +from legenddataflow.scripts.flow.merge_channels import merge_channels + +CHANNELS = ["ch1000000", "ch1000001"] + + +def _channel_file(tmp_path, channel, ext): + return tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{channel}-par_dsp.{ext}" + + +def _run(monkeypatch, args): + monkeypatch.setattr("sys.argv", ["merge-channels", *args]) + merge_channels() + + +def test_merge_channels_yaml(tmp_path, monkeypatch): + infiles = [] + for i, channel in enumerate(CHANNELS): + f = _channel_file(tmp_path, channel, "yaml") + f.write_text(yaml.dump({"a": i})) + infiles.append(str(f)) + out_file = tmp_path / "merged" / "l200-p03-r000-cal-20230101T000000Z-par_dsp.yaml" + + _run(monkeypatch, ["--input", *infiles, "--output", str(out_file)]) + + out_dict = yaml.safe_load(out_file.read_text()) + assert out_dict == {"ch1000000": {"a": 0}, "ch1000001": {"a": 1}} + + +def test_merge_channels_extension_mismatch(tmp_path, monkeypatch): + f = _channel_file(tmp_path, CHANNELS[0], "json") + f.write_text("{}") + out_file = tmp_path / "out.yaml" + + with pytest.raises(RuntimeError, match="does not match the extension"): + _run(monkeypatch, ["--input", str(f), "--output", str(out_file)]) + + +def test_merge_channels_pickle(tmp_path, monkeypatch): + infiles = [] + for i, channel in enumerate(CHANNELS): + f = _channel_file(tmp_path, channel, "pkl") + with f.open("wb") as w: + pkl.dump({"result": i}, w) + infiles.append(str(f)) + out_file = tmp_path / "l200-p03-r000-cal-20230101T000000Z-par_dsp.pkl" + + _run(monkeypatch, ["--input", *infiles, "--output", str(out_file)]) + + with out_file.open("rb") as r: + out_dict = pkl.load(r) + assert out_dict == {"ch1000000": {"result": 0}, "ch1000001": {"result": 1}} + + +def test_merge_channels_shelve(tmp_path, monkeypatch): + infiles = [] + for i, channel in enumerate(CHANNELS): + f = _channel_file(tmp_path, channel, "pkl") + with f.open("wb") as w: + pkl.dump({"result": i, "common": {"shared": True}}, w) + infiles.append(str(f)) + out_file = tmp_path / "l200-p03-r000-cal-20230101T000000Z-plt_dsp.dir" + + _run(monkeypatch, ["--input", *infiles, "--output", str(out_file)]) + + with shelve.open(str(out_file.with_suffix(""))) as shelf: + assert shelf["ch1000000"] == {"result": 0} + assert shelf["ch1000001"] == {"result": 1} + # per-channel "common" blocks are collected into one entry + assert shelf["common"] == { + "ch1000000": {"shared": True}, + "ch1000001": {"shared": True}, + } + + +def test_merge_channels_lh5(tmp_path, monkeypatch): + infiles = [] + for i, channel in enumerate(CHANNELS): + f = _channel_file(tmp_path, channel, "lh5") + lh5.write( + Struct({"data": Array([i, i + 1])}), + name=channel, + lh5_file=str(f), + wo_mode="a", + ) + infiles.append(str(f)) + out_file = tmp_path / "l200-p03-r000-cal-20230101T000000Z-par_dsp.lh5" + + in_db = tmp_path / "in_db.yaml" + in_db.write_text( + yaml.dump( + {channel: {"file": f} for channel, f in zip(CHANNELS, infiles, strict=True)} + ) + ) + out_db = tmp_path / "out_db.yaml" + + _run( + monkeypatch, + [ + "--input", + *infiles, + "--output", + str(out_file), + "--in-db", + str(in_db), + "--out-db", + str(out_db), + ], + ) + + for i, channel in enumerate(CHANNELS): + assert list(lh5.read(channel, str(out_file))["data"]) == [i, i + 1] + # per-channel db file references are rewritten to the merged file + out_db_dict = yaml.safe_load(out_db.read_text()) + assert out_db_dict == { + channel: {"file": f"$_/{out_file.name}"} for channel in CHANNELS + } diff --git a/tests/test_merge_in_channel.py b/tests/test_merge_in_channel.py new file mode 100644 index 00000000..6d630891 --- /dev/null +++ b/tests/test_merge_in_channel.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import pickle as pkl + +import lh5 +import pytest +import yaml +from lgdo.types import Array, Struct + +from legenddataflow.scripts.flow.merge_in_channel import merge_in_channel + +CHANNEL = "ch1000000" + + +def _infile(tmp_path, step, ext): + return tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{CHANNEL}-{step}.{ext}" + + +@pytest.fixture +def run(monkeypatch): + def _run(args): + monkeypatch.setattr("sys.argv", ["merge-in-channels", *args]) + merge_in_channel() + + return _run + + +def test_merge_in_channel_yaml(tmp_path, run): + f1 = _infile(tmp_path, "par_dsp_eopt", "yaml") + f1.write_text(yaml.dump({"eopt": {"a": 1}})) + f2 = _infile(tmp_path, "par_dsp_dplms", "yaml") + f2.write_text(yaml.dump({"dplms": {"b": 2}})) + out_file = tmp_path / "merged" / f"l200-p03-r000-cal-{CHANNEL}-par_dsp.yaml" + + run(["--input", str(f1), str(f2), "--output", str(out_file)]) + + assert yaml.safe_load(out_file.read_text()) == { + "eopt": {"a": 1}, + "dplms": {"b": 2}, + } + + +def test_merge_in_channel_pickle(tmp_path, run): + f1 = _infile(tmp_path, "par_dsp_eopt", "pkl") + with f1.open("wb") as w: + pkl.dump({"eopt": 1}, w) + f2 = _infile(tmp_path, "par_dsp_dplms", "pkl") + with f2.open("wb") as w: + pkl.dump({"dplms": 2}, w) + out_file = tmp_path / f"l200-p03-r000-cal-{CHANNEL}-par_dsp.pkl" + + run(["--input", str(f1), str(f2), "--output", str(out_file)]) + + with out_file.open("rb") as r: + assert pkl.load(r) == {"eopt": 1, "dplms": 2} + + +def test_merge_in_channel_lh5(tmp_path, run): + f1 = _infile(tmp_path, "par_dsp_eopt", "lh5") + lh5.write( + Struct({"eopt": Array([1, 2, 3])}), name=CHANNEL, lh5_file=str(f1), wo_mode="a" + ) + f2 = _infile(tmp_path, "par_dsp_dplms", "lh5") + lh5.write( + Struct({"dplms": Array([4, 5])}), name=CHANNEL, lh5_file=str(f2), wo_mode="a" + ) + out_file = tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{CHANNEL}-par_dsp.lh5" + + run(["--input", str(f1), str(f2), "--output", str(out_file)]) + + res = lh5.read(CHANNEL, str(out_file)) + # the identifier of each input's processing step becomes a struct field + assert set(res.keys()) == {"eopt", "dplms"} + assert list(res["eopt"]) == [1, 2, 3] + assert list(res["dplms"]) == [4, 5] + + +def test_merge_in_channel_lh5_multipart_identifier(tmp_path, run): + # identifiers may contain underscores and must be kept whole + f = _infile(tmp_path, "par_dsp_svm_train", "lh5") + lh5.write( + Struct({"svm_train": Array([1, 2])}), name=CHANNEL, lh5_file=str(f), wo_mode="a" + ) + out_file = tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{CHANNEL}-par_dsp.lh5" + + run(["--input", str(f), "--output", str(out_file)]) + + res = lh5.read(CHANNEL, str(out_file)) + assert set(res.keys()) == {"svm_train"} + + +def test_merge_in_channel_lh5_no_identifier(tmp_path, run): + f = _infile(tmp_path, "par_dsp", "lh5") + lh5.write(Struct({"a": Array([1])}), name=CHANNEL, lh5_file=str(f), wo_mode="a") + out_file = tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{CHANNEL}-par_dsp.lh5" + + with pytest.raises(ValueError, match="has no identifier part"): + run(["--input", str(f), "--output", str(out_file)]) + + +def test_merge_in_channel_lh5_mixed_channels(tmp_path, run): + f1 = _infile(tmp_path, "par_dsp_eopt", "lh5") + lh5.write(Struct({"eopt": Array([1])}), name=CHANNEL, lh5_file=str(f1), wo_mode="a") + other = "ch1000001" + f2 = tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{other}-par_dsp_dplms.lh5" + lh5.write(Struct({"dplms": Array([2])}), name=other, lh5_file=str(f2), wo_mode="a") + out_file = tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{CHANNEL}-par_dsp.lh5" + + with pytest.raises(ValueError, match="span multiple channels"): + run(["--input", str(f1), str(f2), "--output", str(out_file)]) + + +def test_merge_in_channel_lh5_missing_field(tmp_path, run): + f = _infile(tmp_path, "par_dsp_eopt", "lh5") + lh5.write(Struct({"wrong": Array([1])}), name=CHANNEL, lh5_file=str(f), wo_mode="a") + out_file = tmp_path / f"l200-p03-r000-cal-20230101T000000Z-{CHANNEL}-par_dsp.lh5" + + with pytest.raises(KeyError, match="has no field 'eopt'"): + run(["--input", str(f), "--output", str(out_file)]) diff --git a/tests/test_pars_loading.py b/tests/test_pars_loading.py index b8f73992..9c10340d 100644 --- a/tests/test_pars_loading.py +++ b/tests/test_pars_loading.py @@ -4,6 +4,8 @@ from pathlib import Path from unittest.mock import MagicMock, patch +import pytest + from legenddataflow.methods import ParsCatalog log = logging.getLogger(__name__) @@ -28,16 +30,20 @@ def test_match_pars_files(): assert result1 == ["file1_step1_typeA", "file3_step1_typeA"] -def test_get_par_file(): +def test_get_par_file(tmp_path): + overwrite_path = tmp_path / "overwrite" setup = { "paths": { "par": "/pars", - "overwrite": "/overwrite", + "overwrite": str(overwrite_path), "det_status": "/det_status", } } timestamp = "20230101T000000Z" tier = "test_tier" + # the overwrite validity file must exist for get_par_file to read it + (overwrite_path / tier).mkdir(parents=True) + (overwrite_path / tier / "validity.yaml").touch() catalog = ParsCatalog.get( [ @@ -81,7 +87,30 @@ def test_get_par_file(): expected_result = [ Path("/pars/file1.yaml"), Path("/pars/file2.yaml"), - Path("/overwrite/test_tier/file3.yaml"), + overwrite_path / tier / "file3.yaml", ] assert result == expected_result + + +def test_get_par_file_missing_overwrite_validity(tmp_path): + setup = {"paths": {"par": "/pars", "overwrite": str(tmp_path / "overwrite")}} + catalog = ParsCatalog.get( + [{"valid_from": "20230101T000000Z", "apply": ["file1.yaml"]}] + ) + + with ( + patch("legenddataflow.methods.pars_loading.pars_path") as mock_pars_path, + patch( + "legenddataflow.methods.pars_loading.get_pars_path" + ) as mock_get_pars_path, + patch( + "legenddataflow.methods.pars_loading.par_overwrite_path" + ) as mock_par_overwrite_path, + ): + mock_pars_path.return_value = setup["paths"]["par"] + mock_get_pars_path.return_value = setup["paths"]["par"] + mock_par_overwrite_path.return_value = setup["paths"]["overwrite"] + + with pytest.raises(FileNotFoundError, match="par-overwrite validity file"): + ParsCatalog.get_par_file(catalog, setup, "20230101T000000Z", "test_tier") diff --git a/tests/test_paths.py b/tests/test_paths.py new file mode 100644 index 00000000..f6bd331f --- /dev/null +++ b/tests/test_paths.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import logging +from pathlib import Path + +import pytest +import yaml +from legenddataflowscripts.workflow import subst_vars + +from legenddataflow.methods import paths + +testprod = Path(__file__).parent / "dummy_cycle" + +with (testprod / "config.yaml").open() as r: + setup = yaml.safe_load(r) +subst_vars(setup, var_values={"_": str(testprod)}) + + +def test_simple_resolvers(): + assert paths.tier_path(setup) == f"{testprod}/generated/tier" + assert paths.pars_path(setup) == f"{testprod}/generated/par" + assert paths.par_overwrite_path(setup) == f"{testprod}/inputs/dataprod/overrides" + assert paths.metadata_path(setup) == f"{testprod}/inputs" + assert paths.filelist_path(setup) == f"{testprod}/generated/tmp/filelists" + + +def test_sandbox_path(): + assert paths.sandbox_path(setup) == "" + assert paths.sandbox_path({"paths": {}}) is None + + +def test_missing_path_key(): + with pytest.raises(KeyError, match=r"dataflow-config paths\.tier_daq is not set"): + paths.tier_daq_path({"paths": {}}) + with pytest.raises(KeyError, match=r"dataflow-config paths\.tier_dsp is not set"): + paths.get_tier_path({"paths": {}}, "dsp") + + +def test_get_tier_path(): + assert paths.get_tier_path(setup, "dsp") == f"{testprod}/generated/tier/dsp" + assert paths.get_tier_path(setup, "skm") == f"{testprod}/generated/tier/skm" + with pytest.raises(ValueError, match="no tier matching"): + paths.get_tier_path(setup, "daq") + with pytest.raises(ValueError, match="no tier matching"): + paths.get_tier_path(setup, "not_a_tier") + + +def test_get_pars_path(): + assert paths.get_pars_path(setup, "hit") == f"{testprod}/generated/par/hit" + with pytest.raises(ValueError, match="no tier matching"): + paths.get_pars_path(setup, "skm") + + +def test_link_external_paths_creates_symlink(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + external = tmp_path / "other-prod" / "generated" / "tier" / "dsp" + external.mkdir(parents=True) + + paths.link_external_paths({"paths": {"tier_dsp": str(external)}}) + + default = tmp_path / "generated" / "tier" / "dsp" + assert default.is_symlink() + assert default.resolve() == external.resolve() + + # a second call with the same config is a no-op + paths.link_external_paths({"paths": {"tier_dsp": str(external)}}) + assert default.resolve() == external.resolve() + + +def test_link_external_paths_removes_stale_symlink(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + default = tmp_path / "generated" / "tier" / "dsp" + default.parent.mkdir(parents=True) + default.symlink_to(tmp_path / "somewhere-else") + + paths.link_external_paths({"paths": {"tier_dsp": str(default)}}) + + assert not default.is_symlink() + assert not default.exists() + + +def test_link_external_paths_keeps_real_directory(tmp_path, monkeypatch, caplog): + monkeypatch.chdir(tmp_path) + default = tmp_path / "generated" / "tier" / "dsp" + default.mkdir(parents=True) + external = tmp_path / "other-prod" / "dsp" + external.mkdir(parents=True) + + with caplog.at_level(logging.WARNING): + paths.link_external_paths({"paths": {"tier_dsp": str(external)}}) + + assert not default.is_symlink() + assert "ignoring override" in caplog.text diff --git a/tests/test_patterns.py b/tests/test_patterns.py new file mode 100644 index 00000000..8f4b1fa8 --- /dev/null +++ b/tests/test_patterns.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import copy +from pathlib import Path + +import pytest +import yaml +from legenddataflowscripts.workflow import subst_vars + +from legenddataflow.methods import patterns + +testprod = Path(__file__).parent / "dummy_cycle" + +with (testprod / "config.yaml").open() as r: + setup = yaml.safe_load(r) +subst_vars(setup, var_values={"_": str(testprod)}) + + +def test_key_patterns(): + assert ( + patterns.key_pattern() == "{experiment}-{period}-{run}-{datatype}-{timestamp}" + ) + assert patterns.processing_pattern().endswith("-{processing_step}.{ext}") + assert patterns.par_validity_pattern().startswith("{datatype}/{period}/{run}/") + assert "{channel}" in patterns.full_channel_pattern() + + +def test_get_pattern_tier(): + assert ( + str(patterns.get_pattern_tier(setup, "dsp")) + == f"{testprod}/generated/tier/dsp/{{datatype}}/{{period}}/{{run}}/" + "{experiment}-{period}-{run}-{datatype}-{timestamp}-tier_dsp.lh5" + ) + # concat patterns have no timestamp and no period/run directories + assert ( + str(patterns.get_pattern_tier(setup, "evt_concat")) + == f"{testprod}/generated/tier/evt/{{datatype}}/" + "{experiment}-{period}-{run}-{datatype}-tier_evt.lh5" + ) + # skm is phy-only and not per-run + assert ( + str(patterns.get_pattern_tier(setup, "skm")) + == f"{testprod}/generated/tier/skm/phy/" + "{experiment}-{period}-{run}-{datatype}-tier_skm.lh5" + ) + with pytest.raises(ValueError, match="invalid tier"): + patterns.get_pattern_tier(setup, "not_a_tier") + + +def test_get_pattern_tier_out_of_cycle(): + # the dummy config points tier_raw outside the production cycle + out = patterns.get_pattern_tier(setup, "raw") + assert str(out).startswith("/tmp/") + # with the check disabled the configured path is returned unchanged + kept = patterns.get_pattern_tier(setup, "raw", check_in_cycle=False) + assert str(kept).startswith("/data2/") + + +def test_get_pattern_pars(): + assert ( + str(patterns.get_pattern_pars(setup, "dsp")) + == f"{testprod}/generated/par/dsp/cal/{{period}}/{{run}}/" + "{experiment}-{period}-{run}-cal-{timestamp}-par_dsp.yaml" + ) + assert str(patterns.get_pattern_pars(setup, "dsp", name="eopt")).endswith( + "par_dsp_eopt.yaml" + ) + assert ( + str(patterns.get_pattern_pars(setup, "dsp", datatype=None)).count("{datatype}") + == 2 + ) + with pytest.raises(ValueError, match="invalid tier"): + patterns.get_pattern_pars(setup, "skm") + + +def test_get_pattern_pars_out_of_cycle(): + modified = copy.deepcopy(setup) + modified["paths"]["par_dsp"] = "/somewhere/else/par/dsp" + out = patterns.get_pattern_pars(modified, "dsp", name="eopt", extension="pkl") + assert str(out) == ( + "/tmp/{experiment}-{period}-{run}-cal-{timestamp}-par_dsp_eopt.pkl" + ) + + +def test_get_pattern_pars_tmp_channel(): + assert ( + str(patterns.get_pattern_pars_tmp_channel(setup, "dsp")) + == f"{testprod}/generated/tmp/par/" + "{experiment}-{period}-{run}-cal-{timestamp}-{channel}-par_dsp.yaml" + ) + assert str( + patterns.get_pattern_pars_tmp_channel(setup, "dsp", name="eopt") + ).endswith("-{channel}-par_dsp_eopt.yaml") + + +def test_get_pattern_log(): + assert ( + str(patterns.get_pattern_log(setup, "build_dsp", "20230101T000000Z")) + == f"{testprod}/generated/tmp/log/20230101T000000Z/build_dsp/" + "{experiment}-{period}-{run}-{datatype}-{timestamp}-build_dsp.log" + ) + assert str( + patterns.get_pattern_log_channel(setup, "par_dsp", "20230101T000000Z") + ).endswith("-cal-{timestamp}-{channel}-par_dsp.log") diff --git a/tests/test_write_filelist.py b/tests/test_write_filelist.py new file mode 100644 index 00000000..91cacd65 --- /dev/null +++ b/tests/test_write_filelist.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import runpy +from types import SimpleNamespace + +import snakemake.script + +import legenddataflow.scripts.flow.write_filelist as write_filelist_module +from legenddataflow.scripts.flow.write_filelist import write_filelist + + +def test_write_filelist(tmp_path): + out = tmp_path / "all-l200-p03-r000-phy-dsp.filelist" + write_filelist(["/data/a.lh5", "/data/b.lh5"], out) + assert out.read_text() == "/data/a.lh5\n/data/b.lh5\n" + + +def test_write_filelist_empty(tmp_path, capsys): + out = tmp_path / "empty.filelist" + write_filelist([], out, label="all-l200-p99-r999-phy") + assert out.read_text() == "" + captured = capsys.readouterr() + assert "WARNING: no files found" in captured.out + assert "all-l200-p99-r999-phy" in captured.out + + +def test_snakemake_glue(tmp_path, monkeypatch): + # emulate a snakemake script: execution by injecting the snakemake object + out = tmp_path / "out.filelist" + fake = SimpleNamespace( + input=["/data/a.lh5", "/data/b.lh5"], + output=[str(out)], + wildcards=SimpleNamespace(label="all-l200-p03-r000-phy"), + ) + monkeypatch.setattr(snakemake.script, "snakemake", fake, raising=False) + + runpy.run_path(write_filelist_module.__file__) + + assert out.read_text() == "/data/a.lh5\n/data/b.lh5\n" diff --git a/workflow/src/legenddataflow/methods/FileKey.py b/workflow/src/legenddataflow/methods/FileKey.py index cbfa1f79..872065c2 100644 --- a/workflow/src/legenddataflow/methods/FileKey.py +++ b/workflow/src/legenddataflow/methods/FileKey.py @@ -24,6 +24,8 @@ def regex_from_filepattern(filepattern): + """Compile a Snakemake-style pattern (``{wildcard}`` placeholders) into an + anchored regex with named groups; repeated wildcards become backreferences.""" f = [] wildcards = [] last = 0 @@ -49,6 +51,10 @@ def regex_from_filepattern(filepattern): class FileKey( namedtuple("FileKey", ["experiment", "period", "run", "datatype", "timestamp"]) ): + """A ``{experiment}-{period}-{run}-{datatype}-{timestamp}`` key identifying + one DAQ cycle, with conversions between keys, filenames and path patterns. + Unknown components are represented by ``*``.""" + __slots__ = () re_pattern = "(-(?P[^-]+)(\\-(?P[^-]+)(\\-(?P[^-]+)(\\-(?P[^-]+)(\\-(?P[^-]+))?)?)?)?)?$" @@ -79,16 +85,24 @@ def get_filekey_from_filename(cls, filename): @classmethod def get_filekey_from_pattern(cls, filename, pattern=None): + """Match ``filename`` against ``pattern`` (default: the class key + pattern) and build a key from the named groups; fields not present in + the pattern become ``*``. Raises :class:`ValueError` if the filename + doesn't match.""" if isinstance(pattern, Path): pattern = pattern.as_posix() filename = str(filename) - key_pattern_rx = re.compile( - regex_from_filepattern(cls.key_pattern if pattern is None else pattern) - ) - - if key_pattern_rx.match(filename) is None: - return None - d = key_pattern_rx.match(filename).groupdict() + used_pattern = cls.key_pattern if pattern is None else pattern + key_pattern_rx = re.compile(regex_from_filepattern(used_pattern)) + + match = key_pattern_rx.match(filename) + if match is None: + msg = ( + f"'{filename}' does not match the {cls.__name__} " + f"pattern '{used_pattern}'" + ) + raise ValueError(msg) + d = match.groupdict() for entry in list(d): if entry not in cls._fields: d.pop(entry) @@ -107,8 +121,14 @@ def get_unix_timestamp(self): @classmethod def parse_keypart(cls, keypart): + """Parse a possibly partial key (e.g. ``-l200-p00``); missing trailing + components become ``*``.""" keypart_rx = re.compile(cls.re_pattern) - d = keypart_rx.match(keypart).groupdict() + match = keypart_rx.match(keypart) + if match is None: + msg = f"'{keypart}' cannot be parsed as a {cls.__name__} keypart" + raise ValueError(msg) + d = match.groupdict() for key in d: if d[key] is None: d[key] = "*" @@ -118,6 +138,9 @@ def parse_keypart(cls, keypart): return cls(**d) def expand(self, file_pattern, **kwargs): + """Substitute this key's fields (overridden/extended by ``kwargs``, + whose values may be lists) into ``file_pattern``, returning the paths + for all combinations.""" file_pattern = str(file_pattern) wildcard_dict = self._asdict() if kwargs is not None: @@ -136,18 +159,22 @@ def expand(self, file_pattern, **kwargs): result.append(formatter.vformat(file_pattern, (), substitution)) return result - def get_path_from_filekey(self, pattern, **kwargs): - if kwargs is None: - return self.expand(pattern, **kwargs) - for entry, value in kwargs.items(): + def _resolve_dict_kwargs(self, kwargs): + """Resolve dict-valued kwargs by looking up the first of their keys + matching one of this key's fields; drop entries with no match.""" + for entry, value in list(kwargs.items()): if isinstance(value, dict): - if len(next(iter(set(value).intersection(self._list())))) > 0: - kwargs[entry] = value[ - next(iter(set(value).intersection(self._list()))) - ] + matches = set(value).intersection(self._list()) + if matches: + kwargs[entry] = value[next(iter(matches))] else: kwargs.pop(entry) - return self.expand(pattern, **kwargs) + return kwargs + + def get_path_from_filekey(self, pattern, **kwargs): + """Expand ``pattern`` with this key; dict-valued kwargs are resolved by + looking up the first of their keys matching one of this key's fields.""" + return self.expand(pattern, **self._resolve_dict_kwargs(kwargs)) # get_path_from_key @classmethod @@ -158,6 +185,7 @@ def get_full_path_from_filename(cls, filename, pattern, path_pattern): @staticmethod def tier_files(setup, keys, tier): + """Expand each key string in ``keys`` into its file path for ``tier``.""" fn_pattern = get_pattern_tier(setup, tier) files = [] for line in keys: @@ -169,6 +197,9 @@ def tier_files(setup, keys, tier): class ProcessingFileKey(FileKey): + """A :class:`FileKey` extended with a ``processing_step`` component of the + form ``{type}_{tier}[_{identifier}]`` (e.g. ``par_dsp_eopt``).""" + _fields = (*FileKey._fields, "processing_step") key_pattern = processing_pattern() @@ -212,20 +243,12 @@ def get_path_from_filekey(self, pattern, **kwargs): pattern = pattern.as_posix() if not isinstance(pattern, str): pattern = pattern(self.tier, self.identifier) - if kwargs is None: - return self.expand(pattern, **kwargs) - for entry, value in kwargs.items(): - if isinstance(value, dict): - if len(next(iter(set(value).intersection(self._list())))) > 0: - kwargs[entry] = value[ - next(iter(set(value).intersection(self._list()))) - ] - else: - kwargs.pop(entry) - return self.expand(pattern, **kwargs) + return self.expand(pattern, **self._resolve_dict_kwargs(kwargs)) class ChannelProcKey(ProcessingFileKey): + """A :class:`ProcessingFileKey` extended with a ``channel`` component.""" + re_pattern = "all(-(?P[^-]+)(\\-(?P[^-]+)(\\-(?P[^-]+)(\\-(?P[^-]+)(\\-(?P[^-]+)(\\-(?P[^-]+))?)?)?)?)?)?$" key_pattern = full_channel_pattern_with_extension() _fields = (*ProcessingFileKey._fields, "channel") @@ -251,6 +274,8 @@ def _asdict(self): @staticmethod def get_channel_files(keypart, par_pattern, chan_list): + """Expand ``par_pattern`` once per channel in ``chan_list``, using the + components parsed from ``keypart``.""" if isinstance(par_pattern, Path): par_pattern = par_pattern.as_posix() d = ChannelProcKey.parse_keypart(keypart) @@ -266,10 +291,7 @@ def get_channel_files(keypart, par_pattern, chan_list): def per_grouper(files): - """ - Returns list containing lists of each run - """ - + """Group files into one list per ``experiment-period``.""" pers = [] per_files = [] for file in files: @@ -284,7 +306,7 @@ def per_grouper(files): def run_grouper(files): - """Returns list containing lists of each run""" + """Group files into one list per ``experiment-period-run``.""" runs = [] run_files = [] for file in files: @@ -299,10 +321,7 @@ def run_grouper(files): def run_splitter(files): - """ - Returns list containing lists of each run - """ - + """Group files into one list per ``period-run``.""" runs = [] run_files = [] for file in files: diff --git a/workflow/src/legenddataflow/methods/cal_grouping.py b/workflow/src/legenddataflow/methods/cal_grouping.py index 2065cc83..0eb0f0e4 100644 --- a/workflow/src/legenddataflow/methods/cal_grouping.py +++ b/workflow/src/legenddataflow/methods/cal_grouping.py @@ -4,9 +4,11 @@ from __future__ import annotations +import logging from pathlib import Path from dbetto import Props +from legendmeta.utils import expand_runs as _expand_runs from .FileKey import ChannelProcKey, ProcessingFileKey from .paths import filelist_path @@ -17,42 +19,47 @@ get_pattern_plts_tmp_channel, ) +log = logging.getLogger(__name__) + class CalGrouping: + """Resolve partition-level calibration groupings (from + ``cal_groupings.yaml``): which runs belong to each partition per channel, + and the corresponding filelists, parameter/plot files, log paths and + wildcard constraints.""" + def __init__(self, setup, input_file): self.datasets = Props.read_from(input_file) self.expand_runs() self.setup = setup def expand_runs(self): + """Expand ``r000..r005`` run-range strings into explicit run lists.""" for channel, chan_dict in self.datasets.items(): for part, part_dict in chan_dict.items(): for per, runs in part_dict.items(): - if isinstance(runs, str) and ".." in runs: - start, end = runs.split("..") - self.datasets[channel][part][per] = [ - f"r{x:03}" for x in range(int(start[1:]), int(end[1:]) + 1) - ] - if isinstance(runs, list): - final_runs = [] - for run in runs: - if ".." in run: - start, end = run.split("..") - final_runs += [ - f"r{x:03}" - for x in range(int(start[1:]), int(end[1:]) + 1) - ] - else: - final_runs.append(run) - self.datasets[channel][part][per] = final_runs + if runs != "all": + self.datasets[channel][part][per] = _expand_runs(runs) def get_dataset(self, dataset, channel): + """Return the ``{period: runs}`` dict for a partition, with + channel-specific overrides applied on top of ``default``.""" + if "default" not in self.datasets: + msg = "cal-groupings config has no 'default' section" + raise KeyError(msg) partition_dict = self.datasets["default"].copy() if channel in self.datasets: partition_dict.update(self.datasets[channel]) + if dataset not in partition_dict: + msg = ( + f"partition {dataset!r} not defined for channel {channel!r} " + f"(or 'default'): available partitions are {sorted(partition_dict)}" + ) + raise KeyError(msg) return partition_dict[dataset] def get_filelists(self, dataset, channel, tier, experiment="l200", datatype="cal"): + """Return the filelist paths covering all runs in the partition.""" dataset = self.get_dataset(dataset, channel) files = [] for per in dataset: @@ -81,6 +88,9 @@ def get_par_files( extension="yaml", pattern_func=get_pattern_pars_tmp_channel, ): + """Select from ``catalog`` the per-channel parameter files whose keys + fall inside the partition's periods and runs, expanded with + ``pattern_func``.""" dataset = self.get_dataset(dataset, channel) all_par_files = [] for item in catalog.entries["all"]: @@ -134,6 +144,7 @@ def get_plt_files( datatype="cal", name=None, ): + """Like :meth:`get_par_files` but for the pickled plot outputs.""" return self.get_par_files( catalog, dataset, @@ -157,6 +168,8 @@ def get_log_file( datatype="cal", name=None, ): + """Build the log-file path for the partition job, derived from the + first matching parameter file.""" par_files = self.get_par_files( catalog, dataset, @@ -175,11 +188,19 @@ def get_log_file( return fk.get_path_from_filekey( get_pattern_log_channel(self.setup, name, processing_timestamp) )[0] + log.warning( + "partition %r channel %r tier %r resolved to no par files; " + "using placeholder log file name", + dataset, + channel, + tier, + ) return "log.log" def get_timestamp( self, catalog, dataset, channel, tier, experiment="l200", datatype="cal" ): + """Return the timestamp of the partition's first parameter file.""" par_files = self.get_par_files( catalog, dataset, @@ -192,9 +213,18 @@ def get_timestamp( if len(par_files) > 0: fk = ChannelProcKey.get_filekey_from_pattern(Path(par_files[0]).name) return fk.timestamp + log.warning( + "partition %r channel %r tier %r resolved to no par files; " + "using placeholder timestamp", + dataset, + channel, + tier, + ) return "20000101T000000Z" def get_wildcard_constraints(self, dataset, channel): + """Build the channel wildcard regex; for ``default`` it excludes + channels that have their own override for the partition's runs.""" if channel == "default": exclude_chans = [] default_runs = self.get_dataset(dataset, channel) diff --git a/workflow/src/legenddataflow/methods/create_pars_keylist.py b/workflow/src/legenddataflow/methods/create_pars_keylist.py index bf7cd930..35966122 100644 --- a/workflow/src/legenddataflow/methods/create_pars_keylist.py +++ b/workflow/src/legenddataflow/methods/create_pars_keylist.py @@ -111,6 +111,16 @@ def get_keys(keypart, search_pattern, ignore_keys=None): def apply_run_override(cls, hit_par_catalog, name_dict, run_overwrite_validity): run_overwrite_catalog = ParsCatalog.read_from(run_overwrite_validity) + if "all" not in run_overwrite_catalog.entries: + msg = ( + f"run-override validity file {run_overwrite_validity} has no " + "'all' system entry" + ) + raise ValueError(msg) + if "all" not in hit_par_catalog.entries: + msg = "par catalog passed to apply_run_override has no 'all' system entry" + raise ValueError(msg) + for i, entry in enumerate(run_overwrite_catalog.entries["all"]): if len(entry.file) != 0: # check if common timestamp and remove old if this is the case diff --git a/workflow/src/legenddataflow/methods/pars_loading.py b/workflow/src/legenddataflow/methods/pars_loading.py index 11b3d9e8..7e66fdd2 100644 --- a/workflow/src/legenddataflow/methods/pars_loading.py +++ b/workflow/src/legenddataflow/methods/pars_loading.py @@ -79,6 +79,9 @@ def get_par_file(self, setup: dict, timestamp: str, tier: str) -> list: allow_none = setup.get("allow_none_par", False) if pars_path(setup) not in get_pars_path(setup, tier): par_file = Path(get_pars_path(setup, tier)) / "validity.yaml" + if not par_file.is_file(): + msg = f"validity file {par_file} for tier {tier!r} does not exist" + raise FileNotFoundError(msg) catalog = ParsCatalog.read_from(par_file) pars_files = catalog.valid_for(timestamp, allow_none=allow_none) else: @@ -87,6 +90,12 @@ def get_par_file(self, setup: dict, timestamp: str, tier: str) -> list: pars_files = [] par_overwrite_file = Path(par_overwrite_path(setup)) / tier / "validity.yaml" + if not par_overwrite_file.is_file(): + msg = ( + f"par-overwrite validity file {par_overwrite_file} for tier " + f"{tier!r} does not exist" + ) + raise FileNotFoundError(msg) overwrite_catalog = ParsCatalog.read_from(par_overwrite_file) pars_files_overwrite = overwrite_catalog.valid_for( timestamp, allow_none=allow_none diff --git a/workflow/src/legenddataflow/methods/paths.py b/workflow/src/legenddataflow/methods/paths.py index e0822031..e57ff796 100644 --- a/workflow/src/legenddataflow/methods/paths.py +++ b/workflow/src/legenddataflow/methods/paths.py @@ -11,20 +11,29 @@ def sandbox_path(setup): if "sandbox_path" in setup["paths"]: - return setup["paths"]["sandbox_path"] + return _get_path(setup, "sandbox_path") return None +def _get_path(setup, key): + """Look up ``paths.`` in the dataflow config, with a named error.""" + try: + return setup["paths"][key] + except KeyError: + msg = f"dataflow-config paths.{key} is not set" + raise KeyError(msg) from None + + def tier_daq_path(setup): - return setup["paths"]["tier_daq"] + return _get_path(setup, "tier_daq") def tier_raw_blind_path(setup): - return setup["paths"]["tier_raw_blind"] + return _get_path(setup, "tier_raw_blind") def tier_path(setup): - return setup["paths"]["tier"] + return _get_path(setup, "tier") def get_tier_path(setup, tier): @@ -41,72 +50,72 @@ def get_tier_path(setup, tier): "pet", "skm", ]: - return setup["paths"][f"tier_{tier}"] + return _get_path(setup, f"tier_{tier}") msg = f"no tier matching:{tier}" raise ValueError(msg) def pars_path(setup): - return setup["paths"]["par"] + return _get_path(setup, "par") def get_pars_path(setup, tier): if tier in ["raw", "tcm", "dsp", "hit", "evt", "psp", "pht", "pet"]: - return setup["paths"][f"par_{tier}"] + return _get_path(setup, f"par_{tier}") msg = f"no tier matching:{tier}" raise ValueError(msg) def tmp_par_path(setup): - return setup["paths"]["tmp_par"] + return _get_path(setup, "tmp_par") def tmp_plts_path(setup): - return setup["paths"]["tmp_plt"] + return _get_path(setup, "tmp_plt") def plts_path(setup): - return setup["paths"]["plt"] + return _get_path(setup, "plt") def par_overwrite_path(setup): - return setup["paths"]["par_overwrite"] + return _get_path(setup, "par_overwrite") def config_path(setup): - return setup["paths"]["config"] + return _get_path(setup, "config") def chan_map_path(setup): - return setup["paths"]["chan_map"] + return _get_path(setup, "chan_map") def det_status_path(setup): - return setup["paths"]["detector_status"] + return _get_path(setup, "detector_status") def metadata_path(setup): - return setup["paths"]["metadata"] + return _get_path(setup, "metadata") def detector_db_path(setup): - return setup["paths"]["detector_db"] + return _get_path(setup, "detector_db") def log_path(setup): - return setup["paths"]["log"] + return _get_path(setup, "log") def tmp_log_path(setup): - return setup["paths"]["tmp_log"] + return _get_path(setup, "tmp_log") def tmp_benchmark_path(setup): - return setup["paths"]["tmp_benchmark"] + return _get_path(setup, "tmp_benchmark") def filelist_path(setup): - return setup["paths"]["tmp_filelists"] + return _get_path(setup, "tmp_filelists") # Maps each paths. candidate to its canonical default location relative to diff --git a/workflow/src/legenddataflow/methods/patterns.py b/workflow/src/legenddataflow/methods/patterns.py index ea0d5a2d..e0b93157 100644 --- a/workflow/src/legenddataflow/methods/patterns.py +++ b/workflow/src/legenddataflow/methods/patterns.py @@ -118,8 +118,8 @@ def get_pattern_tier(setup, tier, check_in_cycle=True): / "{experiment}-{period}-{run}-{datatype}-tier_skm.lh5" ) else: - msg = "invalid tier" - raise Exception(msg) + msg = f"invalid tier {tier!r}" + raise ValueError(msg) if ( tier_path(setup) not in str(file_pattern.resolve(strict=False)) and check_in_cycle is True @@ -161,8 +161,8 @@ def get_pattern_pars( ) ) else: - msg = "invalid tier" - raise Exception(msg) + msg = f"invalid tier {tier!r}" + raise ValueError(msg) if ( pars_path(setup) not in str(Path(file_pattern).resolve(strict=False)) and check_in_cycle is True diff --git a/workflow/src/legenddataflow/scripts/flow/build_chanlist.py b/workflow/src/legenddataflow/scripts/flow/build_chanlist.py index daccdefd..849d7cff 100644 --- a/workflow/src/legenddataflow/scripts/flow/build_chanlist.py +++ b/workflow/src/legenddataflow/scripts/flow/build_chanlist.py @@ -1,5 +1,8 @@ # ruff: noqa: I002, T201 +"""Determine the channels to process for a given timestamp and build the +corresponding per-channel parameter and plot output filenames.""" + from pathlib import Path from dbetto import TextDB @@ -19,13 +22,13 @@ def get_chanlist(timestamp, datatype, det_status, channelmap, system): channelmap = TextDB(channelmap, lazy=True) if isinstance(det_status, TextDB): - status_map = det_status.statuses.on(timestamp, system=datatype) + status_map = det_status.statuses.on(timestamp, category=datatype) else: - status_map = det_status.valid_for(timestamp, system=datatype) + status_map = det_status.valid_for(timestamp, category=datatype) if isinstance(channelmap, TextDB): - chmap = channelmap.channelmaps.on(timestamp, system=datatype) + chmap = channelmap.channelmaps.on(timestamp, category=datatype) else: - chmap = channelmap.valid_for(timestamp, system=datatype) + chmap = channelmap.valid_for(timestamp, category=datatype) # only restrict to a certain system (geds, spms, ...) channels = [] diff --git a/workflow/src/legenddataflow/scripts/flow/build_filelist.py b/workflow/src/legenddataflow/scripts/flow/build_filelist.py index c7ee719b..5dbcf536 100644 --- a/workflow/src/legenddataflow/scripts/flow/build_filelist.py +++ b/workflow/src/legenddataflow/scripts/flow/build_filelist.py @@ -1,5 +1,8 @@ # ruff: noqa: I002 +"""Build the lists of input files matching a ``.gen`` production target, +grouping keys by run and handling the per-tier concatenation rules.""" + import glob from pathlib import Path @@ -21,23 +24,23 @@ def expand_runs(in_dict): "p01": "r001..r005" } """ + + def _expand_range(run_range): + start, end = run_range.split("..") + return [f"r{x:03}" for x in range(int(start[1:]), int(end[1:]) + 1)] + for datatype, datalist in in_dict.items(): for per, run_list in datalist.items(): if isinstance(run_list, str) and ".." in run_list: - start, end = run_list.split("..") - in_dict[datatype][per] = [ - f"r{x:03}" for x in range(int(start[1:]), int(end[1:]) + 1) - ] + in_dict[datatype][per] = _expand_range(run_list) if isinstance(run_list, list): - for i, run in enumerate(run_list): + expanded_runs = [] + for run in run_list: if isinstance(run, str) and ".." in run: - start, end = run.split("..") - run_list.pop(i) - expanded_runs = [ - f"r{x:03}" for x in range(int(start[1:]), int(end[1:]) + 1) - ] - in_dict[datatype][per] += expanded_runs - in_dict[datatype][per] = sorted(run_list) + expanded_runs += _expand_range(run) + else: + expanded_runs.append(run) + in_dict[datatype][per] = sorted(expanded_runs) return in_dict @@ -63,29 +66,41 @@ def get_analysis_runs( ] else: - msg = "ignore_keys_file file not in json, yaml or keylist format" + msg = ( + f"unsupported ignore_keys file extension " + f"{Path(ignore_keys_file).suffix!r} ({ignore_keys_file}): " + "expected .json, .yaml, .yml or .keylist" + ) raise ValueError(msg) else: - msg = f"no ignore_keys file found: {ignore_keys_file}" + msg = f"ignore_keys file not found: {ignore_keys_file}" raise ValueError(msg) if analysis_runs_file is not None and file_selection != "all": if Path(analysis_runs_file).is_file(): - if Path(ignore_keys_file).suffix in (".json", ".yaml", ".yml"): + if Path(analysis_runs_file).suffix in (".json", ".yaml", ".yml"): analysis_runs = Props.read_from(analysis_runs_file) else: - msg = f"analysis_runs file not in json or yaml format: {analysis_runs_file}" + msg = ( + f"unsupported analysis_runs file extension " + f"{Path(analysis_runs_file).suffix!r} ({analysis_runs_file}): " + "expected .json, .yaml or .yml" + ) raise ValueError(msg) if file_selection in analysis_runs: analysis_runs = expand_runs( analysis_runs[file_selection] ) # select the file_selection and expand out the runs else: - msg = f"Unknown file selection: {file_selection} not in {list(analysis_runs)}" + msg = ( + f"unknown file selection {file_selection!r}: " + f"available selections in {analysis_runs_file} are " + f"{list(analysis_runs)}" + ) raise ValueError(msg) else: - msg = f"no analysis_runs file found: {analysis_runs_file}" + msg = f"analysis_runs file not found: {analysis_runs_file}" raise ValueError(msg) return analysis_runs, ignore_keys @@ -168,12 +183,14 @@ def build_filelist( the keys specified in the analysis_runs dict. """ # the ignore_keys dictionary organizes keys in sections, gather all the - # section contents in a single list - if ignore_keys is not None: + # section contents in a single list; keylist files yield a flat list + if isinstance(ignore_keys, dict): _ignore_keys = ignore_keys.get("unprocessable", []) if remove_level == "removed": _ignore_keys += ignore_keys.get("removed", []) ignore_keys = sorted(_ignore_keys) + elif ignore_keys is not None: + ignore_keys = sorted(ignore_keys) else: ignore_keys = [] diff --git a/workflow/src/legenddataflow/scripts/flow/complete_run.py b/workflow/src/legenddataflow/scripts/flow/complete_run.py index 7dfc3633..81e048ee 100644 --- a/workflow/src/legenddataflow/scripts/flow/complete_run.py +++ b/workflow/src/legenddataflow/scripts/flow/complete_run.py @@ -1,5 +1,9 @@ # ruff: noqa: T201, I002 +"""Run-completion script executed by the ``autogen_output`` rule: checks log +files for warnings and errors, builds ``pygama`` FileDB databases and +collects the produced run information.""" + import datetime import json import os @@ -9,80 +13,57 @@ from legenddataflowscripts.workflow import as_ro, execenv_pyexe from legenddataflowscripts.workflow.execenv import _execenv2str -from snakemake.script import snakemake from legenddataflow.methods import paths as pat from legenddataflow.methods import patterns -print("INFO: dataflow ran successfully, now few final checks and scripts") - - -def as_ro_path(path): - return as_ro(snakemake.params.setup, path) +try: + from snakemake.script import snakemake +except ImportError: # not running inside a snakemake job (e.g. unit tests) + snakemake = None def check_log_files(log_path, output_file, gen_output, warning_file=None): now = datetime.datetime.now(datetime.UTC).strftime("%d/%m/%y %H:%M") - Path(output_file).parent.mkdir(parents=True, exist_ok=True) - if warning_file is not None: - Path(warning_file).parent.mkdir(parents=True, exist_ok=True) - with Path(warning_file).open("w") as w, Path(output_file).open("w") as f: - n_errors = 0 - n_warnings = 0 - for file in Path(log_path).rglob("*.log"): - with Path(file).open() as r: - text = r.read() - if "ERROR" in text or "WARNING" in text: - for line in text.splitlines(): - if "ERROR" in line: - if n_errors == 0: - f.write( - f"{gen_output} successfully generated at {now} with errors \n" - ) - if n_warnings == 0: - w.write( - f"{gen_output} successfully generated at {now} with warnings \n" - ) - f.write(f"{Path(file).name} : {line}\n") - n_errors += 1 - elif "WARNING" in line: - w.write(f"{Path(file).name} : {line}\n") - n_warnings += 1 - else: - pass - Path(file).unlink() - text = None - if n_errors == 0: - f.write( - f"{gen_output} successfully generated at {now} with no errors \n" - ) - if n_warnings == 0: - w.write( - f"{gen_output} successfully generated at {now} with no warnings \n" - ) - else: - with Path(output_file).open("w") as f: - n_errors = 0 - for file in Path(log_path).rglob("*.log"): - with Path(file).open() as r: - text = r.read() - if "ERROR" in text: - for line in text.splitlines(): - if "ERROR" in line: - if n_errors == 0: - f.write( - f"{gen_output} successfully generated at {now} with errors \n" - ) - f.write(f"{Path(file).name} : {line}\n") - n_errors += 1 - else: - pass - Path(file).unlink() - text = None - if n_errors == 0: + + # collect ERROR/WARNING lines from all log files, consuming the logs; + # uncaught tracebacks carry no "ERROR" token, so track traceback blocks + # and report their final exception line as an error + error_lines = [] + warning_lines = [] + for file in Path(log_path).rglob("*.log"): + in_traceback = False + with Path(file).open() as r: + for line in r: + if line.startswith("Traceback (most recent call last):"): + in_traceback = True + elif in_traceback: + if line[:1] in (" ", "\t") or not line.strip(): + continue + # first unindented line is the exception itself + error_lines.append(f"{Path(file).name} : {line.rstrip()}\n") + in_traceback = False + elif "ERROR" in line: + error_lines.append(f"{Path(file).name} : {line.rstrip()}\n") + elif "WARNING" in line: + warning_lines.append(f"{Path(file).name} : {line.rstrip()}\n") + Path(file).unlink() + + def _write_summary(path, lines, label): + Path(path).parent.mkdir(parents=True, exist_ok=True) + with Path(path).open("w") as f: + if lines: + f.write(f"{gen_output} successfully generated at {now} with {label} \n") + f.writelines(lines) + else: f.write( - f"{gen_output} successfully generated at {now} with no errors \n" + f"{gen_output} successfully generated at {now} with no {label} \n" ) + + _write_summary(output_file, error_lines, "errors") + if warning_file is not None: + _write_summary(warning_file, warning_lines, "warnings") + walk = list(os.walk(log_path)) for path, _, _ in walk[::-1]: if len(list(Path(path).iterdir())) == 0: @@ -97,18 +78,18 @@ def find_gen_runs(gen_tier_path): # then look for concat tiers (use filenames now) paths_concat = gen_tier_path.glob("*/*/*.lh5") - # use the directories to build a datatype/period/run string + # {experiment}-{period}-{run}-{datatype}-... -> datatype/period/run runs_concat = { - "/".join([str(p).split("-")[3], *str(p).split("-")[1:3]]) for p in paths_concat + "/".join([p.name.split("-")[3], *p.name.split("-")[1:3]]) for p in paths_concat } return runs | runs_concat -def build_file_dbs(gen_tier_path, outdir, ignore_keys_file=None): +def build_file_dbs(setup, gen_tier_path, outdir, ignore_keys_file=None, threads=1): tic = time.time() - gen_tier_path = Path(as_ro_path(gen_tier_path)) + gen_tier_path = Path(as_ro(setup, gen_tier_path)) outdir = Path(outdir) # find generated directories @@ -124,15 +105,11 @@ def build_file_dbs(gen_tier_path, outdir, ignore_keys_file=None): # TODO: replace l200 with {experiment} outfile = outdir / f"l200-{speck[1]}-{speck[2]}-{speck[0]}-filedb.h5" logfile = ( - Path(pat.tmp_log_path(snakemake.params.setup)) - / "filedb" - / outfile.with_suffix(".log").name + Path(pat.tmp_log_path(setup)) / "filedb" / outfile.with_suffix(".log").name ) print(f"INFO: ......building {outfile}") - pre_cmdline, cmdenv = execenv_pyexe( - snakemake.params.setup, "build-filedb", as_string=False - ) + pre_cmdline, cmdenv = execenv_pyexe(setup, "build-filedb", as_string=False) cmdline = [ *pre_cmdline, @@ -154,7 +131,7 @@ def build_file_dbs(gen_tier_path, outdir, ignore_keys_file=None): # TODO: forward stdout to log file processes.add(subprocess.Popen(cmdline, env=cmdenv)) - if len(processes) >= snakemake.threads: + if len(processes) >= threads: os.wait() processes.difference_update([p for p in processes if p.poll() is not None]) @@ -172,25 +149,20 @@ def build_file_dbs(gen_tier_path, outdir, ignore_keys_file=None): print(f"INFO: ...took {dt}") -def fformat(tier): - abs_path = patterns.get_pattern_tier( - snakemake.params.setup, tier, check_in_cycle=False - ) - return str(abs_path).replace(pat.get_tier_path(snakemake.params.setup, tier), "") +def fformat(setup, tier): + abs_path = patterns.get_pattern_tier(setup, tier, check_in_cycle=False) + return str(abs_path).replace(pat.get_tier_path(setup, tier), "") -if snakemake.params.setup.get("build_file_dbs", True): +def build_file_db_config(setup, filedb_path): file_db_config = {} + prodenv_var = os.getenv("PRODENV") - if os.getenv("PRODENV") is not None and os.getenv("PRODENV") in str( - snakemake.params.filedb_path - ): - prodenv = as_ro_path(os.getenv("PRODENV")) + if prodenv_var is not None and prodenv_var in str(filedb_path): + prodenv = as_ro(setup, prodenv_var) def tdirs(tier): - return as_ro_path(pat.get_tier_path(snakemake.params.setup, tier)).replace( - prodenv, "" - ) + return as_ro(setup, pat.get_tier_path(setup, tier)).replace(prodenv, "") file_db_config["data_dir"] = "$PRODENV" @@ -198,47 +170,56 @@ def tdirs(tier): print("WARNING: $PRODENV not set, the FileDB will not be relocatable") def tdirs(tier): - return as_ro_path(pat.get_tier_path(snakemake.params.setup, tier)) + return as_ro(setup, pat.get_tier_path(setup, tier)) file_db_config["data_dir"] = "/" - file_db_config["tier_dirs"] = { - k: tdirs(k) for k in snakemake.params.setup["table_format"] - } + file_db_config["tier_dirs"] = {k: tdirs(k) for k in setup["table_format"]} file_db_config |= { - "file_format": {k: fformat(k) for k in snakemake.params.setup["table_format"]}, - "table_format": snakemake.params.setup["table_format"], + "file_format": {k: fformat(setup, k) for k in setup["table_format"]}, + "table_format": setup["table_format"], } + return file_db_config + + +if snakemake is not None: + print("INFO: dataflow ran successfully, now few final checks and scripts") + setup = snakemake.params.setup + + if (snakemake.wildcards.tier not in ("daq", "daq_compress")) and setup.get( + "build_file_dbs", True + ): + print(f"INFO: ...building FileDBs with {snakemake.threads} threads") + + file_db_config = build_file_db_config(setup, snakemake.params.filedb_path) + + Path(snakemake.params.filedb_path).mkdir(parents=True, exist_ok=True) + with (Path(snakemake.params.filedb_path) / "file_db_config.json").open( + "w" + ) as f: + json.dump(file_db_config, f, indent=2) + + build_file_dbs( + setup, + pat.tier_path(setup), + snakemake.params.filedb_path, + snakemake.params.ignore_keys_file, + threads=snakemake.threads, + ) + (Path(snakemake.params.filedb_path) / "file_db_config.json").unlink() + + if setup.get("check_log_files", True): + print("INFO: ...checking log files") + check_log_files( + pat.tmp_log_path(setup), + snakemake.output.summary_log, + snakemake.output.gen_output, + warning_file=snakemake.output.warning_log, + ) + else: + Path(snakemake.output.summary_log).parent.mkdir(parents=True, exist_ok=True) + Path(snakemake.output.summary_log).touch() + Path(snakemake.output.warning_log).touch() -if ( - snakemake.wildcards.tier not in ("daq", "daq_compress") -) and snakemake.params.setup.get("build_file_dbs", True): - print(f"INFO: ...building FileDBs with {snakemake.threads} threads") - - Path(snakemake.params.filedb_path).mkdir(parents=True, exist_ok=True) - - with (Path(snakemake.params.filedb_path) / "file_db_config.json").open("w") as f: - json.dump(file_db_config, f, indent=2) - - build_file_dbs( - pat.tier_path(snakemake.params.setup), - snakemake.params.filedb_path, - snakemake.params.ignore_keys_file, - ) - (Path(snakemake.params.filedb_path) / "file_db_config.json").unlink() - -if snakemake.params.setup.get("check_log_files", True): - print("INFO: ...checking log files") - check_log_files( - pat.tmp_log_path(snakemake.params.setup), - snakemake.output.summary_log, - snakemake.output.gen_output, - warning_file=snakemake.output.warning_log, - ) -else: - Path(snakemake.output.summary_log).parent.mkdir(parents=True, exist_ok=True) - Path(snakemake.output.summary_log).touch() - Path(snakemake.output.warning_log).touch() - -Path(snakemake.output.gen_output).touch() + Path(snakemake.output.gen_output).touch() diff --git a/workflow/src/legenddataflow/scripts/flow/merge_channels.py b/workflow/src/legenddataflow/scripts/flow/merge_channels.py index 4eee07a2..754494b8 100644 --- a/workflow/src/legenddataflow/scripts/flow/merge_channels.py +++ b/workflow/src/legenddataflow/scripts/flow/merge_channels.py @@ -1,3 +1,7 @@ +"""Console script ``merge-channels``: merge per-channel parameter or plot +outputs for one timestamp into a single per-tier file (yaml/json, +pickle/shelve or lh5).""" + from __future__ import annotations import argparse @@ -9,6 +13,7 @@ import lh5 from dbetto import AttrsDict from dbetto.catalog import Props +from legenddataflowscripts.utils import check_input_files from legenddataflow.methods import ChannelProcKey @@ -40,6 +45,15 @@ def merge_channels() -> None: channel_files = args.input.infiles if hasattr(args.input, "infiles") else args.input + if len(channel_files) == 0: + msg = "no input files provided, cannot merge anything" + raise ValueError(msg) + check_input_files(channel_files, "--input") + + if args.out_db and not args.in_db: + msg = "--out-db requires --in-db: there is no input database to update" + raise ValueError(msg) + file_extension = Path(args.output).suffix if file_extension in (".dat", ".dir"): @@ -58,7 +72,10 @@ def merge_channels() -> None: fkey = ChannelProcKey.get_filekey_from_pattern(Path(channel).name) out_dict[fkey.channel] = AttrsDict(channel_dict).to_dict() else: - msg = "Output file extension does not match input file extension" + msg = ( + f"input file {channel} does not match the extension " + f"{file_extension!r} of the output file {args.output}" + ) raise RuntimeError(msg) Props.write_to(out_file, out_dict) @@ -74,8 +91,6 @@ def merge_channels() -> None: with Path(out_file).open("wb") as w: pkl.dump(out_dict, w, protocol=pkl.HIGHEST_PROTOCOL) - Path(out_file).rename(out_file) - elif file_extension in (".dat", ".dir"): common_dict = {} # Open the 'dumb' database directly. @@ -112,11 +127,27 @@ def merge_channels() -> None: wo_mode="a", ) if args.in_db: + if fkey.channel not in db_dict: + msg = ( + f"channel {fkey.channel} (from {channel}) not found in " + f"the input database {args.in_db}" + ) + raise KeyError(msg) db_dict[fkey.channel] = replace_path( db_dict[fkey.channel], channel, args.output ) else: - msg = "Output file extension does not match input file extension" + msg = ( + f"input file {channel} does not match the extension " + f"{file_extension!r} of the output file {args.output}" + ) raise RuntimeError(msg) if args.out_db: Props.write_to(args.out_db, AttrsDict(db_dict).to_dict()) + + else: + msg = ( + f"unsupported output file extension {file_extension!r} for {args.output}: " + "expected one of .json, .yaml, .yml, .pkl, .dat, .dir or .lh5" + ) + raise ValueError(msg) diff --git a/workflow/src/legenddataflow/scripts/flow/merge_in_channel.py b/workflow/src/legenddataflow/scripts/flow/merge_in_channel.py index fdb2ab97..672665ae 100644 --- a/workflow/src/legenddataflow/scripts/flow/merge_in_channel.py +++ b/workflow/src/legenddataflow/scripts/flow/merge_in_channel.py @@ -1,3 +1,6 @@ +"""Console script ``merge-in-channels``: merge per-channel parameter outputs +into a single channel-keyed output file.""" + from __future__ import annotations import argparse @@ -7,6 +10,7 @@ import lh5 from dbetto import AttrsDict from dbetto.catalog import Props +from legenddataflowscripts.utils import check_input_files from lgdo.types import Struct from legenddataflow.methods import ChannelProcKey @@ -39,6 +43,15 @@ def merge_in_channel() -> None: input_files = args.input.infiles if hasattr(args.input, "infiles") else args.input + if len(input_files) == 0: + msg = "no input files provided, cannot merge anything" + raise ValueError(msg) + check_input_files(input_files, "--input") + + if args.out_db and not args.in_db: + msg = "--out-db requires --in-db: there is no input database to update" + raise ValueError(msg) + file_extension = Path(args.output).suffix out_file = args.output @@ -61,23 +74,49 @@ def merge_in_channel() -> None: if args.in_db: db_dict = Props.read_from(args.in_db) objects = {} + channel = None for infile in input_files: fkey = ChannelProcKey.get_filekey_from_pattern(Path(infile).name) - tb_in = lh5.read(f"{fkey.channel}", infile) - - objects[f"{fkey.processing_step}".split("_", maxsplit=3)[2]] = tb_in[ - f"{fkey.processing_step}".split("_", maxsplit=3)[2] - ] + if fkey.identifier is None: + msg = ( + f"cannot derive the output group for {infile}: processing " + f"step {fkey.processing_step!r} has no identifier part " + "(expected {type}_{tier}_{identifier})" + ) + raise ValueError(msg) + if channel is not None and fkey.channel != channel: + msg = ( + f"input files span multiple channels ({channel} and " + f"{fkey.channel}): merge-in-channels merges processing " + "steps of a single channel" + ) + raise ValueError(msg) + channel = fkey.channel + + tb_in = lh5.read(fkey.channel, infile) + if fkey.identifier not in tb_in: + msg = ( + f"table {fkey.channel!r} in {infile} has no field " + f"{fkey.identifier!r}: available fields are {list(tb_in)}" + ) + raise KeyError(msg) + objects[fkey.identifier] = tb_in[fkey.identifier] lh5.write( Struct(objects), - name=fkey.channel, + name=channel, lh5_file=out_file, wo_mode="a", ) if args.out_db: - if args.in_db: + for infile in input_files: db_dict = replace_path(db_dict, infile, args.output) - Props.write_to(args.out_db, AttrsDict(db_dict).to_dict()) + + else: + msg = ( + f"unsupported output file extension {file_extension!r} for {args.output}: " + "expected one of .json, .yaml, .yml, .pkl or .lh5" + ) + raise ValueError(msg) diff --git a/workflow/src/legenddataflow/scripts/flow/utils.py b/workflow/src/legenddataflow/scripts/flow/utils.py index 7169a6fd..c80e4ae1 100644 --- a/workflow/src/legenddataflow/scripts/flow/utils.py +++ b/workflow/src/legenddataflow/scripts/flow/utils.py @@ -1,3 +1,5 @@ +"""Shared helpers for the flow scripts.""" + from __future__ import annotations from pathlib import Path diff --git a/workflow/src/legenddataflow/scripts/flow/write_filelist.py b/workflow/src/legenddataflow/scripts/flow/write_filelist.py index b363b78d..47abca9e 100644 --- a/workflow/src/legenddataflow/scripts/flow/write_filelist.py +++ b/workflow/src/legenddataflow/scripts/flow/write_filelist.py @@ -1,16 +1,30 @@ # ruff: noqa: T201, I002 +"""Snakemake ``script:`` for the ``gen_filelist`` checkpoint: write the input +files matched by a ``.gen`` target into a filelist.""" + from pathlib import Path -from snakemake.script import snakemake # snakemake > 8.16 +try: + from snakemake.script import snakemake # snakemake > 8.16 +except ImportError: # not running inside a snakemake job (e.g. unit tests) + snakemake = None + + +def write_filelist(files, output_file, label=""): + print(f"INFO: found {len(files)} files") + if len(files) == 0: + print( + f"WARNING: no files found for the given pattern: {label}. " + "make sure patterns follows the format: " + "all-{experiment}-{period}-{run}-{datatype}-{timestamp}-{tier}.gen" + ) + with Path(output_file).open("w") as f: + for fn in files: + f.write(f"{fn}\n") + -print(f"INFO: found {len(snakemake.input)} files") -if len(snakemake.input) == 0: - print( - f"WARNING: no files found for the given pattern: {snakemake.wildcards.label}. " - "make sure patterns follows the format: " - "all-{experiment}-{period}-{run}-{datatype}-{timestamp}-{tier}.gen" +if snakemake is not None: + write_filelist( + list(snakemake.input), snakemake.output[0], label=snakemake.wildcards.label ) -with Path(snakemake.output[0]).open("w") as f: - for fn in snakemake.input: - f.write(f"{fn}\n") diff --git a/workflow/src/legenddataflow/scripts/par/geds/pht/aoe.py b/workflow/src/legenddataflow/scripts/par/geds/pht/aoe.py index c595f1f4..621aac4e 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/pht/aoe.py +++ b/workflow/src/legenddataflow/scripts/par/geds/pht/aoe.py @@ -1,3 +1,6 @@ +"""Console script ``par-geds-pht-aoe``: derive the partition-level A/E +calibration for a HPGe channel.""" + from __future__ import annotations import argparse @@ -5,9 +8,18 @@ import numpy as np import pandas as pd -from dbetto import Props, TextDB +from dbetto import Props from legenddataflowscripts.par.geds.hit.aoe import run_aoe_calibration -from legenddataflowscripts.utils import build_log, get_pulser_mask +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + check_pulser_mask, + get_channel_config, + get_pulser_mask, + get_rule_config, + prepare_output_paths, + require_config_keys, +) from pygama.pargen.AoE_cal import * # noqa: F403 from pygama.pargen.utils import load_data @@ -56,12 +68,22 @@ def par_geds_pht_aoe() -> None: argparser.add_argument("-d", "--debug", help="debug_mode", action="store_true") args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["pars_pht_aoecal"] + config_dict = get_rule_config( + args.configs, "pars_pht_aoecal", args.timestamp, args.datatype + ) - log = build_log(config_dict, args.log) + build_log(config_dict, args.log) - channel_dict = config_dict["inputs"]["par_pht_aoecal_config"][args.channel] + check_input_files(args.pulser_files, "--pulser-files") + prepare_output_paths( + *(args.hit_pars or []), *(args.aoe_results or []), *(args.plot_file or []) + ) + + channel_dict = get_channel_config( + config_dict["inputs"]["par_pht_aoecal_config"], + args.channel, + name="par_pht_aoecal_config", + ) kwarg_dict = Props.read_from(channel_dict) # par files @@ -78,10 +100,15 @@ def par_geds_pht_aoe() -> None: inplots_dict = get_run_dict(args.inplots) # get files split by run - final_dict, _ = split_files_by_run(args.files) + final_dict, _ = split_files_by_run(args.input_files) # run aoe cal if kwarg_dict.get("run_aoe", True) is True: + require_config_keys( + kwarg_dict, + ["final_cut_field", "current_param", "energy_param", "cal_energy_param"], + f"channel {args.channel} aoecal config ({channel_dict})", + ) params = [ kwarg_dict["final_cut_field"], kwarg_dict["current_param"], @@ -107,6 +134,7 @@ def par_geds_pht_aoe() -> None: ) mask = get_pulser_mask(pulser_file=args.pulser_files) + check_pulser_mask(mask, threshold_mask, args.channel) data["is_pulser"] = mask[threshold_mask] for tstamp in cal_dicts: @@ -124,8 +152,8 @@ def par_geds_pht_aoe() -> None: cal_dicts, results_dicts, object_dicts, + inplots_dict, config=channel_dict, - log=log, debug_mode=args.debug, # gen_plots=bool(args.plot_file), ) diff --git a/workflow/src/legenddataflow/scripts/par/geds/pht/ecal_part.py b/workflow/src/legenddataflow/scripts/par/geds/pht/ecal_part.py index 6316d755..266c6abb 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/pht/ecal_part.py +++ b/workflow/src/legenddataflow/scripts/par/geds/pht/ecal_part.py @@ -1,3 +1,6 @@ +"""Console script ``par-geds-pht-ecal-part``: derive the partition-level +energy calibration for a HPGe channel.""" + from __future__ import annotations import argparse @@ -9,7 +12,16 @@ import pygama.math.distributions as pgf import pygama.math.histogram as pgh from dbetto import Props, TextDB -from legenddataflowscripts.utils import build_log, get_pulser_mask +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + check_pulser_mask, + get_channel_config, + get_pulser_mask, + get_rule_config, + prepare_output_paths, + require_config_keys, +) from pygama.math.distributions import nb_poly from pygama.pargen.energy_cal import FWHMLinear, FWHMQuadratic, HPGeCalibration from pygama.pargen.utils import load_data @@ -253,7 +265,11 @@ def calibrate_partition( bin_width_kev=0.2 if nruns > 3 else 0.5, ) else: - err = f"2614.511 peak not found in {cal_energy_param} fit, reduced csqr {csqr[0] / csqr[1]} not below 10, check fit" + err = ( + f"2614.511 keV peak not found in the {cal_energy_param} fit and " + f"its reduced chi-square {csqr[0] / csqr[1]:.1f} is not below " + "100, check the fit" + ) raise ValueError(err) full_object_dict[cal_energy_param].get_energy_res_curve( @@ -382,7 +398,7 @@ def calibrate_partition( out_plot_dicts = {} for tstamp, plot_dict in plot_dicts.items(): if "common" in list(plot_dict) and common_dict is not None: - plot_dict["common"].update(partcal_plot_dict["common"]) + plot_dict["common"].update(common_dict) elif common_dict is not None: plot_dict["common"] = common_dict plot_dict.update({"partition_ecal": partcal_plot_dict}) @@ -427,13 +443,19 @@ def par_geds_pht_ecal_part() -> None: argparser.add_argument("-d", "--debug", help="debug_mode", action="store_true") args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["pars_pht_partcal"] + config_dict = get_rule_config( + args.configs, "pars_pht_partcal", args.timestamp, args.datatype + ) build_log(config_dict, args.log) + check_input_files(args.pulser_files, "--pulser-files") + prepare_output_paths( + *(args.hit_pars or []), *(args.fit_results or []), *(args.plot_file or []) + ) + chmap = TextDB(path=args.metadata, lazy=True).on( - args.timestamp, system=args.datatype + args.timestamp, category=args.datatype ) # par files @@ -450,10 +472,19 @@ def par_geds_pht_ecal_part() -> None: inplots_dict = get_run_dict(args.inplots) # get files split by run - final_dict, _ = split_files_by_run(args.files) + final_dict, _ = split_files_by_run(args.input_files) - channel_dict = config_dict["inputs"]["pars_pht_partcal_config"][args.channel] + channel_dict = get_channel_config( + config_dict["inputs"]["pars_pht_partcal_config"], + args.channel, + name="pars_pht_partcal_config", + ) kwarg_dict = Props.read_from(channel_dict) + require_config_keys( + kwarg_dict, + ["final_cut_field", "energy_params", "threshold"], + f"channel {args.channel} partcal config ({channel_dict})", + ) params = [ kwarg_dict["final_cut_field"], @@ -476,6 +507,7 @@ def par_geds_pht_ecal_part() -> None: if "pulser_multiplicity_threshold" in kwarg_dict: kwarg_dict.pop("pulser_multiplicity_threshold") + check_pulser_mask(mask, threshold_mask, args.channel) data["is_pulser"] = mask[threshold_mask] for tstamp in cal_dicts: diff --git a/workflow/src/legenddataflow/scripts/par/geds/pht/fast.py b/workflow/src/legenddataflow/scripts/par/geds/pht/fast.py index b547f101..48583aef 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/pht/fast.py +++ b/workflow/src/legenddataflow/scripts/par/geds/pht/fast.py @@ -1,3 +1,6 @@ +"""Console script ``par-geds-pht-fast``: run the partition-level calibrations +(energy, A/E, LQ) for a HPGe channel in a single pass.""" + from __future__ import annotations import argparse @@ -8,11 +11,19 @@ import numpy as np import pandas as pd -from dbetto import TextDB from dbetto.catalog import Props from legenddataflowscripts.par.geds.hit.aoe import run_aoe_calibration from legenddataflowscripts.par.geds.hit.lq import run_lq_calibration -from legenddataflowscripts.utils import build_log, get_pulser_mask +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + check_pulser_mask, + get_channel_config, + get_pulser_mask, + get_rule_config, + prepare_output_paths, + require_config_keys, +) from legendmeta import LegendMetadata from pygama.pargen.utils import load_data @@ -38,9 +49,6 @@ def par_geds_pht_fast() -> None: argparser.add_argument( "--pulser-files", help="pulser_file", nargs="*", type=str, required=False ) - argparser.add_argument( - "--tcm-filelist", help="tcm_filelist", type=str, nargs="*", required=False - ) argparser.add_argument( "--ecal-file", help="ecal_file", type=str, nargs="*", required=True ) @@ -69,10 +77,22 @@ def par_geds_pht_fast() -> None: argparser.add_argument("-d", "--debug", help="debug_mode", action="store_true") args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"] + partcal_dict = get_rule_config( + args.configs, "pars_pht_partcal", args.timestamp, args.datatype + ) + aoecal_dict = get_rule_config( + args.configs, "pars_pht_aoecal", args.timestamp, args.datatype + ) + lqcal_dict = get_rule_config( + args.configs, "pars_pht_lqcal", args.timestamp, args.datatype + ) - log = build_log(config_dict["pars_pht_partcal"], args.log) + log = build_log(partcal_dict, args.log) + + check_input_files(args.pulser_files, "--pulser-files") + prepare_output_paths( + *(args.hit_pars or []), *(args.fit_results or []), *(args.plot_file or []) + ) chmap = LegendMetadata(args.metadata).channelmap( args.timestamp, system=args.datatype @@ -111,6 +131,7 @@ def par_geds_pht_fast() -> None: files = sorted( np.unique(files) ) # need this as sometimes files get double counted as it somehow puts in the p%-* filelist and individual runs also + check_input_files(files, "--input-files") final_dict = {} all_file = run_splitter(sorted(files)) @@ -119,16 +140,39 @@ def par_geds_pht_fast() -> None: timestamp = fk.timestamp final_dict[timestamp] = sorted(filelist) - kwarg_dict = Props.read_from( - config_dict["pars_pht_partcal"]["inputs"]["pars_pht_partcal_config"][ - args.channel - ] + partcal_config = get_channel_config( + partcal_dict["inputs"]["pars_pht_partcal_config"], + args.channel, + name="pars_pht_partcal_config", ) - aoe_kwarg_dict = Props.read_from( - config_dict["pars_pht_aoecal"]["inputs"]["par_pht_aoecal_config"][args.channel] + aoecal_config = get_channel_config( + aoecal_dict["inputs"]["par_pht_aoecal_config"], + args.channel, + name="par_pht_aoecal_config", + ) + lqcal_config = get_channel_config( + lqcal_dict["inputs"]["lqcal_config"], + args.channel, + name="lqcal_config", ) - lq_kwarg_dict = Props.read_from( - config_dict["pars_pht_lqcal"]["inputs"]["lqcal_config"][args.channel] + kwarg_dict = Props.read_from(partcal_config) + aoe_kwarg_dict = Props.read_from(aoecal_config) + lq_kwarg_dict = Props.read_from(lqcal_config) + + require_config_keys( + kwarg_dict, + ["final_cut_field", "energy_params", "threshold"], + f"channel {args.channel} partcal config ({partcal_config})", + ) + require_config_keys( + aoe_kwarg_dict, + ["run_aoe"], + f"channel {args.channel} aoecal config ({aoecal_config})", + ) + require_config_keys( + lq_kwarg_dict, + ["run_lq"], + f"channel {args.channel} lqcal config ({lqcal_config})", ) params = [ @@ -138,6 +182,11 @@ def par_geds_pht_fast() -> None: params += kwarg_dict["energy_params"] if aoe_kwarg_dict["run_aoe"] is True: + require_config_keys( + aoe_kwarg_dict, + ["final_cut_field", "current_param", "energy_param", "cal_energy_param"], + f"channel {args.channel} aoecal config ({aoecal_config})", + ) aoe_params = [ aoe_kwarg_dict["final_cut_field"], aoe_kwarg_dict["current_param"], @@ -155,6 +204,11 @@ def par_geds_pht_fast() -> None: params += aoe_params if lq_kwarg_dict["run_lq"] is True: + require_config_keys( + lq_kwarg_dict, + ["energy_param", "cal_energy_param", "cut_field"], + f"channel {args.channel} lqcal config ({lqcal_config})", + ) params += [ "lq80", "dt_eff", @@ -181,6 +235,7 @@ def par_geds_pht_fast() -> None: if "pulser_multiplicity_threshold" in kwarg_dict: kwarg_dict.pop("pulser_multiplicity_threshold") + check_pulser_mask(mask, threshold_mask, args.channel) data["is_pulser"] = mask[threshold_mask] msg = f"{len(data.query('~is_pulser'))} non pulser events" @@ -195,10 +250,6 @@ def par_geds_pht_fast() -> None: row = pd.DataFrame(row) data = pd.concat([data, row]) - configs = TextDB(path=args.configs, lazy=True).on( - args.timestamp, system=args.datatype - )["snakemake_rules"] - start = time.time() cal_dicts, results_dicts, object_dicts, plot_dicts = calibrate_partition( @@ -209,9 +260,7 @@ def par_geds_pht_fast() -> None: inplots_dict, args.channel, chmap, - configs=configs["pars_pht_partcal"]["inputs"]["pars_pht_partcal_config"][ - args.channel - ], + configs=partcal_config, gen_plots=bool(args.plot_file), debug_mode=args.debug, ) @@ -225,9 +274,7 @@ def par_geds_pht_fast() -> None: results_dicts, object_dicts, plot_dicts, - config=configs["pars_pht_aoecal"]["inputs"]["par_pht_aoecal_config"][ - args.channel - ], + config=aoecal_config, debug_mode=args.debug, # gen_plots=bool(args.plot_file), ) @@ -242,7 +289,7 @@ def par_geds_pht_fast() -> None: results_dicts, object_dicts, plot_dicts, - configs=configs["pars_pht_lqcal"]["inputs"]["lqcal_config"][args.channel], + configs=lqcal_config, debug_mode=args.debug, # gen_plots=bool(args.plot_file), ) @@ -252,7 +299,8 @@ def par_geds_pht_fast() -> None: msg = f"Total calibration took {time.time() - start:.2f} seconds" log.info(msg) - save_dict_to_files(args.plot_file, plot_dicts) + if args.plot_file: + save_dict_to_files(args.plot_file, plot_dicts) save_dict_to_files( sorted(args.hit_pars), diff --git a/workflow/src/legenddataflow/scripts/par/geds/pht/lq.py b/workflow/src/legenddataflow/scripts/par/geds/pht/lq.py index 5caf1d90..91cd2a81 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/pht/lq.py +++ b/workflow/src/legenddataflow/scripts/par/geds/pht/lq.py @@ -1,144 +1,176 @@ -from __future__ import annotations - -import argparse -import warnings - -import numpy as np -import pandas as pd -from dbetto import Props, TextDB -from legenddataflowscripts.par.geds.hit.lq import run_lq_calibration -from legenddataflowscripts.utils import build_log, get_pulser_mask -from pygama.pargen.AoE_cal import * # noqa: F403 -from pygama.pargen.lq_cal import * # noqa: F403 -from pygama.pargen.utils import load_data - -from .util import ( - get_run_dict, - save_dict_to_files, - split_files_by_run, -) - -warnings.filterwarnings(action="ignore", category=RuntimeWarning) - - -def par_geds_pht_lq() -> None: - argparser = argparse.ArgumentParser() - argparser.add_argument( - "--input-files", help="files", type=str, nargs="*", required=True - ) - argparser.add_argument( - "--pulser-files", help="pulser_file", type=str, nargs="*", required=False - ) - argparser.add_argument( - "--ecal-file", help="ecal_file", type=str, nargs="*", required=True - ) - argparser.add_argument( - "--eres-file", help="eres_file", type=str, nargs="*", required=True - ) - argparser.add_argument( - "--inplots", help="eres_file", type=str, nargs="*", required=True - ) - - argparser.add_argument("--configs", help="configs", type=str, required=True) - argparser.add_argument("--metadata", help="metadata path", type=str, required=True) - argparser.add_argument("--log", help="log_file", type=str) - - argparser.add_argument("--timestamp", help="Datatype", type=str, required=True) - argparser.add_argument("--datatype", help="Datatype", type=str, required=True) - argparser.add_argument("--channel", help="Channel", type=str, required=True) - argparser.add_argument("--table-name", help="table name", type=str, required=True) - - argparser.add_argument( - "--plot-file", help="plot_file", type=str, nargs="*", required=False - ) - argparser.add_argument("--hit-pars", help="hit_pars", nargs="*", type=str) - argparser.add_argument("--lq-results", help="lq_results", nargs="*", type=str) - - argparser.add_argument("-d", "--debug", help="debug_mode", action="store_true") - args = argparser.parse_args() - - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["pars_pht_lqcal"] - - log = build_log(config_dict, args.log) - - channel_dict = config_dict["inputs"]["lqcal_config"][args.channel] - kwarg_dict = Props.read_from(channel_dict) - - # par files - par_dict = get_run_dict(args.ecal_file) - cal_dicts = {tstamp: val["pars"] for tstamp, val in par_dict.items()} - results_dicts = {tstamp: val["results"] for tstamp, val in par_dict.items()} - - # obj files - object_dicts = get_run_dict(args.eres_file) - - # plot files - inplots_dict = {} - if args.inplots: - inplots_dict = get_run_dict(args.inplots) - - # get files split by run - final_dict, _ = split_files_by_run(args.files) - - # run lq cal - if kwarg_dict.pop("run_lq") is True: - params = [ - "lq80", - "dt_eff", - kwarg_dict["energy_param"], - kwarg_dict["cal_energy_param"], - kwarg_dict["cut_field"], - ] - - # load data in - data, threshold_mask = load_data( - final_dict, - args.table_name, - cal_dicts, - params=params, - threshold=kwarg_dict.pop("threshold"), - return_selection_mask=True, - ) - msg = f"Loaded {len(data)} events from {len(final_dict)} files" - log.info(msg) - - mask = get_pulser_mask(pulser_file=args.pulser_files) - if "pulser_multiplicity_threshold" in kwarg_dict: - kwarg_dict.pop("pulser_multiplicity_threshold") - - data["is_pulser"] = mask[threshold_mask] - msg = f"{len(data.query('~is_pulser'))} non pulser events" - log.info(msg) - - for tstamp in cal_dicts: - if tstamp not in np.unique(data["run_timestamp"]): - row = { - key: [False] if data.dtypes[key] == "bool" else [np.nan] - for key in data - } - row["run_timestamp"] = tstamp - row = pd.DataFrame(row) - data = pd.concat([data, row]) - - cal_dicts, results_dicts, object_dicts, plot_dicts = run_lq_calibration( - data, - cal_dicts, - results_dicts, - object_dicts, - inplots_dict, - kwarg_dict, - # gen_plots=bool(args.plot_file), - ) - - if args.plot_file: - save_dict_to_files(args.plot_file, plot_dicts) - - save_dict_to_files( - args.hit_pars, - { - tstamp: {"pars": cal_dicts[tstamp], "results": results_dicts[tstamp]} - for tstamp in cal_dicts - }, - ) - save_dict_to_files(args.lq_results, object_dicts) +"""Console script ``par-geds-pht-lq``: derive the partition-level LQ +calibration for a HPGe channel.""" + +from __future__ import annotations + +import argparse +import warnings + +import numpy as np +import pandas as pd +from dbetto import Props +from legenddataflowscripts.par.geds.hit.lq import run_lq_calibration +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + check_pulser_mask, + get_channel_config, + get_pulser_mask, + get_rule_config, + prepare_output_paths, + require_config_keys, +) +from pygama.pargen.AoE_cal import * # noqa: F403 +from pygama.pargen.lq_cal import * # noqa: F403 +from pygama.pargen.utils import load_data + +from .util import ( + get_run_dict, + save_dict_to_files, + split_files_by_run, +) + +warnings.filterwarnings(action="ignore", category=RuntimeWarning) + + +def par_geds_pht_lq() -> None: + argparser = argparse.ArgumentParser() + argparser.add_argument( + "--input-files", help="files", type=str, nargs="*", required=True + ) + argparser.add_argument( + "--pulser-files", help="pulser_file", type=str, nargs="*", required=False + ) + argparser.add_argument( + "--ecal-file", help="ecal_file", type=str, nargs="*", required=True + ) + argparser.add_argument( + "--eres-file", help="eres_file", type=str, nargs="*", required=True + ) + argparser.add_argument( + "--inplots", help="eres_file", type=str, nargs="*", required=True + ) + + argparser.add_argument("--configs", help="configs", type=str, required=True) + argparser.add_argument("--metadata", help="metadata path", type=str, required=True) + argparser.add_argument("--log", help="log_file", type=str) + + argparser.add_argument("--timestamp", help="Datatype", type=str, required=True) + argparser.add_argument("--datatype", help="Datatype", type=str, required=True) + argparser.add_argument("--channel", help="Channel", type=str, required=True) + argparser.add_argument("--table-name", help="table name", type=str, required=True) + + argparser.add_argument( + "--plot-file", help="plot_file", type=str, nargs="*", required=False + ) + argparser.add_argument("--hit-pars", help="hit_pars", nargs="*", type=str) + argparser.add_argument("--lq-results", help="lq_results", nargs="*", type=str) + + argparser.add_argument("-d", "--debug", help="debug_mode", action="store_true") + args = argparser.parse_args() + + config_dict = get_rule_config( + args.configs, "pars_pht_lqcal", args.timestamp, args.datatype + ) + + log = build_log(config_dict, args.log) + + check_input_files(args.pulser_files, "--pulser-files") + prepare_output_paths( + *(args.hit_pars or []), *(args.lq_results or []), *(args.plot_file or []) + ) + + channel_dict = get_channel_config( + config_dict["inputs"]["lqcal_config"], args.channel, name="lqcal_config" + ) + kwarg_dict = Props.read_from(channel_dict) + + # par files + par_dict = get_run_dict(args.ecal_file) + cal_dicts = {tstamp: val["pars"] for tstamp, val in par_dict.items()} + results_dicts = {tstamp: val["results"] for tstamp, val in par_dict.items()} + + # obj files + object_dicts = get_run_dict(args.eres_file) + + # plot files + inplots_dict = {} + if args.inplots: + inplots_dict = get_run_dict(args.inplots) + + # get files split by run + final_dict, _ = split_files_by_run(args.input_files) + + if "run_lq" not in kwarg_dict: + msg = f"lqcal config for {args.channel} is missing the required key 'run_lq'" + raise KeyError(msg) + + # run lq cal + if kwarg_dict.pop("run_lq") is True: + require_config_keys( + kwarg_dict, + ["energy_param", "cal_energy_param", "cut_field", "threshold"], + f"channel {args.channel} lqcal config ({channel_dict})", + ) + params = [ + "lq80", + "dt_eff", + kwarg_dict["energy_param"], + kwarg_dict["cal_energy_param"], + kwarg_dict["cut_field"], + ] + + # load data in + data, threshold_mask = load_data( + final_dict, + args.table_name, + cal_dicts, + params=params, + threshold=kwarg_dict.pop("threshold"), + return_selection_mask=True, + ) + msg = f"Loaded {len(data)} events from {len(final_dict)} files" + log.info(msg) + + mask = get_pulser_mask(pulser_file=args.pulser_files) + if "pulser_multiplicity_threshold" in kwarg_dict: + kwarg_dict.pop("pulser_multiplicity_threshold") + + check_pulser_mask(mask, threshold_mask, args.channel) + data["is_pulser"] = mask[threshold_mask] + msg = f"{len(data.query('~is_pulser'))} non pulser events" + log.info(msg) + + for tstamp in cal_dicts: + if tstamp not in np.unique(data["run_timestamp"]): + row = { + key: [False] if data.dtypes[key] == "bool" else [np.nan] + for key in data + } + row["run_timestamp"] = tstamp + row = pd.DataFrame(row) + data = pd.concat([data, row]) + + cal_dicts, results_dicts, object_dicts, plot_dicts = run_lq_calibration( + data, + cal_dicts, + results_dicts, + object_dicts, + inplots_dict, + kwarg_dict, + # gen_plots=bool(args.plot_file), + ) + else: + plot_dicts = inplots_dict + + if args.plot_file: + save_dict_to_files(args.plot_file, plot_dicts) + + save_dict_to_files( + args.hit_pars, + { + tstamp: {"pars": cal_dicts[tstamp], "results": results_dicts[tstamp]} + for tstamp in cal_dicts + }, + ) + save_dict_to_files(args.lq_results, object_dicts) diff --git a/workflow/src/legenddataflow/scripts/par/geds/pht/qc.py b/workflow/src/legenddataflow/scripts/par/geds/pht/qc.py index da006127..da87a106 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/pht/qc.py +++ b/workflow/src/legenddataflow/scripts/par/geds/pht/qc.py @@ -1,3 +1,6 @@ +"""Console script ``par-geds-pht-qc``: derive partition-level quality-control +cuts for a HPGe channel from calibration data.""" + from __future__ import annotations import argparse @@ -6,10 +9,14 @@ from pathlib import Path import numpy as np -from dbetto import Props, TextDB +from dbetto import Props from legenddataflowscripts.par.geds.hit.qc import build_qc from legenddataflowscripts.utils import ( build_log, + check_input_files, + get_channel_config, + get_rule_config, + prepare_output_paths, ) warnings.filterwarnings(action="ignore", category=RuntimeWarning) @@ -47,13 +54,19 @@ def par_geds_pht_qc() -> None: ) args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["pars_pht_qc"] + config_dict = get_rule_config( + args.configs, "pars_pht_qc", args.timestamp, args.datatype + ) build_log(config_dict, args.log) + check_input_files(args.pulser_files, "--pulser-files") + prepare_output_paths(*(args.save_path or []), *(args.plot_path or [])) + # get metadata dictionary - channel_dict = config_dict["inputs"]["qc_config"][args.channel] + channel_dict = get_channel_config( + config_dict["inputs"]["qc_config"], args.channel, name="qc_config" + ) kwarg_dict = Props.read_from(channel_dict) if args.overwrite_files: @@ -78,6 +91,7 @@ def par_geds_pht_qc() -> None: cal_files = sorted( np.unique(cal_files) ) # need this as sometimes files get double counted as it somehow puts in the p%-* filelist and individual runs also + check_input_files(cal_files, "--cal-files") if isinstance(args.fft_files, list): fft_files = [] @@ -91,6 +105,9 @@ def par_geds_pht_qc() -> None: fft_files = sorted( np.unique(fft_files) ) # need this as sometimes files get double counted as it somehow puts in the p%-* filelist and individual runs also + # an empty fft filelist is legitimate; only check the members it lists + if fft_files: + check_input_files(fft_files, "--fft-files") hit_dict, plot_dict = build_qc( config=kwarg_dict, diff --git a/workflow/src/legenddataflow/scripts/par/geds/pht/qc_phy.py b/workflow/src/legenddataflow/scripts/par/geds/pht/qc_phy.py index 6e810399..809c606e 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/pht/qc_phy.py +++ b/workflow/src/legenddataflow/scripts/par/geds/pht/qc_phy.py @@ -1,3 +1,6 @@ +"""Console script ``par-geds-pht-qc-phy``: derive partition-level +quality-control cuts for a HPGe channel from physics data.""" + from __future__ import annotations import argparse @@ -9,9 +12,15 @@ import lh5 import numpy as np -from dbetto import TextDB from dbetto.catalog import Props -from legenddataflowscripts.utils import build_log, convert_dict_np_to_float +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + convert_dict_np_to_float, + get_channel_config, + get_rule_config, + prepare_output_paths, +) from lh5 import ls from pygama.pargen.data_cleaning import ( generate_cut_classifiers, @@ -49,13 +58,18 @@ def par_geds_pht_qc_phy() -> None: ) args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["pars_pht_qc"] + config_dict = get_rule_config( + args.configs, "pars_pht_qc", args.timestamp, args.datatype + ) log = build_log(config_dict, args.log) + prepare_output_paths(*(args.save_path or []), *(args.plot_path or [])) + # get metadata dictionary - channel_dict = config_dict["inputs"]["qc_config"][args.channel] + channel_dict = get_channel_config( + config_dict["inputs"]["qc_config"], args.channel, name="qc_config" + ) kwarg_dict = Props.read_from(channel_dict) # sort files in dictionary where keys are first timestamp from run @@ -68,11 +82,12 @@ def par_geds_pht_qc_phy() -> None: if len(run_files) == 0: continue run_files = sorted(np.unique(run_files)) + check_input_files(run_files, "--phy-files") phy_files += run_files bls = lh5.read( - args.baseline_name, phy_files, field_mask=["wf_max", "bl_mean"] + args.baseline_name, run_files, field_mask=["wf_max", "bl_mean"] ) - puls = lh5.read(args.pulser_name, phy_files, field_mask=["trapTmax"]) + puls = lh5.read(args.pulser_name, run_files, field_mask=["trapTmax"]) bl_idxs = ((bls["wf_max"].nda - bls["bl_mean"].nda) > 1000) & ( puls["trapTmax"].nda < 200 ) @@ -87,6 +102,10 @@ def par_geds_pht_qc_phy() -> None: puls["trapTmax"].nda < 200 ) + if len(phy_files) == 0: + msg = "no files found in the provided phy filelists, cannot derive qc cuts" + raise RuntimeError(msg) + kwarg_dict_fft = kwarg_dict["fft_fields"] search_name = ( @@ -107,7 +126,7 @@ def par_geds_pht_qc_phy() -> None: ) discharges = data["t_sat_lo"] > 0 - discharge_timestamps = np.where(data["timestamp"][discharges])[0] + discharge_timestamps = data["timestamp"][discharges] is_recovering = np.full(len(data), False, dtype=bool) for tstamp in discharge_timestamps: is_recovering = is_recovering | np.where( @@ -145,10 +164,10 @@ def par_geds_pht_qc_phy() -> None: 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: argparser.add_argument("--channel", help="Channel", type=str, required=True) args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) + config_dict = get_rule_config( + args.configs, "pars_psp", args.timestamp, args.datatype + ) merge_config = Props.read_from( - configs["snakemake_rules"]["pars_psp"]["inputs"]["psp_config"][args.channel] + get_channel_config( + config_dict["inputs"]["psp_config"], args.channel, name="psp_config" + ) ) ave_fields = merge_config["average_fields"] + check_input_files(args.input, "--input") + prepare_output_paths( + *(args.output or []), *(args.out_plots or []), *(args.out_obj or []) + ) + # partitions could be different for different channels - do separately for each channel in_dicts = {} for file in args.input: @@ -82,15 +100,25 @@ def par_geds_psp_average() -> None: tmp_dict[key] = {} tmp_dict = tmp_dict[key] if isinstance(vals[0], str): - if "*" in vals[0]: - unit = vals[0].split("*")[1] - rounding = ( - len(val.split("*")[0].split(".")[-1]) if "." in vals[0] else 16 + try: + if "*" in vals[0]: + unit = vals[0].split("*")[1] + rounding = ( + len(vals[0].split("*")[0].split(".")[-1]) + if "." in vals[0] + else 16 + ) + vals = np.array([float(val.split("*")[0]) for val in vals]) + else: + unit = None + rounding = 16 + vals = np.array([float(val) for val in vals]) + except (ValueError, IndexError) as err: + msg = ( + f"cannot average field {field!r}: values {vals} are not " + "numbers or 'number*unit' strings" ) - vals = np.array([float(val.split("*")[0]) for val in vals]) - else: - unit = None - rounding = 16 + raise ValueError(msg) from err else: vals = np.array(vals) unit = None @@ -129,17 +157,23 @@ def par_geds_psp_average() -> None: file, AttrsDict(convert_dict_np_to_float(in_dicts[tstamp])).to_dict() ) + def _find_input_for_timestamp(infiles, tstamp, outfile): + for infile in infiles: + if tstamp in infile: + with Path(infile).open("rb") as f: + return pkl.load(f) + msg = ( + f"no input file matching timestamp {tstamp} (from output {outfile}) " + f"found among {infiles}" + ) + raise ValueError(msg) + if args.out_plots: for file in args.out_plots: tstamp = ChannelProcKey.get_filekey_from_pattern(Path(file).name).timestamp if args.in_plots: - for infile in args.in_plots: - if tstamp in infile: - with Path(infile).open("rb") as f: - old_plot_dict = pkl.load(f) - break - old_plot_dict.update({"psp": plot_dict}) - new_plot_dict = old_plot_dict + new_plot_dict = _find_input_for_timestamp(args.in_plots, tstamp, file) + new_plot_dict.update({"psp": plot_dict}) else: new_plot_dict = {"psp": plot_dict} with Path(file).open("wb") as f: @@ -149,12 +183,7 @@ def par_geds_psp_average() -> None: for file in args.out_obj: tstamp = ChannelProcKey.get_filekey_from_pattern(Path(file).name).timestamp if args.in_obj: - for infile in args.in_obj: - if tstamp in infile: - with Path(infile).open("rb") as f: - old_obj_dict = pkl.load(f) - break - new_obj_dict = old_obj_dict + new_obj_dict = _find_input_for_timestamp(args.in_obj, tstamp, file) else: new_obj_dict = {} with Path(file).open("wb") as f: diff --git a/workflow/src/legenddataflow/scripts/par/geds/psp/dplms.py b/workflow/src/legenddataflow/scripts/par/geds/psp/dplms.py index 0740914d..32e9d0df 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/psp/dplms.py +++ b/workflow/src/legenddataflow/scripts/par/geds/psp/dplms.py @@ -1,3 +1,6 @@ +"""Console script ``par-geds-psp-dplms``: train DPLMS filters from +calibration waveforms and produce partition-level DSP parameters.""" + from __future__ import annotations import argparse @@ -8,7 +11,11 @@ import numpy as np import pygama.math.distributions as pmd # noqa: F401 from dbetto.catalog import Props -from legenddataflowscripts.utils import build_log +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + prepare_output_paths, +) from lgdo import Array, Table, WaveformTable from pygama.evt.build_tcm import _concat_tables from pygama.pargen.data_cleaning import generate_cuts @@ -208,9 +215,15 @@ def par_geds_psp_dplms() -> None: args = argparser.parse_args() - dsp_config = Props.read_from(args.processing_chain) log = build_log(args.log_config, args.log) + check_input_files(args.peak_files, "--peak-files") + check_input_files(args.database, "--database") + check_input_files(args.inplots, "--inplots") + prepare_output_paths(*args.dsp_pars, *args.lh5_path, *(args.plot_path or [])) + + dsp_config = Props.read_from(args.processing_chain) + t0 = time.time() dplms_dict = Props.read_from(args.config_file) @@ -222,6 +235,12 @@ def par_geds_psp_dplms() -> None: fft_runs, _ = split_files_by_run(args.fft_raw_filelists) n_runs = len(fft_runs) + if n_runs == 0: + msg = ( + "no fft raw files found in the provided filelists " + f"{args.fft_raw_filelists}, cannot train the DPLMS filter" + ) + raise RuntimeError(msg) n_per_run = dplms_dict["n_baselines"] // n_runs t0 = time.time() diff --git a/workflow/src/legenddataflow/scripts/par/geds/raw/blindcal.py b/workflow/src/legenddataflow/scripts/par/geds/raw/blindcal.py index f3e880cc..53968755 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/raw/blindcal.py +++ b/workflow/src/legenddataflow/scripts/par/geds/raw/blindcal.py @@ -14,9 +14,13 @@ import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np -from dbetto import TextDB from dbetto.catalog import Props -from legenddataflowscripts.utils import build_log +from legenddataflowscripts.utils import ( + build_log, + expand_filelist, + get_rule_config, + prepare_output_paths, +) from pygama.pargen.energy_cal import HPGeCalibration mpl.use("agg") @@ -43,22 +47,20 @@ def par_geds_raw_blindcal() -> None: argparser.add_argument("-d", "--debug", help="debug_mode", action="store_true") args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["tier_raw_blind_check"] + config_dict = get_rule_config( + args.configs, "tier_raw_blind_check", args.timestamp, args.datatype + ) log = build_log(config_dict, args.log) + prepare_output_paths(args.blind_curve, args.plot_file) + # peaks to search for peaks_keV = np.array( [238, 583.191, 727.330, 860.564, 1592.53, 1620.50, 2103.53, 2614.50] ) - if isinstance(args.files, list) and args.files[0].split(".")[-1] == "filelist": - input_file = args.files[0] - with Path(input_file).open() as f: - input_file = f.read().splitlines() - else: - input_file = args.files + input_file = expand_filelist(args.files) E_uncal = lh5.read_as(f"{args.raw_table_name}/daqenergy", input_file, library="np") E_uncal = E_uncal[E_uncal > 200] diff --git a/workflow/src/legenddataflow/scripts/par/geds/raw/blindcheck.py b/workflow/src/legenddataflow/scripts/par/geds/raw/blindcheck.py index 825f2a2a..ccc96baf 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/raw/blindcheck.py +++ b/workflow/src/legenddataflow/scripts/par/geds/raw/blindcheck.py @@ -17,9 +17,14 @@ import matplotlib.pyplot as plt import numexpr as ne import numpy as np -from dbetto import TextDB from dbetto.catalog import Props -from legenddataflowscripts.utils import build_log +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + expand_filelist, + get_rule_config, + prepare_output_paths, +) from legendmeta import LegendMetadata from pygama.math.histogram import get_hist from pygama.pargen.energy_cal import get_i_local_maxima @@ -47,11 +52,15 @@ def par_geds_raw_blindcheck() -> None: argparser.add_argument("--log", help="log file", type=str) args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["tier_raw_blind_check"] + config_dict = get_rule_config( + args.configs, "tier_raw_blind_check", args.timestamp, args.datatype + ) log = build_log(config_dict, args.log) + check_input_files(args.blind_curve, "--blind-curve") + prepare_output_paths(args.output, args.plot_file) + # get the usability status for this channel chmap = LegendMetadata(args.metadata).channelmap( @@ -62,12 +71,7 @@ def par_geds_raw_blindcheck() -> None: # read in calibration curve for this channel blind_curve = Props.read_from(args.blind_curve)[args.channel]["pars"]["operations"] - if isinstance(args.files, list) and args.files[0].split(".")[-1] == "filelist": - input_file = args.files[0] - with Path(input_file).open() as f: - input_file = f.read().splitlines() - else: - input_file = args.files + input_file = expand_filelist(args.files) # load in the data daqenergy = lh5.read_as( @@ -113,7 +117,6 @@ def par_geds_raw_blindcheck() -> None: if ( np.any(np.abs(maxs - 2614) < 5) and np.any(np.abs(maxs - 583) < 5) ) or det_status is False: - Path(args.output).parent.mkdir(parents=True, exist_ok=True) Props.write_to( args.output, { @@ -122,5 +125,9 @@ def par_geds_raw_blindcheck() -> None: }, ) else: - msg = "peaks not found in daqenergy" + msg = ( + f"blinding check failed for channel {args.channel}: no peaks found " + f"within 5 keV of both 583 and 2614 keV in the calibrated daqenergy " + f"(local maxima at {maxs} keV), the blinding calibration may be stale" + ) raise RuntimeError(msg) diff --git a/workflow/src/legenddataflow/scripts/par/geds/tcm/pulser.py b/workflow/src/legenddataflow/scripts/par/geds/tcm/pulser.py index 5edc73ef..e56270d4 100644 --- a/workflow/src/legenddataflow/scripts/par/geds/tcm/pulser.py +++ b/workflow/src/legenddataflow/scripts/par/geds/tcm/pulser.py @@ -1,12 +1,17 @@ +"""Console script ``par-geds-tcm-pulser``: identify pulser events in the tcm +tier and save their ids.""" + from __future__ import annotations import argparse -from pathlib import Path -import numpy as np -from dbetto import TextDB from dbetto.catalog import Props -from legenddataflowscripts.utils import build_log +from legenddataflowscripts.utils import ( + build_log, + expand_filelist, + get_rule_config, + prepare_output_paths, +) from pygama.pargen.data_cleaning import get_tcm_pulser_ids @@ -21,35 +26,26 @@ def par_geds_tcm_pulser() -> None: argparser.add_argument("--channel", help="Channel", type=str, required=True) argparser.add_argument("--rawid", help="rawid", type=str, required=True) - argparser.add_argument( - "--pulser-file", help="pulser file", type=str, required=False - ) + argparser.add_argument("--pulser-file", help="pulser file", type=str, required=True) argparser.add_argument("--tcm-files", help="tcm_files", nargs="*", type=str) args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["pars_tcm_pulser"] + config_dict = get_rule_config( + args.configs, "pars_tcm_pulser", args.timestamp, args.datatype + ) build_log(config_dict, args.log) + prepare_output_paths(args.pulser_file) + kwarg_dict = config_dict["inputs"]["pulser_config"] kwarg_dict = Props.read_from(kwarg_dict) - if ( - isinstance(args.tcm_files, list) - and args.tcm_files[0].split(".")[-1] == "filelist" - ): - tcm_files = args.tcm_files[0] - with Path(tcm_files).open() as f: - tcm_files = f.read().splitlines() - else: - tcm_files = args.tcm_files # get pulser mask from tcm files - tcm_files = sorted(np.unique(tcm_files)) + tcm_files = expand_filelist(args.tcm_files, "--tcm-files") ids, mask = get_tcm_pulser_ids( tcm_files, args.rawid, kwarg_dict.pop("pulser_multiplicity_threshold") ) - Path(args.pulser_file).parent.mkdir(parents=True, exist_ok=True) Props.write_to(args.pulser_file, {"idxs": ids.tolist(), "mask": mask.tolist()}) diff --git a/workflow/src/legenddataflow/scripts/par/spms/dsp/trigger_threshold.py b/workflow/src/legenddataflow/scripts/par/spms/dsp/trigger_threshold.py index a4b9d03f..1bb3b7bd 100644 --- a/workflow/src/legenddataflow/scripts/par/spms/dsp/trigger_threshold.py +++ b/workflow/src/legenddataflow/scripts/par/spms/dsp/trigger_threshold.py @@ -1,14 +1,23 @@ +"""Console scripts ``par-spms-dsp-trg-thr`` and ``-multi``: determine SiPM +trigger thresholds from the noise distribution of DSP-processed raw +waveforms.""" + from __future__ import annotations import argparse -from pathlib import Path import hist import lh5 import numpy as np -from dbetto import AttrsDict, Props, TextDB, utils +from dbetto import AttrsDict, Props, utils from dspeed import build_dsp -from legenddataflowscripts.utils import build_log, cfgtools +from legenddataflowscripts.utils import ( + build_log, + cfgtools, + check_input_files, + get_rule_config, + prepare_output_paths, +) def get_channel_trg_thr(df_configs, sipm_name, dsp_db, raw_file, raw_table_name, log): @@ -31,7 +40,7 @@ def get_channel_trg_thr(df_configs, sipm_name, dsp_db, raw_file, raw_table_name, if len(lh5.ls(raw_file, f"{raw_table_name}/waveform_bit_drop")) == 0: msg = ( f"could not find waveform '{raw_table_name}/waveform_bit_drop' " - "in {args.raw_file}, returning null pars" + f"in {raw_file}, returning null pars" ) log.warning(msg) @@ -46,14 +55,15 @@ def get_channel_trg_thr(df_configs, sipm_name, dsp_db, raw_file, raw_table_name, if len(data) == 0: msg = ( f"could not find any waveforms '{raw_table_name}/waveform_bit_drop' " - "in {args.raw_file}, returning null pars" + f"in {raw_file}, returning null pars" ) log.warning(msg) elif len(data) < settings.n_events: msg = ( - f"number of waveforms ({len(data)}) in '{raw_table_name}/waveform_bit_drop' < {settings.n_events}" - " in {args.raw_file}, can't build histogram" + f"found only {len(data)} waveforms in " + f"'{raw_table_name}/waveform_bit_drop' of {raw_file}, need at " + f"least {settings.n_events} to build the histogram" ) raise RuntimeError(msg) else: @@ -85,7 +95,10 @@ def get_channel_trg_thr(df_configs, sipm_name, dsp_db, raw_file, raw_table_name, fwhm = edges[idx_over_half[-1]] - edges[idx_over_half[0]] if fwhm <= 0: - msg = f"determined FWHM of baseline derivative distribution is so <= 0: {fwhm:.3f}" + msg = ( + "determined FWHM of the baseline derivative distribution " + f"is not positive: {fwhm:.3f}" + ) raise RuntimeError(msg) return fwhm @@ -107,15 +120,17 @@ def par_spms_dsp_trg_thr() -> None: args = argparser.parse_args() # dataflow configs - df_configs = ( - TextDB(args.config_path, lazy=True) - .on(args.timestamp, system=args.datatype) - .snakemake_rules.pars_spms_dsp_trg_thr + df_configs = get_rule_config( + args.config_path, "pars_spms_dsp_trg_thr", args.timestamp, args.datatype ) # setup logging log = build_log(df_configs, args.logfile) + check_input_files(args.raw_file, "--raw-file") + check_input_files(args.dsp_db, "--dsp-db") + prepare_output_paths(args.output_file) + fwhm = get_channel_trg_thr( df_configs, args.sipm_name, @@ -126,7 +141,6 @@ def par_spms_dsp_trg_thr() -> None: ) msg = f"writing out baseline_curr_fwhm = {fwhm}" log.debug(msg) - Path(args.output_file).parent.mkdir(parents=True, exist_ok=True) Props.write_to( args.output_file, {"baseline_curr_fwhm": float(fwhm) if fwhm is not None else fwhm}, @@ -148,18 +162,28 @@ def par_spms_dsp_trg_thr_multi() -> None: args = argparser.parse_args() # dataflow configs - df_configs = ( - TextDB(args.config_path, lazy=True) - .on(args.timestamp, system=args.datatype) - .snakemake_rules.pars_spms_dsp_trg_thr + df_configs = get_rule_config( + args.config_path, "pars_spms_dsp_trg_thr", args.timestamp, args.datatype ) # setup logging log = build_log(df_configs, args.logfile) + check_input_files(args.raw_file, "--raw-file") + check_input_files(args.dsp_db, "--dsp-db") + prepare_output_paths(args.output_file) + + if len(args.sipm_names) != len(args.raw_table_names): + msg = ( + f"--sipm-names ({len(args.sipm_names)} entries) and " + f"--raw-table-names ({len(args.raw_table_names)} entries) must " + "have the same length" + ) + raise ValueError(msg) + out_dict = {} for sipm_name, raw_table_name in zip( - args.sipm_names, args.raw_table_names, strict=False + args.sipm_names, args.raw_table_names, strict=True ): fwhm = get_channel_trg_thr( df_configs, @@ -175,7 +199,6 @@ def par_spms_dsp_trg_thr_multi() -> None: out_dict[sipm_name] = { "baseline_curr_fwhm": float(fwhm) if fwhm is not None else fwhm } - Path(args.output_file).parent.mkdir(parents=True, exist_ok=True) Props.write_to( args.output_file, out_dict, diff --git a/workflow/src/legenddataflow/scripts/tier/evt.py b/workflow/src/legenddataflow/scripts/tier/evt.py index 2328bce1..abd4cb6c 100644 --- a/workflow/src/legenddataflow/scripts/tier/evt.py +++ b/workflow/src/legenddataflow/scripts/tier/evt.py @@ -1,3 +1,7 @@ +"""Console script ``build-tier-evt``: build the evt tier (event-level +reconstruction and cross-talk correction) from hit, dsp and tcm inputs, +optionally including ANN files.""" + from __future__ import annotations import argparse @@ -6,8 +10,13 @@ import lh5 import numpy as np -from dbetto import AttrsDict, Props, TextDB -from legenddataflowscripts.utils import build_log +from dbetto import AttrsDict, Props +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + get_rule_config, + prepare_output_paths, +) from legendmeta import LegendMetadata from lgdo.types import Array from pygama.evt import build_evt @@ -15,14 +24,37 @@ from legenddataflow.methods import ProcessingFileKey +def _resolve_channels(chmap, field, dic, timestamp): + """Resolve one evt-config ``channels`` entry into a list of ``chXXX`` + names using the channel map valid at ``timestamp``.""" + system_map = chmap.map("system", unique=False) + if dic["system"] not in system_map: + msg = ( + f"channels field {field!r} requests system {dic['system']!r}, " + f"which is not in the channel map for {timestamp}: available " + f"systems are {sorted(system_map)}" + ) + raise ValueError(msg) + chans = system_map[dic["system"]] + if "selectors" in dic: + try: + for k, val in dic["selectors"].items(): + chans = chans.map(k, unique=False)[val] + except KeyError: + chans = None + if chans is not None: + return [f"ch{chan}" for chan in list(chans.map("daq.rawid"))] + return [] + + def build_tier_evt() -> None: argparser = argparse.ArgumentParser() argparser.add_argument("--hit-file") argparser.add_argument("--dsp-file") argparser.add_argument("--tcm-file") - argparser.add_argument("--ann-file", nargs="*") - argparser.add_argument("--xtc-file", nargs="*") - argparser.add_argument("--par-files", nargs="*") + argparser.add_argument("--ann-file", nargs="*", default=[]) + argparser.add_argument("--xtc-file", nargs="*", default=[]) + argparser.add_argument("--par-files", nargs="*", default=[]) argparser.add_argument("--datatype", required=True) argparser.add_argument("--timestamp", required=True) @@ -36,17 +68,21 @@ def build_tier_evt() -> None: args = argparser.parse_args() if args.tier not in ("evt", "pet"): - msg = f"unsupported tier {args.tier}" + msg = f"unsupported tier {args.tier!r}: this script only builds 'evt' or 'pet'" raise ValueError(msg) # load in config - df_config = ( - TextDB(args.configs, lazy=True) - .on(args.timestamp, system=args.datatype) - .snakemake_rules.tier_evt - ) + df_config = get_rule_config(args.configs, "tier_evt", args.timestamp, args.datatype) log = build_log(df_config, args.log) + check_input_files(args.hit_file, "--hit-file") + check_input_files(args.dsp_file, "--dsp-file") + check_input_files(args.tcm_file, "--tcm-file") + check_input_files(args.par_files, "--par-files") + check_input_files(args.xtc_file, "--xtc-file") + check_input_files(args.ann_file, "--ann-file") + prepare_output_paths(args.output) + chmap = LegendMetadata(args.metadata, lazy=True).channelmap(on=args.timestamp) evt_config = AttrsDict(Props.read_from(df_config.inputs.evt_config)) @@ -80,18 +116,9 @@ def build_tier_evt() -> None: # block for snakemake to fill in channel lists for field, dic in evt_config.channels.items(): if isinstance(dic, dict): - chans = chmap.map("system", unique=False)[dic["system"]] - if "selectors" in dic: - try: - for k, val in dic["selectors"].items(): - chans = chans.map(k, unique=False)[val] - except KeyError: - chans = None - if chans is not None: - chans = [f"ch{chan}" for chan in list(chans.map("daq.rawid"))] - else: - chans = [] - evt_config.channels[field] = chans + evt_config.channels[field] = _resolve_channels( + chmap, field, dic, args.timestamp + ) log.debug(json.dumps(evt_config.channels, indent=2)) @@ -99,7 +126,12 @@ def build_tier_evt() -> None: f"ch{chan}": dic.name for chan, dic in chmap.map("daq.rawid").items() } - Path(args.output).parent.mkdir(parents=True, exist_ok=True) + if "hardware_tcm_1" not in lh5.ls(args.tcm_file): + msg = ( + f"'hardware_tcm_1' not found in tcm file {args.tcm_file}: " + f"available objects are {lh5.ls(args.tcm_file)}" + ) + raise ValueError(msg) file_table = { "tcm": (args.tcm_file, "hardware_tcm_1", "ch{}"), @@ -125,18 +157,9 @@ def build_tier_evt() -> None: # block for snakemake to fill in channel lists for field, dic in muon_config["channels"].items(): if isinstance(dic, dict): - chans = chmap.map("system", unique=False)[dic["system"]] - if "selectors" in dic: - try: - for k, val in dic["selectors"].items(): - chans = chans.map(k, unique=False)[val] - except KeyError: - chans = None - if chans is not None: - chans = [f"ch{chan}" for chan in list(chans.map("daq.rawid"))] - else: - chans = [] - muon_config["channels"][field] = chans + muon_config["channels"][field] = _resolve_channels( + chmap, field, dic, args.timestamp + ) trigger_timestamp = table[field_config["ged_timestamp"]["table"]][ field_config["ged_timestamp"]["field"] @@ -187,7 +210,7 @@ def _find_matching_values_with_delay(arr1, arr2, jit_delay): matching_values = [] # Create an array with all possible delay values - delays = np.arange(0, int(1e9 * jit_delay)) * jit_delay + delays = np.linspace(0, int(1e9 * jit_delay), 10000) * jit_delay for delay in delays: arr2_delayed = arr2 + delay diff --git a/workflow/src/legenddataflow/scripts/tier/raw_blind.py b/workflow/src/legenddataflow/scripts/tier/raw_blind.py index 0a50b031..a852555a 100644 --- a/workflow/src/legenddataflow/scripts/tier/raw_blind.py +++ b/workflow/src/legenddataflow/scripts/tier/raw_blind.py @@ -20,8 +20,15 @@ import numexpr as ne import numpy as np from dbetto.catalog import Props -from legenddataflowscripts.utils import alias_table, build_log -from legendmeta import LegendMetadata, TextDB +from legenddataflowscripts.utils import ( + alias_table, + build_log, + check_input_files, + get_rule_config, + parse_json_arg, + prepare_output_paths, +) +from legendmeta import LegendMetadata filter_map = { "zstd": hdf5plugin.Zstd(), @@ -49,13 +56,21 @@ def build_tier_raw_blind() -> None: argparser.add_argument("--alias-table", help="Alias table", type=str, default=None) args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True) - config_dict = configs.on(args.timestamp, system=args.datatype)["snakemake_rules"][ - "tier_raw_blind" - ] + config_dict = get_rule_config( + args.configs, "tier_raw_blind", args.timestamp, args.datatype + ) log = build_log(config_dict, args.log) + check_input_files(args.input, "--input") + check_input_files(args.blind_curve, "--blind-curve") + alias_map = ( + parse_json_arg(args.alias_table, "--alias-table") + if args.alias_table is not None + else None + ) + prepare_output_paths(args.output) + hdf_settings = Props.read_from(config_dict["settings"])["hdf5_settings"] if "compression" in hdf_settings: @@ -79,20 +94,26 @@ def build_tier_raw_blind() -> None: for system, chan_dict in chmap.map("system", unique=False).items() } - main_channels = ( - list(chans["geds"]) - + list(chans["spms"]) - + list(chans["auxs"]) - + list(chans["blsns"]) - + list(chans["puls"]) - ) + main_systems = ("geds", "spms", "auxs", "bsln", "puls") + missing_systems = [system for system in main_systems if system not in chans] + if missing_systems: + msg = ( + f"systems {missing_systems} not found in the channel map for " + f"{args.timestamp}: available systems are {sorted(chans)}" + ) + raise RuntimeError(msg) + + main_channels = [chnum for system in main_systems for chnum in chans[system]] # rows that need blinding toblind = np.array([]) + daqenergy = None log.info("Blinding Ge channels") start = time.time() + blind_curves = Props.read_from(args.blind_curve) + # first, loop through the Ge detector channels, calibrate them and look for events that should be blinded for chnum in chans["geds"]: # skip Ge detectors that are anti-coincidence only or not able to be blinded for some other reason @@ -103,9 +124,14 @@ def build_tier_raw_blind() -> None: daqenergy = lh5.read(f"ch{chnum}/raw/daqenergy", args.input) # read in calibration curve for this channel - blind_curve = Props.read_from(args.blind_curve)[chmap.map("daq.rawid").name][ - "pars" - ]["operations"] + channel_name = chans["geds"][chnum].name + if channel_name not in blind_curves: + msg = ( + f"no blinding curve found for channel {channel_name} (rawid {chnum}) " + f"in {args.blind_curve}" + ) + raise RuntimeError(msg) + blind_curve = blind_curves[channel_name]["pars"]["operations"] # calibrate daq energy using pre existing curve daqenergy_cal = ne.evaluate( @@ -121,6 +147,14 @@ def build_tier_raw_blind() -> None: np.nonzero(np.abs(np.asarray(daqenergy_cal) - centroid) <= width)[0], ) + if daqenergy is None: + msg = ( + f"no blindable Ge channels found in the channel map for {args.timestamp} " + "(all channels have analysis.is_blinded set to false), refusing to " + "produce an unblinded output file" + ) + raise RuntimeError(msg) + # remove duplicates toblind = np.unique(toblind) @@ -178,6 +212,6 @@ def build_tier_raw_blind() -> None: log.info("Finished blinding Ge channels") msg = f"Time taken: {time.time() - start:.2f} seconds" log.info(msg) - if args.alias_table is not None: + if alias_map is not None: log.info("Creating alias table") - alias_table(args.output, args.alias_table) + alias_table(args.output, alias_map) diff --git a/workflow/src/legenddataflow/scripts/tier/raw_fcio.py b/workflow/src/legenddataflow/scripts/tier/raw_fcio.py index b85e20e9..0cbeb547 100644 --- a/workflow/src/legenddataflow/scripts/tier/raw_fcio.py +++ b/workflow/src/legenddataflow/scripts/tier/raw_fcio.py @@ -1,14 +1,22 @@ +"""Console script ``build-tier-raw-fcio``: convert FCIO (FlashCam) DAQ files +to the LH5 raw tier with ``daq2lh5``.""" + from __future__ import annotations import argparse import time -from pathlib import Path import hdf5plugin from daq2lh5 import build_raw -from dbetto import TextDB from dbetto.catalog import Props -from legenddataflowscripts.utils import alias_table, build_log +from legenddataflowscripts.utils import ( + alias_table, + build_log, + check_input_files, + get_rule_config, + parse_json_arg, + prepare_output_paths, +) filter_map = { "zstd": hdf5plugin.Zstd(), @@ -35,16 +43,20 @@ def build_tier_raw_fcio() -> None: ) args = argparser.parse_args() - Path(args.output).parent.mkdir(parents=True, exist_ok=True) - - config_dict = ( - TextDB(args.configs, lazy=True) - .on(args.timestamp, system=args.datatype) - .snakemake_rules.tier_raw_fcio + config_dict = get_rule_config( + args.configs, "tier_raw_fcio", args.timestamp, args.datatype ) log = build_log(config_dict, args.log) + check_input_files(args.input, "input") + alias_map = ( + parse_json_arg(args.alias_table, "--alias-table") + if args.alias_table is not None + else None + ) + prepare_output_paths(args.output) + channel_dict = config_dict.inputs settings = Props.read_from(channel_dict.settings) channel_dict = channel_dict.out_spec @@ -83,6 +95,6 @@ def build_tier_raw_fcio() -> None: msg = f"Built raw in {time.time() - start:.2f} seconds" log.info(msg) - if args.alias_table is not None: + if alias_map is not None: log.info("Creating alias table") - alias_table(args.output, args.alias_table) + alias_table(args.output, alias_map) diff --git a/workflow/src/legenddataflow/scripts/tier/raw_orca.py b/workflow/src/legenddataflow/scripts/tier/raw_orca.py index 8c643894..434e9727 100644 --- a/workflow/src/legenddataflow/scripts/tier/raw_orca.py +++ b/workflow/src/legenddataflow/scripts/tier/raw_orca.py @@ -1,15 +1,23 @@ +"""Console script ``build-tier-raw-orca``: convert ORCA DAQ files to the LH5 +raw tier with ``daq2lh5``.""" + from __future__ import annotations import argparse -import logging import time -from pathlib import Path import hdf5plugin from daq2lh5 import build_raw from dbetto import TextDB from dbetto.catalog import Props -from legenddataflowscripts.utils import alias_table, build_log +from legenddataflowscripts.utils import ( + alias_table, + build_log, + check_input_files, + get_rule_config, + parse_json_arg, + prepare_output_paths, +) filter_map = { "zstd": hdf5plugin.Zstd(), @@ -38,18 +46,21 @@ def build_tier_raw_orca() -> None: ) args = argparser.parse_args() - Path(args.log).parent.mkdir(parents=True, exist_ok=True) - logging.basicConfig(level=logging.INFO, filename=args.log, filemode="w") - - Path(args.output).parent.mkdir(parents=True, exist_ok=True) - - configs = TextDB(args.configs, lazy=True) - config_dict = configs.on(args.timestamp, system=args.datatype)["snakemake_rules"][ - "tier_raw_orca" - ] + config_dict = get_rule_config( + args.configs, "tier_raw_orca", args.timestamp, args.datatype + ) log = build_log(config_dict, args.log) + check_input_files(args.input, "input") + check_input_files(args.inl_table, "--inl-table") + alias_map = ( + parse_json_arg(args.alias_table, "--alias-table") + if args.alias_table is not None + else None + ) + prepare_output_paths(args.output) + channel_dict = config_dict["inputs"] settings = Props.read_from(channel_dict["settings"]) channel_dict = channel_dict["out_spec"] @@ -131,6 +142,6 @@ def build_tier_raw_orca() -> None: msg = f"Built raw in {time.time() - start:.2f} seconds" log.info(msg) - if args.alias_table is not None: + if alias_map is not None: log.info("Creating alias table") - alias_table(args.output, args.alias_table) + alias_table(args.output, alias_map) diff --git a/workflow/src/legenddataflow/scripts/tier/skm.py b/workflow/src/legenddataflow/scripts/tier/skm.py index 7f0433f2..b4da0b22 100644 --- a/workflow/src/legenddataflow/scripts/tier/skm.py +++ b/workflow/src/legenddataflow/scripts/tier/skm.py @@ -1,12 +1,19 @@ +"""Console script ``build-tier-skm``: build the skm tier (flattened physics +skim) from concatenated evt-tier data.""" + from __future__ import annotations import argparse import awkward as ak import lh5 -from dbetto import TextDB from dbetto.catalog import Props -from legenddataflowscripts.utils import build_log +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + get_rule_config, + prepare_output_paths, +) from lgdo.types import Array, Struct, Table, VectorOfVectors @@ -32,12 +39,15 @@ def build_tier_skm() -> None: args = argparser.parse_args() # load in config - config_dict = TextDB(args.configs, lazy=True).on( - args.timestamp, system=args.datatype - )["snakemake_rules"]["tier_skm"] + config_dict = get_rule_config( + args.configs, "tier_skm", args.timestamp, args.datatype + ) build_log(config_dict, args.log) + check_input_files(args.evt_file, "--evt-file") + prepare_output_paths(args.output) + skm_config_file = config_dict["inputs"]["skm_config"] evt_filter = Props.read_from(skm_config_file)["evt_filter"] out_fields = Props.read_from(skm_config_file)["keep_fields"] @@ -51,17 +61,24 @@ def build_tier_skm() -> None: # make it rectangular and make an LGDO Table out_table = Table(skm) - for field in out_fields: + # expand fields that refer to whole sub-tables into their leaf fields; + # iterate over a copy since out_fields is modified in the loop + for field in list(out_fields): items = field.split(".") ptr1 = out_table - for item in items[:-1]: - ptr1 = ptr1[item] + try: + for item in items: + ptr1 = ptr1[item] + except KeyError as err: + msg = ( + f"keep_fields entry {field!r} not found in the evt tier data: " + f"missing key {err}" + ) + raise KeyError(msg) from err - if isinstance(ptr1[items[-1]], Table): + if isinstance(ptr1, Table): out_fields.remove(field) - out_fields = get_all_out_fields( - ptr1[items[-1]], out_fields, current_field=field - ) + out_fields = get_all_out_fields(ptr1, out_fields, current_field=field) # remove unwanted columns out_table_skm = Table(size=len(out_table)) diff --git a/workflow/src/legenddataflow/scripts/tier/tcm.py b/workflow/src/legenddataflow/scripts/tier/tcm.py index 2470cc92..b4eb3d0c 100644 --- a/workflow/src/legenddataflow/scripts/tier/tcm.py +++ b/workflow/src/legenddataflow/scripts/tier/tcm.py @@ -1,13 +1,20 @@ +"""Console script ``build-tier-tcm``: build the time-coincidence map (tcm) +tier from raw-tier data.""" + from __future__ import annotations import argparse -from pathlib import Path import lh5 from daq2lh5.orca import orca_flashcam -from dbetto import AttrsDict, TextDB +from dbetto import AttrsDict from dbetto.catalog import Props -from legenddataflowscripts.utils import build_log +from legenddataflowscripts.utils import ( + build_log, + check_input_files, + get_rule_config, + prepare_output_paths, +) from pygama.evt.build_tcm import build_tcm @@ -21,15 +28,23 @@ def build_tier_tcm() -> None: argparser.add_argument("--log", help="log file", type=str) args = argparser.parse_args() - configs = TextDB(args.configs, lazy=True).on(args.timestamp, system=args.datatype) - config_dict = configs["snakemake_rules"]["tier_tcm"] + config_dict = get_rule_config( + args.configs, "tier_tcm", args.timestamp, args.datatype + ) log = build_log(config_dict, args.log) + + check_input_files(args.input, "input") + prepare_output_paths(args.output) + # the config can either be a single file used for all fcids or a mapping + # of per-fcid config files + per_fcid_settings = None if isinstance(config_dict["inputs"]["config"], dict | AttrsDict): - settings = { - key: Props.read_from(val) + per_fcid_settings = { + int(key): Props.read_from(val) for key, val in config_dict["inputs"]["config"].items() } + settings = None else: settings = Props.read_from(config_dict["inputs"]["config"]) @@ -37,7 +52,7 @@ def build_tier_tcm() -> None: ch_list = lh5.ls(args.input, "/ch*") if len(ch_list) == 0: - msg = "no tables matching /ch* found in input file" + msg = f"no tables matching /ch* found in input file {args.input}" raise RuntimeError(msg) log.debug(ch_list) @@ -50,17 +65,26 @@ def build_tier_tcm() -> None: fcid_channels[fcid] = [] fcid_channels[fcid].append(f"/{ch}/raw") - Path(args.output).parent.mkdir(parents=True, exist_ok=True) # make a hardware_tcm_[fcid] for each fcid for fcid, fcid_dict in fcid_channels.items(): msg = f"building tcm for fcid: {fcid}" log.info(msg) + if per_fcid_settings is not None: + if fcid not in per_fcid_settings: + msg = ( + f"no tcm config found for fcid {fcid}: the config mapping " + f"only covers fcids {sorted(per_fcid_settings)}" + ) + raise RuntimeError(msg) + fcid_settings = per_fcid_settings[fcid] + else: + fcid_settings = settings build_tcm( [(args.input, fcid_dict)], out_file=args.output, out_name=f"hardware_tcm_{fcid}", wo_mode="o", - **settings.get(fcid, settings), + **fcid_settings, ) msg = f"built tcm for fcid: {fcid}" log.info(msg)