All Functors in DaVinci

This example shows how to:

  1. create a dictionary with functors of interest for various type of particles

  2. add a tuple with the selected functors.

__author__ = "P. Koppenburg"
__date__ = "2021-11-23"

import Functors as F
import FunTuple.functorcollections as FC
from FunTuple import FunctorCollection
from FunTuple import FunTuple_Particles as Funtuple
from PyConf.reading import get_particles, get_pvs
from DaVinci.algorithms import create_lines_filter
from PyConf.reading import get_odin  # get_decreports,
from DecayTreeFitter import DecayTreeFitter
from DaVinciMCTools import MCTruthAndBkgCat
from PyConf.Algorithms import PrintDecayTree

from DaVinci import Options, make_config

#
# Definition of Sprucing line
#
bd2dsk_line = "SpruceB2OC_BdToDsmK_DsmToHHH_FEST"

_basic = "basic"
_composite = "composite"
_toplevel = "toplevel"


def all_variables(pvs, dtf, mctruth, ptype, candidates=None, ftAlg=None):
    """
    function that returns dictionary of functors that work.

    functors are listed in order of https://lhcbdoc.web.cern.ch/lhcbdoc/moore/master/selection/thor_functors_reference.html#module-Functors
    """
    if ptype not in [_basic, _composite]:
        Exception(f"I want {_basic} or {_composite}. Got {ptype}")
    all_vars = FunctorCollection({})

    comp = _composite == ptype or _toplevel == ptype  # is composite
    basic = _basic == ptype  # is not composite
    top = _toplevel == ptype  # the B

    # First import everything that comes in functorcollections
    all_vars += FC.Kinematics()
    if basic:
        all_vars += FC.ParticleID(extra_info=True)
    all_vars += FC.MCKinematics(mctruth_alg=mctruth)
    all_vars += FC.MCHierarchy(mctruth_alg=mctruth)
    all_vars += FC.MCPrimaryVertexInfo(mctruth_alg=mctruth)
    Hlt1_decisions = ["Hlt1TrackMVADecision", "Hlt1TwoTrackMVADecision"]
    if candidates:
        all_vars += FC.HltTisTos(
            selection_type="Hlt", trigger_lines=Hlt1_decisions, data=candidates
        )
    if comp:
        all_vars += FC.MCVertexInfo(mctruth_alg=mctruth)
    if top:
        all_vars += FC.MCPromptDecay(mctruth_alg=mctruth)

    #
    # FTAlg not yet implemented
    # For Track isolation see weightedrelations_trackvariables
    #
    # Now all other functors

    # ALL : Not useful for tupling

    if comp:
        all_vars.update({"ALV": F.ALV(Child1=1, Child2=2)})

        all_vars.update({"BKGCAT": mctruth.BkgCat})

    if comp:  # all these require a vertex
        all_vars.update({"BPVCORRM": F.BPVCORRM(pvs)})
        all_vars.update({"BPVCORRMERR": F.BPVCORRMERR(pvs)})
        all_vars.update({"BPVDIRA": F.BPVDIRA(pvs)})
        all_vars.update({"BPVDLS": F.BPVDLS(pvs)})
        all_vars.update({"BPVETA": F.BPVETA(pvs)})
        all_vars.update({"BPVFD": F.BPVFD(pvs)})
        all_vars.update({"BPVFDCHI2": F.BPVFDCHI2(pvs)})
        all_vars.update({"BPVFDIR": F.BPVFDIR(pvs)})
        all_vars.update({"BPVFDVEC": F.BPVFDVEC(pvs)})

    all_vars.update({"BPVIP": F.BPVIP(pvs)})
    all_vars.update({"BPVIPCHI2": F.BPVIPCHI2(pvs)})
    all_vars.update({"BPVX": F.BPVX(pvs)})
    all_vars.update({"BPVY": F.BPVY(pvs)})
    all_vars.update({"BPVZ": F.BPVZ(pvs)})
    # When storing variable length array one can
    # give a custom branch name for the index.
    # This can be achieved by enclosing custom index
    # name within square brackets (see code below).
    # The branch name ("nPV") will correspond to the
    # index of the PV. If no index branch name given i.e.
    # all_vars.update({ 'ALLPVX'] the default "indx" is used.

    all_vars.update({"ALLPVX[nPVs]": F.ALLPVX(pvs)})
    all_vars.update({"ALLPVY[nPVs]": F.ALLPVY(pvs)})
    all_vars.update({"ALLPVZ[nPVs]": F.ALLPVZ(pvs)})

    if comp:  # all these require a vertex
        all_vars.update({"ALLPV_FD[nPVs]": F.ALLPV_FD(pvs)})
        all_vars.update({"ALLPV_IP[nPVs]": F.ALLPV_IP(pvs)})
        all_vars.update({"BPVLTIME": F.BPVLTIME(pvs)})
        all_vars.update({"BPVVDRHO": F.BPVVDRHO(pvs)})
        all_vars.update({"BPVVDX": F.BPVVDX(pvs)})
        all_vars.update({"BPVVDY": F.BPVVDY(pvs)})
        all_vars.update({"BPVVDZ": F.BPVVDZ(pvs)})

    all_vars.update({"CHARGE": F.CHARGE})
    all_vars.update({"CHI2": F.CHI2})
    all_vars.update({"CHI2DOF": F.CHI2DOF})
    if top:  # apply this only to B
        all_vars.update({"CHILD1_PT": F.CHILD(1, F.PT)})  # example of CHILD
        all_vars.update({"Ds_END_VZ": F.CHILD(1, F.END_VZ)})
        all_vars.update({"Delta_END_VZ_DsB0": F.CHILD(1, F.END_VZ) - F.END_VZ})

    # if basic: all_vars.update({ 'CLOSESTTOBEAM' : F.CLOSESTTOBEAM # 'Track__ClosestToBeamState' object has no attribute 'to_json'
    # COMB
    # if basic: all_vars.update({ 'COV' : F.COV # 'Track__Covariance' object has no attribute 'to_json'

    if comp:
        all_vars.update({"DOCA": F.SDOCA(Child1=1, Child2=2)})
        all_vars.update({"DOCACHI2": F.SDOCACHI2(Child1=1, Child2=2)})
        all_vars.update({"END_VRHO": F.END_VRHO})
        all_vars.update({"END_VX": F.END_VX})
        all_vars.update({"END_VY": F.END_VY})
        all_vars.update({"END_VZ": F.END_VZ})

    # duplicated from FC   all_vars.update({ 'ENERGY' : F.ENERGY})
    all_vars.update({"ETA": F.ETA})
    all_vars.update({"FOURMOMENTUM": F.FOURMOMENTUM})
    all_vars.update({"ISBASIC": F.ISBASICPARTICLE})

    if basic:
        all_vars.update({"GHOSTPROB": F.GHOSTPROB})
        all_vars.update({"ISMUON": F.ISMUON})
        all_vars.update({"INMUON": F.INMUON})
        all_vars.update({"INECAL": F.INECAL})
        all_vars.update({"INHCAL": F.INHCAL})
        all_vars.update({"HASBREM": F.HASBREM})
        all_vars.update({"HASBREMADDED": F.HASBREMADDED})
        all_vars.update({"BREMENERGY": F.BREMENERGY})
        all_vars.update({"BREMBENDCORR": F.BREMBENDCORR})
        all_vars.update({"BREMPIDE": F.BREMPIDE})
        all_vars.update({"ECALPIDE": F.ECALPIDE})
        all_vars.update({"ECALPIDMU": F.ECALPIDMU})
        all_vars.update({"HCALPIDE": F.HCALPIDE})
        all_vars.update({"HCALPIDMU": F.HCALPIDMU})
        all_vars.update({"ELECTRONSHOWEREOP": F.ELECTRONSHOWEREOP})
        all_vars.update({"ELECTRONSHOWERDLL": F.ELECTRONSHOWERDLL})
        all_vars.update({"CLUSTERID": F.CLUSTERID})
        all_vars.update({"CLUSTERMATCH_CHI2": F.CLUSTERMATCH_CHI2})
        all_vars.update({"ELECTRONMATCH_CHI2": F.ELECTRONMATCH_CHI2})
        all_vars.update({"BREMHYPOMATCH_CHI2": F.BREMHYPOMATCH_CHI2})
        all_vars.update({"ELECTRONENERGY": F.ELECTRONENERGY})
        all_vars.update({"BREMHYPOENERGY": F.BREMHYPOENERGY})
        all_vars.update({"BREMHYPODELTAX": F.BREMHYPODELTAX})
        all_vars.update({"BREMTRACKBASEDENERGY": F.BREMTRACKBASEDENERGY})
        all_vars.update({"ELECTRONID": F.ELECTRONID})
        all_vars.update({"HCALEOP": F.HCALEOP})
        # Note: the observables for the two functors below are (TRACK_MOM_X, TRACK_MOM_Y, TRACK_MOM_Z})
        # and (TRACK_POS_CLOSEST_TO_BEAM_X, TRACK_POS_CLOSEST_TO_BEAM_Y, TRACK_POS_CLOSEST_TO_BEAM_Z),
        # which is why the trailing underscore in the name is added i.e. "TRACK_MOM_" and "TRACK_POS_CLOSEST_TO_BEAM_"
        all_vars.update({"TRACK_MOM_": F.TRACK_MOMVEC})
        all_vars.update({"TRACK_POS_CLOSESTTOBEAM_": F.TRACK_POSVEC_CLOSESTTOBEAM})

        all_vars.update({"IS_ABS_ID_pi": F.IS_ABS_ID("pi+")})
        all_vars.update({"IS_ID_pi": F.IS_ID("pi-")})
        all_vars.update({"PDG_MASS_pi": F.PDG_MASS("pi+")})
        all_vars.update({"SIGNED_DELTA_MASS_pi": F.SIGNED_DELTA_MASS("pi+")})
        all_vars.update({"ABS_DELTA_MASS_pi": F.ABS_DELTA_MASS("pi+")})
        all_vars.update({"IS_NOT_H": F.IS_NOT_H})
        all_vars.update({"IS_PHOTON": F.IS_PHOTON})

    all_vars.update({"DTF_PT": dtf(F.PT)})
    all_vars.update({"DTF_BPVIPCHI2": dtf(F.BPVIPCHI2(pvs))})

    all_vars.update(
        {
            "DTF_MASS_SmallestDELTAPT": F.MASS
            @ F.TO
            @ F.ENTRY_WITH_MIN_REL_VALUE_OF(
                F.PT @ F.TO @ F.FORWARDARG0 - F.PT @ F.FORWARDARG1
            ).bind(
                F.RELATIONS.bind(F.TES(dtf.OutputRelations), F.FORWARDARGS),
                F.FORWARDARGS,
            )
        }
    )
    all_vars.update(
        {
            "DTF_MASS_BiggestDELTAPT": F.MASS
            @ F.TO
            @ F.ENTRY_WITH_MAX_REL_VALUE_OF(
                F.PT @ F.TO @ F.FORWARDARG0 - F.PT @ F.FORWARDARG1
            ).bind(
                F.RELATIONS.bind(F.TES(dtf.OutputRelations), F.FORWARDARGS),
                F.FORWARDARGS,
            )
        }
    )

    if top:
        all_vars.update({"DTF_NITER": dtf.NITER})
        all_vars.update({"DTF_CHI2": dtf.CHI2})
        all_vars.update({"DTF_NDOF": dtf.NDOF})
        all_vars.update({"DTF_CHI2DOF": dtf.CHI2DOF})

    if comp:
        all_vars.update({"DTF_MASS": dtf.MASS})
        all_vars.update({"DTF_MASSERR": dtf.MASSERR})
        all_vars.update({"DTF_P": dtf.P})
        all_vars.update({"DTF_PERR": dtf.PERR})
        all_vars.update({"DTF_CTAU": dtf.CTAU})
        all_vars.update({"DTF_CTAUERR": dtf.CTAUERR})
        all_vars.update({"DTF_FD": dtf.FD})
        all_vars.update({"DTF_FDERR": dtf.FDERR})

    all_vars.update({"MASS": F.MASS})
    if top:  # B
        all_vars.update({"MASSWITHHYPOTHESES": F.MASSWITHHYPOTHESES((939.0, 939.0))})
    elif comp:  # Ds
        all_vars.update(
            {"MASSWITHHYPOTHESES": F.MASSWITHHYPOTHESES((493.7, 493.7, 139.6))}
        )
    if comp:
        all_vars.update({"MAXPT": F.MAX(F.PT)})
        all_vars.update({"MAXDOCA": F.MAXSDOCA})
        all_vars.update({"MAXDOCACHI2": F.MAXSDOCACHI2})
        # the above in cut versions.

    # duplicated from FC    all_vars.update({ 'MC_MOTHER_ID' : F.VALUE_OR(0) @ mctruth(
    # duplicated from FC        F.MC_MOTHER(1, F.PARTICLE_ID))})

    if comp:
        all_vars.update({"MINPT": F.MIN(F.PT)})
    all_vars.update({"MINIP": F.MINIP(pvs)})
    all_vars.update({"MINIPCHI2": F.MINIPCHI2(pvs)})

    if basic:
        all_vars.update({"TRACKPT": F.TRACK_PT})
        all_vars.update({"TRACKHISTORY": F.VALUE_OR(-1) @ F.TRACKHISTORY @ F.TRACK})
        all_vars.update({"QOVERP": F.QOVERP @ F.TRACK})
        all_vars.update({"NDOF": F.VALUE_OR(-1) @ F.NDOF @ F.TRACK})
        all_vars.update({"NFTHITS": F.VALUE_OR(-1) @ F.NFTHITS @ F.TRACK})
        all_vars.update({"NHITS": F.VALUE_OR(-1) @ F.NHITS @ F.TRACK})
        all_vars.update({"NUTHITS": F.VALUE_OR(-1) @ F.NUTHITS @ F.TRACK})
        all_vars.update({"NVPHITS": F.VALUE_OR(-1) @ F.NVPHITS @ F.TRACK})
        all_vars.update({"TRACKHASVELO": F.VALUE_OR(-1) @ F.TRACKHASVELO @ F.TRACK})
        all_vars.update({"TRACKHASUT": F.VALUE_OR(-1) @ F.TRACKHASUT @ F.TRACK})

    all_vars.update({"OBJECT_KEY": F.OBJECT_KEY})

    # duplicated from FC        all_vars.update({ 'ORIGIN_VX' : mctruth(F.ORIGIN_VX)})
    # duplicated from FC        all_vars.update({ 'ORIGIN_VY' : mctruth(F.ORIGIN_VY)})
    # duplicated from FC        all_vars.update({ 'ORIGIN_VZ' : mctruth(F.ORIGIN_VZ)})

    # duplicated from FC    all_vars.update({ 'P' : F.P})
    # duplicated from FC    all_vars.update({ 'PARTICLE_ID' : F.PARTICLE_ID})
    all_vars.update({"PHI": F.PHI})

    # duplicated from FC    if basic:
    # duplicated from FC        all_vars.update({ 'PID_E' : F.PID_E})
    # duplicated from FC        all_vars.update({ 'PID_K' : F.PID_K})
    # duplicated from FC        all_vars.update({ 'PID_MU' : F.PID_MU})
    # duplicated from FC        all_vars.update({ 'PID_P' : F.PID_P})
    # duplicated from FC        all_vars.update({ 'PID_PI' : F.PID_PI})
    # POD
    # duplicated from FC        all_vars.update({ 'PROBNN_D' : F.PROBNN_D})
    # duplicated from FC        all_vars.update({ 'PROBNN_E' : F.PROBNN_E})
    # duplicated from FC        all_vars.update({ 'PROBNN_GHOST' : F.PROBNN_GHOST})
    # duplicated from FC        all_vars.update({ 'PROBNN_K' : F.PROBNN_K})
    # duplicated from FC        all_vars.update({ 'PROBNN_MU' : F.PROBNN_MU})
    # duplicated from FC        all_vars.update({ 'PROBNN_P' : F.PROBNN_P})
    # duplicated from FC        all_vars.update({ 'PROBNN_PI' : F.PROBNN_PI})

    # duplicated from FC    all_vars.update({ 'PT' : F.PT})
    # duplicated from FC    all_vars.update({ 'PX' : F.PX})
    # duplicated from FC    all_vars.update({ 'PY' : F.PY})
    # duplicated from FC    all_vars.update({ 'PZ' : F.PZ})
    all_vars.update({"ABS_PX": F.ABS @ F.PX})

    all_vars.update({"REFERENCEPOINT_X": F.REFERENCEPOINT_X})
    all_vars.update({"REFERENCEPOINT_Y": F.REFERENCEPOINT_Y})
    all_vars.update({"REFERENCEPOINT_Z": F.REFERENCEPOINT_Z})

    if comp:
        all_vars.update({"SDOCA": F.SDOCA(1, 2)})
        all_vars.update({"SDOCACHI2": F.SDOCACHI2(1, 2)})
    if basic:
        all_vars.update({"SHOWER_SHAPE": F.CALO_NEUTRAL_SHOWER_SHAPE})

    if comp:
        all_vars.update({"SUBCOMB12_MM": F.SUBCOMB(Functor=F.MASS, Indices=(1, 2))})
        all_vars.update({"SUMPT": F.SUM(F.PT)})

    if basic:
        all_vars.update({"TX": F.TX})
        all_vars.update({"TY": F.TY})

    print(f"### For {ptype} returning variables {all_vars.functor_dict.keys()}")
    return all_vars


