Skip to content
Snippets Groups Projects
Commit 206d84a2 authored by Abhijit Mathad's avatar Abhijit Mathad Committed by Sebastien Ponce
Browse files

Align the helper classes used in DaVinci

parent 1064c488
No related branches found
No related tags found
No related merge requests found
###############################################################################
# (c) Copyright 2022 CERN for the benefit of the LHCb Collaboration #
# (c) Copyright 2022-2023 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". #
......@@ -15,6 +15,7 @@ from Functors.grammar import BoundFunctor
from PyConf.Algorithms import SubstitutePIDAlg_Particles
import Functors as F
import re
from typing import List
class SubstitutePID:
......@@ -26,22 +27,22 @@ class SubstitutePID:
You can configure the mass hypothesis subsitution using the following syntax:
' Old_PID{{New_PID}} '
The '[]cc' syntax is not yet supported in the SubstitutePID tool, you have to
specify the substitution rule explicitly for each cc case.
The '[]CC' syntax is not yet supported in the SubstitutePID tool, you have to
specify the substitution rule explicitly for each CC case.
e.g.
'[B+ -> K+ K+ K-{{pi-}}]cc' => ['B+ -> K+ K+ K-{{pi-}}', 'B- -> K- K- K+{{pi+}}']
'[B+ -> K+ K+ K-{{pi-}}]CC' => ['B+ -> K+ K+ K-{{pi-}}', 'B- -> K- K- K+{{pi+}}']
Args:
name (str): name of the substitution algorithm, will printed in the Gaudi.
substitutions (list): substitution rules, using the substitution syntax.
input (DataHandle): the input data handle (TES)
input_particles (DataHandle): the input_particles data handle (TES)
output_level (int, optional): standard Gaudi Algorithm OutputLevel.
Example:
# Create the tool
Subs_PhiG = SubstitutePID(
name = 'Subs_PhiG',
input = B_Data,
input_particles = B_Data,
substitutions = [
'B0{{B_s0}} -> ( K*(892)0{{phi(1020)}} -> K+ K- ) gamma'
]
......@@ -49,7 +50,7 @@ class SubstitutePID:
Subs_KstG = SubstitutePID(
name = 'Subs_KstG',
input = B_Data,
input_particles = B_Data,
substitutions = [
'B0 -> ( K*(892)0 -> K+ K-{{pi-}} ) gamma'
]
......@@ -57,15 +58,15 @@ class SubstitutePID:
# Get the energy of the substituted particle
allvariables['ENERGY'] = F.ENERGY
allvariables['PhiG_ENERGY'] = Subs_PhiG.get_info( F.ENERGY )
allvariables['KstG_ENERGY'] = Subs_KstG.get_info( F.ENERGY )
allvariables['PhiG_ENERGY'] = Subs_PhiG( F.ENERGY )
allvariables['KstG_ENERGY'] = Subs_KstG( F.ENERGY )
"""
def __init__(self,
name: str,
input: DataHandle,
substitutions: list,
output_level=INFO):
input_particles: DataHandle,
substitutions: List[str],
output_level: int = INFO):
# Check
if not substitutions:
raise ValueError("The 'substitutions' rules can't be empty.")
......@@ -97,7 +98,7 @@ class SubstitutePID:
# Create gaudi algorithm
self.Algorithm = SubstitutePIDAlg_Particles(
name=name,
Input=input,
Input=input_particles,
Decays=decays,
Substitute=subsitution_map,
OutputLevel=output_level)
......@@ -105,11 +106,12 @@ class SubstitutePID:
# Store the algorithm result
self.AllParticles = self.Algorithm.AllParticles
self.Particles = self.Algorithm.Particles
self.Relation = self.Algorithm.Relation
self.Relation = self.Algorithm.Relation # Particle -> PID Substituted Particle
def get_info(self, Functor: BoundFunctor):
def __call__(self, Functor: BoundFunctor):
"""
Apply a specified functor to the result particle of substitution
Apply a specified functor to the resultant particle
obtained from substituting the PID.
Args:
Functor (BoundFunctor): the functor to be applied
......@@ -118,11 +120,11 @@ class SubstitutePID:
BoundFunctor: result functor
Example:
variables['SUBS_CHI2DOF'] = Subs.get_info(F.CHI2DOF)
variables['SUBS_CHI2DOF'] = Subs(F.CHI2DOF)
"""
return F.MAP_INPUT(Functor, self.Relation)
def _get_subsitution_map(self, substitutions: list, output_level):
def _get_subsitution_map(self, substitutions: list, output_level: int):
if output_level in {DEBUG, VERBOSE, ALL}:
print("==============================")
......
###############################################################################
# (c) Copyright 2022 CERN for the benefit of the LHCb Collaboration #
# (c) Copyright 2022-2023 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". #
......@@ -15,6 +15,7 @@ from Functors.grammar import BoundFunctor
import Functors as F
from PyConf.Algorithms import DecayTreeFitterAlg, DecayTreeFitterAlgWithPV
from DaVinciTools import SubstitutePID
from typing import List
class DecayTreeFitter:
......@@ -33,8 +34,8 @@ class DecayTreeFitter:
Args:
name (str): name of the substitution algorithm, will printed in the Gaudi.
input (DataHandle): the input data handle (TES)
input_pvs (DataHandle or None, optional): location of primary vertices to which to constrain.
input_particles (DataHandle): TES location of the reconstructed particles.
input_pvs (DataHandle or None, optional): TES location of primary vertices to which to constrain.
Defaults to None.
mass_constraints (list of str, optional): list of mass constraints to apply. Defaults to [].
substitutions (list): substitution rules, using the substitution syntax.
......@@ -45,7 +46,7 @@ class DecayTreeFitter:
# Create the tool
DTF_PhiG = DecayTreeFitter(
name = 'DTF_PhiG',
input = B_Data,
input_particles = B_Data,
substitutions = [
'B0{{B_s0}} -> ( K*(892)0{{phi(1020)}} -> K+ K- ) gamma'
],
......@@ -55,7 +56,7 @@ class DecayTreeFitter:
DTF_KstG = SubstitutePID(
name = 'DTF_KstG',
input = B_Data,
input_particles = B_Data,
substitutions = [
'B0 -> ( K*(892)0 -> K+ K-{{pi-}} ) gamma'
],
......@@ -65,29 +66,29 @@ class DecayTreeFitter:
# Get the CHI2DOF of subtituted particle
allvariables['CHI2DOF'] = F.CHI2DOF
allvariables['PhiG_CHI2DOF'] = DTF_PhiG.get_info( F.CHI2DOF )
allvariables['KstG_CHI2DOF'] = DTF_KstG.get_info( F.CHI2DOF )
allvariables['PhiG_CHI2DOF'] = DTF_PhiG( F.CHI2DOF )
allvariables['KstG_CHI2DOF'] = DTF_KstG( F.CHI2DOF )
"""
def __init__(self,
name: str,
input: DataHandle,
input_pvs=None,
mass_constraints=[],
substitutions=[],
input_particles: DataHandle,
input_pvs: DataHandle = None,
mass_constraints: List[str] = [],
substitutions: List[str] = [],
output_level=INFO):
# Config subsitution if needed
if substitutions:
self.SubstitutePID = SubstitutePID(
name='PIDSubstitution_' + name,
input=input,
input_particles=input_particles,
substitutions=substitutions,
output_level=output_level)
DTF_input = self.SubstitutePID.Particles
else:
self.SubstitutePID = None
DTF_input = input
DTF_input = input_particles
# Config algorithm
if input_pvs:
......@@ -105,11 +106,12 @@ class DecayTreeFitter:
OutputLevel=output_level)
self.Output = self.Algorithm.Output
self.OutputRelations = self.Algorithm.OutputRelations
self.OutputRelations = self.Algorithm.OutputRelations #Reco particle -> DTF particle
def get_info(self, Functor: BoundFunctor):
def __call__(self, Functor: BoundFunctor):
"""
Apply a specified functor to the result particle of DTF
Apply a specified functor to the resultant particle
obtained from the decay tree fit.
Args:
Functor (BoundFunctor): the functor to be applied
......@@ -118,15 +120,19 @@ class DecayTreeFitter:
BoundFunctor: result functor
Example:
variables['DTF_CHI2DOF'] = DTF.get_info(F.CHI2DOF)
variables['DTF_CHI2DOF'] = DTF(F.CHI2DOF)
"""
if self.SubstitutePID is not None:
return self.SubstitutePID.get_info(
#The following maps the particle twice i.e. MAP_INPUT(MAP_INPUT(functor, P_2_DTFP), P_2_PIDSUBSTITUTEDP).
#Go from reco particle to PID substituted particle, then to DTF particle, then apply the functor
return self.SubstitutePID(
F.MAP_INPUT(Functor, self.OutputRelations))
else:
return F.MAP_INPUT(Functor, self.OutputRelations)
def apply_functors(self, functors=[F.MASS], head='DTF_'):
def apply_functors(self,
functors: List[BoundFunctor] = [F.MASS],
head: str = 'DTF_'):
"""
Helper function returning a dictionary of functors to apply to DecayTreeFitted chain
......@@ -135,20 +141,19 @@ class DecayTreeFitter:
head: header string. The default is "DTF_" and so F.PT of a refitted B will be decoded to "B_DTF_PT"
Returns:
Dictionary of names. By default it will return { 'DTF_Mass' : DTF.get_info(F.Mass) }
Dictionary of names. By default it will return { 'DTF_Mass' : DTF(F.Mass) }
Examples:
>>>
>>> from DecayTreeFitter import DecayTreeFitter
>>> DTF_pv = DecayTreeFitter(
>>> name='DTF_dimuons',
>>> input=dimuons,
>>> input_particles=dimuons,
>>> input_pvs=pvs,
>>> mass_constraints=["J/psi(1S)"])
>>> variables_jpsi.update(DTF_pv.apply_functors(functors=[F.PT,F.MASS], head='DTF_PV_'))
>>>
"""
outdict = {}
for functor in functors:
outdict[head + functor.name()] = self.get_info(functor)
outdict[head + functor.name()] = self.__call__(functor)
return outdict
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment