diff --git a/Phys/FunTuple/python/FunTuple/functorcollections.py b/Phys/FunTuple/python/FunTuple/functorcollections.py index 7e559e01f9743791ed7ac73dd5cf8db2fdf5cca7..fd863ccd8dc477086af098496df369c3ce9da492 100644 --- a/Phys/FunTuple/python/FunTuple/functorcollections.py +++ b/Phys/FunTuple/python/FunTuple/functorcollections.py @@ -26,7 +26,7 @@ import io, sys import Functors as F from Functors.grammar import BoundFunctor from PyConf.application import make_data_with_FetchDataFromFile -from PyConf.Algorithms import MCTruthAndBkgCatAlg +from PyConf.Algorithms import MCTruthAndBkgCatAlg, WeightedRelTableAlg from .FunctorCollection import FunctorCollection __all__ = ( @@ -36,6 +36,7 @@ __all__ = ( "MCHierarchy", "MCKinematics", "MCVertexInfo", + "TrackIsolation", ) def EventInfo(odin, extra_info = False): @@ -268,6 +269,42 @@ def MCVertexInfo(mctruth): return FunctorCollection(MCVertexInfo) +def TrackIsolation(iso_rel_table): + """ + Collection of functors on track isolation, using information related to the particles + inside the cone around a given reference particle. + + Args: + iso_rel_table (WeightedRelTableAlg): algorithm instance returning the required relations table + that maps a reference paricle (forming the head of the cone) + to associated particles inside the cone itself. + + Todo: + Exclude particles that are in the decay descriptor and therefore belong to the decay that is being looked for. + At the moment WeightedRelTableAlg is making relations between the reference particle (head of the cone) + and related particles (particles within the cone) without checking whenever the particles in the cone + belong to the same decay of the head. + """ + SUMCONE = lambda func: F.SUMCONE(Functor=func, Relations=iso_rel_table.OutputRelations) + ASYM = lambda func: F.ASYM(Functor=func, Relations=iso_rel_table.OutputRelations) + TrackIsolationVariables = { + "HEAD_CMULT": F.VALUE_OR(0) @ F.NINCONE(Relations=iso_rel_table.OutputRelations), + "HEAD_CP": SUMCONE(F.P), + "HEAD_CPT": SUMCONE(F.PT), + "HEAD_CPX": SUMCONE(F.PX), + "HEAD_CPY": SUMCONE(F.PY), + "HEAD_CPZ": SUMCONE(F.PZ), + "HEAD_PASY": ASYM(F.P), + "HEAD_PTASY": ASYM(F.PT), + "HEAD_PXASY": ASYM(F.PX), + "HEAD_PYASY": ASYM(F.PY), + "HEAD_PZASY": ASYM(F.PZ), + "HEAD_DETA": F.MAP( F.ETA @ F.FORWARDARG1 - F.ETA @ F.TO @ F.FORWARDARG0 ).bind(F.RELATIONS.bind(F.TES(iso_rel_table.OutputRelations), F.FORWARDARGS), F.FORWARDARGS), + "HEAD_DPHI": F.MAP( F.ADJUST_ANGLE @ (F.PHI @ F.FORWARDARG1 - F.PHI @ F.TO @ F.FORWARDARG0) ).bind(F.RELATIONS.bind(F.TES(iso_rel_table.OutputRelations), F.FORWARDARGS), F.FORWARDARGS) + } + return FunctorCollection(TrackIsolationVariables) + + def __print_collection(coll_name): """Nicely print out the contents of a collection.""" print(f"{coll_name}:") @@ -285,6 +322,13 @@ def __print_collection(coll_name): dummy_dh = make_data_with_FetchDataFromFile("dummy3") dummy_alg = MCTruthAndBkgCatAlg(Input=dummy_dh) f_dict = coll(dummy_alg).functor_dict + elif coll_name == 'TrackIsolation': + dummy_dh1 = make_data_with_FetchDataFromFile("dummy4") + dummy_dh2 = make_data_with_FetchDataFromFile("dummy5") + dummy_alg = WeightedRelTableAlg( + ReferenceParticles=dummy_dh1, + InputCandidates=dummy_dh2) + f_dict = coll(dummy_alg).functor_dict else: f_dict = coll().functor_dict