Skip to content
Snippets Groups Projects

Lb2D0Lz Run3

Merged Dong Ao requested to merge aodong/Lb2DL into master
Files
3
+ 164
0
##############################################################################
# (c) Copyright 2024 CERN for the benefit of the LHCb Collaboration #
# #
# This software is distributed under the terms of the GNU General Public #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING". #
# #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization #
# or submit itself to any jurisdiction. #
###############################################################################
"""Configuration functions for DTF
"""
import Functors as F
from Functors.math import log
from FunTuple import FunctorCollection
from FunTuple.functorcollections import Kinematics
from DecayTreeFitter import DecayTreeFitter
def make_basic_dtf_variables(pvs, data, DTF=None, pv_constraint=False, mass_constraint=False, particle_name=""):
variables = (
FunctorCollection(
{
"TRCHI2DOF" : F.CHI2DOF @ F.TRACK,
"ETA" : F.ETA,
"PHI" : F.PHI,
"TRGHOSTPROB" : F.GHOSTPROB,
"BPVIPCHI2" : F.BPVIPCHI2(pvs),
"BPVIP" : F.BPVIP(pvs),
"BPVX" : F.BPVX(pvs),
"BPVY" : F.BPVY(pvs),
"BPVZ" : F.BPVZ(pvs),
"TX" : F.TX,
"TY" : F.TY,
"MINIPCHI2" : F.MINIPCHI2(pvs),
"MINIP" : F.MINIP(pvs),
"KEY" : F.VALUE_OR(-1) @ F.OBJECT_KEY @ F.TRACK,
#"CTB" : F.POSITION @ F.CLOSESTTOBEAM @ F.TRACK,
"TRACKPT" : F.TRACK_PT,
"TRACKHISTORY" : F.VALUE_OR(-1) @ F.TRACKHISTORY @ F.TRACK,
"QOVERP" : F.QOVERP @ F.TRACK,
"NDOF" : F.VALUE_OR(-1) @ F.NDOF @ F.TRACK,
}
)
+ Kinematics()
)
if(mass_constraint):
if(pv_constraint): # MASS + PV
dtf_variables_mass_pv = FunctorCollection({ 'DTF_PV_Fix'+ particle_name + '_' + k: DTF(v) for k, v in variables.get_thor_functors().items() })
return dtf_variables_mass_pv
else: # MASS
dtf_variables_mass = FunctorCollection( {'DTF_Fix'+ particle_name + '_' + k: DTF(v) for k, v in variables.get_thor_functors().items()})
return dtf_variables_mass
elif(pv_constraint): # PV
dtf_variables_pv = FunctorCollection({ 'DTF_PV_' + k: DTF(v) for k, v in variables.get_thor_functors().items() })
return dtf_variables_pv
else: # NO MASS/PV
dtf_variables = FunctorCollection( {'DTF_' + k: DTF(v) for k, v in variables.get_thor_functors().items()})
return dtf_variables
def make_composite_dtf_variables(pvs, data, DTF=None, pv_constraint=False, mass_constraint=False, particle_name=""):
variables = (
FunctorCollection(
{
"MAXPT": F.MAX(F.PT),
"MINPT": F.MIN(F.PT),
"SUMPT": F.SUM(F.PT),
"MAXP": F.MAX(F.P),
"MINP": F.MIN(F.P),
"BPVDIRA": F.BPVDIRA(pvs),
"VCHI2DOF": F.CHI2DOF, #CHI2VXNDOF
"BPVFDCHI2": F.BPVFDCHI2(pvs),
"BPVFD": F.BPVFD(pvs),
"BPVVDRHO": F.BPVVDRHO(pvs),
"BPVVDZ": F.BPVVDZ(pvs),
"BPVIPCHI2": F.BPVIPCHI2(pvs),
"BPVIP": F.BPVIP(pvs),
# "LOGBPVIPCHI2": log(F.BPVIPCHI2(pvs)),
"BPVLTIME": F.BPVLTIME(pvs),
"MAXBPVIPCHI2": F.MAX(F.BPVIPCHI2(pvs)), #MAX_
"MINBPVIPCHI2": F.MIN(F.BPVIPCHI2(pvs)),
"MAXBPVIP": F.MAX(F.BPVIP(pvs)),
"MINBPVIP": F.MIN(F.BPVIP(pvs)),
"ETA": F.ETA,
"PHI": F.PHI,
"END_VX": F.END_VX, #END_
"END_VY": F.END_VY,
"END_VZ": F.END_VZ,
"BPVX": F.BPVX(pvs),
"BPVY": F.BPVY(pvs),
"BPVZ": F.BPVZ(pvs),
"ALLPVFD" : F.ALLPV_FD(pvs),
"ALLPVIP" : F.ALLPV_IP(pvs),
}
)
+ Kinematics()
)
addstring = "DTF"
if(pv_constraint):
addstring += '_PV'
if(mass_constraint):
addstring += '_Fix'
addstring += particle_name
DTF_chi2ndof = FunctorCollection(
{
addstring+"_DTFCHI2": DTF.CHI2,
addstring+"_DTFNDOF": DTF.NDOF,
addstring+"_CTAU": DTF.CTAU,
addstring+"_CTAUERR": DTF.CTAUERR,
addstring+"_MERR": DTF.MASSERR,
}
)
if(mass_constraint):
if(pv_constraint): # MASS + PV
dtf_variables_mass_pv = FunctorCollection({ 'DTF_PV_Fix'+ particle_name + '_' + k: DTF(v) for k, v in variables.get_thor_functors().items() })
return dtf_variables_mass_pv+DTF_chi2ndof
else: # MASS
dtf_variables_mass = FunctorCollection( {'DTF_Fix'+ particle_name + '_' + k: DTF(v) for k, v in variables.get_thor_functors().items()})
return dtf_variables_mass+DTF_chi2ndof
elif(pv_constraint): # PV
dtf_variables_pv = FunctorCollection({ 'DTF_PV_' + k: DTF(v) for k, v in variables.get_thor_functors().items() })
return dtf_variables_pv+DTF_chi2ndof
else: # NO MASS/PV
dtf_variables = FunctorCollection( {'DTF_' + k: DTF(v) for k, v in variables.get_thor_functors().items()})
return dtf_variables+DTF_chi2ndof
def make_dtf_variables_lb2d0lz(pvs, input_data, ptype):
if ptype not in ["basic", "composite"]:
Exception(f"I want \'basic\' or \'composite\'. Got {ptype}")
DTF = DecayTreeFitter( name=f'DTF_{{hash}}' , input_particles=input_data)
DTF_PV = DecayTreeFitter( name=f'DTF_PV_{{hash}}' , input_particles=input_data, input_pvs=pvs)
DTF_PV_FixD = DecayTreeFitter( name=f'DTF_PV_FixD_{{hash}}' , input_particles=input_data, input_pvs=pvs, mass_constraints=["D0"])
DTF_PV_FixLz = DecayTreeFitter( name=f'DTF_PV_FixLz_{{hash}}' , input_particles=input_data, input_pvs=pvs, mass_constraints=["Lambda0"])
DTF_PV_FixDLz = DecayTreeFitter( name=f'DTF_PV_FixDLz_{{hash}}' , input_particles=input_data, input_pvs=pvs, mass_constraints=[ "D0", "Lambda0"])
DTF_PV_FixB = DecayTreeFitter( name=f'DTF_PV_FixB_{{hash}}' , input_particles=input_data, input_pvs=pvs, mass_constraints=["Lambda_b0", "D0", "Lambda0"])
if ptype == "basic":
dtf_vars = make_basic_dtf_variables(pvs, input_data, DTF=DTF, pv_constraint=False, mass_constraint=False)
dtf_vars += make_basic_dtf_variables(pvs, input_data, DTF=DTF_PV, pv_constraint=True, mass_constraint=False)
dtf_vars += make_basic_dtf_variables(pvs, input_data, DTF=DTF_PV_FixD, pv_constraint=True, mass_constraint=True, particle_name="D")
dtf_vars += make_basic_dtf_variables(pvs, input_data, DTF=DTF_PV_FixLz, pv_constraint=True, mass_constraint=True, particle_name="Lz")
dtf_vars += make_basic_dtf_variables(pvs, input_data, DTF=DTF_PV_FixDLz, pv_constraint=True, mass_constraint=True, particle_name="DLz")
dtf_vars += make_basic_dtf_variables(pvs, input_data, DTF=DTF_PV_FixB, pv_constraint=True, mass_constraint=True, particle_name="B")
return dtf_vars
if ptype == "composite":
dtf_vars = make_composite_dtf_variables(pvs, input_data, DTF=DTF, pv_constraint=False, mass_constraint=False)
dtf_vars += make_composite_dtf_variables(pvs, input_data, DTF=DTF_PV, pv_constraint=True, mass_constraint=False)
dtf_vars += make_composite_dtf_variables(pvs, input_data, DTF=DTF_PV_FixD, pv_constraint=True, mass_constraint=True, particle_name="D")
dtf_vars += make_composite_dtf_variables(pvs, input_data, DTF=DTF_PV_FixLz, pv_constraint=True, mass_constraint=True, particle_name="Lz")
dtf_vars += make_composite_dtf_variables(pvs, input_data, DTF=DTF_PV_FixDLz, pv_constraint=True, mass_constraint=True, particle_name="DLz")
dtf_vars += make_composite_dtf_variables(pvs, input_data, DTF=DTF_PV_FixB, pv_constraint=True, mass_constraint=True, particle_name="B")
return dtf_vars
Loading