Skip to content
Snippets Groups Projects

Draft: Upload code for Lc2pKs_Xic2L0Ks_Omgc2L0Ks AP version1

Closed Xuanjia Peng requested to merge Xuanjia/LcTopKs_XicToL0Ks_OmgcToL0Ks into master
1 unresolved thread
4 files
+ 877
0
Compare changes
  • Side-by-side
  • Inline
Files
4
+ 224
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(pvs, input_data, ptype, particle_to_constrain):
if ptype not in ["basic", "composite"]:
Exception(f"I want \'basic\' or \'composite\'. Got {ptype}")
#particle_to_constrain = ["KS0"]
DTF = DecayTreeFitter(
name=f'DTF_{{hash}}', input_particles=input_data, output_level=5)
DTF_PV = DecayTreeFitter(
name=f'DTF_PV_{{hash}}',
input_particles=input_data,
input_pvs=pvs,
output_level=5)
DTF_PV_mass = DecayTreeFitter(
name=f'DTF_PV_mass_{{hash}}',
input_particles=input_data,
input_pvs=pvs,
mass_constraints=particle_to_constrain,
output_level=5)
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_mass,
pv_constraint=True,
mass_constraint=True,
particle_name=particle_to_constrain[0])
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_mass,
pv_constraint=True,
mass_constraint=True,
particle_name=particle_to_constrain[0])
return dtf_vars
Loading