Merge branch 'tfulghes-trackisolation-collection' into 'master'

Add TrackIsolation functor collection

See merge request !906

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