From a4b69df6e4a6b08f4ec07c57e2dc7db8f1ef4a13 Mon Sep 17 00:00:00 2001 From: juvanden Date: Wed, 24 Jun 2026 11:13:01 +0200 Subject: [PATCH 1/7] add claude context --- CLAUDE.md | 23 + .../01_framework_structure.md | 278 ++++++++++++ .../02_coding_style.md | 403 ++++++++++++++++++ .../03_task_commands.md | 358 ++++++++++++++++ .../04_awkward_reference.md | 305 +++++++++++++ .../05_dos_and_donts.md | 310 ++++++++++++++ 6 files changed, 1677 insertions(+) create mode 100644 CLAUDE.md create mode 100644 docs/columnflow_claude_guide/01_framework_structure.md create mode 100644 docs/columnflow_claude_guide/02_coding_style.md create mode 100644 docs/columnflow_claude_guide/03_task_commands.md create mode 100644 docs/columnflow_claude_guide/04_awkward_reference.md create mode 100644 docs/columnflow_claude_guide/05_dos_and_donts.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..2a136f3cd --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,23 @@ +# Columnflow Project — Claude Instructions + +This is a **columnflow** analysis project. Columnflow is a fully orchestrated columnar HEP analysis framework built on [law](https://github.com/riga/law) (workflow orchestration) and [order](https://github.com/riga/order) (metadata management), using [awkward-array](https://awkward-array.org) for columnar event data. + +Read all five guide files before writing or modifying any code: + +| Guide | Contents | +|---|---| +| [01 Framework Structure](docs/columnflow_claude_guide/01_framework_structure.md) | Analysis pipeline, task graph, config objects (Analysis/Campaign/Config), Order objects | +| [02 Coding Style](docs/columnflow_claude_guide/02_coding_style.md) | TAF decorators, uses/produces, imports, module registration, naming conventions | +| [03 Task Commands](docs/columnflow_claude_guide/03_task_commands.md) | `law run` commands with all standard parameters for every task | +| [04 Awkward Reference](docs/columnflow_claude_guide/04_awkward_reference.md) | Key `ak.*` functions used in columnflow code | +| [05 Dos and Don'ts](docs/columnflow_claude_guide/05_dos_and_donts.md) | Explicit rules Claude must follow when writing columnflow code | + +## Quick orientation + +- All data operations act on **chunks** of events encoded as `ak.Array` named `events`. +- Task array functions (TAFs): `Calibrator`, `Selector`, `Producer`, `Reducer`, `HistProducer`. +- Each TAF declares `uses` (columns to read) and `produces` (columns to write) in its decorator. +- Events are **never modified in-place** — always use `set_ak_column(events, "FieldName", value)` and reassign `events`. +- The standard pipeline order: `GetDatasetLFNs → CalibrateEvents → SelectEvents → ReduceEvents → ProduceColumns → CreateHistograms → PlotVariables1D`. +- All tasks are run with `law run cf. --version [--config ] [...]`. +- New Python modules must be registered in `law.cfg` under the correct `*_modules` key. diff --git a/docs/columnflow_claude_guide/01_framework_structure.md b/docs/columnflow_claude_guide/01_framework_structure.md new file mode 100644 index 000000000..f7c23b96d --- /dev/null +++ b/docs/columnflow_claude_guide/01_framework_structure.md @@ -0,0 +1,278 @@ +# 01 — Columnflow Framework Structure + +## Overview + +Columnflow is a backend for columnar, fully orchestrated HEP analyses in pure Python. + +| Layer | Package | Role | +|---|---|---| +| Workflow orchestration | [law](https://github.com/riga/law) | Task graph, CLI, remote execution | +| Metadata / bookkeeping | [order](https://github.com/riga/order) | Analysis, Campaign, Config, Dataset, Process, Variable, Shift | +| Columnar data | [awkward-array](https://awkward-array.org) | Event arrays; all physics data | +| Histogramming | [Hist](https://hist.readthedocs.io) | Histogram objects produced by `CreateHistograms` | +| Physics objects | [coffea](https://coffeateam.github.io/coffea/) | Lorentz vector behaviour on awkward arrays | + +--- + +## Analysis Pipeline (Task Graph) + +The standard linear pipeline from raw data to plots: + +``` +GetDatasetLFNs + │ +CalibrateEvents ← applies calibrations (e.g. JEC) + │ +SelectEvents ← creates event and object masks; produces stats.json + │ +ReduceEvents ← applies masks; writes reduced parquet files + │ +MergeReducedEvents ← merges per-file reduced files into one file per dataset + │ +ProduceColumns ← creates additional high-level columns + │ +CreateHistograms ← fills Hist histograms per variable / category / shift + │ +MergeHistograms ← merges per-dataset histograms + │ +MergeShiftedHistograms ← merges nominal + shifted histograms for inference + │ +PlotVariables1D ← produces matplotlib/mplhep plots +CreateDatacards ← produces CMS combine datacards + ROOT shape files +``` + +Tasks with `cf.` prefix are columnflow built-ins; user-defined tasks live in `/tasks/`. + +### Key properties + +- **Chunking**: events are processed in chunks of ≤100 000 events (set in `law.cfg` under `chunked_io_chunk_size`). +- **Parallelism**: each file / chunk is a separate law task branch; submit to HTCondor/Slurm with `--workflow htcondor`. +- **Reproducibility**: `--version` tags all intermediate outputs; different versions coexist on disk. +- **Columnar storage**: intermediate results are stored as **Parquet** files (events) or **pickle** files (histograms). + +--- + +## Five Task Array Function (TAF) Types + +| TAF type | Class | CLI parameter | Quantity | Task | +|---|---|---|---|---| +| Calibrator | `Calibrator` | `--calibrators` | 0..N | `CalibrateEvents` | +| Selector | `Selector` | `--selector` | exactly 1 | `SelectEvents` | +| Reducer | `Reducer` | `--reducer` | exactly 1 | `ReduceEvents` | +| Producer | `Producer` | `--producers` | 0..N | `ProduceColumns` | +| HistProducer | `HistProducer` | `--hist-producer` | exactly 1 | `CreateHistograms` | + +All TAFs share the same decorator pattern — see [02 Coding Style](02_coding_style.md). + +--- + +## Configuration Objects (order package) + +### Hierarchy + +``` +Analysis + └── Config (links Analysis + Campaign) + └── Campaign + └── Dataset → DatasetInfo (files, events, keys) +``` + +### Analysis + +Top-level container; rarely holds analysis logic itself. + +```python +import order as od +analysis = od.Analysis(name="my_analysis", id=1) +``` + +### Campaign + +Experiment-period-specific information (year, energy, tier, file locations). + +```python +cpn = od.Campaign( + name="run3_2022", + id=1, + ecm=13.6, + aux={ + "tier": "NanoAOD", + "year": 2022, + "location": "root://...", + }, +) +``` + +### Config + +Analysis + campaign combination; carries all per-config parameters. + +```python +cfg = analysis.add_config(campaign, name="run3_2022", id=1) +``` + +### Process + +Physical process with cross-section and plot metadata. + +```python +from scinum import Number +proc = od.Process( + name="tt", + id=1000, + label=r"$t\bar{t}$", + color=(128, 76, 153), + xsecs={13.6: Number(831.76, {"scale": (19.77, 29.20)})}, +) +``` + +### Dataset + +A sample linked to a Campaign and one or more Processes. + +```python +cpn.add_dataset( + name="tt_dl_powheg", + id=1, + processes=[procs.tt], + info={ + "nominal": od.DatasetInfo( + keys=["/TT.../NANOAODSIM"], + n_files=242, + n_events=276079127, + ), + }, +) +``` + +### Variable + +Defines the column expression and histogram binning for `CreateHistograms`. + +```python +cfg.add_variable( + name="jet1_pt", + expression="Jet.pt[:,0]", + null_value=EMPTY_FLOAT, + binning=(40, 0.0, 400.0), + unit="GeV", + x_title=r"Jet $p_{T}$", +) +``` + +### Shift + +Systematic uncertainty variant (rate-only, weight-based, or dedicated dataset). + +```python +cfg.add_shift(name="nominal", id=0) +cfg.add_shift(name="mu_up", id=1, type=od.Shift.SHAPE) +cfg.add_shift(name="mu_down", id=2, type=od.Shift.SHAPE) +``` + +### Category / Categorizer + +Analysis phase-space regions used in histogramming and plotting. + +```python +from columnflow.config_util import add_category +add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") +``` + +--- + +## Important Config Auxiliaries + +These `cfg.x.*` entries have special meaning in columnflow: + +| Key | Type | Purpose | +|---|---|---| +| `cfg.x.keep_columns` | `DotDict` of sets | Which columns survive `ReduceEvents` | +| `cfg.x.luminosity` | `scinum.Number` | Luminosity with uncertainties | +| `cfg.x.external_files` | `DotDict` | Paths/URLs to external scale-factor files | +| `cfg.x.get_dataset_lfns` | callable | Custom LFN-retrieval function | +| `cfg.x.default_calibrator` | str | Default `--calibrator` value | +| `cfg.x.default_selector` | str | Default `--selector` value | +| `cfg.x.default_producer` | str/tuple | Default `--producer` value | +| `cfg.x.default_variables` | tuple | Default `--variables` value | +| `cfg.x.reduced_file_size` | float | Target merged-file size in MB | +| `cfg.x.versions` | dict | Pinned task versions | + +### keep_columns example + +```python +from columnflow.util import DotDict +from columnflow.columnar_util import ColumnCollection + +cfg.x.keep_columns = DotDict.wrap({ + "cf.ReduceEvents": { + "{Jet,FatJet}.{pt,eta,phi,mass,btagDeepFlavB}", + "Electron.{pt,eta,phi,mass,charge}", + "Muon.{pt,eta,phi,mass,charge}", + "MET.{pt,phi}", + "event", "run", "luminosityBlock", + ColumnCollection.ALL_FROM_SELECTOR, + }, + "cf.ProduceColumns": { + "ht", "n_jet", + }, +}) +``` + +--- + +## law.cfg Structure + +The `law.cfg` file drives the workflow. Critical sections: + +```ini +[analysis] +default_analysis: myanalysis.analysis.my_analysis.my_analysis +default_config: run3_2022 +default_dataset: tt_dl_powheg + +# Register all Python modules columnflow should know about +calibration_modules: columnflow.calibration.cms.{jets,met}, myanalysis.calibration.default +selection_modules: columnflow.selection.empty, columnflow.selection.cms.{json_filter,met_filters}, myanalysis.selection.default +reduction_modules: columnflow.reduction.default +production_modules: columnflow.production.{categories,normalization,processes}, myanalysis.production.default +categorization_modules: myanalysis.categorization.example +hist_production_modules: columnflow.histogramming.default, myanalysis.histogramming.example +inference_modules: columnflow.inference, myanalysis.inference.example + +[outputs] +# Map tasks to storage locations +cf.CalibrateEvents: local, /path/to/store +cf.SelectEvents: local, /path/to/store +cf.ReduceEvents: local, /path/to/store +cf.ProduceColumns: local, /path/to/store +cf.CreateHistograms: local, /path/to/store +``` + +**Critical**: after adding any new Python file (calibrator, selector, producer, etc.) you **must** register it in `law.cfg` under the correct `*_modules` key. No spaces after commas inside `{}` brace expansions. + +--- + +## Systematic Uncertainties (Shifts) + +Three classes, ordered by complexity: + +### 1. Rate-only uncertainties +Only affect yields, not selection. Defined entirely in the inference model. No additional workflow steps needed. + +### 2. Weight-based (shape) uncertainties +Applied via event weights. Require: +1. `cfg.add_shift(...)` for up/down variants +2. `add_shift_aliases(cfg, "source_name", ...)` to map column names +3. The `WeightProducer` / `HistProducer` must declare the shift in its `shifts` set +4. Considered only from `CreateHistograms` onwards + +### 3. Selection-modifying uncertainties (e.g. JEC) +Propagate through the entire pipeline. Require: +1. `cfg.add_shift(...)` for up/down variants with `tags={"selection_dependent"}` +2. Column aliases for varied kinematic columns +3. The `Selector` (and upstream `Calibrator`) must declare the shifts in their `shifts` set +4. Separate task branches are run for each shift + +### 4. Dedicated-dataset uncertainties (e.g. tune, hdamp) +A completely separate dataset is processed for each variation. The dataset `info` dict key must match the shift name. diff --git a/docs/columnflow_claude_guide/02_coding_style.md b/docs/columnflow_claude_guide/02_coding_style.md new file mode 100644 index 000000000..233036555 --- /dev/null +++ b/docs/columnflow_claude_guide/02_coding_style.md @@ -0,0 +1,403 @@ +# 02 — Columnflow Coding Style + +## Module Layout + +A typical analysis module tree mirrors columnflow's own structure: + +``` +myanalysis/ +├── analysis/ +│ └── my_analysis.py # Analysis + Campaign + Config definitions +├── config/ +│ ├── processes.py # order.Process definitions +│ ├── datasets.py # Campaign + DatasetInfo definitions +│ ├── variables.py # order.Variable definitions +│ └── categories.py # Category + Categorizer definitions +├── calibration/ +│ └── default.py # Calibrator definitions +├── selection/ +│ ├── default.py # Exposed Selector +│ ├── objects.py # Internal object-selection Selectors +│ ├── trigger.py # Trigger Selector +│ └── stats.py # increment_stats Selector +├── production/ +│ ├── default.py # Main exposed Producer +│ ├── weights.py # Weight Producers +│ └── features.py # High-level variable Producers +├── categorization/ +│ └── example.py # Categorizer definitions +├── histogramming/ +│ └── example.py # HistProducer definitions +├── inference/ +│ └── example.py # InferenceModel definitions +└── tasks/ + └── ... # Custom law tasks +``` + +--- + +## Import Conventions + +Always use `maybe_import` for packages not available in the default (non-columnar) sandbox. Only `numpy` and `awkward` are safe to import at module level via `maybe_import`. Never import `coffea` at module level. + +```python +# Standard columnflow imports — safe at module level +from columnflow.production import Producer, producer +from columnflow.selection import Selector, SelectionResult, selector +from columnflow.calibration import Calibrator, calibrator +from columnflow.util import maybe_import, four_vec, DotDict +from columnflow.columnar_util import set_ak_column, optional_column, has_ak_column, EMPTY_FLOAT, EMPTY_INT + +# Deferred heavy imports — use maybe_import +np = maybe_import("numpy") +ak = maybe_import("awkward") + +# coffea must ONLY be imported inside the function body, never at module level +def my_function(...): + import coffea + import coffea.nanoevents.methods.nanoaod +``` + +--- + +## TAF Decorator Pattern + +All five TAF types share the same decorator pattern. The decorator registers `uses`, `produces`, and optional metadata on the class. + +### Producer + +```python +from columnflow.production import Producer, producer +from columnflow.util import maybe_import +from columnflow.columnar_util import set_ak_column + +ak = maybe_import("awkward") +np = maybe_import("numpy") + +@producer( + uses={"Jet.pt", "Jet.eta"}, # columns to read from parquet + produces={"ht", "n_jet"}, # columns to write to parquet +) +def my_producer(self: Producer, events: ak.Array, **kwargs) -> ak.Array: + events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) + events = set_ak_column(events, "n_jet", ak.sum(events.Jet.pt > 0, axis=1)) + return events +``` + +### Selector + +```python +from columnflow.selection import Selector, SelectionResult, selector +from columnflow.util import maybe_import + +ak = maybe_import("awkward") + +@selector( + uses={"Jet.pt", "Jet.eta"}, + produces=set(), + exposed=True, # True = reachable from CLI --selector +) +def my_selector( + self: Selector, + events: ak.Array, + stats: dict, + **kwargs, +) -> tuple[ak.Array, SelectionResult]: + jet_mask = (events.Jet.pt > 25.0) & (abs(events.Jet.eta) < 2.4) + jet_sel = ak.sum(jet_mask, axis=1) >= 2 + jet_indices = ak.local_index(events.Jet.pt)[jet_mask] + + return events, SelectionResult( + steps={"jet": jet_sel}, + objects={"Jet": {"Jet": jet_indices}}, + aux={"jet_mask": jet_mask}, + ) +``` + +### Calibrator + +```python +from columnflow.calibration import Calibrator, calibrator +from columnflow.util import maybe_import +from columnflow.columnar_util import set_ak_column + +ak = maybe_import("awkward") + +@calibrator( + uses={"Jet.pt"}, + produces={"Jet.pt"}, +) +def my_calibrator(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array: + corrected_pt = events.Jet.pt * 1.02 # placeholder correction + events = set_ak_column(events, "Jet.pt", corrected_pt) + return events +``` + +--- + +## uses and produces Syntax + +`uses` and `produces` are sets of strings, other TAF instances, or brace-expanded patterns. + +```python +# Dot-notation for nested fields +uses={"Jet.pt", "Jet.eta", "Jet.phi", "Jet.mass"} + +# Brace expansion (bash-style) — equivalent to the above +uses={"Jet.{pt,eta,phi,mass}"} + +# Wildcard — read all columns of a collection +uses={"Jet.*"} + +# Include all uses/produces from another TAF (nested calls) +uses={other_producer} +produces={other_producer, "my_new_column"} + +# four_vec helper — expands to pt, eta, phi, mass for each collection +from columnflow.util import four_vec +uses=four_vec({"Jet", "Electron"}, {"btagDeepFlavB"}) +# equivalent to: Jet.{pt,eta,phi,mass,btagDeepFlavB}, Electron.{pt,eta,phi,mass} +``` + +--- + +## Nested TAF Calls + +To call another TAF from within a TAF: + +```python +events = self[other_producer](events, **kwargs) +events, sub_result = self[sub_selector](events, **kwargs) +``` + +The called TAF must be listed in the parent's `uses` and `produces` sets. + +--- + +## SelectionResult Structure + +```python +SelectionResult( + steps={ + "step_name": bool_mask_1d, # per-event boolean array + "other_step": bool_mask_1d, + }, + objects={ + "SourceField": { + "DestField": index_array, # ak.local_index(...)[mask] + }, + }, + aux={ + "any_key": any_value, # discarded after ReduceEvents + }, +) +``` + +The **exposed** selector must set `results.event` to the combined per-event boolean: + +```python +from operator import and_ +from functools import reduce + +results.event = reduce(and_, results.steps.values()) +``` + +Combine partial results with `+=`: + +```python +results = SelectionResult() +events, jet_results = self[jet_selector](events, **kwargs) +results += jet_results +``` + +--- + +## TAF Lifecycle Hooks + +Register additional behaviour on an existing TAF using decorator methods: + +```python +@my_producer.pre_init +def my_producer_pre_init(self: Producer) -> None: + # called before dependency tree; can set deps_kwargs + ... + +@my_producer.init +def my_producer_init(self: Producer) -> None: + # dynamic uses/produces/shifts registration + if self.dataset_inst.is_mc: + self.uses.add("mc_weight") + +@my_producer.post_init +def my_producer_post_init(self: Producer, task) -> None: + # first hook with access to `task` + ... + +@my_producer.requires +def my_producer_requires(self: Producer, task, reqs: dict) -> None: + # add extra law task requirements + from columnflow.tasks.external import BundleExternalFiles + reqs["ext_files"] = BundleExternalFiles.req(task) + +@my_producer.setup +def my_producer_setup(self: Producer, task, reqs: dict, inputs: dict, reader_targets: dict) -> None: + # load external resources (e.g. scale factor files) onto self + bundle = inputs["ext_files"] + self.sf_file = bundle["collection"][0]["muon_sf"].load() + +@my_producer.teardown +def my_producer_teardown(self: Producer, task) -> None: + # free memory + del self.sf_file +``` + +--- + +## Accessing Config, Dataset, Analysis in a TAF + +Inside any TAF function or hook, three objects are always available as `self` attributes: + +```python +self.config_inst # order.Config +self.dataset_inst # order.Dataset +self.analysis_inst # order.Analysis + +# Common patterns +if self.dataset_inst.is_mc: + ... +year = self.config_inst.campaign.x.year +lumi = self.config_inst.x.luminosity.nominal +``` + +--- + +## set_ak_column — The Only Way to Modify Events + +Never modify `events` fields directly. Always use `set_ak_column` and reassign: + +```python +from columnflow.columnar_util import set_ak_column +import numpy as np + +# Scalar column +events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) + +# Nested column +events = set_ak_column(events, "Jet.pt_corr", events.Jet.pt * 1.02) +``` + +--- + +## Coffea Behavior (Lorentz Vectors) + +To perform four-vector arithmetic, attach coffea behaviour first: + +```python +from columnflow.production.util import attach_coffea_behavior + +events = self[attach_coffea_behavior](events, **kwargs) +# Now events.Jet, events.Electron etc. support .mass, .delta_r(), etc. + +# Custom collection type +events = self[attach_coffea_behavior]( + events, + collections={"MyJets": {"type_name": "Jet"}}, + **kwargs, +) +``` + +`attach_coffea_behavior` must be in `uses` of the TAF that calls it. + +--- + +## Handling Missing / Optional Values + +```python +from columnflow.columnar_util import EMPTY_FLOAT, EMPTY_INT, optional_column + +# Access Jet.pt[:,2] safely (EMPTY_FLOAT for events with < 3 jets) +from columnflow.columnar_util import Route +jet3_pt = Route("Jet.pt[:,2]").apply(events, null_value=EMPTY_FLOAT) + +# Pad to fixed length, fill with dict of defaults +padded = ak.fill_none(ak.pad_none(events.Jet, 4, axis=1), {"pt": -999.0, "eta": -999.0}) + +# Optional column in uses (present in some datasets only) +uses={optional_column("veto")} + +# Check before using +from columnflow.columnar_util import has_ak_column +if has_ak_column(events, "veto"): + events = events[~events.veto] +``` + +--- + +## DeferredColumn — Campaign-Dependent Column Requirements + +```python +from columnflow.columnar_util import deferred_column + +@deferred_column +def IF_RUN3(self, func) -> object: + if func.config_inst.campaign.x.year >= 2022: + return super(IF_RUN3, self).__call__(func) + return None + +@producer( + uses={"common_col", IF_RUN3("run3_only_col")}, +) +def my_prod(self, events, **kwargs): + ... +``` + +--- + +## Naming Conventions + +| Object | Convention | Example | +|---|---|---| +| TAF function | `snake_case` | `jet_selection`, `event_weights` | +| Exposed TAF | `snake_case` | `default` (standard name for the main TAF) | +| Column names | `snake_case` | `"ht"`, `"n_jet"`, `"Jet.pt"` | +| Config name | `snake_case` with year | `"run3_2022"`, `"l18"` (local 2018) | +| Version strings | `snake_case` or `v1` | `"dev1"`, `"prod1"`, `"selection_v2"` | +| Dataset names | `snake_case` with tag | `"tt_dl_powheg"`, `"data_mu_b"` | +| Process names | `snake_case` | `"tt"`, `"dy_lep"`, `"wjets"` | + +--- + +## stats.json — Selection Statistics + +The exposed Selector updates a `stats: dict` (or `defaultdict(float)`) in place. These keys are printed by `cf.SelectEvents`: + +```python +stats["num_events"] += len(events) +stats["num_events_selected"] += ak.sum(results.event, axis=0) +stats["sum_mc_weight"] += ak.sum(events.mc_weight) +stats["sum_mc_weight_selected"] += ak.sum(events.mc_weight[results.event]) +``` + +Use columnflow's built-in helper `increment_stats` where possible (from `columnflow.selection.stats`). + +--- + +## Data vs. Monte Carlo Branching + +Always guard MC-only operations: + +```python +if self.dataset_inst.is_mc: + events = self[mc_weight](events, **kwargs) + events = self[pileup_weight](events, **kwargs) +``` + +For producers that should only run on MC, use the `init` hook to conditionally add columns: + +```python +@my_producer.init +def my_producer_init(self: Producer) -> None: + if self.dataset_inst.is_mc: + self.uses.add("GenJet.pt") + self.produces.add("n_gen_jet") +``` diff --git a/docs/columnflow_claude_guide/03_task_commands.md b/docs/columnflow_claude_guide/03_task_commands.md new file mode 100644 index 000000000..9174c1fd0 --- /dev/null +++ b/docs/columnflow_claude_guide/03_task_commands.md @@ -0,0 +1,358 @@ +# 03 — Task Commands + +All tasks are invoked with `law run cf.`. The `--version` flag is always required. `--config` and `--analysis` can be set in `law.cfg` as defaults. + +Use `--help` on any task for the full parameter list: +```bash +law run cf.SelectEvents --help +``` + +--- + +## Universal Flags + +| Flag | Purpose | +|---|---| +| `--version ` | Version tag for output paths (required) | +| `--config ` | Config name (default: `law.cfg default_config`) | +| `--analysis ` | Analysis module path (default: `law.cfg default_analysis`) | +| `--branch ` | Run a single branch (file) instead of all | +| `--print-status -1` | Show task tree + output existence (recursive) | +| `--print-output ` | Show output file paths at given depth | +| `--remove-output ,[,y]` | Remove outputs; modes: `a` all, `i` interactive, `d` dry-run | +| `--workflow htcondor` | Submit to HTCondor | +| `--workflow slurm` | Submit to Slurm | +| `--pilot` | Run one branch to check before launching all | + +--- + +## Step-by-step Commands + +### GetDatasetLFNs + +Resolves logical file names for a dataset (requires GRID proxy for CMS data). + +```bash +law run cf.GetDatasetLFNs \ + --version dev1 \ + --dataset tt_dl_powheg +``` + +--- + +### CalibrateEvents + +Applies calibrations (e.g. jet energy corrections). Runs per file/branch. + +```bash +law run cf.CalibrateEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --calibrators default \ + --shift nominal + +# Single branch test +law run cf.CalibrateEvents --version dev1 --dataset tt_dl_powheg --branch 0 + +# With a systematic shift +law run cf.CalibrateEvents --version dev1 --dataset tt_dl_powheg --shift jec_up +``` + +--- + +### SelectEvents + +Runs the Selector; produces event/object masks (parquet) and selection statistics (json). + +```bash +law run cf.SelectEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --selector default \ + --calibrators default + +# Verify outputs recursively +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --print-status -1 + +# Cutflow plot requires this first (SelectEvents at cutflow step) +law run cf.SelectEvents --version dev1 --config l18 +``` + +--- + +### ReduceEvents + +Applies selection masks to columns; writes the reduced parquet files. + +```bash +law run cf.ReduceEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --selector default \ + --calibrators default + +# Single branch (first file) — quick local test +law run cf.ReduceEvents --version dev1 --branch 0 +``` + +--- + +### MergeReducedEvents + +Merges per-file reduced parquet files into per-dataset merged files. + +```bash +law run cf.MergeReducedEvents \ + --version dev1 \ + --dataset tt_dl_powheg +``` + +--- + +### ProduceColumns + +Runs Producer(s) on the merged reduced events; writes additional column parquet files. + +```bash +law run cf.ProduceColumns \ + --version dev1 \ + --dataset tt_dl_powheg \ + --producers default + +# Multiple producers (comma-separated, order matters for column overwriting) +law run cf.ProduceColumns \ + --version dev1 \ + --dataset tt_dl_powheg \ + --producers default,extra_features +``` + +--- + +### CreateHistograms + +Fills histograms for the given variables, categories, and shift. + +```bash +law run cf.CreateHistograms \ + --version dev1 \ + --dataset tt_dl_powheg \ + --variables jet1_pt,n_jet \ + --categories incl \ + --shift nominal \ + --producers default \ + --hist-producer default + +# All shifts at once (nominal + all registered shifts) +law run cf.CreateHistograms \ + --version dev1 \ + --datasets tt_dl_powheg,wjets_madgraph \ + --variables jet1_pt \ + --shifts nominal,jec_up,jec_down,mu_up,mu_down +``` + +--- + +### MergeHistograms + +Merges per-dataset histogram files. + +```bash +law run cf.MergeHistograms \ + --version dev1 \ + --variables jet1_pt \ + --processes tt,wjets +``` + +--- + +### PlotVariables1D + +Creates 1D variable plots per process, stacked or overlaid. + +```bash +law run cf.PlotVariables1D \ + --version dev1 \ + --processes tt,wjets,data \ + --variables jet1_pt,n_jet \ + --categories incl \ + --shift nominal + +# Save to PDF and PNG +law run cf.PlotVariables1D \ + --version dev1 \ + --processes tt \ + --variables jet1_pt \ + --file-types pdf,png + +# With custom plot function +law run cf.PlotVariables1D \ + --version dev1 \ + --processes tt \ + --variables jet1_pt \ + --plot-function myanalysis.plotting.my_plot_func + +# Open plots automatically +law run cf.PlotVariables1D \ + --version dev1 \ + --processes tt \ + --variables jet1_pt \ + --view-cmd evince +``` + +--- + +### PlotVariables2D + +Creates 2D plots. + +```bash +law run cf.PlotVariables2D \ + --version dev1 \ + --processes tt \ + --variables jet1_pt__jet1_eta \ + --categories incl +``` + +--- + +### PlotShiftedVariables1D + +Shows nominal + up/down shift variations for a variable. + +```bash +law run cf.PlotShiftedVariables1D \ + --version dev1 \ + --processes tt \ + --variables jet1_pt \ + --shift-sources jec,mu \ + --categories incl +``` + +--- + +### PlotCutflow + +Plots total event yield at each selection step. + +```bash +law run cf.PlotCutflow \ + --version dev1 \ + --datasets tt_dl_powheg \ + --selector-steps muon,jet,bjet +``` + +--- + +### PlotCutflowVariables1D + +Plots a variable distribution at each selector step. + +```bash +law run cf.PlotCutflowVariables1D \ + --version dev1 \ + --datasets tt_dl_powheg \ + --processes tt \ + --variables genTop_pt \ + --categories incl \ + --skip-ratio +``` + +--- + +### CreateDatacards + +Produces CMS combine-compatible datacards and ROOT shape files. + +```bash +law run cf.CreateDatacards \ + --version dev1 \ + --inference-model example \ + --variables jet1_pt \ + --categories incl +``` + +--- + +## Useful Utility Commands + +```bash +# Check if a task's output already exists +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --print-status 0 + +# Remove just the histogram task output and re-run it +law run cf.CreateHistograms --version dev1 --variables jet1_pt --remove-output 0,a,y + +# List all available law tasks +law index --verbose + +# Get help on a specific task +law run cf.ProduceColumns --help + +# Run inside the columnar sandbox manually (for debugging) +cf_sandbox venv_columnar_dev bash +``` + +--- + +## Remote Execution (HTCondor) + +```bash +# Submit SelectEvents for all files of a dataset to HTCondor +law run cf.SelectEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --workflow htcondor + +# Monitor running jobs +law run cf.SelectEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --workflow htcondor \ + --print-status -1 +``` + +--- + +## Pinning Upstream Task Versions + +Use task-family-specific version flags to pin upstream outputs: + +```bash +# Use v1 outputs from CalibrateEvents while running SelectEvents as v2 +law run cf.SelectEvents \ + --version v2 \ + --cf.CalibrateEvents-version v1 \ + --dataset tt_dl_powheg +``` + +--- + +## Dataset / Process Wildcards + +Many tasks accept glob patterns: + +```bash +# All tt datasets +law run cf.PlotVariables1D --version dev1 --datasets "tt*" --variables jet1_pt + +# Multiple processes +law run cf.PlotVariables1D --version dev1 --processes "tt,wjets,dy*" --variables n_jet +``` + +--- + +## Full Pipeline — Single Command Chain + +Running the plotting task automatically triggers all upstream tasks: + +```bash +# This single command triggers the full pipeline for one dataset + one variable +law run cf.PlotVariables1D \ + --version dev1 \ + --datasets tt_dl_powheg \ + --variables jet1_pt \ + --processes tt \ + --categories incl \ + --shift nominal \ + --branch 0 # single branch for a quick test +``` diff --git a/docs/columnflow_claude_guide/04_awkward_reference.md b/docs/columnflow_claude_guide/04_awkward_reference.md new file mode 100644 index 000000000..2dfb82fc2 --- /dev/null +++ b/docs/columnflow_claude_guide/04_awkward_reference.md @@ -0,0 +1,305 @@ +# 04 — Awkward Array Reference for Columnflow + +Columnflow event data is stored as `ak.Array` named `events`. Fields like `events.Jet` are collections (variable-length sub-arrays); scalar fields like `events.MET.pt` are uniform arrays over events. + +Always import awkward via `maybe_import`: +```python +from columnflow.util import maybe_import +ak = maybe_import("awkward") +np = maybe_import("numpy") +``` + +--- + +## Accessing Fields + +```python +# Scalar per-event field +events.event # event number, shape (N,) +events.MET.pt # MET pT, shape (N,) + +# Collection field — variable-length inner dimension +events.Jet.pt # shape (N, var) +events.Electron.eta # shape (N, var) + +# Brace expansion works in uses/produces but NOT in ak access — use dot notation +events.Jet.pt # correct +events["Jet", "pt"] # also valid +``` + +--- + +## Indexing and Slicing + +```python +# First jet in every event (raises if any event has 0 jets) +events.Jet.pt[:, 0] + +# Safe nth-object access with fill for missing +from columnflow.columnar_util import Route, EMPTY_FLOAT +jet3_pt = Route("Jet.pt[:, 2]").apply(events, null_value=EMPTY_FLOAT) + +# Select specific events +selected_events = events[event_mask] # event_mask is bool, shape (N,) + +# Select objects within events using an index array +jet_indices = ak.local_index(events.Jet.pt)[jet_mask] +selected_jets = events.Jet[jet_indices] +``` + +--- + +## Boolean Masks + +```python +# Object-level mask (same shape as the collection) +jet_mask = (events.Jet.pt > 25.0) & (abs(events.Jet.eta) < 2.4) # shape (N, var) + +# Event-level mask derived from object counts +event_mask = ak.sum(jet_mask, axis=1) >= 2 # shape (N,) + +# Multi-condition +mask = ( + (events.Electron.pt > 15.0) & + (abs(events.Electron.eta) < 2.5) & + (events.Electron.cutBased >= 3) +) + +# Invert +not_mask = ~mask +``` + +--- + +## Reduction Operations (axis=1 for per-event sums) + +```python +# Sum over objects in each event +ht = ak.sum(events.Jet.pt, axis=1) # shape (N,) +n_jet = ak.sum(events.Jet.pt > 0, axis=1) # count jets (bool sum) +n_bjet = ak.sum(events.Jet.btagDeepFlavB > 0.5, axis=1) + +# Min / max per event +max_jet_pt = ak.max(events.Jet.pt, axis=1) +min_jet_eta = ak.min(abs(events.Jet.eta), axis=1) + +# Any / all per event +has_forward_jet = ak.any(abs(events.Jet.eta) > 2.4, axis=1) +all_jets_central = ak.all(abs(events.Jet.eta) < 2.4, axis=1) + +# Mean per event +mean_jet_pt = ak.mean(events.Jet.pt, axis=1) + +# axis=0 for global sums (used in stats) +total_mc_weight = ak.sum(events.mc_weight) +``` + +--- + +## Padding and Filling None + +```python +# Pad to at least N objects, filling with None +padded = ak.pad_none(events.Jet, 4, axis=1) # shape (N, >=4) with None + +# Fill None with scalar +filled = ak.fill_none(padded.pt, EMPTY_FLOAT) + +# Fill None with dict (for record arrays) +fill_values = {"pt": -999.0, "eta": -999.0, "phi": -999.0, "mass": -999.0} +filled = ak.fill_none(ak.pad_none(events.Lepton, 2, axis=1), fill_values) + +# Access padded element safely +lepton0_pt = ak.fill_none(ak.pad_none(events.Lepton.pt, 1, axis=1)[:, 0], EMPTY_FLOAT) + +# Check for None +is_missing = ak.is_none(events.Jet.pt, axis=1) +``` + +--- + +## Sorting and Argsort + +```python +# Sort jets by pT descending +sorted_jets = events.Jet[ak.argsort(events.Jet.pt, axis=1, ascending=False)] + +# Sort and take first N +leading_2_jets = events.Jet[ak.argsort(events.Jet.pt, axis=1, ascending=False)][:, :2] +``` + +--- + +## Concatenation and Combination + +```python +# Merge two collections (e.g. electrons and muons into leptons) +leptons = ak.concatenate([events.Electron, events.Muon], axis=1) +# Sort by pt descending +leptons = leptons[ak.argsort(leptons.pt, axis=1, ascending=False)] + +# Zip fields into a record array +dijet = ak.zip({ + "pt": events.Jet.pt[:, 0] + events.Jet.pt[:, 1], + "eta": events.Jet.eta[:, 0], +}) +``` + +--- + +## Flattening + +```python +# Flatten all jets from all events into a 1D array (for histogram filling) +all_jet_pts = ak.flatten(events.Jet.pt) # shape (sum_of_jet_counts,) + +# Flatten with axis=None (full recursive flatten) +flat = ak.flatten(events.Jet.pt, axis=None) +``` + +--- + +## Type Conversions + +```python +# Convert to numpy (only works on regular arrays — use after flatten/reduction) +arr_np = ak.to_numpy(ak.sum(events.Jet.pt, axis=1)) + +# Cast types +events_int = ak.values_astype(events.Jet.nConstituents, np.int32) +events_float = ak.values_astype(events.Jet.pt, np.float32) +``` + +--- + +## Conditional Selection (where) + +```python +# Per-event conditional +selected_pt = ak.where(event_mask, events.Jet.pt[:, 0], EMPTY_FLOAT) + +# Per-object conditional +corrected_pt = ak.where(events.Jet.pt > 50, events.Jet.pt * 1.02, events.Jet.pt) +``` + +--- + +## local_index — Creating Index Arrays for SelectionResult.objects + +The standard way to create index arrays for `SelectionResult.objects`: + +```python +# Create indices of jets that pass the selection +jet_mask = (events.Jet.pt > 25.0) & (abs(events.Jet.eta) < 2.4) +jet_indices = ak.local_index(events.Jet.pt, axis=1)[jet_mask] + +# Use in SelectionResult +return events, SelectionResult( + objects={"Jet": {"Jet": jet_indices}}, +) +``` + +--- + +## set_ak_column — Writing Columns + +Never mutate `events` fields directly. Always use `set_ak_column`: + +```python +from columnflow.columnar_util import set_ak_column + +# Add a new scalar column +events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) + +# Add a nested column (inside an existing collection) +events = set_ak_column(events, "Jet.pt_corr", events.Jet.pt * 1.02) + +# Replace an existing collection field +events = set_ak_column(events, "Jet.btagDeepFlavB", corrected_btag) +``` + +--- + +## Lorentz Vectors (coffea behaviour) + +After attaching coffea behaviour, awkward records expose four-vector methods: + +```python +from columnflow.production.util import attach_coffea_behavior + +events = self[attach_coffea_behavior](events, **kwargs) + +# Four-momentum operations (only after attach_coffea_behavior) +jet1 = events.Jet[:, 0] +jet2 = events.Jet[:, 1] +dijet_mass = (jet1 + jet2).mass +delta_r = jet1.delta_r(jet2) +dijet_pt = (jet1 + jet2).pt +``` + +To build a Lorentz vector manually: + +```python +# Inside a function (NOT at module level) +import coffea.nanoevents.methods.vector + +vec = ak.zip( + {"pt": arr.pt, "eta": arr.eta, "phi": arr.phi, "mass": arr.mass}, + with_name="PtEtaPhiMLorentzVector", + behavior=coffea.nanoevents.methods.vector.behavior, +) +mass = (vec[:, 0] + vec[:, 1]).mass +``` + +--- + +## Common Patterns in Columnflow + +### Count objects after selection + +```python +jet_mask = (events.Jet.pt > 25) & (abs(events.Jet.eta) < 2.4) +n_jet = ak.sum(jet_mask, axis=1) # per-event count +events = set_ak_column(events, "n_jet", n_jet) +``` + +### Per-process event weights + +```python +# Sum of MC weights per process (used in stats) +for pid in np.unique(ak.to_numpy(events.process_id)): + process_mask = events.process_id == pid + stats[f"sum_mc_weight_per_process"][int(pid)] += ak.sum(events.mc_weight[process_mask]) +``` + +### HT (scalar sum of jet pT) + +```python +ht = ak.sum(events.Jet.pt, axis=1) +events = set_ak_column(events, "ht", ht, value_type=np.float32) +``` + +### Leading lepton pT + +```python +lep_pt = ak.fill_none(ak.pad_none(events.Lepton.pt, 1, axis=1)[:, 0], EMPTY_FLOAT) +events = set_ak_column(events, "lep1_pt", lep_pt, value_type=np.float32) +``` + +### Invariant mass of two objects + +```python +# Requires coffea behaviour to be attached +mll = (events.Lepton[:, 0] + events.Lepton[:, 1]).mass +z_veto = abs(mll - 91.2) < 15.0 +``` + +--- + +## Gotchas + +- `axis=1` operates over the **object** (inner) dimension; `axis=0` operates over **events** (outer). +- Arithmetic between collections with different lengths requires `ak.broadcast_arrays` or `ak.zip`. +- Never iterate over `ak.Array` events in Python — always use vectorized operations. +- `ak.Array` returned from a TAF is a copy (with shared underlying data buffers), not the same object as the input; this is intentional and efficient. +- `ak.to_numpy()` fails on variable-length (jagged) arrays — flatten first. diff --git a/docs/columnflow_claude_guide/05_dos_and_donts.md b/docs/columnflow_claude_guide/05_dos_and_donts.md new file mode 100644 index 000000000..4b8b3d3a8 --- /dev/null +++ b/docs/columnflow_claude_guide/05_dos_and_donts.md @@ -0,0 +1,310 @@ +# 05 — Do's and Don'ts for Claude in Columnflow Projects + +## Critical Rules + +### DO always use `set_ak_column` to modify event data + +```python +# CORRECT +events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1)) + +# WRONG — direct mutation is not supported and may silently fail +events.ht = ak.sum(events.Jet.pt, axis=1) +events["ht"] = ak.sum(events.Jet.pt, axis=1) +``` + +### DO always declare `uses` and `produces` in TAF decorators + +Columnflow uses these sets to track which columns are loaded from disk and which are written. Missing a column in `uses` means it won't be available. Missing it in `produces` means it won't be saved. + +```python +# CORRECT — both sets explicitly declared +@producer( + uses={"Jet.pt", "Jet.eta"}, + produces={"ht", "n_jet"}, +) +def my_producer(self, events, **kwargs): ... + +# WRONG — omitting uses/produces causes runtime failures +@producer() +def my_producer(self, events, **kwargs): ... +``` + +### DO use `maybe_import` for heavy packages + +```python +# CORRECT +ak = maybe_import("awkward") +np = maybe_import("numpy") + +# WRONG — direct import fails when columnflow loads the module in a non-columnar sandbox +import awkward as ak +import numpy as np +``` + +### DO import coffea inside the function, never at module level + +```python +# CORRECT +def my_function(self, events, **kwargs): + import coffea.nanoevents.methods.vector + vec = ak.zip({"pt": ...}, with_name="PtEtaPhiMLorentzVector", + behavior=coffea.nanoevents.methods.vector.behavior) + +# WRONG — coffea at module level breaks non-coffea sandboxes +import coffea +``` + +### DO register new Python files in `law.cfg` + +Every new file containing a Calibrator, Selector, Producer, Categorizer, or HistProducer must be added to the appropriate `*_modules` key in `law.cfg`. + +```ini +# CORRECT — after adding myanalysis/selection/jets.py +selection_modules: ..., myanalysis.selection.{default,jets} + +# No spaces after commas inside braces! +# WRONG +selection_modules: ..., myanalysis.selection.{default, jets} +``` + +### DO include sub-TAF's `uses`/`produces` by passing the TAF itself + +```python +# CORRECT — propagates all uses/produces from sub_producer +@producer( + uses={sub_producer, "extra_col"}, + produces={sub_producer, "my_new_col"}, +) +def parent(self, events, **kwargs): ... + +# WRONG — forgets to re-declare sub_producer columns +@producer( + uses={"extra_col"}, + produces={"my_new_col"}, +) +def parent(self, events, **kwargs): + events = self[sub_producer](events, **kwargs) # sub_producer's columns not declared! +``` + +### DO gate Monte Carlo logic on `self.dataset_inst.is_mc` + +```python +# CORRECT +if self.dataset_inst.is_mc: + events = self[mc_weight](events, **kwargs) + +# WRONG — mc_weight column does not exist in data +events = self[mc_weight](events, **kwargs) +``` + +### DO return `events` from every TAF function + +```python +# CORRECT — Producer returns events +def my_producer(self, events, **kwargs) -> ak.Array: + events = set_ak_column(events, "ht", ...) + return events + +# WRONG — missing return +def my_producer(self, events, **kwargs): + events = set_ak_column(events, "ht", ...) +``` + +### DO return `(events, SelectionResult)` from Selectors + +```python +# CORRECT +def my_selector(self, events, stats, **kwargs) -> tuple[ak.Array, SelectionResult]: + ... + return events, SelectionResult(steps={...}) + +# WRONG — wrong return type +def my_selector(self, events, stats, **kwargs): + ... + return events +``` + +### DO set `results.event` in the exposed Selector + +The `ReduceEvents` task reads `results.event` to filter events. Missing it causes no events to be selected (or an error). + +```python +from operator import and_ +from functools import reduce + +results.event = reduce(and_, results.steps.values()) +``` + +### DO add `keep_columns` in the config for columns created before ReduceEvents + +Any column produced by a Calibrator or Selector that should survive `ReduceEvents` must be explicitly listed in `cfg.x.keep_columns` under `"cf.ReduceEvents"`. + +```python +cfg.x.keep_columns = DotDict.wrap({ + "cf.ReduceEvents": { + "Jet.{pt,eta,phi,mass}", + "my_new_column", # REQUIRED if produced in Selector + ColumnCollection.ALL_FROM_SELECTOR, # convenient catch-all for selector columns + }, +}) +``` + +--- + +## Architecture Don'ts + +### DON'T call `law run` inside Python code + +Use the standard `law run cf.` CLI. If triggering a task programmatically from a script is truly necessary, use the pattern in `best_practices.md` with `task.law_run()`. + +### DON'T put analysis logic in `CLAUDE.md` or config files + +Analysis code (Selectors, Producers, etc.) belongs in the appropriate subdirectory, not in configuration files or documentation. + +### DON'T create a new exposed Selector unless truly needed + +One exposed (top-level) Selector should be the entry point, composed of many internal Selectors. Creating multiple exposed Selectors for the same purpose leads to maintenance burden. + +### DON'T mix column-level and event-level operations without explicit axis + +```python +# WRONG — ambiguous without axis +n_jet = ak.sum(events.Jet.pt > 0) # sums over EVERYTHING + +# CORRECT +n_jet = ak.sum(events.Jet.pt > 0, axis=1) # per-event count +``` + +### DON'T use Python loops over events + +```python +# WRONG — extremely slow; defeats the columnar paradigm +for event in events: + ht = sum(jet.pt for jet in event.Jet) + +# CORRECT +ht = ak.sum(events.Jet.pt, axis=1) +``` + +### DON'T modify `events` returned from a sub-TAF call without reassigning + +```python +# WRONG — changes are discarded +self[other_producer](events, **kwargs) + +# CORRECT +events = self[other_producer](events, **kwargs) +``` + +--- + +## Shift / Systematics Don'ts + +### DON'T hardcode weight column names in Producers + +Use the column alias mechanism so that shifted variants are used automatically when a shift is active: + +```python +# WRONG — always uses the nominal weight +weight = events.muon_weight + +# CORRECT — alias resolves to muon_weight_up/down under the corresponding shift +weight = events[self.config_inst.x.column_aliases.get("muon_weight", "muon_weight")] +``` + +Or rely on the `WeightProducer`/`HistProducer` shift mechanism rather than doing manual name resolution. + +### DON'T forget to declare shifts in the TAF's `shifts` set + +For weight-based and selection-modifying uncertainties, the TAF must declare which shifts it is sensitive to: + +```python +@my_producer.init +def my_producer_init(self): + from columnflow.config_util import get_shifts_from_sources + self.shifts |= get_shifts_from_sources(self.config_inst, "mu") +``` + +--- + +## Config Don'ts + +### DON'T add whitespace after commas inside `{}` brace expansions in law.cfg + +```ini +# CORRECT +selection_modules: myanalysis.selection.{default,jets,trigger} + +# WRONG — spaces break the brace expansion parser +selection_modules: myanalysis.selection.{default, jets, trigger} +``` + +### DON'T duplicate config entries across Analysis and Config + +Put analysis-independent information in the `Campaign` object (year, ecm, tier). Put analysis-specific information in `Config`. Do not copy the same value into both. + +### DON'T put large data arrays into Config auxiliaries + +`cfg.x.*` is for lightweight metadata (names, thresholds, lookup tables). Heavy data (scale factor histograms, b-tag efficiency arrays) should be loaded from external files at TAF setup time via the `requires`/`setup` hooks. + +--- + +## Style Don'ts + +### DON'T import at the function call site what should be a module-level `maybe_import` + +```python +# WRONG — hides the deferred import intent +def my_func(self, events, **kwargs): + import awkward as ak + ... + +# CORRECT — at module level +ak = maybe_import("awkward") +``` + +### DON'T use `ak.Array` type annotations without the `maybe_import` guard + +If `awkward` is not available at import time, type annotations that reference `ak.Array` will fail. Use string annotations or guard them: + +```python +# SAFE with maybe_import at module level +ak = maybe_import("awkward") + +def my_producer(self: Producer, events: ak.Array, **kwargs) -> ak.Array: ... +``` + +### DON'T return early from a Producer without returning `events` + +```python +# WRONG +if condition: + return # forgot to return events! +events = set_ak_column(events, "col", value) +return events + +# CORRECT +if condition: + return events +events = set_ak_column(events, "col", value) +return events +``` + +--- + +## Summary Checklist + +When writing or reviewing columnflow code, verify: + +- [ ] All heavy packages use `maybe_import` at module level +- [ ] `coffea` is imported inside function bodies only +- [ ] Every column read is in `uses`; every column written is in `produces` +- [ ] `set_ak_column` is used (never direct field assignment) +- [ ] `events` is reassigned after every `set_ak_column` or sub-TAF call +- [ ] Selectors return `(events, SelectionResult)` with `results.event` set +- [ ] MC-only logic is guarded by `self.dataset_inst.is_mc` +- [ ] New files are registered in `law.cfg` under the correct `*_modules` key +- [ ] Columns created before `ReduceEvents` are in `cfg.x.keep_columns` +- [ ] No brace-expansion whitespace in `law.cfg` +- [ ] No Python loops over event arrays (use vectorized `ak.*` operations) From 3326816f6f939cc94657e71ddce5b7e76add3524 Mon Sep 17 00:00:00 2001 From: juvanden Date: Wed, 24 Jun 2026 11:30:37 +0200 Subject: [PATCH 2/7] docs: expand columnflow Claude guide with user-facing documentation Rewrites and extends the docs/columnflow_claude_guide/ to serve as both persistent Claude instructions and a proper user guide: - 01_framework_structure.md: adds "big picture" section explaining why columnflow exists, the three pillars (law/order/awkward), a plain-English walkthrough of every pipeline step, and annotated config-object examples including systematic uncertainty classes - 03_task_commands.md: rewritten as a per-task user guide with explanations of what each task does, its outputs, key parameters, multiple command examples, and practical tips for debugging and re-running - 05_dos_and_donts.md: rewritten for new users with correct/wrong code pairs covering set_ak_column usage, vectorized operations, imports, Selector return types, keep_columns, law.cfg registration, versioning, and systematics - 06_custom_tasks.md: new file with full worked examples of writing Selectors (object-level, event-level, composed exposed), Producers (features, weights), Calibrators with lifecycle hooks, Categorizers, HistProducers, InferenceModels, and custom law Tasks; includes common patterns from the hh2bbww analysis - CLAUDE.md: updated table of contents to include the new 06 guide Co-Authored-By: Claude Sonnet 4.6 --- CLAUDE.md | 7 +- .../01_framework_structure.md | 422 ++++++----- .../03_task_commands.md | 397 ++++++---- .../05_dos_and_donts.md | 383 +++++----- .../06_custom_tasks.md | 701 ++++++++++++++++++ 5 files changed, 1417 insertions(+), 493 deletions(-) create mode 100644 docs/columnflow_claude_guide/06_custom_tasks.md diff --git a/CLAUDE.md b/CLAUDE.md index 2a136f3cd..93bbd4e90 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,11 +6,12 @@ Read all five guide files before writing or modifying any code: | Guide | Contents | |---|---| -| [01 Framework Structure](docs/columnflow_claude_guide/01_framework_structure.md) | Analysis pipeline, task graph, config objects (Analysis/Campaign/Config), Order objects | +| [01 Framework Structure](docs/columnflow_claude_guide/01_framework_structure.md) | Big picture, analysis pipeline, task graph, config objects (Analysis/Campaign/Config), systematics | | [02 Coding Style](docs/columnflow_claude_guide/02_coding_style.md) | TAF decorators, uses/produces, imports, module registration, naming conventions | -| [03 Task Commands](docs/columnflow_claude_guide/03_task_commands.md) | `law run` commands with all standard parameters for every task | +| [03 Task User Guide](docs/columnflow_claude_guide/03_task_commands.md) | Per-task explanation, `law run` commands, key parameters, tips | | [04 Awkward Reference](docs/columnflow_claude_guide/04_awkward_reference.md) | Key `ak.*` functions used in columnflow code | -| [05 Dos and Don'ts](docs/columnflow_claude_guide/05_dos_and_donts.md) | Explicit rules Claude must follow when writing columnflow code | +| [05 Dos and Don'ts](docs/columnflow_claude_guide/05_dos_and_donts.md) | Common mistakes and best practices for new users | +| [06 Custom Tasks](docs/columnflow_claude_guide/06_custom_tasks.md) | Writing Selectors, Producers, Calibrators, HistProducers, Categorizers, custom law tasks | ## Quick orientation diff --git a/docs/columnflow_claude_guide/01_framework_structure.md b/docs/columnflow_claude_guide/01_framework_structure.md index f7c23b96d..381509465 100644 --- a/docs/columnflow_claude_guide/01_framework_structure.md +++ b/docs/columnflow_claude_guide/01_framework_structure.md @@ -1,135 +1,207 @@ -# 01 — Columnflow Framework Structure +# 01 — Columnflow Framework: The Big Picture -## Overview +## What is columnflow? -Columnflow is a backend for columnar, fully orchestrated HEP analyses in pure Python. +Columnflow is a Python framework for **fully orchestrated, columnar High Energy Physics analyses**. It sits between raw NanoAOD ROOT files and final statistical results (plots, datacards), handling the entire processing chain automatically. -| Layer | Package | Role | -|---|---|---| -| Workflow orchestration | [law](https://github.com/riga/law) | Task graph, CLI, remote execution | -| Metadata / bookkeeping | [order](https://github.com/riga/order) | Analysis, Campaign, Config, Dataset, Process, Variable, Shift | -| Columnar data | [awkward-array](https://awkward-array.org) | Event arrays; all physics data | -| Histogramming | [Hist](https://hist.readthedocs.io) | Histogram objects produced by `CreateHistograms` | -| Physics objects | [coffea](https://coffeateam.github.io/coffea/) | Lorentz vector behaviour on awkward arrays | +It solves three practical problems every HEP analyst faces: + +| Problem | How columnflow solves it | +|---|---| +| **Scalability** — datasets contain billions of events that do not fit in memory | Processes events in chunks; submits to HTCondor/Slurm automatically | +| **Reproducibility** — analysts need to re-run with different settings and track intermediate results | Every task output is versioned and cached; re-running only re-processes what changed | +| **Systematic uncertainties** — dozens of weight variations and kinematic corrections must be propagated consistently | First-class support for shifts; the task graph branches automatically per systematic | + +--- + +## The Three Pillars + +### 1. Law — Workflow orchestration + +[Law](https://github.com/riga/law) (built on [Luigi](https://luigi.readthedocs.io)) defines the **task graph**: every processing step is a `Task` object with declared outputs. When you run a downstream task (e.g. `PlotVariables1D`), law checks which upstream tasks are missing and runs them first — automatically. + +Key properties: +- Tasks are **idempotent**: if the output already exists, the task is skipped. +- Tasks run **locally** or on **HTCondor / Slurm** with a single flag change (`--workflow htcondor`). +- The `--version` tag separates parallel analysis branches on disk — you can have `dev1` and `prod1` coexist. + +### 2. Order — Metadata management + +[Order](https://github.com/riga/order) provides Python objects for HEP bookkeeping: `Analysis`, `Campaign`, `Config`, `Dataset`, `Process`, `Variable`, `Shift`, `Category`. These objects replace scattered config files with a typed, queryable in-memory structure. Every analysis object (process colour, cross-section, dataset file count) lives here. + +### 3. Awkward-array — Columnar event data + +All event data is stored as [awkward arrays](https://awkward-array.org): irregular (jagged) arrays where each event can have a different number of jets, leptons, etc. Operations are vectorized — you never write a Python loop over events. The `events` array is the central object passed through the entire pipeline. --- -## Analysis Pipeline (Task Graph) +## The Analysis Pipeline -The standard linear pipeline from raw data to plots: +The standard pipeline processes data from ROOT files to plots and datacards: ``` -GetDatasetLFNs - │ -CalibrateEvents ← applies calibrations (e.g. JEC) - │ -SelectEvents ← creates event and object masks; produces stats.json - │ -ReduceEvents ← applies masks; writes reduced parquet files - │ -MergeReducedEvents ← merges per-file reduced files into one file per dataset - │ -ProduceColumns ← creates additional high-level columns - │ -CreateHistograms ← fills Hist histograms per variable / category / shift - │ -MergeHistograms ← merges per-dataset histograms - │ -MergeShiftedHistograms ← merges nominal + shifted histograms for inference - │ -PlotVariables1D ← produces matplotlib/mplhep plots -CreateDatacards ← produces CMS combine datacards + ROOT shape files +Raw NanoAOD ROOT files + │ + ▼ +GetDatasetLFNs Resolves file paths (LFNs) from DAS or custom source + │ + ▼ +CalibrateEvents Applies corrections (JEC, MET corrections, tau ES, ...) + │ + ▼ +SelectEvents Defines event and object masks; saves selection statistics + │ + ▼ +ReduceEvents Applies masks; writes reduced Parquet files (~100× smaller) + │ + ▼ +MergeReducedEvents Merges per-file outputs into one file per dataset + │ + ├──────────────────────────────────────────────────┐ + ▼ ▼ +ProduceColumns CreateHistograms (directly) +(creates extra columns) + │ + ▼ +CreateHistograms Fills Hist histograms for all variables / categories / shifts + │ + ▼ +MergeHistograms Merges histograms across branches (files) of a dataset + │ + ├────────────────────────┐ + ▼ ▼ + PlotVariables1D CreateDatacards + (matplotlib plots) (CMS combine format) ``` -Tasks with `cf.` prefix are columnflow built-ins; user-defined tasks live in `/tasks/`. +### What happens at each step + +**GetDatasetLFNs** — Queries CMS DAS (or a custom function) to find the ROOT file paths for each dataset. Saves them as a JSON file used by all downstream tasks. Requires a valid GRID proxy for CMS data. + +**CalibrateEvents** — Runs user-defined `Calibrator` objects on raw events. Calibrators add corrected columns (e.g. `Jet.pt` after JEC) without removing the originals. Multiple calibrators can run sequentially. For systematic variations that modify kinematics (e.g. JEC up/down), separate task branches are created. -### Key properties +**SelectEvents** — Runs one `Selector` on the calibrated events. The selector produces: +- Boolean event masks (one per selection step, e.g. `"trigger"`, `"muon"`, `"jet"`) +- Index arrays for object collections (e.g. which jets pass `pt > 25 GeV && |eta| < 2.4`) +- A `stats.json` file with event counts and MC weight sums (needed for normalization later) -- **Chunking**: events are processed in chunks of ≤100 000 events (set in `law.cfg` under `chunked_io_chunk_size`). -- **Parallelism**: each file / chunk is a separate law task branch; submit to HTCondor/Slurm with `--workflow htcondor`. -- **Reproducibility**: `--version` tags all intermediate outputs; different versions coexist on disk. -- **Columnar storage**: intermediate results are stored as **Parquet** files (events) or **pickle** files (histograms). +Masks are saved to Parquet — they are **not yet applied** here. + +**ReduceEvents** — Applies the masks from `SelectEvents` to all columns. Writes a compact Parquet file containing only selected events and selected objects. Columns not listed in `cfg.x.keep_columns["cf.ReduceEvents"]` are dropped here permanently. + +**MergeReducedEvents** — Merges the many small per-file Parquet outputs into larger files per dataset, targeting a configurable size (default ~512 MB, set via `cfg.x.reduced_file_size`). + +**ProduceColumns** — Runs `Producer` objects on the merged reduced events to create new high-level columns: `ht`, `n_bjet`, `mll`, category IDs, event weights, ML scores, etc. These are saved in separate Parquet files that are transparently merged with the reduced events by downstream tasks. + +**CreateHistograms** — Fills `Hist` histograms for all requested variables, categories, and systematic shifts. The `HistProducer` controls event weighting and histogram filling logic. Each histogram is labelled by dataset, shift, and category. + +**MergeHistograms** — Merges histograms from all branches (files) of a dataset. A second merge step (`MergeShiftedHistograms`) further combines nominal and shifted histograms across all datasets for a given process, ready for plotting or inference. + +**PlotVariables1D / 2D** — Creates matplotlib/mplhep-styled plots from merged histograms. Supports stacked MC + data overlays, ratio panels, and shifted variable comparisons. + +**CreateDatacards** — Produces CMS `combine`-compatible datacards (`.txt`) and shape ROOT files from the merged histograms, driven by an `InferenceModel` object. --- ## Five Task Array Function (TAF) Types -| TAF type | Class | CLI parameter | Quantity | Task | +User-defined code hooks into the pipeline through **Task Array Functions**. Each TAF type slots into one specific task: + +| TAF | Class | Task | CLI parameter | Quantity | |---|---|---|---|---| -| Calibrator | `Calibrator` | `--calibrators` | 0..N | `CalibrateEvents` | -| Selector | `Selector` | `--selector` | exactly 1 | `SelectEvents` | -| Reducer | `Reducer` | `--reducer` | exactly 1 | `ReduceEvents` | -| Producer | `Producer` | `--producers` | 0..N | `ProduceColumns` | -| HistProducer | `HistProducer` | `--hist-producer` | exactly 1 | `CreateHistograms` | +| Calibrator | `Calibrator` | `CalibrateEvents` | `--calibrators` | 0 or more | +| Selector | `Selector` | `SelectEvents` | `--selector` | exactly 1 | +| Reducer | `Reducer` | `ReduceEvents` | `--reducer` | exactly 1 | +| Producer | `Producer` | `ProduceColumns` | `--producers` | 0 or more | +| HistProducer | `HistProducer` | `CreateHistograms` | `--hist-producer` | exactly 1 | + +All TAFs share the same decorator pattern and lifecycle. See [02 Coding Style](02_coding_style.md) for the full pattern. + +--- -All TAFs share the same decorator pattern — see [02 Coding Style](02_coding_style.md). +## Directory Structure of an Analysis + +``` +myanalysis/ +├── analysis/ +│ └── my_analysis.py # Creates Analysis, Campaign, Config objects +├── config/ +│ ├── processes.py # Process definitions (name, xsec, colour) +│ ├── datasets.py # Dataset definitions (files, keys) +│ ├── variables.py # Variable definitions (binning, expression) +│ └── categories.py # Category and Categorizer definitions +├── calibration/ +│ └── jets.py # Custom Calibrator(s) +├── selection/ +│ ├── default.py # Exposed (top-level) Selector +│ ├── objects.py # Internal object-selection Selectors +│ └── stats.py # Stats-increment helper +├── production/ +│ ├── default.py # Main exposed Producer (called from CLI) +│ ├── weights.py # Event weight Producers +│ └── features.py # High-level variable Producers +├── categorization/ +│ └── categories.py # Categorizer definitions +├── histogramming/ +│ └── default.py # HistProducer +├── inference/ +│ └── default.py # InferenceModel for datacards +├── tasks/ # Custom analysis-specific law tasks +├── law.cfg # Workflow configuration (must register all modules) +└── setup.sh # Environment setup +``` --- -## Configuration Objects (order package) +## Configuration Objects (order) ### Hierarchy ``` -Analysis - └── Config (links Analysis + Campaign) - └── Campaign - └── Dataset → DatasetInfo (files, events, keys) +Analysis ─── top-level container (rarely holds analysis logic itself) + └── Config ─── analysis + campaign combination; carries all per-config settings + └── Campaign ─── experimental period (year, energy, tier, file locations) + └── Dataset ─── one Monte Carlo or data sample ``` ### Analysis -Top-level container; rarely holds analysis logic itself. - ```python import order as od -analysis = od.Analysis(name="my_analysis", id=1) +analysis = od.Analysis(name="hbw", id=1) ``` ### Campaign -Experiment-period-specific information (year, energy, tier, file locations). - ```python cpn = od.Campaign( - name="run3_2022", - id=1, - ecm=13.6, + name="run2_2018", + id=4, + ecm=13, aux={ "tier": "NanoAOD", - "year": 2022, - "location": "root://...", + "year": 2018, + "location": "root://xrootd-cms.infn.it//", }, ) ``` -### Config - -Analysis + campaign combination; carries all per-config parameters. - -```python -cfg = analysis.add_config(campaign, name="run3_2022", id=1) -``` - ### Process -Physical process with cross-section and plot metadata. - ```python from scinum import Number -proc = od.Process( + +tt = od.Process( name="tt", id=1000, label=r"$t\bar{t}$", - color=(128, 76, 153), - xsecs={13.6: Number(831.76, {"scale": (19.77, 29.20)})}, + color=(205, 0, 9), + xsecs={13: Number(831.76, {"scale": (19.77, 29.20), "pdf": 35.06})}, ) ``` ### Dataset -A sample linked to a Campaign and one or more Processes. - ```python cpn.add_dataset( name="tt_dl_powheg", @@ -137,142 +209,160 @@ cpn.add_dataset( processes=[procs.tt], info={ "nominal": od.DatasetInfo( - keys=["/TT.../NANOAODSIM"], + keys=["/TTTo2L2Nu.../RunIISummer20UL18NanoAODv9.../NANOAODSIM"], n_files=242, n_events=276079127, ), + # For dedicated-dataset systematics, add extra info entries: + "tune_up": od.DatasetInfo( + keys=["/TTTo2L2Nu_TuneUp.../NANOAODSIM"], + n_files=30, + n_events=10000000, + ), }, ) ``` -### Variable +### Config -Defines the column expression and histogram binning for `CreateHistograms`. +The `Config` object links an Analysis and a Campaign and carries all per-config settings: ```python +cfg = analysis.add_config(campaign, name="run2_2018", id=4) + +# --- Required order objects --- +cfg.add_process(procs.tt) +cfg.add_dataset(campaign.get_dataset("tt_dl_powheg")) +cfg.add_shift(name="nominal", id=0) + +# Variable definition cfg.add_variable( name="jet1_pt", expression="Jet.pt[:,0]", null_value=EMPTY_FLOAT, binning=(40, 0.0, 400.0), unit="GeV", - x_title=r"Jet $p_{T}$", + x_title=r"Leading jet $p_T$", ) -``` -### Shift - -Systematic uncertainty variant (rate-only, weight-based, or dedicated dataset). +# Category +from columnflow.config_util import add_category +add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") -```python -cfg.add_shift(name="nominal", id=0) -cfg.add_shift(name="mu_up", id=1, type=od.Shift.SHAPE) -cfg.add_shift(name="mu_down", id=2, type=od.Shift.SHAPE) -``` +# --- Auxiliary configuration (cfg.x.*) --- +from columnflow.util import DotDict +from columnflow.columnar_util import ColumnCollection -### Category / Categorizer +cfg.x.luminosity = Number(59740, {"lumi_13TeV_2018": 0.025j}) -Analysis phase-space regions used in histogramming and plotting. +cfg.x.keep_columns = DotDict.wrap({ + "cf.ReduceEvents": { + "{Jet,FatJet}.{pt,eta,phi,mass,btagDeepFlavB}", + "Electron.{pt,eta,phi,mass,charge}", + "Muon.{pt,eta,phi,mass,charge}", + "MET.{pt,phi}", + "event", "run", "luminosityBlock", + ColumnCollection.ALL_FROM_SELECTOR, + }, +}) -```python -from columnflow.config_util import add_category -add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") +cfg.x.default_calibrator = "default" +cfg.x.default_selector = "default" +cfg.x.default_producer = "default" +cfg.x.default_variables = ("n_jet", "jet1_pt") ``` ---- - -## Important Config Auxiliaries - -These `cfg.x.*` entries have special meaning in columnflow: +### Important Config Auxiliaries | Key | Type | Purpose | |---|---|---| -| `cfg.x.keep_columns` | `DotDict` of sets | Which columns survive `ReduceEvents` | +| `cfg.x.keep_columns` | `DotDict` of sets | Columns that survive `ReduceEvents` | | `cfg.x.luminosity` | `scinum.Number` | Luminosity with uncertainties | | `cfg.x.external_files` | `DotDict` | Paths/URLs to external scale-factor files | | `cfg.x.get_dataset_lfns` | callable | Custom LFN-retrieval function | -| `cfg.x.default_calibrator` | str | Default `--calibrator` value | +| `cfg.x.default_calibrator` | str | Default `--calibrators` value | | `cfg.x.default_selector` | str | Default `--selector` value | -| `cfg.x.default_producer` | str/tuple | Default `--producer` value | +| `cfg.x.default_producer` | str/tuple | Default `--producers` value | | `cfg.x.default_variables` | tuple | Default `--variables` value | | `cfg.x.reduced_file_size` | float | Target merged-file size in MB | -| `cfg.x.versions` | dict | Pinned task versions | +| `cfg.x.versions` | dict | Pinned task versions (see best_practices) | -### keep_columns example +--- -```python -from columnflow.util import DotDict -from columnflow.columnar_util import ColumnCollection +## Systematic Uncertainties (Shifts) -cfg.x.keep_columns = DotDict.wrap({ - "cf.ReduceEvents": { - "{Jet,FatJet}.{pt,eta,phi,mass,btagDeepFlavB}", - "Electron.{pt,eta,phi,mass,charge}", - "Muon.{pt,eta,phi,mass,charge}", - "MET.{pt,phi}", - "event", "run", "luminosityBlock", - ColumnCollection.ALL_FROM_SELECTOR, - }, - "cf.ProduceColumns": { - "ht", "n_jet", - }, -}) +Columnflow has first-class support for three classes of systematics: + +### Rate-only uncertainties + +Affect the overall yield only (e.g. luminosity uncertainty). Defined entirely in the inference model — no extra task branches needed. + +```python +# In the inference model: +model.add_parameter("lumi", type="lnN", effect=1.025) ``` ---- +### Weight-based uncertainties (shape) + +Applied as event weight variations (e.g. muon scale factors, pile-up weights). They produce separate histograms for up/down variations without re-running the selection. + +```python +# In config: +cfg.add_shift(name="mu_up", id=1, type=od.Shift.SHAPE) +cfg.add_shift(name="mu_down", id=2, type=od.Shift.SHAPE) -## law.cfg Structure - -The `law.cfg` file drives the workflow. Critical sections: - -```ini -[analysis] -default_analysis: myanalysis.analysis.my_analysis.my_analysis -default_config: run3_2022 -default_dataset: tt_dl_powheg - -# Register all Python modules columnflow should know about -calibration_modules: columnflow.calibration.cms.{jets,met}, myanalysis.calibration.default -selection_modules: columnflow.selection.empty, columnflow.selection.cms.{json_filter,met_filters}, myanalysis.selection.default -reduction_modules: columnflow.reduction.default -production_modules: columnflow.production.{categories,normalization,processes}, myanalysis.production.default -categorization_modules: myanalysis.categorization.example -hist_production_modules: columnflow.histogramming.default, myanalysis.histogramming.example -inference_modules: columnflow.inference, myanalysis.inference.example - -[outputs] -# Map tasks to storage locations -cf.CalibrateEvents: local, /path/to/store -cf.SelectEvents: local, /path/to/store -cf.ReduceEvents: local, /path/to/store -cf.ProduceColumns: local, /path/to/store -cf.CreateHistograms: local, /path/to/store +from columnflow.config_util import add_shift_aliases +add_shift_aliases(cfg, "mu", {"muon_weight": "muon_weight_{direction}"}) + +# In the weight Producer's init hook: +from columnflow.config_util import get_shifts_from_sources +self.shifts |= get_shifts_from_sources(self.config_inst, "mu") ``` -**Critical**: after adding any new Python file (calibrator, selector, producer, etc.) you **must** register it in `law.cfg` under the correct `*_modules` key. No spaces after commas inside `{}` brace expansions. +### Selection-modifying uncertainties (e.g. JEC) ---- +Kinematic corrections that affect object selection. A separate task branch runs the complete pipeline (from `CalibrateEvents` onwards) for each variation. -## Systematic Uncertainties (Shifts) +```python +# In config: +cfg.add_shift(name="jec_up", id=10, type=od.Shift.SHAPE, tags={"selection_dependent"}) +cfg.add_shift(name="jec_down", id=11, type=od.Shift.SHAPE, tags={"selection_dependent"}) + +# Column aliases map "Jet.pt" → "Jet.pt_jec_up" when jec_up shift is active: +add_shift_aliases(cfg, "jec", {"Jet.pt": "Jet.pt_{name}", "Jet.eta": "Jet.eta_{name}"}) + +# In the Selector's init hook: +from columnflow.config_util import get_shifts_from_sources +self.shifts |= get_shifts_from_sources(self.config_inst, "jec", tags={"selection_dependent"}) +``` -Three classes, ordered by complexity: +### Dedicated-dataset uncertainties -### 1. Rate-only uncertainties -Only affect yields, not selection. Defined entirely in the inference model. No additional workflow steps needed. +A completely separate Monte Carlo sample is processed (e.g. tune, hdamp variations). The dataset `info` dict key must match the shift name exactly. -### 2. Weight-based (shape) uncertainties -Applied via event weights. Require: -1. `cfg.add_shift(...)` for up/down variants -2. `add_shift_aliases(cfg, "source_name", ...)` to map column names -3. The `WeightProducer` / `HistProducer` must declare the shift in its `shifts` set -4. Considered only from `CreateHistograms` onwards +```python +cfg.add_shift(name="tune_up", id=20, type=od.Shift.SHAPE, tags={"disjoint_from_nominal"}) +# The dataset info key "tune_up" automatically links this shift to the correct dataset. +``` + +--- -### 3. Selection-modifying uncertainties (e.g. JEC) -Propagate through the entire pipeline. Require: -1. `cfg.add_shift(...)` for up/down variants with `tags={"selection_dependent"}` -2. Column aliases for varied kinematic columns -3. The `Selector` (and upstream `Calibrator`) must declare the shifts in their `shifts` set -4. Separate task branches are run for each shift +## Data Flow Summary + +``` +ROOT files → chunks of ak.Array named "events" + │ + [Calibrator] adds corrected columns (never removes) + │ + [Selector] creates boolean masks → stats.json + masks.parquet + │ + [Reducer] applies masks → reduced_events.parquet + │ + [Producer] adds columns → extra_columns.parquet + │ + [HistProducer] event weights + fills Hist → histograms.pkl + │ + [Plots / Datacards] +``` -### 4. Dedicated-dataset uncertainties (e.g. tune, hdamp) -A completely separate dataset is processed for each variation. The dataset `info` dict key must match the shift name. +All intermediate data is stored in a directory tree under the path configured in `law.cfg`, organised by task family, version, config, dataset, and shift. diff --git a/docs/columnflow_claude_guide/03_task_commands.md b/docs/columnflow_claude_guide/03_task_commands.md index 9174c1fd0..605ac5313 100644 --- a/docs/columnflow_claude_guide/03_task_commands.md +++ b/docs/columnflow_claude_guide/03_task_commands.md @@ -1,253 +1,373 @@ -# 03 — Task Commands +# 03 — User Guide: Working with Columnflow Tasks -All tasks are invoked with `law run cf.`. The `--version` flag is always required. `--config` and `--analysis` can be set in `law.cfg` as defaults. - -Use `--help` on any task for the full parameter list: -```bash -law run cf.SelectEvents --help -``` - ---- +All tasks are run with `law run cf.`. Law automatically determines and runs any missing upstream tasks before the requested one. ## Universal Flags -| Flag | Purpose | -|---|---| -| `--version ` | Version tag for output paths (required) | -| `--config ` | Config name (default: `law.cfg default_config`) | -| `--analysis ` | Analysis module path (default: `law.cfg default_analysis`) | -| `--branch ` | Run a single branch (file) instead of all | -| `--print-status -1` | Show task tree + output existence (recursive) | -| `--print-output ` | Show output file paths at given depth | -| `--remove-output ,[,y]` | Remove outputs; modes: `a` all, `i` interactive, `d` dry-run | -| `--workflow htcondor` | Submit to HTCondor | -| `--workflow slurm` | Submit to Slurm | -| `--pilot` | Run one branch to check before launching all | +| Flag | Purpose | Notes | +|---|---|---| +| `--version ` | Version tag for output paths | Required for every task | +| `--config ` | Which config to use | Default set in `law.cfg` | +| `--analysis ` | Analysis module path | Default set in `law.cfg` | +| `--branch ` | Process a single file/branch | Great for quick local tests | +| `--print-status -1` | Show task tree + output existence | Use `-1` for full recursion | +| `--print-output ` | Show output file paths at depth | `0` = the task itself | +| `--remove-output ,[,y]` | Delete outputs | `a`=all, `i`=interactive, `d`=dry | +| `--workflow htcondor` | Submit to HTCondor | Requires `[job]` config in `law.cfg` | +| `--workflow slurm` | Submit to Slurm | Requires `[job]` config in `law.cfg` | +| `--help` | Full parameter list for a task | Use before every new task | --- -## Step-by-step Commands +## GetDatasetLFNs + +**What it does:** Resolves the list of ROOT file paths (Logical File Names) for a dataset. For CMS data it queries DAS using `dasgoclient`. Custom retrieval functions can be configured via `cfg.x.get_dataset_lfns`. -### GetDatasetLFNs +**Requires:** A valid GRID proxy (`voms-proxy-init -rfc -valid 196:00`) for CMS datasets, unless a custom LFN source is configured. -Resolves logical file names for a dataset (requires GRID proxy for CMS data). +**Output:** A JSON file mapping file indices to LFN paths. ```bash +# Resolve LFNs for one dataset law run cf.GetDatasetLFNs \ --version dev1 \ --dataset tt_dl_powheg + +# Check which files were found +law run cf.GetDatasetLFNs --version dev1 --dataset tt_dl_powheg --print-output 0 ``` +**Tips:** +- Run this first for all datasets before starting the pipeline. +- If your datasets are stored locally or on a custom SE, configure `cfg.x.get_dataset_lfns` to bypass DAS. +- For debugging, add `--branch 0` to retrieve only the first file entry. + --- -### CalibrateEvents +## CalibrateEvents -Applies calibrations (e.g. jet energy corrections). Runs per file/branch. +**What it does:** Applies one or more `Calibrator` objects to raw events, adding corrected column versions (e.g. corrected `Jet.pt` after JEC). The original columns are preserved. + +**Outputs:** A Parquet file containing any new/modified columns declared in the calibrators' `produces` sets. + +**Key parameters:** +- `--calibrators` — comma-separated list of calibrator names (order matters) +- `--shift` — which systematic shift variant to run (`nominal`, `jec_up`, `jec_down`, ...) ```bash +# Run default calibration on all files (triggers GetDatasetLFNs first) +law run cf.CalibrateEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --calibrators default + +# Test on a single file first law run cf.CalibrateEvents \ --version dev1 \ --dataset tt_dl_powheg \ --calibrators default \ - --shift nominal + --branch 0 -# Single branch test -law run cf.CalibrateEvents --version dev1 --dataset tt_dl_powheg --branch 0 +# Run with a JEC up-shift +law run cf.CalibrateEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --calibrators default \ + --shift jec_up -# With a systematic shift -law run cf.CalibrateEvents --version dev1 --dataset tt_dl_powheg --shift jec_up +# Check if all branches completed +law run cf.CalibrateEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --print-status 0 ``` +**Tips:** +- For JEC uncertainties, separate task branches run for `nominal`, `jec_up`, and `jec_down` — this is handled automatically when you request a downstream task with multiple shifts. +- You rarely need to run `CalibrateEvents` standalone; it is triggered automatically by `SelectEvents`. + --- -### SelectEvents +## SelectEvents + +**What it does:** Runs the analysis `Selector` on calibrated events. Produces: +1. A Parquet file with event/object selection masks. +2. A `stats.json` file with event counts and MC weight sums (used for normalization in histogramming). +3. Optionally, selection-step histograms for cutflow plots. -Runs the Selector; produces event/object masks (parquet) and selection statistics (json). +Masks are **not yet applied** here — they are passed to `ReduceEvents`. + +**Key parameters:** +- `--selector` — name of the exposed Selector to run +- `--calibrators` — calibrators to use upstream ```bash +# Run selector on all files of a dataset law run cf.SelectEvents \ --version dev1 \ --dataset tt_dl_powheg \ --selector default \ --calibrators default -# Verify outputs recursively -law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --print-status -1 +# Inspect the full upstream task tree +law run cf.SelectEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --print-status -1 -# Cutflow plot requires this first (SelectEvents at cutflow step) -law run cf.SelectEvents --version dev1 --config l18 +# Quick test: single file, no remote submission +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --branch 0 + +# Run for a systematic shift (runs separate CalibrateEvents branch for jec_up) +law run cf.SelectEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --shift jec_up ``` +**Tips:** +- After any change to the `Selector` code, bump the version to avoid reusing stale cached outputs. +- Check `stats.json` to verify event counts and selection efficiencies. +- The selector step names (keys in `SelectionResult.steps`) appear in cutflow plots. + --- -### ReduceEvents +## ReduceEvents + +**What it does:** Applies the event and object masks from `SelectEvents` to all columns, writing a compact Parquet file with only the selected events. Columns not in `cfg.x.keep_columns["cf.ReduceEvents"]` are permanently dropped here. -Applies selection masks to columns; writes the reduced parquet files. +**Output:** A Parquet file per input file with selected events and selected object collections. ```bash +# Reduce events for one dataset law run cf.ReduceEvents \ --version dev1 \ - --dataset tt_dl_powheg \ - --selector default \ - --calibrators default + --dataset tt_dl_powheg -# Single branch (first file) — quick local test -law run cf.ReduceEvents --version dev1 --branch 0 +# Single-file test +law run cf.ReduceEvents --version dev1 --dataset tt_dl_powheg --branch 0 + +# Check output size +law run cf.ReduceEvents --version dev1 --dataset tt_dl_powheg --print-output 0 ``` +**Tips:** +- Verify that all columns you need downstream are listed in `cfg.x.keep_columns["cf.ReduceEvents"]`. +- Use `ColumnCollection.ALL_FROM_SELECTOR` to automatically include all columns produced by your Selector. +- After changing `keep_columns`, remove and re-run `ReduceEvents`: `law run cf.ReduceEvents --version dev1 --remove-output 0,a,y` + --- -### MergeReducedEvents +## MergeReducedEvents -Merges per-file reduced parquet files into per-dataset merged files. +**What it does:** Merges the many per-file reduced Parquet files into larger files, targeting a configurable size (default 512 MB). This avoids having thousands of small files in downstream tasks. ```bash law run cf.MergeReducedEvents \ --version dev1 \ --dataset tt_dl_powheg + +# Adjust target file size (also configurable in cfg.x.reduced_file_size) +law run cf.MergeReducedEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --merged-size 256 ``` +**Tips:** +- You rarely run this standalone; it is triggered automatically by `ProduceColumns` or `CreateHistograms`. + --- -### ProduceColumns +## ProduceColumns + +**What it does:** Runs one or more `Producer` objects on the merged reduced events, adding new columns (high-level variables, weights, ML scores, category IDs). Each producer's output is stored in a separate Parquet file alongside the reduced events. + +**Output:** A Parquet file per producer per dataset, containing the new columns. -Runs Producer(s) on the merged reduced events; writes additional column parquet files. +**Key parameters:** +- `--producers` — comma-separated list of producer names (order matters: later producers can overwrite earlier columns of the same name) ```bash +# Run the default producer law run cf.ProduceColumns \ --version dev1 \ --dataset tt_dl_powheg \ --producers default -# Multiple producers (comma-separated, order matters for column overwriting) +# Multiple producers (run in order; last one wins on name conflicts) law run cf.ProduceColumns \ --version dev1 \ --dataset tt_dl_powheg \ --producers default,extra_features + +# Single-file test +law run cf.ProduceColumns \ + --version dev1 \ + --dataset tt_dl_powheg \ + --producers default \ + --branch 0 ``` +**Tips:** +- Producers can call other Producers; declare the sub-producers in `uses` and `produces`. +- Columns created in `ProduceColumns` do **not** need to be in `cfg.x.keep_columns`; they are automatically available downstream. +- If multiple producers write the same column name, the last producer in the `--producers` list wins. + --- -### CreateHistograms +## CreateHistograms + +**What it does:** Fills `Hist` histograms for all requested variables, categories, and shifts. The `HistProducer` controls event weighting. One histogram file is produced per dataset branch. -Fills histograms for the given variables, categories, and shift. +**Key parameters:** +- `--variables` — comma-separated variable names (defined in config) +- `--categories` — comma-separated category names +- `--shift` — the systematic shift to use +- `--producers` — producers to run before histogramming +- `--hist-producer` — the HistProducer to use ```bash +# Basic: one variable, inclusive category, nominal shift law run cf.CreateHistograms \ --version dev1 \ --dataset tt_dl_powheg \ - --variables jet1_pt,n_jet \ + --variables jet1_pt \ --categories incl \ --shift nominal \ - --producers default \ - --hist-producer default + --producers default -# All shifts at once (nominal + all registered shifts) +# Multiple variables and categories law run cf.CreateHistograms \ --version dev1 \ - --datasets tt_dl_powheg,wjets_madgraph \ + --dataset tt_dl_powheg \ + --variables "jet1_pt,n_jet,ht,lep1_pt" \ + --categories "incl,sr,cr_1b" + +# Run for all registered shifts at once (triggers separate branches per shift) +law run cf.CreateHistograms \ + --version dev1 \ + --datasets "tt_dl_powheg,wjets_madgraph" \ --variables jet1_pt \ - --shifts nominal,jec_up,jec_down,mu_up,mu_down + --shifts "nominal,jec_up,jec_down,mu_up,mu_down" ``` +**Tips:** +- Running with `--datasets "tt*"` processes all datasets matching the glob pattern. +- The `--shifts` flag (plural) allows specifying multiple shifts in one command; the task graph branches automatically. +- Use `--branch 0` to test on one file before submitting all branches. + --- -### MergeHistograms +## MergeHistograms -Merges per-dataset histogram files. +**What it does:** Merges per-branch histogram files into one file per dataset and shift. A second stage (`MergeShiftedHistograms`) further merges shifted histograms across all datasets of a process. ```bash law run cf.MergeHistograms \ --version dev1 \ --variables jet1_pt \ - --processes tt,wjets + --processes tt,wjets,data_mu ``` +**Tips:** +- This task is almost always triggered automatically by plotting or datacard tasks. You rarely call it standalone. + --- -### PlotVariables1D +## PlotVariables1D -Creates 1D variable plots per process, stacked or overlaid. +**What it does:** Creates 1D variable plots (stacked MC + data, or process-overlaid) from merged histograms using matplotlib/mplhep CMS style. + +**Key parameters:** +- `--processes` — which processes to plot (alternative to `--datasets`) +- `--variables` — variables to plot +- `--categories` — categories to plot +- `--shift` — which shift to show +- `--file-types` — output format(s): `pdf`, `png`, `svg` +- `--skip-ratio` — disable the ratio panel below the main plot +- `--density` — normalize to unit area +- `--yscale` — `log` or `linear` +- `--view-cmd` — open plots automatically (e.g. `evince`, `eog`) ```bash +# Standard stacked plot: all tt + wjets + data law run cf.PlotVariables1D \ --version dev1 \ - --processes tt,wjets,data \ - --variables jet1_pt,n_jet \ + --processes "tt,wjets,data_mu" \ + --variables "jet1_pt,n_jet,ht" \ --categories incl \ --shift nominal -# Save to PDF and PNG +# Save as both PDF and PNG law run cf.PlotVariables1D \ --version dev1 \ --processes tt \ --variables jet1_pt \ - --file-types pdf,png + --file-types "pdf,png" -# With custom plot function +# Log scale, skip ratio law run cf.PlotVariables1D \ --version dev1 \ - --processes tt \ + --processes "tt,wjets" \ --variables jet1_pt \ - --plot-function myanalysis.plotting.my_plot_func + --yscale log \ + --skip-ratio -# Open plots automatically +# Find where plots are saved law run cf.PlotVariables1D \ - --version dev1 \ - --processes tt \ - --variables jet1_pt \ - --view-cmd evince + --version dev1 --processes tt --variables jet1_pt --print-output 0 ``` --- -### PlotVariables2D - -Creates 2D plots. +## PlotVariables2D ```bash law run cf.PlotVariables2D \ --version dev1 \ --processes tt \ - --variables jet1_pt__jet1_eta \ + --variables "jet1_pt__jet1_eta" \ --categories incl ``` --- -### PlotShiftedVariables1D +## PlotShiftedVariables1D -Shows nominal + up/down shift variations for a variable. +**What it does:** Shows nominal and up/down shift variations for a variable, overlaid on one plot. Useful for visualizing the impact of a systematic on a distribution. ```bash law run cf.PlotShiftedVariables1D \ --version dev1 \ --processes tt \ --variables jet1_pt \ - --shift-sources jec,mu \ + --shift-sources "jec,mu" \ --categories incl ``` --- -### PlotCutflow +## PlotCutflow -Plots total event yield at each selection step. +**What it does:** Bar chart showing the event yield after each selection step defined in the Selector (`SelectionResult.steps`). ```bash law run cf.PlotCutflow \ --version dev1 \ --datasets tt_dl_powheg \ - --selector-steps muon,jet,bjet + --selector-steps "trigger,muon,jet,btag" \ + --categories incl ``` --- -### PlotCutflowVariables1D +## PlotCutflowVariables1D -Plots a variable distribution at each selector step. +**What it does:** Plots a variable distribution at one or more intermediate selection steps (before the full selection is applied). Useful for gen-level checks or debugging. ```bash +# Gen-level top-quark pT at each selector step law run cf.PlotCutflowVariables1D \ --version dev1 \ --datasets tt_dl_powheg \ @@ -259,100 +379,73 @@ law run cf.PlotCutflowVariables1D \ --- -### CreateDatacards +## CreateDatacards -Produces CMS combine-compatible datacards and ROOT shape files. +**What it does:** Produces CMS `combine`-compatible datacards (`.txt`) and shape ROOT files from merged histograms, driven by an `InferenceModel` defined in your analysis. ```bash law run cf.CreateDatacards \ --version dev1 \ - --inference-model example \ + --inference-model default \ --variables jet1_pt \ - --categories incl + --categories sr ``` --- -## Useful Utility Commands +## Useful Workflow Commands ```bash -# Check if a task's output already exists -law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --print-status 0 +# --- Inspect --- +# Full task tree with output existence for a plot task +law run cf.PlotVariables1D --version dev1 --processes tt --variables jet1_pt --print-status -1 -# Remove just the histogram task output and re-run it -law run cf.CreateHistograms --version dev1 --variables jet1_pt --remove-output 0,a,y - -# List all available law tasks -law index --verbose - -# Get help on a specific task -law run cf.ProduceColumns --help - -# Run inside the columnar sandbox manually (for debugging) -cf_sandbox venv_columnar_dev bash -``` - ---- - -## Remote Execution (HTCondor) - -```bash -# Submit SelectEvents for all files of a dataset to HTCondor -law run cf.SelectEvents \ - --version dev1 \ - --dataset tt_dl_powheg \ - --workflow htcondor +# Find output paths at depth 0 (the plot itself) +law run cf.PlotVariables1D --version dev1 --processes tt --variables jet1_pt --print-output 0 -# Monitor running jobs -law run cf.SelectEvents \ - --version dev1 \ - --dataset tt_dl_powheg \ - --workflow htcondor \ - --print-status -1 -``` - ---- - -## Pinning Upstream Task Versions - -Use task-family-specific version flags to pin upstream outputs: +# --- Re-run --- +# Remove histogram task output and re-run it immediately (mode a = all) +law run cf.CreateHistograms --version dev1 --variables jet1_pt --remove-output 0,a,y -```bash -# Use v1 outputs from CalibrateEvents while running SelectEvents as v2 +# --- Remote --- +# Submit SelectEvents to HTCondor law run cf.SelectEvents \ - --version v2 \ - --cf.CalibrateEvents-version v1 \ - --dataset tt_dl_powheg -``` - ---- + --version dev1 --dataset tt_dl_powheg --workflow htcondor -## Dataset / Process Wildcards - -Many tasks accept glob patterns: - -```bash -# All tt datasets -law run cf.PlotVariables1D --version dev1 --datasets "tt*" --variables jet1_pt +# --- Pin upstream versions --- +# Use prod1 outputs from ReduceEvents while creating new histograms as dev2 +law run cf.CreateHistograms \ + --version dev2 \ + --cf.MergeReducedEvents-version prod1 \ + --variables jet1_pt -# Multiple processes -law run cf.PlotVariables1D --version dev1 --processes "tt,wjets,dy*" --variables n_jet +# --- Open columnar sandbox for ad-hoc scripting --- +cf_sandbox venv_columnar_dev bash ``` --- -## Full Pipeline — Single Command Chain +## Full Pipeline in One Command -Running the plotting task automatically triggers all upstream tasks: +Law resolves all upstream dependencies automatically. Running a plot task triggers the entire chain: ```bash -# This single command triggers the full pipeline for one dataset + one variable +# This single command runs the full pipeline for one file of one dataset law run cf.PlotVariables1D \ --version dev1 \ --datasets tt_dl_powheg \ - --variables jet1_pt \ --processes tt \ + --variables jet1_pt \ --categories incl \ - --shift nominal \ - --branch 0 # single branch for a quick test + --branch 0 +``` + +For production over all files and datasets, use HTCondor: + +```bash +law run cf.MergeHistograms \ + --version prod1 \ + --datasets "tt_dl_powheg,wjets_madgraph,data_mu_b" \ + --variables "jet1_pt,n_jet,ht,lep1_pt" \ + --workflow htcondor ``` diff --git a/docs/columnflow_claude_guide/05_dos_and_donts.md b/docs/columnflow_claude_guide/05_dos_and_donts.md index 4b8b3d3a8..65e659a0c 100644 --- a/docs/columnflow_claude_guide/05_dos_and_donts.md +++ b/docs/columnflow_claude_guide/05_dos_and_donts.md @@ -1,310 +1,349 @@ -# 05 — Do's and Don'ts for Claude in Columnflow Projects +# 05 — Do's and Don'ts for Columnflow Users -## Critical Rules +This is a practical checklist for anyone writing code with columnflow, covering the most common mistakes and best practices. Each rule has a concrete code example showing the wrong and correct approach. -### DO always use `set_ak_column` to modify event data +--- + +## Event Data: the `events` Array + +### DO use `set_ak_column` to add or modify columns + +`events` is an awkward array. Its fields cannot be set by direct attribute assignment. Always use `set_ak_column` and **reassign** the return value. ```python # CORRECT -events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1)) +from columnflow.columnar_util import set_ak_column +events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) -# WRONG — direct mutation is not supported and may silently fail -events.ht = ak.sum(events.Jet.pt, axis=1) +# WRONG — silently ignored or raises events["ht"] = ak.sum(events.Jet.pt, axis=1) +events.ht = ak.sum(events.Jet.pt, axis=1) ``` -### DO always declare `uses` and `produces` in TAF decorators - -Columnflow uses these sets to track which columns are loaded from disk and which are written. Missing a column in `uses` means it won't be available. Missing it in `produces` means it won't be saved. +### DO reassign `events` after every TAF call ```python -# CORRECT — both sets explicitly declared -@producer( - uses={"Jet.pt", "Jet.eta"}, - produces={"ht", "n_jet"}, -) -def my_producer(self, events, **kwargs): ... +# CORRECT +events = self[other_producer](events, **kwargs) +events = set_ak_column(events, "my_col", value) -# WRONG — omitting uses/produces causes runtime failures -@producer() -def my_producer(self, events, **kwargs): ... +# WRONG — the return value is discarded +self[other_producer](events, **kwargs) ``` -### DO use `maybe_import` for heavy packages +### DON'T loop over events in Python + +Columnflow is columnar: all operations must be vectorized with `ak.*`. Python loops over events are thousands of times slower and defeat the purpose of the framework. ```python -# CORRECT -ak = maybe_import("awkward") -np = maybe_import("numpy") +# WRONG — never do this +hts = [] +for event in events: + hts.append(sum(jet.pt for jet in event.Jet)) -# WRONG — direct import fails when columnflow loads the module in a non-columnar sandbox -import awkward as ak -import numpy as np +# CORRECT +ht = ak.sum(events.Jet.pt, axis=1) ``` -### DO import coffea inside the function, never at module level +### DO specify `axis=1` when reducing over objects within an event ```python -# CORRECT -def my_function(self, events, **kwargs): - import coffea.nanoevents.methods.vector - vec = ak.zip({"pt": ...}, with_name="PtEtaPhiMLorentzVector", - behavior=coffea.nanoevents.methods.vector.behavior) +# WRONG — sums over everything (axis=None by default in some contexts) +n_jet = ak.sum(events.Jet.pt > 0) # scalar, not per-event -# WRONG — coffea at module level breaks non-coffea sandboxes -import coffea +# CORRECT — one value per event +n_jet = ak.sum(events.Jet.pt > 0, axis=1) # shape (N_events,) ``` -### DO register new Python files in `law.cfg` +--- -Every new file containing a Calibrator, Selector, Producer, Categorizer, or HistProducer must be added to the appropriate `*_modules` key in `law.cfg`. +## TAF Declarations (`uses` and `produces`) -```ini -# CORRECT — after adding myanalysis/selection/jets.py -selection_modules: ..., myanalysis.selection.{default,jets} +### DO declare every column you read in `uses` -# No spaces after commas inside braces! -# WRONG -selection_modules: ..., myanalysis.selection.{default, jets} +Columnflow only loads columns from disk that are listed in `uses`. Missing a column means `events.Jet.pt` raises `AttributeError` at runtime. + +```python +# CORRECT +@producer( + uses={"Jet.{pt,eta,phi,mass}", "Electron.pt"}, + produces={"ht", "n_jet"}, +) +def my_producer(self, events, **kwargs): ... + +# WRONG — Jet.eta read without being declared +@producer( + uses={"Jet.pt"}, # missing Jet.eta! + produces={"ht"}, +) +def my_producer(self, events, **kwargs): + mask = abs(events.Jet.eta) < 2.4 # KeyError at runtime ``` -### DO include sub-TAF's `uses`/`produces` by passing the TAF itself +### DO declare every column you write in `produces` + +Columns not in `produces` are computed but never saved to disk, and are invisible to downstream tasks. + +### DO pass sub-TAF objects in `uses`/`produces` to propagate their columns ```python -# CORRECT — propagates all uses/produces from sub_producer +# CORRECT — sub_producer's columns are automatically included @producer( uses={sub_producer, "extra_col"}, - produces={sub_producer, "my_new_col"}, + produces={sub_producer, "my_output"}, ) -def parent(self, events, **kwargs): ... +def parent(self, events, **kwargs): + events = self[sub_producer](events, **kwargs) + ... -# WRONG — forgets to re-declare sub_producer columns +# WRONG — sub_producer reads "Jet.pt" but it's not declared here @producer( uses={"extra_col"}, - produces={"my_new_col"}, + produces={"my_output"}, ) def parent(self, events, **kwargs): - events = self[sub_producer](events, **kwargs) # sub_producer's columns not declared! + events = self[sub_producer](events, **kwargs) # Jet.pt not loaded! ``` -### DO gate Monte Carlo logic on `self.dataset_inst.is_mc` +--- -```python -# CORRECT -if self.dataset_inst.is_mc: - events = self[mc_weight](events, **kwargs) +## Imports -# WRONG — mc_weight column does not exist in data -events = self[mc_weight](events, **kwargs) -``` +### DO use `maybe_import` for heavy packages at module level -### DO return `events` from every TAF function +Columnflow modules are imported in the default (non-columnar) sandbox at setup time. Packages like `awkward`, `numpy`, and `coffea` are not available there. ```python -# CORRECT — Producer returns events -def my_producer(self, events, **kwargs) -> ak.Array: - events = set_ak_column(events, "ht", ...) - return events +# CORRECT — deferred import +from columnflow.util import maybe_import +ak = maybe_import("awkward") +np = maybe_import("numpy") -# WRONG — missing return -def my_producer(self, events, **kwargs): - events = set_ak_column(events, "ht", ...) +# WRONG — fails at import time in the default sandbox +import awkward as ak +import numpy as np ``` -### DO return `(events, SelectionResult)` from Selectors +### DO import coffea inside the function body, never at module level ```python # CORRECT -def my_selector(self, events, stats, **kwargs) -> tuple[ak.Array, SelectionResult]: - ... - return events, SelectionResult(steps={...}) +def my_selector(self, events, **kwargs): + import coffea.nanoevents.methods.vector + vec = ak.zip({"pt": events.Jet.pt}, with_name="PtEtaPhiMLorentzVector", + behavior=coffea.nanoevents.methods.vector.behavior) -# WRONG — wrong return type -def my_selector(self, events, stats, **kwargs): - ... - return events +# WRONG +import coffea # breaks any non-coffea sandbox ``` +--- + +## Selectors + ### DO set `results.event` in the exposed Selector -The `ReduceEvents` task reads `results.event` to filter events. Missing it causes no events to be selected (or an error). +`ReduceEvents` reads `results.event` to determine which events to keep. If it is not set, all events (or no events) may be kept depending on the default. ```python from operator import and_ from functools import reduce +# After combining all SelectionResult objects: results.event = reduce(and_, results.steps.values()) ``` -### DO add `keep_columns` in the config for columns created before ReduceEvents - -Any column produced by a Calibrator or Selector that should survive `ReduceEvents` must be explicitly listed in `cfg.x.keep_columns` under `"cf.ReduceEvents"`. +### DO return `(events, SelectionResult)` from every Selector ```python -cfg.x.keep_columns = DotDict.wrap({ - "cf.ReduceEvents": { - "Jet.{pt,eta,phi,mass}", - "my_new_column", # REQUIRED if produced in Selector - ColumnCollection.ALL_FROM_SELECTOR, # convenient catch-all for selector columns - }, -}) -``` +# CORRECT +def my_selector(self, events, stats, **kwargs) -> tuple[ak.Array, SelectionResult]: + ... + return events, SelectionResult(steps={"jet": mask}) ---- +# WRONG — wrong return type breaks ReduceEvents +def my_selector(self, events, stats, **kwargs): + ... + return events # missing SelectionResult +``` -## Architecture Don'ts +### DO use `ak.local_index` to create index arrays for `objects` -### DON'T call `law run` inside Python code +```python +# CORRECT — index arrays for SelectionResult.objects +jet_mask = (events.Jet.pt > 25) & (abs(events.Jet.eta) < 2.4) +jet_indices = ak.local_index(events.Jet.pt, axis=1)[jet_mask] -Use the standard `law run cf.` CLI. If triggering a task programmatically from a script is truly necessary, use the pattern in `best_practices.md` with `task.law_run()`. +return events, SelectionResult( + objects={"Jet": {"Jet": jet_indices}}, +) -### DON'T put analysis logic in `CLAUDE.md` or config files +# WRONG — using a boolean mask directly in objects (must be indices) +return events, SelectionResult( + objects={"Jet": {"Jet": jet_mask}}, # wrong type! +) +``` -Analysis code (Selectors, Producers, etc.) belongs in the appropriate subdirectory, not in configuration files or documentation. +--- -### DON'T create a new exposed Selector unless truly needed +## Configuration -One exposed (top-level) Selector should be the entry point, composed of many internal Selectors. Creating multiple exposed Selectors for the same purpose leads to maintenance burden. +### DO list columns in `cfg.x.keep_columns` that must survive `ReduceEvents` -### DON'T mix column-level and event-level operations without explicit axis +Any column produced in a Calibrator or Selector that is needed downstream must be explicitly listed. Columns from Producers (run after `ReduceEvents`) are automatically kept. ```python -# WRONG — ambiguous without axis -n_jet = ak.sum(events.Jet.pt > 0) # sums over EVERYTHING +from columnflow.util import DotDict +from columnflow.columnar_util import ColumnCollection -# CORRECT -n_jet = ak.sum(events.Jet.pt > 0, axis=1) # per-event count +cfg.x.keep_columns = DotDict.wrap({ + "cf.ReduceEvents": { + "Jet.{pt,eta,phi,mass,btagDeepFlavB}", + "Electron.{pt,eta,phi,mass,charge}", + "Muon.{pt,eta,phi,mass,charge}", + "MET.{pt,phi}", + "event", "run", "luminosityBlock", + ColumnCollection.ALL_FROM_SELECTOR, # all columns produced in the Selector + }, +}) ``` -### DON'T use Python loops over events +### DO register new files in `law.cfg` under the correct `*_modules` key -```python -# WRONG — extremely slow; defeats the columnar paradigm -for event in events: - ht = sum(jet.pt for jet in event.Jet) +After adding a new Python file containing a Selector, Producer, Calibrator, Categorizer, or HistProducer, add it to `law.cfg`. Without this, columnflow cannot find the TAF. -# CORRECT -ht = ak.sum(events.Jet.pt, axis=1) -``` +```ini +# After creating myanalysis/selection/jets.py: +selection_modules: ..., myanalysis.selection.{default,jets} -### DON'T modify `events` returned from a sub-TAF call without reassigning +# After creating myanalysis/production/weights.py: +production_modules: ..., myanalysis.production.{default,weights} +``` -```python -# WRONG — changes are discarded -self[other_producer](events, **kwargs) +### DON'T put spaces after commas inside `{}` brace expansions in `law.cfg` +```ini # CORRECT -events = self[other_producer](events, **kwargs) +selection_modules: myanalysis.selection.{default,jets,trigger} + +# WRONG — spaces break the brace expansion +selection_modules: myanalysis.selection.{default, jets, trigger} ``` --- -## Shift / Systematics Don'ts +## Systematic Uncertainties + +### DO guard MC-only operations with `self.dataset_inst.is_mc` + +```python +if self.dataset_inst.is_mc: + events = self[mc_weight](events, **kwargs) + events = self[pileup_weight](events, **kwargs) +``` -### DON'T hardcode weight column names in Producers +### DON'T hardcode shifted column names — use column aliases -Use the column alias mechanism so that shifted variants are used automatically when a shift is active: +When a shift is active, columnflow swaps column names according to `cfg.x.column_aliases`. Your code should read the nominal name and let the alias mechanism do the substitution. ```python -# WRONG — always uses the nominal weight +# WRONG — hardcoded nominal name, broken under mu_up shift weight = events.muon_weight -# CORRECT — alias resolves to muon_weight_up/down under the corresponding shift -weight = events[self.config_inst.x.column_aliases.get("muon_weight", "muon_weight")] +# CORRECT — the HistProducer's shift mechanism resolves the alias automatically +# Just declare "muon_weight" in uses and let the framework handle the rest ``` -Or rely on the `WeightProducer`/`HistProducer` shift mechanism rather than doing manual name resolution. - -### DON'T forget to declare shifts in the TAF's `shifts` set - -For weight-based and selection-modifying uncertainties, the TAF must declare which shifts it is sensitive to: +### DO declare shifts in the TAF's `shifts` set (for weight-based uncertainties) ```python @my_producer.init -def my_producer_init(self): +def my_producer_init(self: Producer) -> None: from columnflow.config_util import get_shifts_from_sources - self.shifts |= get_shifts_from_sources(self.config_inst, "mu") + self.shifts |= set(get_shifts_from_sources(self.config_inst, "mu")) ``` --- -## Config Don'ts +## Versioning and Caching -### DON'T add whitespace after commas inside `{}` brace expansions in law.cfg +### DO bump `--version` after changing any code in the pipeline -```ini -# CORRECT -selection_modules: myanalysis.selection.{default,jets,trigger} +Law caches outputs based on the version string. If you change a Selector but keep the same version, old cached outputs are reused and your changes have no effect. -# WRONG — spaces break the brace expansion parser -selection_modules: myanalysis.selection.{default, jets, trigger} +```bash +# After changing the Selector: +law run cf.SelectEvents --version v2 --dataset tt_dl_powheg # not v1! ``` -### DON'T duplicate config entries across Analysis and Config +### DO use `--branch 0` for quick local tests before submitting all branches -Put analysis-independent information in the `Campaign` object (year, ecm, tier). Put analysis-specific information in `Config`. Do not copy the same value into both. +```bash +# Test on a single file first to catch bugs early +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --branch 0 -### DON'T put large data arrays into Config auxiliaries +# Then run all files +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg +``` -`cfg.x.*` is for lightweight metadata (names, thresholds, lookup tables). Heavy data (scale factor histograms, b-tag efficiency arrays) should be loaded from external files at TAF setup time via the `requires`/`setup` hooks. +### DO use `--print-status -1` to understand what will run before running it ---- +```bash +law run cf.PlotVariables1D --version dev1 --processes tt --variables jet1_pt --print-status -1 +``` -## Style Don'ts +### DO use `--remove-output 0,a,y` to delete and immediately re-run a task -### DON'T import at the function call site what should be a module-level `maybe_import` +```bash +# Delete CreateHistograms output and re-run it in one command +law run cf.CreateHistograms --version dev1 --variables jet1_pt --remove-output 0,a,y +``` -```python -# WRONG — hides the deferred import intent -def my_func(self, events, **kwargs): - import awkward as ak - ... +--- -# CORRECT — at module level -ak = maybe_import("awkward") -``` +## Storage and Columns -### DON'T use `ak.Array` type annotations without the `maybe_import` guard +### DON'T store large arrays in `cfg.x.*` auxiliaries -If `awkward` is not available at import time, type annotations that reference `ak.Array` will fail. Use string annotations or guard them: +Config auxiliaries are for lightweight metadata (thresholds, names, flags). External data (scale factor histograms, efficiency maps) must be loaded at TAF setup time via the `requires`/`setup` hooks. ```python -# SAFE with maybe_import at module level -ak = maybe_import("awkward") +# WRONG — loading a large array into cfg at config creation time +cfg.x.btag_efficiency = np.load("btag_eff.npy") # loaded before any task runs -def my_producer(self: Producer, events: ak.Array, **kwargs) -> ak.Array: ... +# CORRECT — load it in the setup() hook of the Producer that needs it +@btag_producer.setup +def btag_producer_setup(self, task, reqs, inputs, reader_targets): + self.btag_eff = inputs["ext_files"]["collection"][0]["btag_eff"].load() ``` -### DON'T return early from a Producer without returning `events` +### DO produce columns only after `ReduceEvents` where possible -```python -# WRONG -if condition: - return # forgot to return events! -events = set_ak_column(events, "col", value) -return events +Creating columns in `ProduceColumns` (after reduction) avoids storing them in every file during calibration/selection, saving significant disk space. -# CORRECT -if condition: - return events -events = set_ak_column(events, "col", value) -return events +### DO use `value_type=np.float32` for large float columns + +Float64 uses twice the disk space of float32, with no benefit for typical physics variables. + +```python +events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) ``` --- ## Summary Checklist -When writing or reviewing columnflow code, verify: +Before committing or running new analysis code, verify: -- [ ] All heavy packages use `maybe_import` at module level -- [ ] `coffea` is imported inside function bodies only -- [ ] Every column read is in `uses`; every column written is in `produces` -- [ ] `set_ak_column` is used (never direct field assignment) -- [ ] `events` is reassigned after every `set_ak_column` or sub-TAF call +- [ ] All columns read from `events` are in `uses` +- [ ] All columns written to `events` are in `produces` +- [ ] `set_ak_column` is used (never `events.field = ...`) +- [ ] `events` is reassigned after every `set_ak_column` and sub-TAF call +- [ ] Heavy imports use `maybe_import`; `coffea` is imported inside function bodies - [ ] Selectors return `(events, SelectionResult)` with `results.event` set - [ ] MC-only logic is guarded by `self.dataset_inst.is_mc` -- [ ] New files are registered in `law.cfg` under the correct `*_modules` key -- [ ] Columns created before `ReduceEvents` are in `cfg.x.keep_columns` -- [ ] No brace-expansion whitespace in `law.cfg` -- [ ] No Python loops over event arrays (use vectorized `ak.*` operations) +- [ ] All new Python files are registered in `law.cfg` under `*_modules` +- [ ] Columns needed after `ReduceEvents` are in `cfg.x.keep_columns["cf.ReduceEvents"]` +- [ ] No spaces after commas in `law.cfg` brace expansions +- [ ] No Python loops over the `events` array +- [ ] Version bumped after any code change diff --git a/docs/columnflow_claude_guide/06_custom_tasks.md b/docs/columnflow_claude_guide/06_custom_tasks.md new file mode 100644 index 000000000..bfab8e66f --- /dev/null +++ b/docs/columnflow_claude_guide/06_custom_tasks.md @@ -0,0 +1,701 @@ +# 06 — Writing Custom Analysis Code + +This guide shows how to write the five types of Task Array Functions (TAFs) and custom law tasks for your analysis. Examples are inspired by the [hh2bbww analysis](https://github.com/uhh-cms/hh2bbww) and the Ghent analysis template. + +--- + +## Object Selection Selector + +The most common starting point: a set of internal selectors that define the baseline objects (muons, electrons, jets), then one exposed selector that composes them. + +### Step 1: Object-level Selectors (unexposed, internal) + +```python +# selection/objects.py +from collections import defaultdict +from columnflow.selection import Selector, SelectionResult, selector +from columnflow.columnar_util import set_ak_column, sorted_indices_from_mask +from columnflow.util import maybe_import, four_vec + +ak = maybe_import("awkward") + + +@selector( + uses=four_vec("Muon", {"sip3d", "dxy", "dz", "miniPFRelIso_all", "mediumId"}), +) +def muon_object( + self: Selector, + events: ak.Array, + stats: defaultdict, + **kwargs, +) -> tuple[ak.Array, SelectionResult]: + + # Loose muon selection + mu_mask = ( + (abs(events.Muon.eta) < 2.4) & + (events.Muon.pt > 10.0) & + (events.Muon.miniPFRelIso_all < 0.4) & + (events.Muon.sip3d < 8) & + (abs(events.Muon.dxy) < 0.05) & + (abs(events.Muon.dz) < 0.1) + ) + # Store tight flag as a new column for later use + mu_tight = mu_mask & events.Muon.mediumId + events = set_ak_column(events, "Muon.tight", mu_tight, value_type=bool) + + return events, SelectionResult( + objects={"Muon": {"Muon": sorted_indices_from_mask(mu_mask, events.Muon.pt)}}, + ) + + +@selector( + uses=four_vec("Jet", {"btagDeepFlavB"}), +) +def jet_object( + self: Selector, + events: ak.Array, + results: SelectionResult, + stats: defaultdict, + **kwargs, +) -> tuple[ak.Array, SelectionResult]: + + jet_mask = ( + (events.Jet.pt > 25.0) & + (abs(events.Jet.eta) < 2.4) + ) + jet_indices = ak.local_index(events.Jet.pt, axis=1)[jet_mask] + + return events, SelectionResult( + objects={"Jet": {"Jet": jet_indices}}, + ) +``` + +### Step 2: Event-level Selectors (unexposed, internal) + +```python +# selection/event_selection.py +from columnflow.selection import Selector, SelectionResult, selector +from columnflow.util import maybe_import + +ak = maybe_import("awkward") + + +@selector( + uses=four_vec("Muon"), +) +def lepton_selection( + self: Selector, + events: ak.Array, + results: SelectionResult, + **kwargs, +) -> tuple[ak.Array, SelectionResult]: + # work with the already-selected muons (object indices from object_selection step) + muon = events.Muon[results.objects.Muon.Muon] + n_mu = ak.sum(muon.pt > 0, axis=1) + + return events, SelectionResult(steps={"lepton": n_mu >= 1}) + + +@selector( + uses=four_vec("Jet"), +) +def jet_selection( + self: Selector, + events: ak.Array, + results: SelectionResult, + **kwargs, +) -> tuple[ak.Array, SelectionResult]: + jet = events.Jet[results.objects.Jet.Jet] + n_jet = ak.sum(jet.pt > 0, axis=1) + n_bjet = ak.sum(jet.btagDeepFlavB > 0.5, axis=1) # medium WP + + return events, SelectionResult( + steps={ + "jet": n_jet >= 2, + "btag": n_bjet >= 1, + }, + ) +``` + +### Step 3: The Exposed (Top-Level) Selector + +```python +# selection/default.py +from collections import defaultdict +from operator import and_ +from functools import reduce + +from columnflow.selection import Selector, SelectionResult, selector +from columnflow.production.cms.mc_weight import mc_weight +from columnflow.production.processes import process_ids +from columnflow.production.util import attach_coffea_behavior +from columnflow.selection.stats import increment_stats +from columnflow.util import maybe_import + +from myanalysis.selection.objects import muon_object, jet_object +from myanalysis.selection.event_selection import lepton_selection, jet_selection + +ak = maybe_import("awkward") + + +@selector( + uses={ + attach_coffea_behavior, mc_weight, process_ids, + muon_object, jet_object, + lepton_selection, jet_selection, + increment_stats, + }, + produces={ + attach_coffea_behavior, mc_weight, process_ids, + }, + exposed=True, # reachable from CLI as --selector default +) +def default( + self: Selector, + events: ak.Array, + stats: defaultdict, + **kwargs, +) -> tuple[ak.Array, SelectionResult]: + + # Attach coffea 4-vector behavior + events = self[attach_coffea_behavior](events, **kwargs) + + # MC bookkeeping + if self.dataset_inst.is_mc: + events = self[mc_weight](events, **kwargs) + events = self[process_ids](events, **kwargs) + + results = SelectionResult() + + # Object selection + events, muon_results = self[muon_object](events, stats, **kwargs) + results += muon_results + + events, jet_results = self[jet_object](events, results, stats, **kwargs) + results += jet_results + + # Event-level selection + events, lep_results = self[lepton_selection](events, results, **kwargs) + results += lep_results + + events, jet_ev_results = self[jet_selection](events, results, **kwargs) + results += jet_ev_results + + # Combine all selection steps into the final event mask + results.event = reduce(and_, results.steps.values()) + + # Increment selection statistics (saved in stats.json) + events, results = self[increment_stats]( + events, + results, + stats, + weight_map={ + "mc_weight": (events.mc_weight if self.dataset_inst.is_mc else None, Ellipsis), + "mc_weight_selected": (events.mc_weight if self.dataset_inst.is_mc else None, results.event), + }, + **kwargs, + ) + + return events, results +``` + +--- + +## Producer: High-Level Variables + +```python +# production/features.py +from columnflow.production import Producer, producer +from columnflow.columnar_util import set_ak_column, EMPTY_FLOAT +from columnflow.util import maybe_import, four_vec + +ak = maybe_import("awkward") +np = maybe_import("numpy") + + +@producer( + uses=four_vec("Jet") | four_vec("Electron") | four_vec("Muon"), + produces={"ht", "n_jet", "n_bjet", "n_electron", "n_muon", "jet1_pt", "lep1_pt"}, +) +def kinematic_features(self: Producer, events: ak.Array, **kwargs) -> ak.Array: + + # Scalar sum of jet pT + events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) + + # Object counts + events = set_ak_column(events, "n_jet", ak.num(events.Jet.pt, axis=1)) + events = set_ak_column(events, "n_bjet", ak.sum(events.Jet.btagDeepFlavB > 0.5, axis=1)) + events = set_ak_column(events, "n_electron", ak.num(events.Electron.pt, axis=1)) + events = set_ak_column(events, "n_muon", ak.num(events.Muon.pt, axis=1)) + + # Leading jet pT (EMPTY_FLOAT if no jets) + from columnflow.columnar_util import Route + events = set_ak_column( + events, "jet1_pt", + Route("Jet.pt[:,0]").apply(events, null_value=EMPTY_FLOAT), + value_type=np.float32, + ) + + # Leading lepton pT (combine muons + electrons) + lep_pt = ak.concatenate([events.Muon.pt, events.Electron.pt], axis=1) + lep_pt = lep_pt[ak.argsort(lep_pt, axis=1, ascending=False)] + events = set_ak_column( + events, "lep1_pt", + ak.fill_none(ak.pad_none(lep_pt, 1, axis=1)[:, 0], EMPTY_FLOAT), + value_type=np.float32, + ) + + return events +``` + +--- + +## Producer: Event Weights + +```python +# production/weights.py +from columnflow.production import Producer, producer +from columnflow.columnar_util import set_ak_column +from columnflow.util import maybe_import +from columnflow.config_util import get_shifts_from_sources + +ak = maybe_import("awkward") +np = maybe_import("numpy") + + +@producer( + uses={ + "normalization_weight", # added by columnflow's normalization_weights producer + "muon_weight", # added by columnflow's muon_weights producer + "pu_weight", # added by columnflow's pileup_weights producer + }, + produces={"event_weight"}, +) +def event_weights(self: Producer, events: ak.Array, **kwargs) -> ak.Array: + weight = ak.Array(np.ones(len(events), dtype=np.float32)) + + if self.dataset_inst.is_mc: + weight = weight * events.normalization_weight + weight = weight * events.muon_weight + weight = weight * events.pu_weight + + events = set_ak_column(events, "event_weight", weight, value_type=np.float32) + return events + + +@event_weights.init +def event_weights_init(self: Producer) -> None: + if not getattr(self, "dataset_inst", None) or self.dataset_inst.is_data: + return + # declare that this producer responds to muon and pileup shift sources + self.shifts |= set(get_shifts_from_sources(self.config_inst, "mu", "minbias_xs")) +``` + +--- + +## Calibrator: Custom Jet Energy Correction + +```python +# calibration/jets.py +from columnflow.calibration import Calibrator, calibrator +from columnflow.columnar_util import set_ak_column +from columnflow.util import maybe_import + +ak = maybe_import("awkward") +np = maybe_import("numpy") + + +@calibrator( + uses={"Jet.{pt,eta,phi,mass,rawFactor,area}", "fixedGridRhoFastjetAll"}, + produces={"Jet.{pt,eta,phi,mass}"}, # overwrites nominal; varied columns added in init +) +def default(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array: + # Placeholder: in practice, load JEC from external files in setup() hook + # and apply them here per-event on the Jet collection. + corrected_pt = events.Jet.pt * self.jec_factor # self.jec_factor set in setup() + events = set_ak_column(events, "Jet.pt", corrected_pt) + return events + + +@default.requires +def default_requires(self: Calibrator, task, reqs: dict) -> None: + from columnflow.tasks.external import BundleExternalFiles + reqs["ext"] = BundleExternalFiles.req(task) + + +@default.setup +def default_setup(self: Calibrator, task, reqs, inputs, reader_targets) -> None: + # Load JEC files from the bundled external files + bundle = inputs["ext"]["collection"][0] + # ... parse JEC file from bundle["jec"] ... + self.jec_factor = 1.02 # placeholder + + +@default.init +def default_init(self: Calibrator) -> None: + from columnflow.config_util import get_shifts_from_sources + # declare which shifts this calibrator covers + self.shifts |= set(get_shifts_from_sources(self.config_inst, "jec")) + # for each shift, add the varied column to produces + for shift_inst in self.shifts: + if "jec" in shift_inst.name: + self.produces.add(f"Jet.pt_{shift_inst.name}") +``` + +--- + +## Categorizer + +Categorizers define analysis regions used in histogramming and plotting. + +```python +# categorization/categories.py +from columnflow.categorization import Categorizer, categorizer +from columnflow.util import maybe_import + +ak = maybe_import("awkward") + + +@categorizer(uses={"event"}) +def cat_incl(self: Categorizer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + return events, ak.ones_like(events.event, dtype=bool) + + +@categorizer(uses={"n_jet"}) +def cat_2j(self: Categorizer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + return events, events.n_jet >= 2 + + +@categorizer(uses={"n_jet", "n_bjet"}) +def cat_sr(self: Categorizer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + return events, (events.n_jet >= 4) & (events.n_bjet >= 2) + + +@categorizer(uses={"n_jet", "n_bjet"}) +def cat_cr_1b(self: Categorizer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + return events, (events.n_jet >= 4) & (events.n_bjet == 1) +``` + +Register them in the config: + +```python +from columnflow.config_util import add_category +add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") +add_category(cfg, name="sr", id=2, selection="cat_sr", label="Signal Region") +add_category(cfg, name="cr_1b", id=3, selection="cat_cr_1b", label="CR (1b)") +``` + +And register the file in `law.cfg`: + +```ini +categorization_modules: myanalysis.categorization.categories +``` + +--- + +## HistProducer + +The HistProducer controls event weighting and the histogram filling loop. The most common pattern extends columnflow's `cf_default`: + +```python +# histogramming/default.py +from columnflow.histogramming import HistProducer +from columnflow.histogramming.default import cf_default +from columnflow.columnar_util import Route +from columnflow.util import maybe_import +from columnflow.config_util import get_shifts_from_sources + +ak = maybe_import("awkward") +np = maybe_import("numpy") + + +@cf_default.hist_producer() +def default(self: HistProducer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + weight = ak.Array(np.ones(len(events), dtype=np.float32)) + + if self.dataset_inst.is_mc: + for column in self.weight_columns: + weight = weight * Route(column).apply(events) + + return events, weight + + +@default.init +def default_init(self: HistProducer) -> None: + self.weight_columns = set() + if self.dataset_inst.is_data: + return + + self.weight_columns |= {"normalization_weight", "muon_weight", "pu_weight"} + self.uses |= self.weight_columns + + self.shifts |= set(get_shifts_from_sources(self.config_inst, "mu", "minbias_xs")) +``` + +--- + +## InferenceModel (Datacards) + +```python +# inference/default.py +from columnflow.inference import inference_model, ParameterType, ParameterTransformation + + +@inference_model +def default(self): + + # --- Categories --- + self.add_category( + "sr", + config_category="sr", + config_variable="jet1_pt", + config_data_datasets=["data_mu_b", "data_mu_c", "data_mu_d"], + mc_stats=True, + ) + + # --- Processes --- + self.add_process("TT", is_signal=False, config_process="tt") + self.add_process("ST", is_signal=True, config_process="st") + self.add_process("WJets", is_signal=False, config_process="wjets") + + # --- Parameters --- + # Luminosity (rate uncertainty) + lumi = self.config_inst.x.luminosity + for unc_name in lumi.uncertainties: + self.add_parameter( + unc_name, + type=ParameterType.rate_gauss, + effect=lumi.get(names=unc_name, direction=("down", "up"), factor=True), + transformations=[ParameterTransformation.symmetrize], + ) + + # Muon scale factor (weight-based shape uncertainty) + self.add_parameter( + "mu", + process=["TT", "ST", "WJets"], + type=ParameterType.shape, + config_shift_source="mu", + ) + + # JEC (selection-modifying shape uncertainty) + self.add_parameter( + "jec", + process=["TT", "ST", "WJets"], + type=ParameterType.shape, + config_shift_source="jec", + ) + + # Tune (dedicated dataset) + self.add_parameter( + "tune", + process="TT", + type=ParameterType.shape, + config_shift_source="tune", + ) +``` + +--- + +## Custom Law Task + +For analysis steps that do not fit the standard TAF pattern, you can write a custom law task. The base class `BaseTask` from columnflow provides access to analysis/config/campaign objects. + +```python +# tasks/base.py +from columnflow.tasks.framework.base import BaseTask + + +class MyAnalysisTask(BaseTask): + task_namespace = "my" + + +# tasks/yield_table.py +import law +import order as od + +from columnflow.tasks.framework.base import AnalysisTask +from columnflow.tasks.histograms import MergeHistograms + + +class CreateYieldTable(AnalysisTask, law.LocalWorkflow): + """ + Custom task that reads merged histograms and writes a LaTeX yield table. + """ + + version = luigi.Parameter() + processes = law.CSVParameter(default=("tt", "wjets")) + + def requires(self): + return { + proc: MergeHistograms.req( + self, + datasets=self.config_inst.get_process(proc).datasets, + variables=("n_jet",), + ) + for proc in self.processes + } + + def output(self): + return self.local_target("yield_table.tex") + + def run(self): + import hist + + rows = [] + for proc, inp in self.input().items(): + h = inp.load() + n_events = h.sum().value + rows.append((proc, f"{n_events:.1f}")) + + table = "\\begin{tabular}{lr}\n" + for proc, count in rows: + table += f" {proc} & {count} \\\\\n" + table += "\\end{tabular}\n" + + self.output().dump(table, formatter="text") +``` + +Register in `law.cfg`: + +```ini +[modules] +myanalysis.tasks +``` + +Run it: + +```bash +law run my.CreateYieldTable --version dev1 --processes tt,wjets +``` + +--- + +## Composing the Main Producer + +The main `default` producer (called from `--producers default`) should call sub-producers and set category IDs: + +```python +# production/default.py +from columnflow.production import Producer, producer +from columnflow.production.categories import category_ids +from columnflow.production.normalization import normalization_weights +from columnflow.util import maybe_import + +from myanalysis.production.features import kinematic_features +from myanalysis.production.weights import event_weights + +ak = maybe_import("awkward") + + +@producer( + uses={category_ids, kinematic_features, normalization_weights, event_weights}, + produces={category_ids, kinematic_features, normalization_weights, event_weights}, +) +def default(self: Producer, events: ak.Array, **kwargs) -> ak.Array: + # Reproduce category IDs (needed for histogramming) + events = self[category_ids](events, **kwargs) + + # High-level kinematic features + events = self[kinematic_features](events, **kwargs) + + # Normalization weight (xsec × lumi / sum_mc_weight) + if self.dataset_inst.is_mc: + events = self[normalization_weights](events, **kwargs) + events = self[event_weights](events, **kwargs) + + return events + + +@default.init +def default_init(self: Producer) -> None: + if not getattr(self, "dataset_inst", None) or self.dataset_inst.is_data: + return + # MC-only sub-producers declared dynamically + self.uses.add(event_weights) + self.produces.add(event_weights) +``` + +--- + +## Common Patterns from hh2bbww + +These patterns appear throughout the hh2bbww analysis and are good models to follow: + +### Splitting config setup into helper functions + +```python +# config/config_2018.py +def add_processes_and_datasets(cfg, campaign): + from myanalysis.config import processes as procs + cfg.add_process(procs.tt) + cfg.add_process(procs.wjets) + cfg.add_dataset(campaign.get_dataset("tt_dl_powheg")) + cfg.add_dataset(campaign.get_dataset("wjets_madgraph")) + +def add_variables(cfg): + cfg.add_variable(name="jet1_pt", expression="Jet.pt[:,0]", ...) + cfg.add_variable(name="n_jet", expression="n_jet", ...) + +def add_categories(cfg): + from myanalysis.categorization.categories import cat_incl, cat_sr + add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") + add_category(cfg, name="sr", id=2, selection="cat_sr", label="SR") + +def create_config(analysis, campaign, config_name, config_id): + cfg = analysis.add_config(campaign, name=config_name, id=config_id) + add_processes_and_datasets(cfg, campaign) + add_variables(cfg) + add_categories(cfg) + return cfg +``` + +### Using `four_vec` to declare vector columns compactly + +```python +from columnflow.util import four_vec + +@selector( + uses=four_vec({"Jet", "Electron", "Muon"}, {"btagDeepFlavB"}) | {"MET.pt", "MET.phi"}, +) +def my_selector(self, events, **kwargs): ... +``` + +`four_vec(collections, extra_fields)` expands to `Collection.{pt,eta,phi,mass,*extra_fields}` for each collection. + +### Dynamic `uses`/`produces` in `init` hooks + +```python +@my_producer.init +def my_producer_init(self: Producer) -> None: + # Only add MC-specific columns for MC datasets + if not getattr(self, "dataset_inst", None): + return + if self.dataset_inst.is_mc: + self.uses.add("GenJet.pt") + self.produces.add("n_gen_jet") +``` + +### Composing SelectionResults + +```python +results = SelectionResult() +events, r1 = self[muon_selection](events, stats, **kwargs) +results += r1 +events, r2 = self[jet_selection](events, results, stats, **kwargs) +results += r2 +# Combined event mask +results.event = reduce(and_, results.steps.values()) +``` + +### Accessing config objects inside a TAF + +```python +# self.config_inst → order.Config +# self.dataset_inst → order.Dataset +# self.analysis_inst → order.Analysis +# self.campaign → order.Campaign (via self.config_inst.campaign) + +year = self.config_inst.campaign.x.year +btag_wp = self.config_inst.x.btag_working_points["deepjet"]["medium"] +is_mc = self.dataset_inst.is_mc +``` From 147cfc1eaae4ab655ec403912f4ec58470d0f818 Mon Sep 17 00:00:00 2001 From: juvanden Date: Wed, 24 Jun 2026 11:45:16 +0200 Subject: [PATCH 3/7] docs: add .claude/rules/ context structure for Claude Code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Creates a .claude/rules/ directory following the Claude Code memory spec (https://code.claude.com/docs/en/memory#organize-rules-with-claude/rules/): Always-loaded rules (no paths frontmatter): - 01-pipeline.md: task order, 5 TAF types table, law.cfg module registration - 02-invariants.md: set_ak_column, uses/produces, maybe_import, Selector return type, MC guard, keep_columns, lifecycle hooks, no Python loops over events Path-scoped rules (load only when Claude works in matching directories): - selection.md → **/selection/** - production.md → **/production/** - calibration.md → **/calibration/** - config.md → **/config/** - categorization.md → **/categorization/** - histogramming.md → **/histogramming/** - inference.md → **/inference/** All rule files are under 110 lines; CLAUDE.md trimmed to 24 lines. Rules are condensed from docs/columnflow_claude_guide/ which remains the full user-facing documentation. Co-Authored-By: Claude Sonnet 4.6 --- .claude/rules/01-pipeline.md | 44 +++++++++++++ .claude/rules/02-invariants.md | 105 ++++++++++++++++++++++++++++++++ .claude/rules/calibration.md | 56 +++++++++++++++++ .claude/rules/categorization.md | 42 +++++++++++++ .claude/rules/config.md | 100 ++++++++++++++++++++++++++++++ .claude/rules/histogramming.md | 54 ++++++++++++++++ .claude/rules/inference.md | 80 ++++++++++++++++++++++++ .claude/rules/production.md | 104 +++++++++++++++++++++++++++++++ .claude/rules/selection.md | 88 ++++++++++++++++++++++++++ CLAUDE.md | 48 +++++++-------- 10 files changed, 697 insertions(+), 24 deletions(-) create mode 100644 .claude/rules/01-pipeline.md create mode 100644 .claude/rules/02-invariants.md create mode 100644 .claude/rules/calibration.md create mode 100644 .claude/rules/categorization.md create mode 100644 .claude/rules/config.md create mode 100644 .claude/rules/histogramming.md create mode 100644 .claude/rules/inference.md create mode 100644 .claude/rules/production.md create mode 100644 .claude/rules/selection.md diff --git a/.claude/rules/01-pipeline.md b/.claude/rules/01-pipeline.md new file mode 100644 index 000000000..033d643d3 --- /dev/null +++ b/.claude/rules/01-pipeline.md @@ -0,0 +1,44 @@ +# Columnflow Pipeline + +## Standard task order + +``` +GetDatasetLFNs → CalibrateEvents → SelectEvents → ReduceEvents +→ MergeReducedEvents → ProduceColumns → CreateHistograms +→ MergeHistograms → PlotVariables1D / CreateDatacards +``` + +Law resolves upstream dependencies automatically. Running a downstream task triggers all missing upstream tasks. + +## Five TAF types + +| TAF | Class | Task | CLI flag | Count | +|---|---|---|---|---| +| Calibrator | `Calibrator` | `CalibrateEvents` | `--calibrators` | 0..N | +| Selector | `Selector` | `SelectEvents` | `--selector` | exactly 1 | +| Reducer | `Reducer` | `ReduceEvents` | `--reducer` | exactly 1 | +| Producer | `Producer` | `ProduceColumns` | `--producers` | 0..N | +| HistProducer | `HistProducer` | `CreateHistograms` | `--hist-producer` | exactly 1 | + +## What each task produces + +- **CalibrateEvents** → Parquet with additional/corrected columns +- **SelectEvents** → Parquet with event/object masks + `stats.json` (event counts, MC weight sums) +- **ReduceEvents** → Parquet with selected events only; columns not in `cfg.x.keep_columns["cf.ReduceEvents"]` are permanently dropped +- **ProduceColumns** → Parquet with new columns alongside the reduced events +- **CreateHistograms** → pickle with `Hist` histograms per dataset/shift/category + +## law.cfg: module registration + +Every new Python file containing a TAF must be added to `law.cfg`: + +```ini +calibration_modules: myanalysis.calibration.{default,jets} +selection_modules: myanalysis.selection.{default,objects} +production_modules: myanalysis.production.{default,weights} +categorization_modules: myanalysis.categorization.categories +hist_production_modules: myanalysis.histogramming.default +inference_modules: myanalysis.inference.default +``` + +No spaces after commas inside `{}` brace expansions. diff --git a/.claude/rules/02-invariants.md b/.claude/rules/02-invariants.md new file mode 100644 index 000000000..2f3e5dd78 --- /dev/null +++ b/.claude/rules/02-invariants.md @@ -0,0 +1,105 @@ +# Columnflow Coding Invariants + +## Column writes — always use set_ak_column + +```python +# CORRECT — reassign events +from columnflow.columnar_util import set_ak_column +events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) + +# WRONG — direct field mutation does not work +events["ht"] = ... # wrong +events.ht = ... # wrong +``` + +## uses / produces — declare every column + +- Every column read from `events` must be in `uses`. +- Every column written via `set_ak_column` must be in `produces`. +- Include sub-TAF objects in both sets to propagate their column declarations: + +```python +@producer(uses={sub_producer, "extra"}, produces={sub_producer, "my_col"}) +def parent(self, events, **kwargs): + events = self[sub_producer](events, **kwargs) # reassign! +``` + +## Imports — use maybe_import for heavy packages + +```python +# At module level — deferred import (CORRECT) +from columnflow.util import maybe_import +ak = maybe_import("awkward") +np = maybe_import("numpy") + +# coffea — inside function body only, never at module level +def my_func(self, events, **kwargs): + import coffea.nanoevents.methods.vector +``` + +## Selectors — return signature and event mask + +```python +def my_selector(self, events, stats, **kwargs) -> tuple[ak.Array, SelectionResult]: + ... + return events, SelectionResult(steps={"jet": mask}) + +# Exposed selector must set results.event: +from operator import and_ +from functools import reduce +results.event = reduce(and_, results.steps.values()) +``` + +## Monte Carlo guard + +```python +if self.dataset_inst.is_mc: + events = self[mc_weight](events, **kwargs) +``` + +## Vectorized operations — no Python loops over events + +```python +# WRONG +for event in events: ... + +# CORRECT +n_jet = ak.sum(events.Jet.pt > 25, axis=1) # axis=1 = per-event +total = ak.sum(events.mc_weight) # axis=0 = global scalar +``` + +## keep_columns — required for pre-Reduce columns + +Columns produced in Calibrators/Selectors that are needed downstream must be listed: + +```python +from columnflow.util import DotDict +from columnflow.columnar_util import ColumnCollection + +cfg.x.keep_columns = DotDict.wrap({ + "cf.ReduceEvents": { + "Jet.{pt,eta,phi,mass,btagDeepFlavB}", + "Electron.{pt,eta,phi,mass,charge}", + "MET.{pt,phi}", + "event", "run", "luminosityBlock", + ColumnCollection.ALL_FROM_SELECTOR, + }, +}) +``` + +## TAF lifecycle hooks — when to use each + +```python +@my_taf.init # dynamic uses/produces/shifts; no task access +@my_taf.post_init # first hook with task access (for late registration) +@my_taf.requires # add law task requirements (e.g. BundleExternalFiles) +@my_taf.setup # load external resources (files, scale factors) onto self +@my_taf.teardown # free memory +``` + +## Calling sub-TAFs + +```python +events = self[other_producer](events, **kwargs) +events, sub_result = self[sub_selector](events, **kwargs) +``` diff --git a/.claude/rules/calibration.md b/.claude/rules/calibration.md new file mode 100644 index 000000000..ddff0bccf --- /dev/null +++ b/.claude/rules/calibration.md @@ -0,0 +1,56 @@ +--- +paths: + - "**/calibration/**" +--- + +# Calibrator Rules + +## Decorator pattern + +```python +from columnflow.calibration import Calibrator, calibrator +from columnflow.columnar_util import set_ak_column +from columnflow.util import maybe_import +ak = maybe_import("awkward") + +@calibrator( + uses={"Jet.{pt,eta,phi,mass,rawFactor,area}", "fixedGridRhoFastjetAll"}, + produces={"Jet.{pt,eta,phi,mass}"}, # nominal; add varied columns in init +) +def default(self: Calibrator, events: ak.Array, **kwargs) -> ak.Array: + corrected_pt = events.Jet.pt * self.jec_factor # loaded in setup() + events = set_ak_column(events, "Jet.pt", corrected_pt) + return events +``` + +## Hooks for external files and shifts + +```python +@default.requires +def default_requires(self, task, reqs): + from columnflow.tasks.external import BundleExternalFiles + reqs["ext"] = BundleExternalFiles.req(task) + +@default.setup +def default_setup(self, task, reqs, inputs, reader_targets): + bundle = inputs["ext"]["collection"][0] + # parse JEC file from bundle["jec"] and store correction objects on self + self.jec_factor = 1.02 # placeholder + +@default.init +def default_init(self): + from columnflow.config_util import get_shifts_from_sources + # declare which systematic shifts this calibrator is responsible for + self.shifts |= set(get_shifts_from_sources(self.config_inst, "jec")) + # add varied output columns for each shift + for shift_inst in self.shifts: + self.produces.add(f"Jet.pt_{shift_inst.name}") +``` + +## Key differences from Producers + +- Calibrators run on **raw NanoAOD events** (before selection/reduction). +- They typically **overwrite** kinematic columns (e.g. `Jet.pt`) rather than creating new named columns. +- Separate task branches run for each registered systematic shift. +- Varied columns (e.g. `Jet.pt_jec_up`) must be declared in `produces` dynamically in the `init` hook. +- Column aliases in `cfg` map `"Jet.pt"` → `"Jet.pt_jec_up"` when the `jec_up` shift is active. diff --git a/.claude/rules/categorization.md b/.claude/rules/categorization.md new file mode 100644 index 000000000..652d09ee8 --- /dev/null +++ b/.claude/rules/categorization.md @@ -0,0 +1,42 @@ +--- +paths: + - "**/categorization/**" +--- + +# Categorizer Rules + +## Decorator and return signature + +```python +from columnflow.categorization import Categorizer, categorizer +from columnflow.util import maybe_import +ak = maybe_import("awkward") + +@categorizer(uses={"event"}) +def cat_incl(self: Categorizer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + # Returns (events, boolean_mask_1d) — True = event belongs to this category + return events, ak.ones_like(events.event, dtype=bool) + +@categorizer(uses={"n_jet", "n_bjet"}) +def cat_sr(self: Categorizer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + return events, (events.n_jet >= 4) & (events.n_bjet >= 2) +``` + +## Registering in config + +```python +from columnflow.config_util import add_category + +add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") +add_category(cfg, name="sr", id=2, selection="cat_sr", label="Signal Region") +add_category(cfg, name="cr_1b", id=3, selection="cat_cr_1b", label="CR (1b)") +``` + +The `selection` argument is the **name of the Categorizer function** (a string), not a boolean expression. + +## Key rules + +- Always include `"event"` in `uses` for inclusive categorizers (it's always available). +- Categorizers run inside `CreateHistograms` after `ProduceColumns`; they read columns produced by Producers. +- Categories are mutually exclusive by convention; combine them via the `--categories` CLI flag to run several in parallel. +- Register the file in `law.cfg` under `categorization_modules`. diff --git a/.claude/rules/config.md b/.claude/rules/config.md new file mode 100644 index 000000000..66003b671 --- /dev/null +++ b/.claude/rules/config.md @@ -0,0 +1,100 @@ +--- +paths: + - "**/config/**" +--- + +# Config Object Rules + +## Hierarchy: Analysis → Config → Campaign → Dataset + +```python +import order as od + +analysis = od.Analysis(name="my_analysis", id=1) +cpn = od.Campaign(name="run2_2018", id=4, ecm=13, aux={"tier": "NanoAOD", "year": 2018}) +cfg = analysis.add_config(cpn, name="run2_2018", id=4) +``` + +## Adding order objects to Config + +```python +# Process (must be added before datasets that reference it) +cfg.add_process(procs.tt) + +# Dataset (fetched from Campaign by name) +cfg.add_dataset(cpn.get_dataset("tt_dl_powheg")) + +# Shift +cfg.add_shift(name="nominal", id=0) +cfg.add_shift(name="mu_up", id=1, type=od.Shift.SHAPE) + +# Variable +cfg.add_variable( + name="jet1_pt", + expression="Jet.pt[:,0]", # awkward expression evaluated in HistProducer + null_value=EMPTY_FLOAT, + binning=(40, 0.0, 400.0), + unit="GeV", + x_title=r"Leading jet $p_T$", +) + +# Category +from columnflow.config_util import add_category +add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") +``` + +## Required auxiliaries (cfg.x.*) + +```python +from scinum import Number +cfg.x.luminosity = Number(59740, {"lumi_13TeV_2018": 0.025j}) + +from columnflow.util import DotDict +from columnflow.columnar_util import ColumnCollection + +cfg.x.keep_columns = DotDict.wrap({ + "cf.ReduceEvents": { + "Jet.{pt,eta,phi,mass,btagDeepFlavB}", + "Electron.{pt,eta,phi,mass,charge}", + "Muon.{pt,eta,phi,mass,charge}", + "MET.{pt,phi}", + "event", "run", "luminosityBlock", + ColumnCollection.ALL_FROM_SELECTOR, + }, +}) + +# Default CLI argument values (override on command line) +cfg.x.default_calibrator = "default" +cfg.x.default_selector = "default" +cfg.x.default_producer = "default" +cfg.x.default_variables = ("n_jet", "jet1_pt") +``` + +## Shift aliases (weight-based uncertainties) + +```python +from columnflow.config_util import add_shift_aliases + +# Maps "muon_weight" → "muon_weight_up" when mu_up shift is active +add_shift_aliases(cfg, "mu", { + "muon_weight": "muon_weight_{direction}", +}) +``` + +## External files + +```python +cfg.x.external_files = DotDict.wrap({ + "muon_sf": "/path/to/muon_sf.json", + "btag_sf": ("/path/to/btag_sf.json", "v2"), # with version +}) +``` + +## Variable expression formats + +```python +"Jet.pt[:,0]" # leading jet pT (EMPTY_FLOAT if no jets — set null_value) +"Jet.pt" # all jets flattened (1D histogram across all events and jets) +"n_jet" # scalar column (from ProduceColumns) +"MET.pt" # scalar per-event field +``` diff --git a/.claude/rules/histogramming.md b/.claude/rules/histogramming.md new file mode 100644 index 000000000..aae4f2e99 --- /dev/null +++ b/.claude/rules/histogramming.md @@ -0,0 +1,54 @@ +--- +paths: + - "**/histogramming/**" +--- + +# HistProducer Rules + +## Extending cf_default (standard approach) + +```python +from columnflow.histogramming import HistProducer +from columnflow.histogramming.default import cf_default +from columnflow.columnar_util import Route +from columnflow.util import maybe_import +from columnflow.config_util import get_shifts_from_sources + +ak = maybe_import("awkward") +np = maybe_import("numpy") + + +@cf_default.hist_producer() +def default(self: HistProducer, events: ak.Array, **kwargs) -> tuple[ak.Array, ak.Array]: + """Return (events, per-event weight array).""" + weight = ak.Array(np.ones(len(events), dtype=np.float32)) + + if self.dataset_inst.is_mc and len(events): + for column in self.weight_columns: + weight = weight * Route(column).apply(events) + + return events, weight + + +@default.init +def default_init(self: HistProducer) -> None: + self.weight_columns = set() + if self.dataset_inst.is_data: + return + + # Add weight column names; declared in uses so they are loaded from disk + self.weight_columns |= {"normalization_weight", "muon_weight", "pu_weight"} + self.uses |= self.weight_columns + + # Declare which shift sources affect the event weight (for alias resolution) + self.shifts |= set(get_shifts_from_sources(self.config_inst, "mu", "minbias_xs")) +``` + +## Key rules + +- The `__call__` method returns `(events, weight_array)` — not just events. +- The weight is a 1D float32 array of length `len(events)`. Data events get weight = 1. +- Add weight columns to `self.uses` (not the decorator `uses`) in the `init` hook so they are dynamically included. +- Declare `self.shifts` in the `init` hook: the framework uses this to resolve column aliases when a shift is active (e.g. `muon_weight` → `muon_weight_up` under `mu_up`). +- `Route(column).apply(events)` safely reads a column; use it instead of `events[column]` when the column name is a variable. +- Register in `law.cfg` under `hist_production_modules`. diff --git a/.claude/rules/inference.md b/.claude/rules/inference.md new file mode 100644 index 000000000..1747d8aee --- /dev/null +++ b/.claude/rules/inference.md @@ -0,0 +1,80 @@ +--- +paths: + - "**/inference/**" +--- + +# InferenceModel Rules + +## Decorator and structure + +```python +from columnflow.inference import inference_model, ParameterType, ParameterTransformation + + +@inference_model +def default(self): + + # --- Categories (one per distribution used in the fit) --- + self.add_category( + "sr", + config_category="sr", # name of order Category in config + config_variable="jet1_pt", # variable whose histogram to use + config_data_datasets=["data_mu_b", "data_mu_c"], + mc_stats=True, # add MC stat uncertainty (Barlow-Beeston) + ) + + # --- Processes --- + self.add_process( + "TT", + is_signal=False, + config_process="tt", + config_mc_datasets=["tt_dl_powheg"], + ) + self.add_process( + "ST", + is_signal=True, + config_process="st", + ) + + # --- Parameters (nuisances) --- + + # Rate uncertainty (log-normal) + lumi = self.config_inst.x.luminosity + for unc_name in lumi.uncertainties: + self.add_parameter( + unc_name, + type=ParameterType.rate_gauss, + effect=lumi.get(names=unc_name, direction=("down", "up"), factor=True), + transformations=[ParameterTransformation.symmetrize], + ) + + # Weight-based shape uncertainty + self.add_parameter( + "mu", + process=["TT", "ST"], + type=ParameterType.shape, + config_shift_source="mu", # links to shifts mu_up / mu_down in config + ) + + # Selection-modifying shape uncertainty (JEC) + self.add_parameter( + "jec", + process=["TT", "ST"], + type=ParameterType.shape, + config_shift_source="jec", + ) + + # Dedicated-dataset uncertainty + self.add_parameter( + "tune", + process="TT", + type=ParameterType.shape, + config_shift_source="tune", + ) +``` + +## Key rules + +- `config_shift_source` must match the shift name prefix in the config (e.g. `"mu"` → shifts `mu_up` / `mu_down`). +- Register in `law.cfg` under `inference_modules`. +- Run with `law run cf.CreateDatacards --inference-model default --variables jet1_pt`. diff --git a/.claude/rules/production.md b/.claude/rules/production.md new file mode 100644 index 000000000..e6fd2c4b8 --- /dev/null +++ b/.claude/rules/production.md @@ -0,0 +1,104 @@ +--- +paths: + - "**/production/**" +--- + +# Producer Rules + +## Decorator pattern + +```python +from columnflow.production import Producer, producer +from columnflow.columnar_util import set_ak_column, EMPTY_FLOAT +from columnflow.util import maybe_import, four_vec +ak = maybe_import("awkward") +np = maybe_import("numpy") + +@producer( + uses={"Jet.{pt,eta}", "Electron.{pt,eta,phi,mass}"}, + produces={"ht", "n_jet", "jet1_pt"}, +) +def features(self: Producer, events: ak.Array, **kwargs) -> ak.Array: + events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) + events = set_ak_column(events, "n_jet", ak.num(events.Jet.pt, axis=1)) + from columnflow.columnar_util import Route + events = set_ak_column( + events, "jet1_pt", + Route("Jet.pt[:,0]").apply(events, null_value=EMPTY_FLOAT), + value_type=np.float32, + ) + return events +``` + +## four_vec helper + +```python +from columnflow.util import four_vec + +# Expands to Collection.{pt,eta,phi,mass} (+ extra fields) for each collection: +uses=four_vec({"Jet", "Electron"}, {"btagDeepFlavB"}) +# equivalent to: {"Jet.{pt,eta,phi,mass,btagDeepFlavB}", "Electron.{pt,eta,phi,mass}"} +``` + +## Nesting producers + +```python +@producer( + uses={sub_producer, "extra_col"}, # sub_producer's uses are inherited + produces={sub_producer, "my_col"}, # sub_producer's produces are inherited +) +def parent(self: Producer, events: ak.Array, **kwargs) -> ak.Array: + events = self[sub_producer](events, **kwargs) + events = set_ak_column(events, "my_col", ...) + return events +``` + +## MC-conditional columns — use init hook + +```python +@my_producer.init +def my_producer_init(self: Producer) -> None: + if not getattr(self, "dataset_inst", None) or self.dataset_inst.is_data: + return + self.uses.add("mc_weight") + self.produces.add("event_weight") + +def my_producer(self: Producer, events: ak.Array, **kwargs) -> ak.Array: + if self.dataset_inst.is_mc: + events = set_ak_column(events, "event_weight", events.mc_weight * ...) + return events +``` + +## Loading external files (scale factors, corrections) + +```python +@sf_producer.requires +def sf_producer_requires(self, task, reqs): + from columnflow.tasks.external import BundleExternalFiles + reqs["ext"] = BundleExternalFiles.req(task) + +@sf_producer.setup +def sf_producer_setup(self, task, reqs, inputs, reader_targets): + bundle = inputs["ext"]["collection"][0] + self.sf_file = bundle["muon_sf"].load() # store on self for use in __call__ +``` + +## Common column patterns + +```python +# Scalar sum of jet pT +events = set_ak_column(events, "ht", ak.sum(events.Jet.pt, axis=1), value_type=np.float32) + +# Count objects +events = set_ak_column(events, "n_jet", ak.num(events.Jet.pt, axis=1)) + +# Leading object pT with EMPTY_FLOAT for events with no objects +from columnflow.columnar_util import Route, EMPTY_FLOAT +events = set_ak_column(events, "jet1_pt", + Route("Jet.pt[:,0]").apply(events, null_value=EMPTY_FLOAT), value_type=np.float32) + +# Invariant mass (requires attach_coffea_behavior first) +from columnflow.production.util import attach_coffea_behavior +events = self[attach_coffea_behavior](events, **kwargs) +mll = (events.Lepton[:, 0] + events.Lepton[:, 1]).mass +``` diff --git a/.claude/rules/selection.md b/.claude/rules/selection.md new file mode 100644 index 000000000..633e1d94b --- /dev/null +++ b/.claude/rules/selection.md @@ -0,0 +1,88 @@ +--- +paths: + - "**/selection/**" +--- + +# Selector Rules + +## Decorator pattern + +```python +from columnflow.selection import Selector, SelectionResult, selector +from columnflow.util import maybe_import +ak = maybe_import("awkward") + +@selector( + uses={"Jet.{pt,eta}"}, + produces=set(), + exposed=False, # True only for the top-level CLI-reachable selector +) +def jet_selection(self: Selector, events: ak.Array, **kwargs) -> tuple[ak.Array, SelectionResult]: + jet_mask = (events.Jet.pt > 25) & (abs(events.Jet.eta) < 2.4) + jet_indices = ak.local_index(events.Jet.pt, axis=1)[jet_mask] + return events, SelectionResult( + steps={"jet": ak.sum(jet_mask, axis=1) >= 2}, + objects={"Jet": {"Jet": jet_indices}}, + aux={"jet_mask": jet_mask}, # discarded after ReduceEvents + ) +``` + +## SelectionResult fields + +```python +SelectionResult( + steps={ + "step_name": bool_1d_array, # per-event; applied by ReduceEvents + }, + objects={ + "SourceField": { + "DestField": index_array, # ak.local_index(...)[mask] + }, + }, + aux={"key": value}, # temporary; not persisted after ReduceEvents + event=combined_mask, # required on the exposed selector's final result +) +``` + +## Composing sub-selectors + +```python +results = SelectionResult() +events, jet_result = self[jet_selection](events, **kwargs) +results += jet_result +events, lep_result = self[lepton_selection](events, results, **kwargs) +results += lep_result + +# Combine all steps into the final event mask (required for ReduceEvents) +from operator import and_ +from functools import reduce +results.event = reduce(and_, results.steps.values()) +``` + +## stats.json — increment in the exposed selector + +```python +stats["num_events"] += len(events) +stats["num_events_selected"] += ak.sum(results.event, axis=0) +if self.dataset_inst.is_mc: + stats["sum_mc_weight"] += ak.sum(events.mc_weight) + stats["sum_mc_weight_selected"] += ak.sum(events.mc_weight[results.event]) +``` + +Or use `from columnflow.selection.stats import increment_stats` (built-in helper). + +## Accessing selected objects downstream within selection + +```python +# Objects already reduced by results.objects from a prior sub-selector: +muon = events.Muon[results.objects.Muon.Muon] +jet = events.Jet[results.objects.Jet.Jet] +``` + +## sorted_indices_from_mask helper + +```python +from columnflow.columnar_util import sorted_indices_from_mask +# Sort selected objects by pt descending: +indices = sorted_indices_from_mask(object_mask, events.Jet.pt, ascending=False) +``` diff --git a/CLAUDE.md b/CLAUDE.md index 93bbd4e90..20785a800 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,24 +1,24 @@ -# Columnflow Project — Claude Instructions - -This is a **columnflow** analysis project. Columnflow is a fully orchestrated columnar HEP analysis framework built on [law](https://github.com/riga/law) (workflow orchestration) and [order](https://github.com/riga/order) (metadata management), using [awkward-array](https://awkward-array.org) for columnar event data. - -Read all five guide files before writing or modifying any code: - -| Guide | Contents | -|---|---| -| [01 Framework Structure](docs/columnflow_claude_guide/01_framework_structure.md) | Big picture, analysis pipeline, task graph, config objects (Analysis/Campaign/Config), systematics | -| [02 Coding Style](docs/columnflow_claude_guide/02_coding_style.md) | TAF decorators, uses/produces, imports, module registration, naming conventions | -| [03 Task User Guide](docs/columnflow_claude_guide/03_task_commands.md) | Per-task explanation, `law run` commands, key parameters, tips | -| [04 Awkward Reference](docs/columnflow_claude_guide/04_awkward_reference.md) | Key `ak.*` functions used in columnflow code | -| [05 Dos and Don'ts](docs/columnflow_claude_guide/05_dos_and_donts.md) | Common mistakes and best practices for new users | -| [06 Custom Tasks](docs/columnflow_claude_guide/06_custom_tasks.md) | Writing Selectors, Producers, Calibrators, HistProducers, Categorizers, custom law tasks | - -## Quick orientation - -- All data operations act on **chunks** of events encoded as `ak.Array` named `events`. -- Task array functions (TAFs): `Calibrator`, `Selector`, `Producer`, `Reducer`, `HistProducer`. -- Each TAF declares `uses` (columns to read) and `produces` (columns to write) in its decorator. -- Events are **never modified in-place** — always use `set_ak_column(events, "FieldName", value)` and reassign `events`. -- The standard pipeline order: `GetDatasetLFNs → CalibrateEvents → SelectEvents → ReduceEvents → ProduceColumns → CreateHistograms → PlotVariables1D`. -- All tasks are run with `law run cf. --version [--config ] [...]`. -- New Python modules must be registered in `law.cfg` under the correct `*_modules` key. +# Columnflow Project + +Backend for columnar, fully orchestrated HEP analyses with [law](https://github.com/riga/law) (workflow orchestration), [order](https://github.com/riga/order) (metadata), and [awkward-array](https://awkward-array.org) (columnar event data). + +## Always-loaded rules (in `.claude/rules/`) + +- **01-pipeline.md** — task order, 5 TAF types, law.cfg module registration +- **02-invariants.md** — set_ak_column, uses/produces, imports, Selector return type, MC guard, keep_columns, no loops + +Path-scoped rules load automatically when Claude works with files in the corresponding directory: +`selection/`, `production/`, `calibration/`, `config/`, `categorization/`, `histogramming/`, `inference/` + +## Commands + +```bash +law run cf. --version [--config ] [--dataset ] [--branch 0] +law run cf.PlotVariables1D --version dev1 --processes tt --variables jet1_pt --print-status -1 +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --remove-output 0,a,y +``` + +## User guide + +Full documentation for users is in `docs/columnflow_claude_guide/`: +`01_framework_structure.md` · `02_coding_style.md` · `03_task_commands.md` · `04_awkward_reference.md` · `05_dos_and_donts.md` · `06_custom_tasks.md` From ed91294a441183c5b818d6f23ea9ed39ea3a4e05 Mon Sep 17 00:00:00 2001 From: juvanden Date: Wed, 24 Jun 2026 11:55:11 +0200 Subject: [PATCH 4/7] docs: add custom-tasks and linting rules for Claude Code Adds two new .claude/rules/ files: - custom-tasks.md (path-scoped to tasks/**): patterns for writing custom law tasks that extend the columnflow pipeline, with examples for histogram-consuming tasks, per-category datacard workflows, and custom plot tasks derived from hh2bbww - linting.md (always-loaded): instructs Claude to run flake8 after writing code and fix all errors using the project .flake8 config Updates CLAUDE.md to reference both new rules. Co-Authored-By: Claude Sonnet 4.6 --- .claude/rules/custom-tasks.md | 180 ++++++++++++++++++++++++++++++++++ .claude/rules/linting.md | 37 +++++++ CLAUDE.md | 3 +- 3 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 .claude/rules/custom-tasks.md create mode 100644 .claude/rules/linting.md diff --git a/.claude/rules/custom-tasks.md b/.claude/rules/custom-tasks.md new file mode 100644 index 000000000..435540de2 --- /dev/null +++ b/.claude/rules/custom-tasks.md @@ -0,0 +1,180 @@ +--- +paths: + - "**/tasks/**" +--- + +# Custom Law Tasks + +Custom tasks extend the columnflow pipeline for non-standard outputs (custom plots, derived datacards, scale factors, etc.). + +## Skeleton + +```python +import law +from columnflow.tasks.framework.base import AnalysisTask, ConfigTask, Requirements + +class MyTask(ConfigTask, law.LocalWorkflow): + # class-level task_namespace for the analysis (set once in the analysis base class) + task_namespace = "hbw" + + # Upstream task requirements + def requires(self): + return SomeUpstreamTask.req(self) + + # Branch map: maps branch index → payload (for LocalWorkflow only) + def create_branch_map(self): + return list(self.categories) # one branch per category + + # Output target(s) for this branch + def output(self): + return { + "data": self.local_target("result.json"), + "plot": self.local_target("plot.pdf"), + } + + # Actual computation + def run(self): + inp = self.input() + data = inp["data"].load(formatter="json") + # ... compute ... + self.output()["data"].dump(result, formatter="json") +``` + +## Analysis base class (define once per analysis) + +```python +class HBWTask(AnalysisTask): + task_namespace = "hbw" # prefix for all custom tasks +``` + +## Pattern 1 — Histogram-consuming task (trigger SF, custom plots from histograms) + +Inherits `HistogramsUserSingleShiftBase` which already wires up `MergeHistograms` requirements for all configured datasets and hist_producers. + +```python +from columnflow.tasks.framework.histograms import HistogramsUserSingleShiftBase +from columnflow.tasks.framework.mixins import DatasetsProcessesMixin + +class ComputeMyScaleFactors(HBWTask, HistogramsUserSingleShiftBase, DatasetsProcessesMixin, law.LocalWorkflow): + + def create_branch_map(self): + return list(self.hist_producer_insts) # one branch per hist_producer + + def requires(self): + # HistogramsUserSingleShiftBase provides .requires_histograms() helper + reqs = {} + for dataset_inst in self.dataset_insts: + reqs[dataset_inst.name] = MergeHistograms.req( + self, + dataset=dataset_inst.name, + branch=-1, + ) + return reqs + + def output(self): + return { + "json": self.local_target("corrections.json"), + "plot": self.local_target("efficiency.pdf"), + } + + def run(self): + inputs = self.input() + # load histograms using inherited helper + hist = self.load_histogram(inputs, self.config_inst, dataset_inst, variable) + # compute data/MC ratios, write CorrectionLib JSON, save plots + self.output()["json"].dump(corrections, formatter="json") +``` + +## Pattern 2 — Custom datacard writer (per-category LocalWorkflow) + +Inherits `SerializeInferenceModelBase` which handles inference model + histogram loading. + +```python +from columnflow.tasks.framework.inference import SerializeInferenceModelBase +from columnflow.tasks.histograms import MergeHistograms + +class CreateMultipleDatacards(HBWTask, SerializeInferenceModelBase, law.LocalWorkflow): + + def create_branch_map(self): + # one branch per analysis category + return {i: cat for i, cat in enumerate(self.inference_model_inst.categories)} + + def requires(self): + cat = self.branch_data + reqs = {} + for dataset_inst in self.dataset_insts: + reqs[dataset_inst.name] = MergeHistograms.req( + self, + dataset=dataset_inst.name, + branch=-1, + ) + return reqs + + def output(self): + cat = self.branch_data + return { + "datacard": self.local_target(f"{cat.name}/datacard.txt"), + "shapes": self.local_target(f"{cat.name}/shapes.root"), + } + + def run(self): + inputs = self.input() + cat = self.branch_data + # write datacard for this category + writer = DatacardWriter(self.inference_model_inst, cat) + writer.run(inputs, self.output()) +``` + +## Pattern 3 — Custom plot task (reads external ROOT file) + +```python +from columnflow.plots.util import PlotBase1D +from columnflow.tasks.framework.mixins import ( + DatasetsProcessesMixin, CalibratorClassesMixin, SelectorClassMixin, + ProducerClassesMixin, HistProducerClassMixin, InferenceModelMixin, +) + +class PlotPostfitShapes( + HBWTask, + PlotBase1D, + DatasetsProcessesMixin, + InferenceModelMixin, + HistProducerClassMixin, +): + # No upstream law tasks required — reads an external file directly + def requires(self): + return {} + + def output(self): + return law.LocalDirectoryTarget(self.local_path("plots")) + + def run(self): + fit_file = self.input() # or open an external file + # load histograms from ROOT, call self.call_plot_func(...) + self.output().touch() +``` + +## Registering custom tasks in law.cfg + +```ini +[modules] +# ... existing columnflow modules ... +hbw.tasks.trigger_sf +hbw.tasks.multiple_datacards +hbw.tasks.postfit_plots +``` + +Run with: +```bash +law run hbw.ComputeMyScaleFactors --version dev1 --config run2_2018 --branch 0 +law run hbw.CreateMultipleDatacards --version dev1 --inference-model default +``` + +## Key rules + +- Always call `.req(self, ...)` on upstream tasks instead of constructing them directly — this propagates shared parameters. +- `self.local_target(path)` creates a `law.LocalFileTarget` inside the task's store path. For directories use `law.LocalDirectoryTarget(self.local_path(...))`. +- `law.LocalWorkflow` requires `create_branch_map()` returning a dict or list; each entry becomes one branch. +- Without `LocalWorkflow`, the task is a single-branch task — omit `create_branch_map()`. +- For remote execution add `HTCondorWorkflow` or `SlurmWorkflow` as additional base classes. +- Always register the module path in `law.cfg [modules]` so `law` can discover the task. diff --git a/.claude/rules/linting.md b/.claude/rules/linting.md new file mode 100644 index 000000000..e19038b15 --- /dev/null +++ b/.claude/rules/linting.md @@ -0,0 +1,37 @@ +# Linting + +After writing or editing any Python file, run flake8 and fix all reported errors before considering the task done. + +## Command + +```bash +flake8 --config=.flake8 +``` + +Example: +```bash +flake8 --config=.flake8 hbw/tasks/trigger_sf.py +flake8 --config=.flake8 hbw/ +``` + +## Project settings (from `.flake8`) + +- `max-line-length = 120` +- `ignore = E128, E306, E402, E722, E731, W504, Q003` +- `inline-quotes = double` — use `"..."` not `'...'` for strings + +## Workflow + +1. Write or edit the Python code. +2. Run `flake8 --config=.flake8 `. +3. Fix every reported error. +4. Re-run flake8 to confirm zero errors. +5. Only then consider the code complete. + +## Common fixes + +- **E501** (line too long): break at operator, wrap in parentheses, or shorten names +- **F401** (unused import): remove the import +- **E711** (comparison to None): use `is None` / `is not None`, not `== None` +- **W291/W293** (trailing whitespace): remove trailing spaces +- **Q000** (wrong quote type): replace `'...'` with `"..."` diff --git a/CLAUDE.md b/CLAUDE.md index 20785a800..b276f2d83 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,9 +6,10 @@ Backend for columnar, fully orchestrated HEP analyses with [law](https://github. - **01-pipeline.md** — task order, 5 TAF types, law.cfg module registration - **02-invariants.md** — set_ak_column, uses/produces, imports, Selector return type, MC guard, keep_columns, no loops +- **linting.md** — after writing code run `flake8 --config=.flake8 ` and fix all errors Path-scoped rules load automatically when Claude works with files in the corresponding directory: -`selection/`, `production/`, `calibration/`, `config/`, `categorization/`, `histogramming/`, `inference/` +`selection/`, `production/`, `calibration/`, `config/`, `categorization/`, `histogramming/`, `inference/`, `tasks/` ## Commands From 7fb56d7ac680d5e69beb121c46f90b130c9248c1 Mon Sep 17 00:00:00 2001 From: juvanden Date: Wed, 24 Jun 2026 14:04:40 +0200 Subject: [PATCH 5/7] docs: add github behaviour rule and refine linting workflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - .claude/rules/github.md: new always-loaded rule defining what Claude is and is not allowed to do with git/GitHub — commits and branch creation are allowed freely; push and rebase require confirmation; branch/commit deletion is forbidden; PRs must never go directly from a feature branch to main (intermediate branch required) - .claude/rules/linting.md: restructured into "while coding" (follow .flake8 rules proactively) vs "before pushing" (run run_linting, commit fixes as "linting fixes"); added tests section with sandbox and no-sandbox commands - CLAUDE.md: updated always-loaded rules list to reference both new files Co-Authored-By: Claude Sonnet 4.6 --- .claude/rules/github.md | 31 ++++++++++++++++ .claude/rules/linting.md | 80 ++++++++++++++++++++++++++++++---------- CLAUDE.md | 3 +- 3 files changed, 94 insertions(+), 20 deletions(-) create mode 100644 .claude/rules/github.md diff --git a/.claude/rules/github.md b/.claude/rules/github.md new file mode 100644 index 000000000..b65af97c7 --- /dev/null +++ b/.claude/rules/github.md @@ -0,0 +1,31 @@ +# GitHub / Git Behaviour + +## Allowed (no confirmation needed) + +- Commit to the current branch (any branch, including main) +- Create a new branch for feature development + +## Not allowed — stop and ask the user first + +- **Push to any branch** — always ask for confirmation before pushing, even if the user says "push it" +- **Rebase** — never rebase without explicit permission for that specific rebase +- **Delete a branch or commit** — never delete; warn the user if they request it + +## Pull request rules + +- **Never create a PR from a feature branch directly to main/master.** An intermediate integration branch must always exist between the feature branch and main. +- Branch hierarchy: `feature/* → integration/staging branch → main` +- If asked to open a PR to main directly, refuse and explain that an intermediate branch is required. + +## Summary table + +| Action | Allowed | +|---|---| +| `git commit` | Yes | +| `git checkout -b ` | Yes | +| `git push` | Ask first | +| `git rebase` | Ask first | +| `git branch -d` / `git branch -D` | Never | +| `git reset --hard` / force operations | Never | +| PR: feature → main | Never (use intermediate branch) | +| PR: feature → intermediate | Yes, after push confirmation | diff --git a/.claude/rules/linting.md b/.claude/rules/linting.md index e19038b15..43add585d 100644 --- a/.claude/rules/linting.md +++ b/.claude/rules/linting.md @@ -1,37 +1,79 @@ # Linting -After writing or editing any Python file, run flake8 and fix all reported errors before considering the task done. +## While writing code -## Command +Always respect the `.flake8` settings when writing code — do not produce code that violates them: -```bash -flake8 --config=.flake8 -``` +- `max-line-length = 120` +- `ignore = E128, E306, E402, E722, E731, W504, Q003` +- `inline-quotes = double` — use `"..."` not `'...'` for strings + +## Before pushing to GitHub / GitLab + +Run flake8 and fix all errors before pushing. Commit the fixes separately with the message `linting fixes`. + +### Prerequisites + +`flake8` requires the columnflow environment: -Example: ```bash -flake8 --config=.flake8 hbw/tasks/trigger_sf.py -flake8 --config=.flake8 hbw/ +source setup.sh default ``` -## Project settings (from `.flake8`) +### Commands -- `max-line-length = 120` -- `ignore = E128, E306, E402, E722, E731, W504, Q003` -- `inline-quotes = double` — use `"..."` not `'...'` for strings +```bash +# Full project lint (canonical — mirrors CI) +bash tests/run_linting -## Workflow +# Single file or directory +flake8 --config=.flake8 +``` + +### Workflow -1. Write or edit the Python code. -2. Run `flake8 --config=.flake8 `. -3. Fix every reported error. -4. Re-run flake8 to confirm zero errors. -5. Only then consider the code complete. +1. Finish the feature/fix and commit it. +2. Run `bash tests/run_linting`. +3. If errors are reported: fix them, then commit with message `linting fixes`. +4. Push only when linting is clean. -## Common fixes +### Common fixes - **E501** (line too long): break at operator, wrap in parentheses, or shorten names - **F401** (unused import): remove the import - **E711** (comparison to None): use `is None` / `is not None`, not `== None` - **W291/W293** (trailing whitespace): remove trailing spaces - **Q000** (wrong quote type): replace `'...'` with `"..."` + +--- + +# Tests + +After writing code that touches columnflow core, run the unit tests. + +## Prerequisites + +Same environment: `source setup.sh default` + +## Commands + +```bash +# All tests (some require the columnar sandbox — may be slow) +bash tests/run_tests + +# Tests that run without a sandbox (fast, always available) +python -m unittest tests.test_util tests.test_task_parameters tests.test_base_tasks + +# Tests that need the columnar venv sandbox +bash tests/run_test test_columnar_util sandboxes/venv_columnar.sh +bash tests/run_test test_config_util sandboxes/venv_columnar.sh +bash tests/run_test test_inference sandboxes/venv_columnar.sh +bash tests/run_test test_hist_util sandboxes/venv_columnar.sh +bash tests/run_test test_plotting sandboxes/venv_columnar.sh +``` + +## Workflow + +1. Run `python -m unittest tests.test_util tests.test_task_parameters tests.test_base_tasks` for a quick sanity check (no sandbox needed). +2. If changes touch columnar utilities, histogramming, inference, or plotting, run the full `bash tests/run_tests`. +3. Fix any failures before pushing. diff --git a/CLAUDE.md b/CLAUDE.md index b276f2d83..68bd80ef8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,7 +6,8 @@ Backend for columnar, fully orchestrated HEP analyses with [law](https://github. - **01-pipeline.md** — task order, 5 TAF types, law.cfg module registration - **02-invariants.md** — set_ak_column, uses/produces, imports, Selector return type, MC guard, keep_columns, no loops -- **linting.md** — after writing code run `flake8 --config=.flake8 ` and fix all errors +- **linting.md** — respect `.flake8` rules while coding; run `bash tests/run_linting` before pushing and commit fixes as `linting fixes` +- **github.md** — allowed/forbidden git actions; always ask before pushing or rebasing; never delete branches; never PR directly to main Path-scoped rules load automatically when Claude works with files in the corresponding directory: `selection/`, `production/`, `calibration/`, `config/`, `categorization/`, `histogramming/`, `inference/`, `tasks/` From ea973403b5c55c054610163d742797fad723ce33 Mon Sep 17 00:00:00 2001 From: juvanden Date: Wed, 24 Jun 2026 14:28:21 +0200 Subject: [PATCH 6/7] linting fixes --- docs/columnflow_claude_guide/01_framework_structure.md | 9 +++++---- docs/columnflow_claude_guide/02_coding_style.md | 2 +- docs/columnflow_claude_guide/03_task_commands.md | 1 + docs/columnflow_claude_guide/04_awkward_reference.md | 1 + 4 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/columnflow_claude_guide/01_framework_structure.md b/docs/columnflow_claude_guide/01_framework_structure.md index 381509465..593c73dfb 100644 --- a/docs/columnflow_claude_guide/01_framework_structure.md +++ b/docs/columnflow_claude_guide/01_framework_structure.md @@ -21,6 +21,7 @@ It solves three practical problems every HEP analyst faces: [Law](https://github.com/riga/law) (built on [Luigi](https://luigi.readthedocs.io)) defines the **task graph**: every processing step is a `Task` object with declared outputs. When you run a downstream task (e.g. `PlotVariables1D`), law checks which upstream tasks are missing and runs them first — automatically. Key properties: + - Tasks are **idempotent**: if the output already exists, the task is skipped. - Tasks run **locally** or on **HTCondor / Slurm** with a single flag change (`--workflow htcondor`). - The `--version` tag separates parallel analysis branches on disk — you can have `dev1` and `prod1` coexist. @@ -39,7 +40,7 @@ All event data is stored as [awkward arrays](https://awkward-array.org): irregul The standard pipeline processes data from ROOT files to plots and datacards: -``` +```text Raw NanoAOD ROOT files │ ▼ @@ -121,7 +122,7 @@ All TAFs share the same decorator pattern and lifecycle. See [02 Coding Style](0 ## Directory Structure of an Analysis -``` +```text myanalysis/ ├── analysis/ │ └── my_analysis.py # Creates Analysis, Campaign, Config objects @@ -157,7 +158,7 @@ myanalysis/ ### Hierarchy -``` +```text Analysis ─── top-level container (rarely holds analysis logic itself) └── Config ─── analysis + campaign combination; carries all per-config settings └── Campaign ─── experimental period (year, energy, tier, file locations) @@ -349,7 +350,7 @@ cfg.add_shift(name="tune_up", id=20, type=od.Shift.SHAPE, tags={"disjoint_from_n ## Data Flow Summary -``` +```text ROOT files → chunks of ak.Array named "events" │ [Calibrator] adds corrected columns (never removes) diff --git a/docs/columnflow_claude_guide/02_coding_style.md b/docs/columnflow_claude_guide/02_coding_style.md index 233036555..e42553bde 100644 --- a/docs/columnflow_claude_guide/02_coding_style.md +++ b/docs/columnflow_claude_guide/02_coding_style.md @@ -4,7 +4,7 @@ A typical analysis module tree mirrors columnflow's own structure: -``` +```text myanalysis/ ├── analysis/ │ └── my_analysis.py # Analysis + Campaign + Config definitions diff --git a/docs/columnflow_claude_guide/03_task_commands.md b/docs/columnflow_claude_guide/03_task_commands.md index 605ac5313..0e3ad2e01 100644 --- a/docs/columnflow_claude_guide/03_task_commands.md +++ b/docs/columnflow_claude_guide/03_task_commands.md @@ -38,6 +38,7 @@ law run cf.GetDatasetLFNs --version dev1 --dataset tt_dl_powheg --print-output 0 ``` **Tips:** + - Run this first for all datasets before starting the pipeline. - If your datasets are stored locally or on a custom SE, configure `cfg.x.get_dataset_lfns` to bypass DAS. - For debugging, add `--branch 0` to retrieve only the first file entry. diff --git a/docs/columnflow_claude_guide/04_awkward_reference.md b/docs/columnflow_claude_guide/04_awkward_reference.md index 2dfb82fc2..cebae8110 100644 --- a/docs/columnflow_claude_guide/04_awkward_reference.md +++ b/docs/columnflow_claude_guide/04_awkward_reference.md @@ -3,6 +3,7 @@ Columnflow event data is stored as `ak.Array` named `events`. Fields like `events.Jet` are collections (variable-length sub-arrays); scalar fields like `events.MET.pt` are uniform arrays over events. Always import awkward via `maybe_import`: + ```python from columnflow.util import maybe_import ak = maybe_import("awkward") From ef4eafae6e3d885eae0fb3fbf61adc7f1a1c4252 Mon Sep 17 00:00:00 2001 From: juvanden Date: Wed, 24 Jun 2026 14:32:17 +0200 Subject: [PATCH 7/7] docs: expand linting rule to reflect full CI test suite Updates .claude/rules/linting.md with accurate CI coverage: - adds markdown linting (pymarkdown / run_docs lint) and breakpoint check - adds CI jobs table mapping each job to its script and scope - adds per-module test table with sandbox requirements and what each covers - notes test_selectionresults.py is a manual script, not run by run_tests Updates CLAUDE.md one-liner to mention all three pre-push checks. Co-Authored-By: Claude Sonnet 4.6 --- .claude/rules/linting.md | 91 +++++++++++++++++++++++++--------------- CLAUDE.md | 2 +- 2 files changed, 58 insertions(+), 35 deletions(-) diff --git a/.claude/rules/linting.md b/.claude/rules/linting.md index 43add585d..407aa60d2 100644 --- a/.claude/rules/linting.md +++ b/.claude/rules/linting.md @@ -7,64 +7,72 @@ Always respect the `.flake8` settings when writing code — do not produce code - `max-line-length = 120` - `ignore = E128, E306, E402, E722, E731, W504, Q003` - `inline-quotes = double` — use `"..."` not `'...'` for strings +- Never leave active `breakpoint()` calls (CI scans for them — add `# noqa` if intentional) ## Before pushing to GitHub / GitLab -Run flake8 and fix all errors before pushing. Commit the fixes separately with the message `linting fixes`. +Run all checks below and fix any failures. Commit fixes separately with the message `linting fixes`. -### Prerequisites +All commands require the columnflow environment (`source setup.sh default`). -`flake8` requires the columnflow environment: +### 1. Python linting — `bash tests/run_linting` -```bash -source setup.sh default -``` - -### Commands +Runs `flake8 columnflow tests bin docs setup.py`. Must be clean before pushing. ```bash -# Full project lint (canonical — mirrors CI) bash tests/run_linting - -# Single file or directory -flake8 --config=.flake8 +# or for a single file: +flake8 --config=.flake8 ``` -### Workflow +### 2. Markdown linting — `bash tests/run_docs lint` -1. Finish the feature/fix and commit it. -2. Run `bash tests/run_linting`. -3. If errors are reported: fix them, then commit with message `linting fixes`. -4. Push only when linting is clean. +Runs `pymarkdown` on `README.md` and `docs/`. Rules configured in `.markdownlint`: -### Common fixes +- MD031: fenced code blocks must be surrounded by blank lines +- MD032: lists must be surrounded by blank lines +- MD040: fenced code blocks must have a language tag (use `text` for diagrams/plain text) -- **E501** (line too long): break at operator, wrap in parentheses, or shorten names -- **F401** (unused import): remove the import -- **E711** (comparison to None): use `is None` / `is not None`, not `== None` -- **W291/W293** (trailing whitespace): remove trailing spaces -- **Q000** (wrong quote type): replace `'...'` with `"..."` +```bash +bash tests/run_docs lint +``` + +### 3. Breakpoint check (CI only — no local script) + +CI scans for uncommented `breakpoint()` calls. Do not leave them in pushed code. --- -# Tests +## CI jobs on push / pull request -After writing code that touches columnflow core, run the unit tests. +| Job | Script | What it checks | +|---|---|---| +| `lint` | `tests/run_linting` | flake8 on `columnflow/`, `tests/`, `bin/`, `docs/`, `setup.py` | +| `lint_docs` | `tests/run_docs lint` | pymarkdown on `README.md` and `docs/` | +| `breaks` | inline `git grep` | no active `breakpoint()` calls | +| `test` | `tests/run_tests` | all unit tests (see below) | +| `pypi` | `python setup.py sdist` + `twine check` | package builds and passes PyPI checks | +| `coverage` | `tests/run_coverage` | coverage upload to Codecov | -## Prerequisites +--- -Same environment: `source setup.sh default` +## Unit tests -## Commands +### Run all tests ```bash -# All tests (some require the columnar sandbox — may be slow) bash tests/run_tests +``` + +### Tests without a sandbox (fast) -# Tests that run without a sandbox (fast, always available) +```bash python -m unittest tests.test_util tests.test_task_parameters tests.test_base_tasks +``` + +### Tests requiring the columnar venv sandbox -# Tests that need the columnar venv sandbox +```bash bash tests/run_test test_columnar_util sandboxes/venv_columnar.sh bash tests/run_test test_config_util sandboxes/venv_columnar.sh bash tests/run_test test_inference sandboxes/venv_columnar.sh @@ -72,8 +80,23 @@ bash tests/run_test test_hist_util sandboxes/venv_columnar.sh bash tests/run_test test_plotting sandboxes/venv_columnar.sh ``` -## Workflow +### What each test module covers + +| Module | Sandbox | Covers | +|---|---|---| +| `test_util` | no | `maybe_import`, utility functions (`save_div`, `try_int`, `is_regex`, ...) | +| `test_columnar_util` | yes | `Route` class: join, split, apply, tags | +| `test_config_util` | yes | `get_events_from_categories` and config utilities | +| `test_inference` | yes | `InferenceModel` specs: process, category, parameter, parameter groups | +| `test_hist_util` | yes | `create_hist_from_variables`, `translate_hist_intcat_to_strcat` | +| `test_task_parameters` | no | `SettingsParameter`, `MultiSettingsParameter` | +| `test_base_tasks` | no | `AnalysisTask`: resolve config, categories, variables, datasets, shifts | +| `test_plotting` | yes | `PlotUtil`, confusion matrix, ROC curve plot utilities | + +> `test_selectionresults.py` is a standalone manual script, not a unittest — it is not run by `run_tests`. + +### Workflow -1. Run `python -m unittest tests.test_util tests.test_task_parameters tests.test_base_tasks` for a quick sanity check (no sandbox needed). -2. If changes touch columnar utilities, histogramming, inference, or plotting, run the full `bash tests/run_tests`. +1. Run the no-sandbox tests first for a quick check. +2. If changes touch columnar utilities, histogramming, inference, or plotting, run `bash tests/run_tests`. 3. Fix any failures before pushing. diff --git a/CLAUDE.md b/CLAUDE.md index 68bd80ef8..974b50084 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,7 +6,7 @@ Backend for columnar, fully orchestrated HEP analyses with [law](https://github. - **01-pipeline.md** — task order, 5 TAF types, law.cfg module registration - **02-invariants.md** — set_ak_column, uses/produces, imports, Selector return type, MC guard, keep_columns, no loops -- **linting.md** — respect `.flake8` rules while coding; run `bash tests/run_linting` before pushing and commit fixes as `linting fixes` +- **linting.md** — respect `.flake8` rules while coding; before pushing run `bash tests/run_linting` (Python) + `bash tests/run_docs lint` (markdown) + `bash tests/run_tests` (unit tests); commit fixes as `linting fixes` - **github.md** — allowed/forbidden git actions; always ask before pushing or rebasing; never delete branches; never PR directly to main Path-scoped rules load automatically when Claude works with files in the corresponding directory: