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/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/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/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/linting.md b/.claude/rules/linting.md new file mode 100644 index 000000000..407aa60d2 --- /dev/null +++ b/.claude/rules/linting.md @@ -0,0 +1,102 @@ +# Linting + +## While writing code + +Always respect the `.flake8` settings when writing code — do not produce code that violates them: + +- `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 all checks below and fix any failures. Commit fixes separately with the message `linting fixes`. + +All commands require the columnflow environment (`source setup.sh default`). + +### 1. Python linting — `bash tests/run_linting` + +Runs `flake8 columnflow tests bin docs setup.py`. Must be clean before pushing. + +```bash +bash tests/run_linting +# or for a single file: +flake8 --config=.flake8 +``` + +### 2. Markdown linting — `bash tests/run_docs lint` + +Runs `pymarkdown` on `README.md` and `docs/`. Rules configured in `.markdownlint`: + +- 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) + +```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. + +--- + +## CI jobs on push / pull request + +| 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 | + +--- + +## Unit tests + +### Run all tests + +```bash +bash tests/run_tests +``` + +### Tests without a sandbox (fast) + +```bash +python -m unittest tests.test_util tests.test_task_parameters tests.test_base_tasks +``` + +### Tests requiring 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 +bash tests/run_test test_hist_util sandboxes/venv_columnar.sh +bash tests/run_test test_plotting sandboxes/venv_columnar.sh +``` + +### 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 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/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 new file mode 100644 index 000000000..974b50084 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,26 @@ +# 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 +- **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: +`selection/`, `production/`, `calibration/`, `config/`, `categorization/`, `histogramming/`, `inference/`, `tasks/` + +## 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` 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..593c73dfb --- /dev/null +++ b/docs/columnflow_claude_guide/01_framework_structure.md @@ -0,0 +1,369 @@ +# 01 — Columnflow Framework: The Big Picture + +## What is columnflow? + +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. + +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. + +--- + +## The Analysis Pipeline + +The standard pipeline processes data from ROOT files to plots and datacards: + +```text +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) +``` + +### 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. + +**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) + +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 + +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` | `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. + +--- + +## Directory Structure of an Analysis + +```text +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) + +### 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) + └── Dataset ─── one Monte Carlo or data sample +``` + +### Analysis + +```python +import order as od +analysis = od.Analysis(name="hbw", id=1) +``` + +### Campaign + +```python +cpn = od.Campaign( + name="run2_2018", + id=4, + ecm=13, + aux={ + "tier": "NanoAOD", + "year": 2018, + "location": "root://xrootd-cms.infn.it//", + }, +) +``` + +### Process + +```python +from scinum import Number + +tt = od.Process( + name="tt", + id=1000, + label=r"$t\bar{t}$", + color=(205, 0, 9), + xsecs={13: Number(831.76, {"scale": (19.77, 29.20), "pdf": 35.06})}, +) +``` + +### Dataset + +```python +cpn.add_dataset( + name="tt_dl_powheg", + id=1, + processes=[procs.tt], + info={ + "nominal": od.DatasetInfo( + 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, + ), + }, +) +``` + +### Config + +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"Leading jet $p_T$", +) + +# Category +from columnflow.config_util import add_category +add_category(cfg, name="incl", id=1, selection="cat_incl", label="Inclusive") + +# --- Auxiliary configuration (cfg.x.*) --- +from columnflow.util import DotDict +from columnflow.columnar_util import ColumnCollection + +cfg.x.luminosity = Number(59740, {"lumi_13TeV_2018": 0.025j}) + +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, + }, +}) + +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 + +| Key | Type | Purpose | +|---|---|---| +| `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 `--calibrators` value | +| `cfg.x.default_selector` | str | Default `--selector` 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 (see best_practices) | + +--- + +## Systematic Uncertainties (Shifts) + +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) + +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") +``` + +### 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. + +```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"}) +``` + +### Dedicated-dataset uncertainties + +A completely separate Monte Carlo sample is processed (e.g. tune, hdamp variations). The dataset `info` dict key must match the shift name exactly. + +```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. +``` + +--- + +## Data Flow Summary + +```text +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] +``` + +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/02_coding_style.md b/docs/columnflow_claude_guide/02_coding_style.md new file mode 100644 index 000000000..e42553bde --- /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: + +```text +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..0e3ad2e01 --- /dev/null +++ b/docs/columnflow_claude_guide/03_task_commands.md @@ -0,0 +1,452 @@ +# 03 — User Guide: Working with Columnflow Tasks + +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 | 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 | + +--- + +## 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`. + +**Requires:** A valid GRID proxy (`voms-proxy-init -rfc -valid 196:00`) for CMS datasets, unless a custom LFN source is configured. + +**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 + +**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 \ + --branch 0 + +# Run with a JEC up-shift +law run cf.CalibrateEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --calibrators default \ + --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 + +**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. + +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 + +# Inspect the full upstream task tree +law run cf.SelectEvents \ + --version dev1 \ + --dataset tt_dl_powheg \ + --print-status -1 + +# 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 + +**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. + +**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 + +# 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 + +**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 + +**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. + +**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 (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 + +**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. + +**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 \ + --categories incl \ + --shift nominal \ + --producers default + +# Multiple variables and categories +law run cf.CreateHistograms \ + --version dev1 \ + --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" +``` + +**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 + +**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,data_mu +``` + +**Tips:** +- This task is almost always triggered automatically by plotting or datacard tasks. You rarely call it standalone. + +--- + +## PlotVariables1D + +**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_mu" \ + --variables "jet1_pt,n_jet,ht" \ + --categories incl \ + --shift nominal + +# Save as both PDF and PNG +law run cf.PlotVariables1D \ + --version dev1 \ + --processes tt \ + --variables jet1_pt \ + --file-types "pdf,png" + +# Log scale, skip ratio +law run cf.PlotVariables1D \ + --version dev1 \ + --processes "tt,wjets" \ + --variables jet1_pt \ + --yscale log \ + --skip-ratio + +# Find where plots are saved +law run cf.PlotVariables1D \ + --version dev1 --processes tt --variables jet1_pt --print-output 0 +``` + +--- + +## PlotVariables2D + +```bash +law run cf.PlotVariables2D \ + --version dev1 \ + --processes tt \ + --variables "jet1_pt__jet1_eta" \ + --categories incl +``` + +--- + +## PlotShiftedVariables1D + +**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" \ + --categories incl +``` + +--- + +## PlotCutflow + +**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 "trigger,muon,jet,btag" \ + --categories incl +``` + +--- + +## PlotCutflowVariables1D + +**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 \ + --processes tt \ + --variables genTop_pt \ + --categories incl \ + --skip-ratio +``` + +--- + +## CreateDatacards + +**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 default \ + --variables jet1_pt \ + --categories sr +``` + +--- + +## Useful Workflow Commands + +```bash +# --- 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 + +# Find output paths at depth 0 (the plot itself) +law run cf.PlotVariables1D --version dev1 --processes tt --variables jet1_pt --print-output 0 + +# --- 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 + +# --- Remote --- +# Submit SelectEvents to HTCondor +law run cf.SelectEvents \ + --version dev1 --dataset tt_dl_powheg --workflow htcondor + +# --- 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 + +# --- Open columnar sandbox for ad-hoc scripting --- +cf_sandbox venv_columnar_dev bash +``` + +--- + +## Full Pipeline in One Command + +Law resolves all upstream dependencies automatically. Running a plot task triggers the entire chain: + +```bash +# This single command runs the full pipeline for one file of one dataset +law run cf.PlotVariables1D \ + --version dev1 \ + --datasets tt_dl_powheg \ + --processes tt \ + --variables jet1_pt \ + --categories incl \ + --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/04_awkward_reference.md b/docs/columnflow_claude_guide/04_awkward_reference.md new file mode 100644 index 000000000..cebae8110 --- /dev/null +++ b/docs/columnflow_claude_guide/04_awkward_reference.md @@ -0,0 +1,306 @@ +# 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..65e659a0c --- /dev/null +++ b/docs/columnflow_claude_guide/05_dos_and_donts.md @@ -0,0 +1,349 @@ +# 05 — Do's and Don'ts for Columnflow Users + +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. + +--- + +## 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 +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 — silently ignored or raises +events["ht"] = ak.sum(events.Jet.pt, axis=1) +events.ht = ak.sum(events.Jet.pt, axis=1) +``` + +### DO reassign `events` after every TAF call + +```python +# CORRECT +events = self[other_producer](events, **kwargs) +events = set_ak_column(events, "my_col", value) + +# WRONG — the return value is discarded +self[other_producer](events, **kwargs) +``` + +### 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 +# WRONG — never do this +hts = [] +for event in events: + hts.append(sum(jet.pt for jet in event.Jet)) + +# CORRECT +ht = ak.sum(events.Jet.pt, axis=1) +``` + +### DO specify `axis=1` when reducing over objects within an event + +```python +# WRONG — sums over everything (axis=None by default in some contexts) +n_jet = ak.sum(events.Jet.pt > 0) # scalar, not per-event + +# CORRECT — one value per event +n_jet = ak.sum(events.Jet.pt > 0, axis=1) # shape (N_events,) +``` + +--- + +## TAF Declarations (`uses` and `produces`) + +### DO declare every column you read in `uses` + +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 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 — sub_producer's columns are automatically included +@producer( + uses={sub_producer, "extra_col"}, + produces={sub_producer, "my_output"}, +) +def parent(self, events, **kwargs): + events = self[sub_producer](events, **kwargs) + ... + +# WRONG — sub_producer reads "Jet.pt" but it's not declared here +@producer( + uses={"extra_col"}, + produces={"my_output"}, +) +def parent(self, events, **kwargs): + events = self[sub_producer](events, **kwargs) # Jet.pt not loaded! +``` + +--- + +## Imports + +### DO use `maybe_import` for heavy packages at module level + +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 — deferred import +from columnflow.util import maybe_import +ak = maybe_import("awkward") +np = maybe_import("numpy") + +# WRONG — fails at import time in the default sandbox +import awkward as ak +import numpy as np +``` + +### DO import coffea inside the function body, never at module level + +```python +# CORRECT +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 +import coffea # breaks any non-coffea sandbox +``` + +--- + +## Selectors + +### DO set `results.event` in the exposed Selector + +`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 return `(events, SelectionResult)` from every Selector + +```python +# 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 +``` + +### DO use `ak.local_index` to create index arrays for `objects` + +```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] + +return events, SelectionResult( + objects={"Jet": {"Jet": jet_indices}}, +) + +# WRONG — using a boolean mask directly in objects (must be indices) +return events, SelectionResult( + objects={"Jet": {"Jet": jet_mask}}, # wrong type! +) +``` + +--- + +## Configuration + +### DO list columns in `cfg.x.keep_columns` that must survive `ReduceEvents` + +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 +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, # all columns produced in the Selector + }, +}) +``` + +### DO register new files in `law.cfg` under the correct `*_modules` key + +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. + +```ini +# After creating myanalysis/selection/jets.py: +selection_modules: ..., myanalysis.selection.{default,jets} + +# After creating myanalysis/production/weights.py: +production_modules: ..., myanalysis.production.{default,weights} +``` + +### DON'T put spaces after commas inside `{}` brace expansions in `law.cfg` + +```ini +# CORRECT +selection_modules: myanalysis.selection.{default,jets,trigger} + +# WRONG — spaces break the brace expansion +selection_modules: myanalysis.selection.{default, jets, trigger} +``` + +--- + +## 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 shifted column names — use column aliases + +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 — hardcoded nominal name, broken under mu_up shift +weight = events.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 +``` + +### DO declare shifts in the TAF's `shifts` set (for weight-based uncertainties) + +```python +@my_producer.init +def my_producer_init(self: Producer) -> None: + from columnflow.config_util import get_shifts_from_sources + self.shifts |= set(get_shifts_from_sources(self.config_inst, "mu")) +``` + +--- + +## Versioning and Caching + +### DO bump `--version` after changing any code in the pipeline + +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. + +```bash +# After changing the Selector: +law run cf.SelectEvents --version v2 --dataset tt_dl_powheg # not v1! +``` + +### DO use `--branch 0` for quick local tests before submitting all branches + +```bash +# Test on a single file first to catch bugs early +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg --branch 0 + +# Then run all files +law run cf.SelectEvents --version dev1 --dataset tt_dl_powheg +``` + +### 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 +``` + +### DO use `--remove-output 0,a,y` to delete and immediately re-run a task + +```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 +``` + +--- + +## Storage and Columns + +### DON'T store large arrays in `cfg.x.*` auxiliaries + +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 +# 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 + +# 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() +``` + +### DO produce columns only after `ReduceEvents` where possible + +Creating columns in `ProduceColumns` (after reduction) avoids storing them in every file during calibration/selection, saving significant disk space. + +### 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 + +Before committing or running new analysis code, verify: + +- [ ] 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` +- [ ] 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 +```