Tupling jet info from Spruce data output

This example shows how to produce ntuples with jet ID info using JetInfoRelationTables.

import Functors as F
from DaVinci import Options, make_config
from DaVinci.algorithms import create_lines_filter
from PyConf.reading import get_particles, get_pvs
from FunTuple import FunctorCollection
import FunTuple.functorcollections as FC
from FunTuple import FunTuple_Particles as Funtuple

from PyConf.Algorithms import LHCb__JetInfoRelationTable as JetInfoRelationTable

from PyConf.Algorithms import ThOrParticleSelection


# Copied from rd_isolation.py ...
# To collect the jets in a decay tree
def find_in_decay(input, id):
    """
    Retrieve the selection of particles matching 'id' argument from 'input' decay tree
    Args:
        input: Composite particle
        id: String representation of the ParticleProperty to get
    """
    code = F.FILTER(F.IS_ABS_ID(id)) @ F.GET_ALL_DESCENDANTS()

    return ThOrParticleSelection(InputParticles=input, Functor=code).OutputSelection


def tr(**outputs):
    return {"OutputLocation": [v for (k, v) in outputs.items()]}


def jet_ID_FC(input_data):
    # input_data must be a list of jets locations
    # List of  info (jet ID) to be added
    infos = ["Ntracks", "MTF", "CPF", "MPT", "MNF", "JetWidth", "JetWidthNorm"]
    tables = JetInfoRelationTable(
        InputParticles=input_data,
        JetInfo=infos,
        outputs={prop: None for prop in infos},
        output_transform=tr,
    ).outputs

    jet_vars = {}
    for jetID in infos:
        tables[jetID].force_type("LHCb::Relation1D<LHCb::Particle,double>")
        jet_vars[jetID] = F.MAP_TO_RELATED(tables[jetID])

    return FunctorCollection(jet_vars)


def main(options: Options):

    lines = []

    # Dijet lines
    line = "SpruceQEE_diSVTag1010"
    lines.append(line)

    input_data = get_particles(f"/Event/Spruce/{line}/Particles")

    pvs = get_pvs()

    # Event variables
    evt_vars = FC.EventInfo()
    evt_vars["PV_SIZE"] = F.SIZE(pvs)
    evt_vars["ALLPVX"] = F.ALLPVX(pvs)
    evt_vars["ALLPVY"] = F.ALLPVY(pvs)
    evt_vars["ALLPVZ"] = F.ALLPVZ(pvs)

    # Kinematics
    var_kin = FC.Kinematics()

    all_variables = FunctorCollection(
        {
            "BPVIP": F.BPVIP(pvs),
            "BPVIPCHI2": F.BPVIPCHI2(pvs),
            "BPV_POS": F.BPV_POS(pvs),
        }
    )
    all_variables = var_kin + all_variables

    jet_variables = FunctorCollection(
        {
            "nCons": F.NINGENERATION(F.ALL, 1),
            "nChargedCons": F.NINGENERATION(F.CHARGE != 0, 1),
            "OWNPV_POS": F.OWNPV_POS(),
            "OWNPVIP": F.OWNPVIP(),
            "OWNPVIPCHI2": F.OWNPVIPCHI2(),
        }
    )

    jetID_vars = jet_ID_FC(input_data)

    fields_dijet = {
        "DiJet": "CLUSjet -> CELLjet CELLjet",
        "Jet1": "CLUSjet -> ^CELLjet CELLjet",
        "Jet2": "CLUSjet -> CELLjet ^CELLjet",
    }

    jets = find_in_decay(input_data, id="CELLjet")

    # Functor to the jet ID vars
    jetID_vars = jet_ID_FC(jets)

    variables_dijet = {
        "DiJet": all_variables,
        "Jet1": all_variables + jet_variables + jetID_vars,
        "Jet2": all_variables + jet_variables + jetID_vars,
    }

    # Define instance of FunTuple
    tuple_dijets = Funtuple(
        name=line + "_ntp",
        tuple_name="DecayTree",
        fields=fields_dijet,
        variables=variables_dijet,
        event_variables=evt_vars,
        inputs=input_data,
    )

    # Define a filter
    filter_jets = create_lines_filter("HDRFilter_jets", lines=lines)

    config = make_config(options, [filter_jets, tuple_dijets])
    return config

To run the example:

lbexec DaVinciExamples.tupling.option_davinci_tupling_jets:main $DAVINCIEXAMPLESROOT/example_data/example_tupling_jets.yaml

For reference, these are the options of this example