def event_variables(PVs, ODIN, decreports, lines):
    """
    event variables
    """

    evt_vars = FunctorCollection({})
    evt_vars += FC.EventInfo()

    evt_vars += FC.SelectionInfo(selection_type="Spruce", trigger_lines=lines)
    # duplicated from FC    if ODIN:
    # duplicated from FC        evt_vars.update({ 'BUNCHCROSSING_ID' : F.BUNCHCROSSING_ID(ODIN)})
    # duplicated from FC        evt_vars.update({ 'BUNCHCROSSING_TYPE' : F.BUNCHCROSSING_TYPE(ODIN)})

    if decreports:
        evt_vars.update(
            {
                "DECISIONS": F.DECISIONS(
                    Lines=[bd2dsk_line + "Decision"], DecReports=decreports
                )
            }
        )
        evt_vars.update(
            {
                "DECREPORTS_FILTER": F.DECREPORTS_FILTER(
                    Lines=[bd2dsk_line + "Decision"], DecReports=decreports
                )
            }
        )

    if ODIN:
        evt_vars.update({"EVENTTYPE": F.EVENTTYPE(ODIN)})

    # duplicated from FC        evt_vars.update({ 'GPSTIME' : F.GPSTIME(ODIN)})
    # duplicated from FC        evt_vars.update({ 'ODINTCK' : F.ODINTCK(ODIN)})

    evt_vars.update({"PV_SIZE": F.SIZE(PVs)})
    # duplicated from FC        evt_vars.update({ 'GPSTIME' : F.GPSTIME(ODIN)})
    # duplicated from FC        evt_vars.update({ 'ODINTCK' : F.ODINTCK(ODIN)})

    if decreports:
        evt_vars.update({"TCK": F.TCK(decreports)})

    print(f"### For event returning variables {evt_vars.functor_dict.keys()}")
    return evt_vars


