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

Marcel Rieger developed a function to assign CMS Monte Carlo events to years 2024, 2025, and 2026 based on recorded luminosities. The assignment is pseudo-random, deterministic, and unbiased, preserving event id distributions for downstream analyses.

| """ | | 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 | | | | | | | | constants | | | | | | recorded CMS luminosities | | from https://twiki.cern.ch/twiki/bin/view/CMSPublic/LumiPublicResults?rev=213 | | lumis = { | | 2024: 112.70, | | 2025: 114.85, | | 2026: 30.36, | | } | | | | build fractions relative to their sum, multiply by 1000 for accuracy , then convert to ranges | | 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 | | | | | | | | year assignment function and helpers | | | | | | | | 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 | | | | single event id | | assign year 1234567890 - 2025 | | assign year 1234567891 - 2024 | | ... | | | | array of event ids | | assign year np.array 1234567890, 1234567891 - np.array 2025, 2024 | | """ | | handle single integers | | single input = isinstance event id, int | | | | apply hashing | | h = split mix hash np.array event id , dtype=np.uint64 if single input else event id | | | | take last three digits and cast down | | h = h % 1000 .astype np.uint16 | | | | perform assignment based on lumi ranges | | 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. | | """ | | handle single integers | | single input = isinstance a, int | | | | start from uint64 | | h = np.array a , dtype=np.uint64 if single input else a.copy .astype np.uint64 | | | | start mixing | | 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 | | | | only read events from nano aod file that are assigned to 2025 | | with uproot.open "nano.root" as f: | | tree = f "Events" | | events 2025 = tree.arrays uproot select year 2025 | | """ | | patch uproot uint64 | | | | nested aliases that represent the splitmix64 hash algorithm from :py:func: split mix hash | | 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 expression on lumi range for that year | | 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 | | | | add "uint64" conversion | | py lang functions = uproot.language.python.PythonLanguage.default functions | | if "uint64" not in py lang functions: | | py lang functions "uint64" = np.uint64 | | | | | | | | testing | | | | | | 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 | | | | show lumis and fractions | | 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 specific event ids | | test ids = 1, 9, 10, 16849655 | | for event id in test ids: | | print f"assign {event id} - {assign year event id }" | | print | | | | load event ids from nano file | | input file = os.path.expanduser os.path.expandvars args.input | | with uproot.open input file as f: | | ids = f "Events" "event" .array library="np" | | | | assign years | | years = assign year ids | | | | for each year, plot the distribution of last and last-two digits of event ids | | to check for bias via uniformity | | for digits in range 1, 4 : | | bins = 10 digits | | 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 |