Ntuples in Postprocessing

The CROWN Ntuples can be used by any Postprocessing framework. There are a few things to keep in mind to ensure easy processing. The most important difference is that only quantities affected by a shift are recalculated. This means the postprocessing framework must be able to use a mixture of the original and shifted quantities, when applying shifts. In order to make this step a bit easier, the information which quantities are affected by a shift, is stored in the Ntuple.

Quantity mapping

To read the mapping from an Ntuple, the python function listed below may be used. Two types of mapping are available, depending on the actual use case. In the first, the mapping is sorted by shift; in the second the mapping is sorted by quantity.

import ROOT as r
import glob
import json
import os


def load_crown_mapping(
    filename,
    libdir,
    by_shift=False,
    by_quantity=False,
    load_metadata=False,
    sort_values=True,
):
    """
    Load CROWN quantity/shift mappings from a ROOT file.

    Parameters
    ----------
    filename : str
        Input ROOT file.

    libdir : str
        Directory containing libMyDicts.so, usually "CROWN/.cache/"

    by_shift : bool
        Read shift_quantities_map.

    by_quantity : bool
        Read quantities_shift_map.

    load_metadata : bool
        Also read metadata object from file.

    sort_values : bool
        Sort vectors before returning.

    Returns
    -------
    dict
        Mapping dictionary.

    tuple(dict, dict)
        If load_metadata=True.
    """

    data = {}

    # ------------------------------------------------------------------
    # Validate mapping selection
    # ------------------------------------------------------------------
    if by_shift and not by_quantity:
        name = "shift_quantities_map"
    elif by_quantity and not by_shift:
        name = "quantities_shift_map"
    else:
        raise ValueError(
            "Choose exactly one of by_shift=True or by_quantity=True"
        )

    # ------------------------------------------------------------------
    # Load ROOT dictionary parsing library
    # ------------------------------------------------------------------
    lib_path = os.path.abspath(os.path.join(libdir, "libMyDicts.so"))

    if not os.path.exists(lib_path):
        raise FileNotFoundError(f"Missing library: {lib_path}")

    result = r.gSystem.Load(lib_path)

    if result < 0:
        err_type = (
            "Version mismatch"
            if result == -2
            else "Linker error / Missing dependency"
        )
        raise ImportError(
            f"Failed to load {lib_path} "
            f"(ROOT return code {result}: {err_type})"
        )

    # ------------------------------------------------------------------
    # Open ROOT file
    # ------------------------------------------------------------------
    print(f"Reading {name} from {filename}")

    f = r.TFile.Open(filename)

    if not f or f.IsZombie():
        raise OSError(f"Could not open ROOT file: {filename}")

    # ------------------------------------------------------------------
    # Read mapping
    # ------------------------------------------------------------------
    m = f.Get(name)

    if not m:
        f.Close()
        raise KeyError(f"Object '{name}' not found in file")

    for key, values in m:
        entries = [str(v) for v in values]

        if sort_values:
            entries = sorted(entries)

        data[str(key)] = entries

    # ------------------------------------------------------------------
    # Optional metadata
    # ------------------------------------------------------------------
    metadata = None

    if load_metadata:
        metadata_obj = f.Get("metadata")

        if metadata_obj:
            metadata = json.loads(
                metadata_obj.GetString().Data()
            )["metadata"]

    f.Close()

    print(f"Successfully read {name} from {filename}")

    # ------------------------------------------------------------------
    # Cleanup autogenerated ROOT dictionary artifacts
    # ------------------------------------------------------------------
    print("Cleaning up autogenerated files")

    for autogenerated in glob.glob("AutoDict_*"):
        try:
            os.remove(autogenerated)
        except OSError:
            pass

    if load_metadata:
        return data, metadata

    return data


# Example usage:
#
# data = load_crown_mapping(
#     "test.root",
#     libdir=".cache",
#     by_shift=True,
# )
#
# data, metadata = load_crown_mapping(
#     "test.root",
#     libdir=".cache",
#     by_shift=True,
#     load_metadata=True,
# )