Tupling and DecayTreeFitter with Neutrino

This example demonstrates how to use FunTuple and DecayTreeFitter with a neutrino.

The decay tree involving a neutrino is simplified as ReferenceSystem -> (MotherParticle -> VisibleSystem neutrino) IndependentSystem. Four ntuples are created, each corresponding to a different neutrino strategy: Strategy 0 assumes neutrino is parallel to visible system, forcing the mass constraint of the mother particle. Strategy 1 estimates the flight direction of reference particle using endvertex and primary vertex position, applying the constraint p_perp_visible + p_perp_neutrino + p_perp_independent = 0. Strategy 2 estimates the flight direction of reference particle using endvertex and primary vertex position, calculating the transversal momentum of mother particle from independent system (pt(independent) = - pt(mother)), assuming neutrino has zero contribution in transversal momentum. Strategy 3 determines the direction of flight of the Mother by the end vertex position of the visible system and by the end vertex position of the reference system, then like strategy 2 calculates the transversal momentum of mother particle from independent system.

from PyConf.reading import get_particles
import Functors as F
from FunTuple import FunctorCollection, FunTuple_Particles as Funtuple
from DaVinci.algorithms import create_lines_filter
from DaVinci import Options, make_config
import FunTuple.functorcollections as FC
from DaVinciMCTools import MCTruthAndBkgCat
from PyConf.Algorithms import InsertNeutrinoToDecayAlg
from DecayTreeFitter import DecayTreeFitter
from functools import partial


def chain_functors(fc, chained_functor, prefix=""):
    return FunctorCollection(
        {f"{prefix}{k}": v @ chained_functor for k, v in fc.functor_dict.items()}
    )


def truth_match_functors(truth_match):
    output = FC.MCKinematics(mctruth_alg=truth_match)
    output += FunctorCollection(
        {
            "BKGCAT": truth_match.BkgCat,
            "TRUE": truth_match(F.HAS_VALUE),
        }
    )
    return output


def neutrino_truth_matching(truth_match):

    fetch_nu = F.FIND_MCDECAY("[ tau+ -x> pi+ pi- pi+ ^nu_tau~ ]CC")

    return FunctorCollection(
        {
            "Nu_TRUEP": truth_match(F.FOURMOMENTUM @ fetch_nu),
        }
    )


