scripts/dev/shrink_root_sample.py

"""Create a small ROOT fixture from a Delphes sample."""

from __future__ import annotations

import argparse
from pathlib import Path

import uproot

BRANCHES = (
    "jet_pt",
    "jet_eta",
    "jet_phi",
    "jet_btag",
    "ph_pt",
    "ph_eta",
    "ph_phi",
    "ele_pt",
    "ele_eta",
    "ele_phi",
    "ele_charge",
    "mu_pt",
    "mu_eta",
    "mu_phi",
    "mu_charge",
    "MET_met",
    "MET_phi",
    "weight",
    "Number",
)


def shrink_sample(source: Path, target: Path, entries: int) -> None:
    """Copy the active branches and first ``entries`` events to ``target``."""
    if entries < 1:
        raise ValueError("entries must be positive")
    target.parent.mkdir(parents=True, exist_ok=True)
    with uproot.open(source) as source_file:
        arrays = source_file["output"].arrays(
            BRANCHES, entry_start=0, entry_stop=entries, library="ak"
        )
        with uproot.recreate(target) as target_file:
            target_file["output"] = arrays


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", type=Path)
    parser.add_argument("target", type=Path)
    parser.add_argument("--entries", type=int, default=64)
    args = parser.parse_args()
    shrink_sample(args.source, args.target, args.entries)
    print(f"Wrote {args.entries} events to {args.target}")


if __name__ == "__main__":
    main()