Functor Collections

Read an HLT2 file and create an ntuple using pre-defined Functor collections. This example is also used to test all available functor collections on a sample of \(D^0 \to K^-\pi^+\) decays.

from PyConf.reading import (
    get_particles,
    get_mc_particles,
    get_mc_track_info,
    get_mc_header,
)
from DaVinciMCTools import MCReconstructible as MCRectible, MCReconstructed as MCRected
from FunTuple import (
    FunctorCollection,
    functorcollections as functorcollections_original,
)
from FunTuple import FunTuple_Particles as Funtuple, FunTuple_MCParticles as MCFuntuple
from DaVinci.algorithms import create_lines_filter
from DaVinci import Options, make_config
from DaVinciMCTools import MCTruthAndBkgCat
from DecayTreeFitter import DecayTreeFitter

# define functor collections that need to be excluded.
# Delete the collection "HltTisTos" since the
# sample against which this options file is tested does not contain the Hlt1 selreports
# Delete the collection "ParticleIsolation", "VertexIsolation" and "ConeIsolation" since the
# sample against which this options file is tested does not contain the TES locations from 'extra_outputs'
# For these functorcollections see 'option_trigger_decisions' and 'option_davinci_tupling_relation_isovariables' examples
exclude_functor_collections = [
    "HltTisTos",
    "ParticleIsolation",
    "ConeIsolation",
    "VertexIsolation",
    "FlavourTaggingResults",
]


class check_wrapper:
    def __init__(self, functorcollection):
        self.functorcollection = functorcollection
        self.called = set()

    def __getattr__(self, name: str):
        func = getattr(self.functorcollection, name)
        self.called.add(func.__name__)
        return func

    def check(self):
        sorted_set_all = {
            entry
            for entry in self.functorcollection.__all__
            if entry not in exclude_functor_collections
        }
        sorted_set_called = set(self.called)
        if sorted_set_all != sorted_set_called:
            diff = sorted_set_all - sorted_set_called
            raise Exception(
                f"Oh no! Did you forget to add a new collection to this test? Missing collections: {diff}"
            )


def main(options: Options):
    line_name = "Hlt2Charm_D0ToKmPip_Line"

    d02kpi_data = get_particles(f"/Event/HLT2/{line_name}/Particles")
    mc_data = get_mc_particles("/Event/HLT2/MC/Particles")
    mc_header = get_mc_header()

    # Get variables related to reconstructible information.
    mcrtible = MCRectible(
        input_mctrackinfo=get_mc_track_info()
    )  # tes location of MC track info

    # Get variables related to reconstructed information.
    mcrted_all = MCRected(input_mcparticles=mc_data, use_best_mcmatch=True)

    # get configured "MCTruthAndBkgCatAlg" algorithm for HLT2 output
    MCTRUTH = MCTruthAndBkgCat(d02kpi_data, name="MCTruthAndBkgCat_coll")

    # define DTF
    DTF = DecayTreeFitter(name="DecayTreeFitter", input_particles=d02kpi_data)

    # initialize the checker
    FC = check_wrapper(functorcollections_original)

    # use FC to add variables, please expand when there's new collections :)
    collections = [
        FC.Kinematics(),
        FC.MCHierarchy(mctruth_alg=MCTRUTH),
        FC.MCKinematics(mctruth_alg=MCTRUTH),
        FC.MCVertexInfo(mctruth_alg=MCTRUTH),
        FC.MCPromptDecay(mctruth_alg=MCTRUTH),
        FC.NeutralCaloInfo(extra_info=True),
        FC.ChargedCaloInfo(extra_info=True),
        FC.ParticleID(),
        FC.DecayTreeFitterResults(
            DTF=DTF, decay_origin=False, with_lifetime=False, with_kinematics=True
        ),
    ]

    mc_collections = [
        FC.MCReconstructible(mcreconstructible_alg=mcrtible, extra_info=True),
        FC.MCReconstructed(mcreconstructed_alg=mcrted_all, extra_info=False),
        FC.MCHierarchy(),
        FC.MCVertexInfo(),
        FC.MCPromptDecay(),
    ]

    evt_collections = [
        FC.EventInfo(),
        FC.SelectionInfo(selection_type="Hlt2", trigger_lines=[line_name]),
        FC.RecSummary(),
    ]

    mc_evt_collections = [
        FC.EventInfo(),
        FC.MCPrimaries(mc_header=mc_header),
    ]

    FC.check()

    field_vars = FunctorCollection()
    for coll in collections:
        field_vars += coll

    mc_field_vars = FunctorCollection()
    for coll in mc_collections:
        mc_field_vars += coll

    evt_vars = FunctorCollection()
    for coll in evt_collections:
        evt_vars += coll

    mc_evt_vars = FunctorCollection()
    for coll in mc_evt_collections:
        mc_evt_vars += coll

    fields = {
        "D0": "[D0 -> K- pi+]CC",
        "Kminus": "[D0 -> ^K- pi+]CC",
        "piplus": "[D0 -> K- ^pi+]CC",
    }

    variables = {
        "D0": field_vars,
        "Kminus": field_vars,
        "piplus": field_vars,
    }

    mc_variables = {
        "D0": mc_field_vars,
        "Kminus": mc_field_vars,
        "piplus": mc_field_vars,
    }

    my_tuple = Funtuple(
        name="Tuple",
        tuple_name="DecayTree",
        fields=fields,
        variables=variables,
        event_variables=evt_vars,
        inputs=d02kpi_data,
    )

    my_filter = create_lines_filter(
        name="HDRFilter_D0Kpi", lines=[f"{line_name}Decision"]
    )

    mc_tuple = MCFuntuple(
        name="MCTuple",
        tuple_name="DecayTree",
        fields=fields,
        variables=mc_variables,
        event_variables=mc_evt_vars,
        inputs=mc_data,
    )

    return make_config(options, [my_filter, my_tuple, mc_tuple])

To run the example:

lbexec DaVinciExamples.tupling.option_davinci_tupling_from_collections:main $DAVINCIEXAMPLESROOT/example_data/FEST_November_2021_dst_newPacking.yaml

For reference, these are the options of this example

testfiledb_key: FEST_November_2021
input_manifest_file: root://eoslhcb.cern.ch//eos/lhcb/wg/dpa/wp3/tests/hlt2_D0_Kpi_100evts_newPacking_newDst.tck.json
histo_file: v2_example_histos.root
ntuple_file: v2_example_ntuples.root
input_process: Hlt2
input_stream: default
print_freq: 1000
input_raw_format: 0.5
persistreco_version: 0.0
lumi: False
write_fsr: False