def main(options: Options, isMC=True):

    # Basics
    line_name = "Hlt2RD_BdToKstTauMu_KstToKPi_TauTo3Pi_OS"
    input_particles = get_particles(f"/Event/HLT2/{line_name}/Particles")
    hlt2_filter = create_lines_filter(name="HLT2Filter", lines=[line_name])
    truth_match = MCTruthAndBkgCat(input_particles, name="truth_match")

    # Define nTuple fields
    fields = {
        "B": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        "B_Jpsi": "[B0 -> ^([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        "B_taup": "[B0 -> ([J/psi(1S) -> ^(tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        "B_mum": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) ^mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        "B_Kst": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) ^(K*(892)0 -> K+ pi-)]CC",
        "B_Kp": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> ^K+ pi-)]CC",
        "B_Pim": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ ^pi-)]CC",
        "B_Pi1": "[B0 -> ([J/psi(1S) -> (tau+ -> ^pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        "B_Pi2": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ ^pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        "B_Pi3": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- ^pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        "B_Nu": "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ ^nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
    }

    # Define neutrino algorithm
    insert_neutrinos = partial(
        InsertNeutrinoToDecayAlg,
        InputParticles=input_particles,
        SelectedDecays=[
            "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+) mu-]CC) (K*(892)0 -> K+ pi-)]CC"
        ],
        ParticleInsertionRules={
            "B0 -> (J/psi(1S) -> ^(tau+ -> pi+ pi- pi+) mu-) (K*(892)0 -> K+ pi-)": "nu_tau~",
            "B0 -> (J/psi(1S) -> ^(tau- -> pi- pi+ pi-) mu+) (K*(892)0 -> K+ pi-)": "nu_tau",
            "B~0 -> (J/psi(1S) -> ^(tau- -> pi- pi+ pi-) mu+) (K*(892)~0 -> K- pi+)": "nu_tau",
            "B~0 -> (J/psi(1S) -> ^(tau+ -> pi+ pi- pi+) mu-) (K*(892)~0 -> K- pi+)": "nu_tau~",
        },
        VisibleSystem=[
            "[B0 -> ([J/psi(1S) -> (tau+ -> ^pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
            "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ ^pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
            "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- ^pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        ],
        IndependentSystem=[
            "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) ^mu-]CC) (K*(892)0 -> K+ pi-)]CC",
            "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> ^K+ pi-)]CC",
            "[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ ^pi-)]CC",
        ],
        ReferenceParticle="[B0 -> ([J/psi(1S) -> (tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        MotherParticle="[B0 -> ([J/psi(1S) -> ^(tau+ -> pi+ pi- pi+ nu_tau~) mu-]CC) (K*(892)0 -> K+ pi-)]CC",
        RelativeError=0.1,
        UpdateMomentum=True,
    )

    # Define FunTuple nodes
    davinci_node = {}
    for neutrino_strategy in [0, 1, 2, 3]:

        print(f"Adding FunTuple with neutrino strategy {neutrino_strategy}.")

        # Create neutrino
        neutrino = insert_neutrinos(Strategy=neutrino_strategy)
        relate_to_original_particle = F.MAP_TO_RELATED(neutrino.Relation)

        # Fit neutrino with DTF
        DTF = DecayTreeFitter(
            name=f"DTF_{neutrino_strategy}",
            constrain_to_ownpv=True,
            input_particles=neutrino.OutputDecays,
            mass_constraints=["tau+"],
        )

        # Define functors
        variables_all = FunctorCollection()
        variables_B = FunctorCollection()
        variables_tau = FunctorCollection()

        # Kinematics functors
        variables_all += FC.Kinematics()
        variables_all += chain_functors(
            FC.Kinematics(), relate_to_original_particle, prefix="ORIGINAL_"
        )

        # Truth matching functors
        variables_all += chain_functors(
            truth_match_functors(truth_match), relate_to_original_particle
        )
        variables_tau += chain_functors(
            neutrino_truth_matching(truth_match), relate_to_original_particle
        )

        # DTF functors
        variables_all += FC.DecayTreeFitterResults(
            DTF=DTF,
            decay_origin=False,
            with_lifetime=False,
            with_kinematics=True,
        )
        variables_B += FC.DecayTreeFitterResults(
            DTF=DTF,
            decay_origin=True,
            with_lifetime=True,
            with_kinematics=False,
        )

        # Assign functors
        variables = {
            "ALL": variables_all,
            "B": variables_B,
            "B_taup": variables_tau,
        }

        # Create funtuple
        funtuple = Funtuple(
            name=f"Strategy{neutrino_strategy}",
            tuple_name="DecayTree",
            fields=fields,
            variables=variables,
            inputs=neutrino.OutputDecays,
        )

        # Create FunTuple node
        davinci_node.update({f"Tupling_{neutrino_strategy}": [hlt2_filter, funtuple]})

    return make_config(options, davinci_node)

To run the example:

lbexec DaVinciExamples.tupling.option_davinci_tupling_neutrino:main $DAVINCIEXAMPLESROOT/example_data/test_neutrino.yaml

For reference, these are the options of this example

testfiledb_key: DaVinciTestNeutrino_2024_B2JpsiKst_Jpsi2TauMu_Tau2pipipinu
input_raw_format: 0.5
simulation: true
input_process: Hlt2
input_stream: full
lumi: false
evt_max: 1000
print_freq: 1000
ntuple_file: test_neutrino_ntuple.root
histo_file: test_neutrino_histo.root