Source code for analyzer.modules.common.column_tools

from analyzer.core.analysis_modules import AnalyzerModule
import re

from analyzer.core.columns import addSelection
from analyzer.core.columns import Column
from analyzer.utils.structure_tools import flatten
from analyzer.core.analysis_modules import ParameterSpec, ModuleParameterSpec
import awkward as ak
import itertools as it
from attrs import define, field, evolve
from .axis import RegularAxis
from .histogram_builder import makeHistogram
from analyzer.core.adl import ADLBlock, ADLStatement


import correctionlib
import logging


from analyzer.core.analysis_modules import (
    MetadataExpr,
    MetadataAnd,
    IsRun,
    IsSampleType,
)


[docs]
logger = logging.getLogger("analyzer.modules")



@define

[docs]
class Count(AnalyzerModule):

[docs]
    input_col: Column


[docs]
    output_col: Column



[docs]
    def run(self, columns, params):
        columns[self.output_col] = ak.num(columns[self.input_col], axis=1)
        return columns, []



[docs]
    def inputs(self, metadata):
        return [self.input_col]



[docs]
    def outputs(self, metadata):
        return [self.output_col]




@define

[docs]
class PromoteIndex(AnalyzerModule):
    """
    Promote a fixed index of a nested collection to a top-level column.

    This analyzer selects a single element at a given index from each
    entry of a nested (jagged) input collection and stores it as a
    top-level column. It is commonly used to extract leading or
    sub-leading objects (e.g. leading jet, first lepton) from per-event
    collections.

    Parameters
    ----------
    input_col : Column
        Column containing a nested collection (e.g. ``N × M`` objects).
    output_col : Column
        Column where the selected elements will be stored.
    index : int, optional
        Index of the element to promote from each nested collection,
        by default ``0`` (leading element).

    """


[docs]
    input_col: Column


[docs]
    output_col: Column


[docs]
    index: int = 0



[docs]
    def run(self, columns, params):
        columns[self.output_col] = ak.pad_none(
            columns[self.input_col], self.index + 1, axis=1
        )[:, self.index]
        return columns, []



[docs]
    def inputs(self, metadata):
        return [self.input_col]



[docs]
    def outputs(self, metadata):
        return [self.output_col]




@define

[docs]
class Concatenate(AnalyzerModule):

[docs]
    input_cols: list[Column]


[docs]
    output_col: Column



[docs]
    def run(self, columns, params):
        columns[self.output_col] = ak.concatenate(
            [columns[c] for c in self.input_cols], axis=1
        )
        return columns, []



[docs]
    def inputs(self, metadata):
        return self.input_cols



[docs]
    def outputs(self, metadata):
        return [self.output_col]



[docs]
    def adlExport(self, metadata):
        input_names = [c.adl_name for c in self.input_cols]
        union_expr = "union(" + ", ".join(input_names) + ")"
        return [
            ADLBlock(
                block_type="object",
                name=self.output_col.adl_name,
                statements=[ADLStatement("take", union_expr)]
            )
        ]