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

README.md

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

The snippets below demonstrate how to split MC events originally intended for the 2024 data-taking period into orthogonal parts that can be used to simulate the 2024, 2025 and 2026 data-taking periods. The splitting is done event-by-event and only depends on the event ID. The underlying algorithm is

• fast,
• deterministic (also across different working groups for synchronization purposes),
• unbiased towards the event ID, and
• splits events into three parts with frequencies that reflect the relative recorded luminosities in 2024, 2025 and 2026.

The unbiased behavior is strictly necessary to avoid that algorithms based on event IDs that are used downstream in analyses see a bias towards specific event IDs in any of the three years. For instance, event splitting for ML training / validation purposes, usually done via (e.g.) event % 2 == 0 or alike, should remain statistically unaffected. If there was a bias in the event ID distribution, this would lead to skewed training and validation datasets, which is not desirable.

• Author: Marcel Rieger
• Source: https://gist.github.com/riga/f9476f3b1477f1609683bea68ae64897

Python tools

The following example shows how to filter events

import cms_event_split

# single event test using different event ids
cms_event_split.assign_year(16849654)  # -> 2025
cms_event_split.assign_year(16849655)  # -> 2026

# open a NanoAOD file via uproot and eagerly filter events to only select those assigned to 2025
with uproot.open("nano.root") as f:
    tree = f["Events"]
    events = tree.arrays(..., **cms_event_split.uproot_select_year(2025))

Cpp tools

#include "cms_event_split.h"

// single event test using different event ids
cms_event_split::assign_year(16849654);  // -> 2025
cms_event_split::assign_year(16849655);  // -> 2026

// add a filter node to a RDataFrame to continue processing only events assigned to 2025
auto df = df.Filter(
    [](uint64_t event) {
        return cms_event_split::assign_year(event) == 2025;
    },
    {"event"},
    "event_id_filter_2025"
);

Details

Internally, the assignment logic is based on a hash function (splitmix64) with the goal that a small change in the input event id causes a large change in the output hash. Using the last three digits of the hash, events are then mapped to one of the years 2024, 2025 and 2026. This mapping is done via "binning", where the width of each of the three bins corresponds to the relative luminosity recorded in that year.

Testing

Output of python cms_event_split.py /some/ttbar/nano.root:

cms_event_split.h

/*
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 assign_year() below for more details.

For implementations in C and Python (with tests), see:
https://gist.github.com/riga/f9476f3b1477f1609683bea68ae64897

Author: Marcel Rieger
*/

#ifndef CMS_EVENT_SPLIT_H
#define CMS_EVENT_SPLIT_H

#include <vector>
#include <concepts>
#include <stdexcept>

namespace cms_event_split {

// taken from python version linked in gist above
inline const std::vector<std::pair<int, uint16_t>> lumi_ranges = {
    {2024, 437},  // 2024
    {2025, 882},  // 2025
    {2026, 1000},  // 2026
};

template <std::integral T>
requires (!std::same_as<T, bool> && !std::same_as<T, char>)
uint64_t split_mix_hash(T i) {
    uint64_t h = static_cast<uint64_t>(i);
    h += uint64_t(0x9e3779b97f4a7c15);
    h = (h ^ (h >> 30)) * uint64_t(0xbf58476d1ce4e5b9);
    h = (h ^ (h >> 27)) * uint64_t(0x94d049bb133111eb);
    h = (h ^ (h >> 31));
    return h;
}

template <std::integral T>
requires (!std::same_as<T, bool> && !std::same_as<T, char>)
int assign_year(T event_id) {
    // apply splitmix64 for hashing
    uint64_t h = split_mix_hash(event_id);

    // take last three digits and cast down
    uint16_t d = h % 1000;

    // assign to year
    for (const auto& [year, upper_bound] : lumi_ranges) {
        if (d < upper_bound) {
            return year;
        }
    }

    // should never get here if lumi_ranges is set up correctly
    throw std::runtime_error("failed to assign year for event_id: " + std::to_string(event_id));

    return 0;
}

}  // namespace cms_event_split

#endif  // CMS_EVENT_SPLIT_H

// testing only
// #include <iostream>
// int main() {
//     std::vector<int> test_ids = {1, 9, 10, 16849655};
//     for (int event_id : test_ids) {
//         int year = cms_event_split::assign_year(event_id);
//         std::cout << "assign " << event_id << " -> " << year << std::endl;
//     }
//     return 0;
// }

cms_event_split.py

"""
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()