Event ID based splitting of 2024 MC samples into 2024, 2025 and 2026 parts

| """ | | Functions to assign CMS MC events to years 2024, 2025, and 2026 based on the recorded luminosities | | in each year. The assignment is pseudo-random, deterministic and unbiased, i.e., it does not skew | | the distribution of event ids for any of the years. This means that analysis downstream that need | | split events (e.g. for ML datasets) can still rely on the event id and modulo-based splittings. | | | | See :py:func:assign_year below for more details. | | | | For implementations in C and Python (with tests), see: | | https://gist.github.com/riga/f9476f3b1477f1609683bea68ae64897 | | | | Author: Marcel Rieger | | """ | | | | from future import annotations | | | | from typing import TypeVar, Any | | | | import numpy as np | | from numpy.typing import NDArray | | | | | | IntTypes = TypeVar("IntTypes", int, NDArray) | | | | | | | | | | | | | lumis = { | | 2024: 112.70, | | 2025: 114.85, | | 2026: 30.36, | | } | | | | | lumi_sum = sum(lumis.values()) | | lumi_fractions = {year: int(round(1000 * lumi / lumi_sum)) for year, lumi in lumis.items()} | | if sum(lumi_fractions.values()) != 1000: | | raise RuntimeError(f"fractions do not sum up to 1000: {lumi_fractions}") | | lumi_ranges = {} | | offset = 0 | | for year, fraction in lumi_fractions.items(): | | lumi_ranges[year] = (offset, offset + fraction) | | offset = lumi_ranges[year][1] | | | | | | | | | | | | | def assign_year( | | event_id: IntTypes, | | lumi_ranges: dict[int, tuple[int, int]] = lumi_ranges, | | ) -> IntTypes: | | """ | | Takes a single event_id or an array of ids and assigns them to one of the years defined in lumi_ranges. The | | assignment is pseudo-random and deterministic, i.e. the same event id will always be assigned to the same year | | without any bias towards specific event id. See :py:func:split_mix_hash for details on the hashing algorithm used. | | | | lumi_ranges is a dictionary mapping years to a tuple of (min, max) values that define the range of hashed values | | that will be assigned to that year. The ranges must be non-overlapping and cover the entire range of possible hashed | | values, 0 (incl.) to 1000 (excl.). Thereby, the frequency of assigned years will be proportional to the | | luminosities. | | | | Examples: | | | | .. code-block:: python | | | | | assign_year(1234567890) # -> 2025 | | assign_year(1234567891) # -> 2024 | | ... | | | | | assign_year(np.array([1234567890, 1234567891])) # -> np.array([2025, 2024]) | | """ | | | single_input = isinstance(event_id, int) | | | | | h = split_mix_hash(np.array([event_id], dtype=np.uint64) if single_input else event_id) | | | | | h = (h % 1000).astype(np.uint16) | | | | | years = np.zeros(len(h), dtype=np.uint16) | | for year, (min_val, max_val) in lumi_ranges.items(): | | years[(h >= min_val) & (h < max_val)] = year | | | | return years[0] if single_input else years | | | | | | def split_mix_hash(a: IntTypes) -> IntTypes: | | """ | | Deterministic integer mixing using the splitmix64 algorithm with the goal that a small change in the input causes a | | large change in the output hash. See https://rosettacode.org/wiki/Pseudo-random_numbers/Splitmix64. | | """ | | | single_input = isinstance(a, int) | | | | | h = np.array([a], dtype=np.uint64) if single_input else a.copy().astype(np.uint64) | | | | | h += np.uint64(0x9e3779b97f4a7c15) | | h = (h ^ (h >> 30)) * np.uint64(0xbf58476d1ce4e5b9) | | h = (h ^ (h >> 27)) * np.uint64(0x94d049bb133111eb) | | h = (h ^ (h >> 31)) | | | | return h[0] if single_input else h | | | | | | def uproot_select_year(year: int) -> dict[str, Any]: | | """ | | Returns a dictionary with the necessary arguments to pass to :py:meth:uproot.TTree.arrays (and alike) to read only | | events that are assigned to the given year. The returned dictionary contains aliases and cut expressions. | | | | Example: | | | | .. code-block:: python | | | | | with uproot.open("nano.root") as f: | | tree = f["Events"] | | events_2025 = tree.arrays(uproot_select_year(2025)) | | """ | | _patch_uproot_uint64() | | | | | aliases = { | | "split_mix_hash_1": "event + uint64(11400714819323198485)", | | "split_mix_hash_2": "(split_mix_hash_1 ^ (split_mix_hash_1 >> 30)) * uint64(13787848793156543929)", | | "split_mix_hash_3": "(split_mix_hash_2 ^ (split_mix_hash_2 >> 27)) * uint64(10723151780598845931)", | | "split_mix_hash": "split_mix_hash_3 ^ (split_mix_hash_3 >> 31)", | | "event_split_id": "split_mix_hash % 1000", | | } | | | | | cut = f"(event_split_id >= {lumi_ranges[year][0]}) & (event_split_id < {lumi_ranges[year][1]})" | | | | return {"aliases": aliases, "cut": cut} | | | | | | def _patch_uproot_uint64() -> None: | | try: | | import uproot | | except ImportError: | | return | | | | | py_lang_functions = uproot.language.python.PythonLanguage.default_functions | | if "uint64" not in py_lang_functions: | | py_lang_functions["uint64"] = np.uint64 | | | | | | | | | | | def main() -> None: | | import os | | import argparse | | import uproot | | import uniplot | | | | parser = argparse.ArgumentParser(description=doc.strip().splitlines()[0].strip()) | | parser.add_argument("input", help="input NanoAOD file") | | args = parser.parse_args() | | | | | print("luminosities and fractions:") | | for year, lumi in lumis.items(): | | print(f" - {year}: {lumi:.2f}/fb, 1k-fraction: {lumi_fractions[year]}, range: {lumi_ranges[year]}") | | print() | | | | | test_ids = [1, 9, 10, 16849655] | | for event_id in test_ids: | | print(f"assign {event_id} -> {assign_year(event_id)}") | | print() | | | | | input_file = os.path.expanduser(os.path.expandvars(args.input)) | | with uproot.open(input_file) as f: | | ids = f["Events"]["event"].array(library="np") | | | | | years = assign_year(ids) | | | | | | for digits in range(1, 4): | | bins = 10digits | | data = [ | | (ids[years == year] % bins) | | for year in lumis | | ] | | uniplot.histogram( | | data, | | legend_labels=list(map(str, lumis)), | | color=True, | | width=100, | | height=25, | | bins=min((bins := 10**digits), 100), | | bins_min=-0.5, | | bins_max=bins - 0.5, | | title=f"Distribution of last {digits} digit(s) of event ids after splitting", | | ) | | print() | | | | | | if name == "main": | | main() |