def alg_config(options: Options):
    """
    Algorithm configuration function called from the command line
    """
    # get the particles from line
    bd2dsk_data = get_particles(f"/Event/Spruce/{bd2dsk_line}/Particles")

    #
    # DecayTreeFitter Algorithm
    #
    v2_pvs = get_pvs()

    #
    # DecayTreeFitter Algorithm with "name" and "input_particles" as arguments
    #
    DTF = DecayTreeFitter(name="DTF_Bd2DsK", input_particles=bd2dsk_data)
    #
    # MC truth with "input_particles" as argument
    #
    MCTRUTH = MCTruthAndBkgCat(
        input_particles=bd2dsk_data, name="MCTruthAndBkgCat_functor"
    )
    #
    # Definition of fields (branches) and functors
    #
    fields_dsk = {
        "B0": "[B0 -> (D_s- -> pi+ pi- pi-) K+]CC",
        "Kaon": "[B0 -> (D_s- -> pi+ pi- pi-) ^K+]CC",
        "Ds": "[B0 -> ^(D_s- -> pi+ pi- pi-) K+]CC",
        "pip": "[B0 -> (D_s- -> ^pi+ pi- pi-) K+]CC",
    }
    variables_dsk = {
        "B0": all_variables(v2_pvs, DTF, MCTRUTH, _toplevel),
        "Kaon": all_variables(v2_pvs, DTF, MCTRUTH, _basic),
        "Ds": all_variables(v2_pvs, DTF, MCTRUTH, _composite),
        "pip": all_variables(v2_pvs, DTF, MCTRUTH, _basic),
    }

    #
    # event variables
    odin = get_odin()
    #    decreports = get_decreports('Spruce')
    decreports = None
    evt_vars = event_variables(v2_pvs, odin, decreports, [bd2dsk_line + "Decision"])

    #
    # Sprucing filter
    #
    my_filter = create_lines_filter(name="HDRFilter_B0DsK", lines=[f"{bd2dsk_line}"])

    #
    # FunTuple
    #
    my_tuple = Funtuple(
        name="B0DsK_Tuple",
        tuple_name="DecayTree",
        fields=fields_dsk,
        variables=variables_dsk,
        event_variables=evt_vars,
        loki_preamble=[],
        inputs=bd2dsk_data,
    )

    #
    # Algorithms to be run
    #
    return make_config(
        options, [my_filter, PrintDecayTree(Input=bd2dsk_data), my_tuple]
    )

To run the example:

lbexec DaVinciExamples.tupling.AllFunctors:alg_config $DAVINCIEXAMPLESROOT/example_options/example_tupling_allfunctors.yaml

For reference, these are the options of this example

    input_files:
      - root://eoslhcb.cern.ch//eos/lhcb/wg/dpa/wp3/tests/spruce_all_lines_realtimereco_newPacking_newDst.dst
    input_manifest_file: 'root://eoslhcb.cern.ch//eos/lhcb/wg/dpa/wp3/tests/spruce_all_lines_realtime_newPacking_newDst.tck.json'
    input_type: ROOT
    simulation: true
    conddb_tag: sim-20171127-vc-md100
    dddb_tag: dddb-20171126
    conditions_version: master
    geometry_version: run3/trunk
    input_raw_format: 0.3
    lumi: false
    ntuple_file: 'DV_example_allFunctors_ntp.root'
    print_freq: 1
    input_process: Spruce
    input_stream: default