File size: 1,406 Bytes
97f7eaf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """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()
|