Update the Run3 DTF: Add extra output and v2 PVs support

Phys/FunTuple/python/FunTuple/functorcollections.py

@@ -34,12 +34,14 @@ from PyConf.dataflow import DataHandle
 from PyConf.components import Algorithm
 from PyConf.Algorithms import HltTisTosAlg
 import DaVinciMCTools  # for MCReconstructible, MCReconstructed and MCTruthAndBkgCat helper classes
+from DecayTreeFitter import DecayTreeFitter
 from .FunctorCollection import FunctorCollection
 
 __all__ = ("EventInfo", "SelectionInfo", "HltTisTos", "Kinematics",
            "MCHierarchy", "MCKinematics", "MCVertexInfo", "MCPromptDecay",
            "MCReconstructible", "MCReconstructed", "TrackIsolation",
-           "NeutralCaloInfo", "ParticleID", "MCPrimaries")
+           "NeutralCaloInfo", "DecayTreeFitterResults", "ParticleID",
+           "MCPrimaries")
 
 
 def EventInfo():
@@ -252,6 +254,113 @@ def MCKinematics(
     return FunctorCollection(mckinematics)
 
 
+def DecayTreeFitterResults(*,
+                           DTF: DecayTreeFitter,
+                           decay_origin: bool = True,
+                           with_lifetime: bool = False,
+                           with_kinematics: bool = True,
+                           prefix: str = 'DTF') -> FunctorCollection:
+    """
+    Candidate-level collection of functors that access DecayTreeFitter fit results.
+
+    Args:
+        DTF (DecayTreeFitter.DecayTreeFitter): An instance of the `DecayTreeFitter` helper class
+            that builds the necessary functors.
+        decay_origin (bool): returns functorcollection for decay origin, defaults to True.
+        with_lifetime (bool): The returned functors comprise lifetime and flight distance, defaults to False.
+        with_kinematics (bool): The returned functors comprise kinematics, defaults to True.
+        prefix (str): Additional prefix to output functors. Defaults to 'DTF'
+
+    Example:
+        from FunTuple import FunTuple_Particles as Funtuple
+        import FunTuple.functorcollections as FC
+        from DecayTreeFitter import DecayTreeFitter
+
+        DTF = DecayTreeFitter(
+            name = 'DTF',
+            input_particles = get_particles(...),
+            input_pvs = get_pvs()
+        )
+
+        variable_decay_origin = FC.DecayTreeFitter(DTF, decay_origin=True, with_lifetime=True, with_kinematics=True)
+        variable_kinematics = FC.DecayTreeFitter(DTF, decay_origin=False, with_lifetime=False, with_kinematics=True)
+
+        variables = {
+            'B': variable_decay_origin,
+            'phi_1020': variable_kinematics,
+            'Kplus': variable_kinematics,
+            'Kminus': variable_kinematics,
+        }
+
+        tuple = Funtuple(name="MyTuple",
+                         fields=...,
+                         variables=variables,
+                         ...
+                         )
+    """
+    suffix = f'[{prefix}_IDX]' if DTF.FitAllPVs else ''
+
+    output_collection = FunctorCollection()
+
+    if decay_origin:
+        output_collection += FunctorCollection({
+            f'{prefix}_CHI2{suffix}':
+            DTF.CHI2,
+            f'{prefix}_NDOF{suffix}':
+            DTF.NDOF,
+            f'{prefix}_CHI2DOF{suffix}':
+            DTF.CHI2DOF,
+            f'{prefix}_NITER{suffix}':
+            DTF.NITER,
+            f'{prefix}_MASS{suffix}':
+            DTF.MASS,
+            f'{prefix}_MASSERR{suffix}':
+            DTF.MASSERR,
+            f'{prefix}_P{suffix}':
+            DTF.P,
+            f'{prefix}_PERR{suffix}':
+            DTF.PERR,
+        })
+        if DTF.HasPVConstraint:
+            output_collection += FunctorCollection({
+                f'{prefix}_PV_KEY{suffix}':
+                F.VALUE_OR(-1) @ DTF(
+                    F.VALUE_OR(-1) @ F.OBJECT_KEY @ F.ENDVERTEX),
+                f'{prefix}_PV_X{suffix}':
+                DTF(F.END_VX),
+                f'{prefix}_PV_Y{suffix}':
+                DTF(F.END_VY),
+                f'{prefix}_PV_Z{suffix}':
+                DTF(F.END_VZ),
+            })
+
+    if with_lifetime:
+        output_collection += FunctorCollection({
+            f'{prefix}_TAU{suffix}':
+            DTF.TAU,
+            f'{prefix}_TAUERR{suffix}':
+            DTF.TAUERR,
+            f'{prefix}_FD{suffix}':
+            DTF.FD,
+            f'{prefix}_FDERR{suffix}':
+            DTF.FDERR,
+        })
+
+    if with_kinematics:
+        output_collection += FunctorCollection({
+            f'{prefix}_PX{suffix}':
+            DTF(F.PX),
+            f'{prefix}_PY{suffix}':
+            DTF(F.PY),
+            f'{prefix}_PZ{suffix}':
+            DTF(F.PZ),
+            f'{prefix}_PE{suffix}':
+            DTF(F.ENERGY),
+        })
+
+    return output_collection
+
+
 def MCPromptDecay(
         *,
         mctruth_alg: DaVinciMCTools.MCTruthAndBkgCat = None,
@@ -914,6 +1023,19 @@ def __print_collection(coll_name):
         dummy_dh = make_data_with_FetchDataFromFile("dummy9")
         dummy_dh.force_type('LHCb::MCHeader')
         f_dict = coll(mc_header=dummy_dh).functor_dict
+    elif coll_name == 'DecayTreeFitterResults':
+        dummy_particles = make_data_with_FetchDataFromFile("dummy10")
+        dummy_pvs = make_data_with_FetchDataFromFile("dummy11")
+        dummy_pvs.force_type('LHCb::Event::PV::PrimaryVertexContainer')
+        DTF = DecayTreeFitter(
+            'DecayTreeFitter_{hash}',
+            input_particles=dummy_particles,
+            input_pvs=dummy_pvs)
+        f_dict = coll(
+            DTF=DTF,
+            decay_origin=True,
+            with_lifetime=True,
+            with_kinematics=True).functor_dict
     else:
         f_dict = coll().functor_dict