Make necesary changes in the front-end

Phys/DecayTreeFitter/include/DecayTreeFitter/Fitter.h

@@ -10,6 +10,7 @@
 \*****************************************************************************/
 #pragma once
 
+#include "Event/RecVertex_v2.h"
 #include "Kernel/DecayTree.h"
 #include "LHCbMath/ParticleParams.h"
 #include "LHCbMath/ValueWithError.h"
@@ -17,6 +18,7 @@
 
 #include "GaudiKernel/SmartIF.h"
 
+#include <memory>
 #include <vector>
 
 #include "DetDesc/IGeometryInfo.h"
@@ -47,9 +49,15 @@ namespace DecayTreeFitter {
     /// constructor from the particle (decay head) and primary vertex
     Fitter( const LHCb::Particle& bc, const LHCb::VertexBase& pv, const bool forceFitAll = true,
             const ITrackStateProvider* extrapolator = 0 );
+    /// constructor from the particle (decay head) and v2 primary vertex
+    Fitter( const LHCb::Particle& bc, const LHCb::Event::v2::RecVertex& pv, const bool forceFitAll = true,
+            const ITrackStateProvider* extrapolator = 0 );
     /// constructor from the particle (decay head) and primary vertex
     Fitter( const LHCb::Particle& bc, const LHCb::VertexBase& pv, const ITrackStateProvider* extrapolator,
             const bool forceFitAll = true );
+    /// constructor from the particle (decay head) and v2 primary vertex
+    Fitter( const LHCb::Particle& bc, const LHCb::Event::v2::RecVertex& pv, const ITrackStateProvider* extrapolator,
+            const bool forceFitAll = true );
     /// destructor
     ~Fitter(); // destructor
 
@@ -109,6 +117,8 @@ namespace DecayTreeFitter {
     bool updateCand( LHCb::Particle& cand ) const;
     /// update a particlular candidate in the tree
     bool updateTree( LHCb::Particle& cand ) const;
+    /// update decay tree
+    StatusCode UpdateDecayTree( LHCb::DecayTree& tree ) const;
     /// error code
     int errCode() { return m_errCode; }
     /// set the verbosity level (for debugging only)
@@ -127,10 +137,10 @@ namespace DecayTreeFitter {
 
     Gaudi::Math::ValueWithError decayLengthSum( const ParticleBase&, const ParticleBase& ) const;
 
-    DecayChain*       decaychain() { return m_decaychain; }
-    FitParams*        fitparams() { return m_fitparams; }
-    const DecayChain* decaychain() const { return m_decaychain; }
-    const FitParams*  fitparams() const { return m_fitparams; }
+    DecayChain*       decaychain() { return m_decaychain.get(); }
+    FitParams*        fitparams() { return m_fitparams.get(); }
+    const DecayChain* decaychain() const { return m_decaychain.get(); }
+    const FitParams*  fitparams() const { return m_fitparams.get(); }
 
     double globalChiSquare( IGeometryInfo const& geometry ) const;
 
@@ -159,13 +169,13 @@ namespace DecayTreeFitter {
     /// Get current active algorithm
     const IAlgorithm* getAlg() const;
 
-    const LHCb::Particle* m_particle;
-    DecayChain*           m_decaychain;
-    FitParams*            m_fitparams;
-    FitStatus             m_status;
-    double                m_chiSquare;
-    int                   m_niter;
-    int                   m_errCode;
+    const LHCb::Particle*       m_particle;
+    std::unique_ptr<DecayChain> m_decaychain;
+    std::unique_ptr<FitParams>  m_fitparams;
+    FitStatus                   m_status;
+    double                      m_chiSquare;
+    int                         m_niter;
+    int                         m_errCode;
 
     typedef std::map<const LHCb::Particle*, Gaudi::Math::ParticleParams> Map;
     mutable Map                                                          m_map;

Phys/DecayTreeFitter/python/DecayTreeFitter/DTFHelpers.py

@@ -9,13 +9,13 @@
 # or submit itself to any jurisdiction.                                       #
 ###############################################################################
 
-from PyConf.Algorithms import DecayTreeFitterAlg, DecayTreeFitterAlgWithPV
+from PyConf.Algorithms import DecayTreeFitterAlg_v1Particle, DecayTreeFitterAlg_v1Particle_BestPV, RecV1ToPVConverter
 import Functors as F
 
 
 def DTFAlg(Input, InputPVs=None, MassConstraints=[], OutputLevel=3):
     """
-    Wrapper function for the DecayTreeFitterAlg and DecayTreeFitterAlgWithPV algorithms.
+    Wrapper function for the DecayTreeFitterAlg_v1Particle and DecayTreeFitterAlg_v1Particle_v1Vertex algorithms.
 
     Args:
         Input: Location of input particles (to be refitted).
@@ -25,17 +25,20 @@ def DTFAlg(Input, InputPVs=None, MassConstraints=[], OutputLevel=3):
         OutputLevel (int, optional): OutputLevel of algorithm. Defaults to INFO=3.
 
     Returns:
-        Configured DecayTreeFitterAlg or DecayTreeFitterAlgWithPV instance.
+        Configured DecayTreeFitterAlg_v1Particle or DecayTreeFitterAlg_v1Particle_v1Vertex instance.
     """
-
+    print(
+        '#Warning: DecayTreeFitter.DTFAlg deprecated. Use DecayTreeFitter.DecayTreeFitter'
+    )
     if InputPVs:
-        DTF = DecayTreeFitterAlgWithPV(
+        gaudi_pvs = RecV1ToPVConverter(InputVertices=InputPVs).OutputVertices
+        DTF = DecayTreeFitterAlg_v1Particle_BestPV(
             Input=Input,
-            InputPVs=InputPVs,
+            InputPVs=gaudi_pvs,
             MassConstraints=MassConstraints,
             OutputLevel=OutputLevel)
     else:
-        DTF = DecayTreeFitterAlg(
+        DTF = DecayTreeFitterAlg_v1Particle(
             Input=Input,
             MassConstraints=MassConstraints,
             OutputLevel=OutputLevel)
@@ -65,6 +68,9 @@ def DTF_functors(DTF, functors=[F.MASS], head='DTF_'):
     >>>
 
     """
+    print(
+        '#Warning: DecayTreeFitter.DTF_functors deprecated. Use DecayTreeFitter.DecayTreeFitter'
+    )
     outdict = {}
     for fct in functors:
         outdict[head + fct.name()] = F.MAP_INPUT(

Phys/DecayTreeFitter/python/DecayTreeFitter/DecayTreeFitter.py

@@ -13,7 +13,9 @@ from PyConf.dataflow import DataHandle
 from Gaudi.Configuration import INFO
 from Functors.grammar import BoundFunctor
 import Functors as F
-from PyConf.Algorithms import DecayTreeFitterAlg, DecayTreeFitterAlgWithPV
+from PyConf.Algorithms import DecayTreeFitterAlg_v1Particle, DecayTreeFitterAlg_v1Particle_BestPV, DecayTreeFitterAlg_v1Particle_AllPVs, RecV1ToPVConverter
+
+# return RecV1ToPVConverter(InputVertices=get_pvs_v1()).OutputVertices
 from DaVinciTools import SubstitutePID
 from typing import List
 
@@ -41,7 +43,7 @@ class DecayTreeFitter:
         substitutions (list): substitution rules, using the substitution syntax.
         output_level (int, optional): standard Gaudi Algorithm OutputLevel.
     Example:
-        pvs = get_pvs_v1() # DTF needs v1 pvs
+        pvs = get_pvs()
 
         # Create the tool
         DTF_PhiG = DecayTreeFitter(
@@ -76,7 +78,15 @@ class DecayTreeFitter:
                  input_pvs: DataHandle = None,
                  mass_constraints: List[str] = [],
                  substitutions: List[str] = [],
+                 fit_all_pvs=False,
                  output_level=INFO):
+        if fit_all_pvs and input_pvs is None:
+            raise ValueError(
+                "The 'input_pvs' parameter must not be set to None when fitting all PVs."
+            )
+
+        # Status
+        self.FitAllPVs = fit_all_pvs
 
         # Config subsitution if needed
         if substitutions:
@@ -92,14 +102,27 @@ class DecayTreeFitter:
 
         # Config algorithm
         if input_pvs:
-            self.Algorithm = DecayTreeFitterAlgWithPV(
+            self.HasPVConstrain = True
+            gaudi_algorithm = DecayTreeFitterAlg_v1Particle_AllPVs if fit_all_pvs else DecayTreeFitterAlg_v1Particle_BestPV
+            # If is v2 vertex
+            if 'LHCb::Event::PV::PrimaryVertexContainer' in input_pvs.type:
+                gaudi_pvs = input_pvs
+            elif 'LHCb::RecVertex' in input_pvs.type:
+                # Convert the v1 PVs to v2 PVs
+                gaudi_pvs = RecV1ToPVConverter(
+                    InputVertices=input_pvs).OutputVertices
+            else:
+                raise ValueError(
+                    'Invalid input_pvs type, input_pvs = {input_pvs.type}.')
+            self.Algorithm = gaudi_algorithm(
                 name=name,
                 Input=DTF_input,
-                InputPVs=input_pvs,
+                InputPVs=gaudi_pvs,
                 MassConstraints=mass_constraints,
                 OutputLevel=output_level)
         else:
-            self.Algorithm = DecayTreeFitterAlg(
+            self.HasPVConstrain = False
+            self.Algorithm = DecayTreeFitterAlg_v1Particle(
                 name=name,
                 Input=DTF_input,
                 MassConstraints=mass_constraints,
@@ -107,14 +130,78 @@ class DecayTreeFitter:
 
         self.Output = self.Algorithm.Output
         self.OutputRelations = self.Algorithm.OutputRelations  #Reco particle -> DTF particle
+        self.OutputParticleParams = self.Algorithm.ParticleParams
+        self.OutputNIter = self.Algorithm.NIter
+
+        # Functor to access fit result directly
+        self.NITER = F.VALUE_OR(-1) @ self._apply_functor(
+            functor=F.CAST_TO_INT, relation=self.OutputNIter)
+        self.CHI2 = self.__call__(
+            functor=F.CHI2 @ F.ENDVERTEX, apply_to_particle_params=False)
+        self.NDOF = F.VALUE_OR(-1) @ self.__call__(
+            functor=F.VALUE_OR(-1) @ F.NDOF @ F.ENDVERTEX,
+            apply_to_particle_params=False)
+        self.CHI2DOF = self.__call__(
+            functor=F.CHI2DOF @ F.ENDVERTEX, apply_to_particle_params=False)
+        self.MASS = self.__call__(
+            functor=F.MATH_VALUE @ F.MATH_INVARIANT_MASS @ F.FOURMOMENTUM,
+            apply_to_particle_params=True)
+        self.MASSERR = self.__call__(
+            functor=F.MATH_ERROR @ F.MATH_INVARIANT_MASS @ F.FOURMOMENTUM,
+            apply_to_particle_params=True)
+        self.P = self.__call__(
+            functor=F.MATH_VALUE @ F.MATH_SCALAR_MOMENTUM @ F.FOURMOMENTUM,
+            apply_to_particle_params=True)
+        self.PERR = self.__call__(
+            functor=F.MATH_ERROR @ F.MATH_SCALAR_MOMENTUM @ F.FOURMOMENTUM,
+            apply_to_particle_params=True)
+        self.TAU = self.__call__(
+            functor=F.MATH_VALUE @ F.PARTICLE_PARAMS_CTAU,
+            apply_to_particle_params=True)
+        self.TAUERR = self.__call__(
+            functor=F.MATH_ERROR @ F.PARTICLE_PARAMS_CTAU,
+            apply_to_particle_params=True)
+        self.FD = self.__call__(
+            functor=F.MATH_VALUE @ F.PARTICLE_PARAMS_DECAY_LENGTH,
+            apply_to_particle_params=True)
+        self.FDERR = self.__call__(
+            functor=F.MATH_ERROR @ F.PARTICLE_PARAMS_DECAY_LENGTH,
+            apply_to_particle_params=True)
+
+    def _apply_functor(self, functor: BoundFunctor, relation: DataHandle):
+        '''
+        Internal function that apply functor with certain relation
+        '''
+
+        if self.FitAllPVs and functor is None:
+            DTF_functor = F.MAP(F.TO) @ F.RELATIONS.bind(
+                F.TES(relation), F.FORWARDARGS)
+        elif self.FitAllPVs and functor is not None:
+            DTF_functor = F.MAP_INPUT_ARRAY(
+                Functor=functor, Relations=relation)
+        elif not self.FitAllPVs and functor is None:
+            DTF_functor = F.MAP_TO_RELATED(Relations=relation)
+        else:
+            DTF_functor = F.MAP_INPUT(functor, relation)
+
+        if self.SubstitutePID is not None:
+            #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(DTF_functor)
+        else:
+            return DTF_functor
 
-    def __call__(self, Functor: BoundFunctor):
+    def __call__(self,
+                 functor: BoundFunctor,
+                 apply_to_particle_params: bool = False):
         """
         Apply a specified functor to the resultant particle
-        obtained from the decay tree fit.
+        obtained from the decay tree fit or the resultant particle
+        parameters.
 
         Args:
-            Functor (BoundFunctor): the functor to be applied
+            functor (BoundFunctor): the functor to be applied
+            apply_to_particle_params (bool, optional): apply to the resultant particle parameters. Defaults to False
 
         Return:
             BoundFunctor: result functor
@@ -122,22 +209,23 @@ class DecayTreeFitter:
         Example:
             variables['DTF_CHI2DOF'] = DTF(F.CHI2DOF)
         """
-        if self.SubstitutePID is not None:
-            #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))
+        if apply_to_particle_params:
+            relation = self.OutputParticleParams
         else:
-            return F.MAP_INPUT(Functor, self.OutputRelations)
+            relation = self.OutputRelations
+
+        return self._apply_functor(functor=functor, relation=relation)
 
     def apply_functors(self,
                        functors: List[BoundFunctor] = [F.MASS],
+                       apply_to_particle_params: bool = False,
                        head: str = 'DTF_'):
         """
         Helper function returning a dictionary of functors to apply to DecayTreeFitted chain
 
         Args:
             functors: list of Functors (default: mass)
+            apply_to_particle_params (bool, optional): apply to the resultant particle parameters. Defaults to False
             head: header string. The default is "DTF_" and so F.PT of a refitted B will be decoded to "B_DTF_PT"
 
         Returns:
@@ -155,5 +243,6 @@ class DecayTreeFitter:
         """
         outdict = {}
         for functor in functors:
-            outdict[head + functor.name()] = self.__call__(functor)
+            outdict[head + functor.name()] = self.__call__(
+                functor, apply_to_particle_params=apply_to_particle_params)
         return outdict