tool.py

Go to the documentation of this file.

# Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
#
# @author Giordon Stark
 
# High-level Tool wrapper around PythonToolHandle.
# Automates buffer extraction, column setting, calling, and output
# reconstruction using ColumnInfo metadata.
 
import awkward as ak
import numpy as np
 
from ColumnarToolWrapperPython.buffers import (
    _branch_name_for_column,
    allocate_outputs,
    classify_columns,
    extract_buffers,
    reconstruct_output,
    resolve_optional_columns,
)
from ColumnarToolWrapperPython.python_tool_handle import PythonToolHandle
 
 
class Tool:
    """High-level wrapper around PythonToolHandle.
 
    Handles the boilerplate of classifying columns, extracting awkward-array
    buffers, allocating outputs, setting columns on the handle, calling the
    tool, and reconstructing the output awkward array.
 
    Parameters
    ----------
    type_and_name:
        Tool type and instance name, e.g. "CP::MuonEfficiencyScaleFactors/myTool".
    properties:
        Optional dict of tool properties to set before initialization.
    rename_containers:
        Optional mapping from canonical container names (e.g. "Muons") to the
        branch-name prefix used in the input arrays (e.g. "AnalysisMuonsAuxDyn").
        This is passed to PythonToolHandle.rename_containers *after* initialize
        so that ColumnInfo.name values match input array fields.
    """
 
    def __init__(self, type_and_name, properties=None, rename_containers=None):
        self._handle = PythonToolHandle()
        self._handle.set_type_and_name(type_and_name)
        self._properties = dict(properties) if properties else {}
 
        for key, value in self._properties.items():
            self._handle.set_property(key, value)
 
        self._handle.initialize()
 
        if rename_containers:
            self._handle.rename_containers(rename_containers)
 
        # Cache classified column topology once after initialization
        self._classified = classify_columns(self._handle.columns)
 
    @property
    def properties(self):
        """Dict of properties set on this tool at construction time."""
        return self._properties
 
    @property
    def columns(self):
        """All ColumnInfo objects reported by the tool."""
        return self._handle.columns
 
    @property
    def input_columns(self):
        """Input (non-offset) ColumnInfo objects, including nested-vector data columns."""
        cols = []
        for info in self._classified.values():
            cols.extend(info["inputs"])
            for nested in info["nested_offsets"].values():
                cols.extend(nested["inputs"])
        return cols
 
    @property
    def output_columns(self):
        """Output ColumnInfo objects, including nested-vector outputs (currently unsupported)."""
        cols = []
        for info in self._classified.values():
            cols.extend(info["outputs"])
            for nested in info["nested_offsets"].values():
                cols.extend(nested["outputs"])
        return cols
 
    @property
    def required_input_branch_names(self):
        """Deduplicated uproot branch names needed to fill non-optional inputs.
 
        For nested-vector inputs, the trailing ``.data`` suffix is stripped so
        the returned name matches the actual ROOT branch (e.g.
        ``"Particles.NumTrkPt500"`` instead of ``"Particles.NumTrkPt500.data"``).
        Optional inputs are excluded.
        """
        seen = []
        seen_set = set()
        for col in self.input_columns:
            if col.is_optional:
                continue
            branch = _branch_name_for_column(col.name)
            if branch not in seen_set:
                seen_set.add(branch)
                seen.append(branch)
        return seen
 
    @property
    def recommended_systematics(self):
        """List of recommended systematic variation names."""
        return self._handle.get_recommended_systematics()
 
    def apply_systematic_variation(self, sys_name):
        """Apply a systematic variation by name.
 
        Parameters
        ----------
        sys_name:
            Systematic variation name, e.g. ``"MUON_EFF_RECO_SYS__1up"``.
            Pass ``""`` to reset to the nominal.
        """
        self._handle.apply_systematic_variation(sys_name)
 
    def __call__(self, events, systematic=None):
        """Run the tool on events and return output columns as an ak.Array.
 
        Parameters
        ----------
        events:
            An ak.Array with fields matching the tool's input column names
            (after any rename_containers mapping).
        systematic:
            Optional systematic variation name (e.g. "MUON_EFF_RECO_SYS__1up").
            If provided, applied before running the tool and reset to nominal
            after execution.
 
        Returns
        -------
        ak.Array
            Record array with one field per output column, each a
            variable-length list over the per-particle values.
        """
        if systematic is not None:
            self.apply_systematic_variation(systematic)
 
        try:
            num_events = int(ak.num(events, axis=0))
 
            # Resolve optional columns against the actual fields present
            effective = resolve_optional_columns(self._classified, events)
 
            # Extract flat buffers from the awkward array
            buffer_dict = extract_buffers(events, effective)
 
            # Allocate zero-filled output arrays (added into buffer_dict in-place)
            allocate_outputs(effective, buffer_dict)
 
            # Set all columns on the handle
            for container_name, info in effective.items():
                # Container offset (always immutable)
                self._handle[container_name] = np.asarray(buffer_dict[container_name])
 
                # Nested-vector offsets and their data inputs/outputs
                for nested_offset_name, nested in info["nested_offsets"].items():
                    self._handle[nested_offset_name] = np.asarray(
                        buffer_dict[nested_offset_name]
                    )
                    for col in nested["inputs"]:
                        self._handle[col.name] = np.asarray(buffer_dict[col.name])
                    for col in nested["outputs"]:
                        self._handle.set_column_void(col.name, buffer_dict[col.name], False)
 
                # Flat input data columns (immutable)
                for col in info["inputs"]:
                    self._handle[col.name] = np.asarray(buffer_dict[col.name])
 
                # Flat output data columns (mutable)
                for col in info["outputs"]:
                    self._handle.set_column_void(col.name, buffer_dict[col.name], False)
 
            self._handle.call()
 
            return reconstruct_output(effective, buffer_dict, num_events)
        finally:
            # Reset to nominal if systematic was applied
            if systematic is not None:
                self.apply_systematic_variation("")