PWGHF: Add number of primary-vertex contributors in the 2-prong candidate creator (#3954)

Author: DelloStritto

PWGHF/DataModel/CandidateReconstructionTables.h

@@ -521,7 +521,7 @@ DECLARE_SOA_TABLE(HfCand2ProngBase, "AOD", "HFCAND2PBASE", //!
                   hf_cand::PxProng1, hf_cand::PyProng1, hf_cand::PzProng1,
                   hf_cand::ImpactParameter0, hf_cand::ImpactParameter1,
                   hf_cand::ErrorImpactParameter0, hf_cand::ErrorImpactParameter1,
-                  hf_track_index::Prong0Id, hf_track_index::Prong1Id,
+                  hf_track_index::Prong0Id, hf_track_index::Prong1Id, hf_cand::NProngsContributorsPV,
                   hf_track_index::HFflag,
                   /* dynamic columns */
                   hf_cand_2prong::M<hf_cand::PxProng0, hf_cand::PyProng0, hf_cand::PzProng0, hf_cand::PxProng1, hf_cand::PyProng1, hf_cand::PzProng1>,

PWGHF/TableProducer/candidateCreator2Prong.cxx

@@ -212,8 +212,17 @@ struct HfCandidateCreator2Prong {
       auto errorDecayLength = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, theta) + getRotatedCovMatrixXX(covMatrixPCA, phi, theta));
       auto errorDecayLengthXY = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, 0.) + getRotatedCovMatrixXX(covMatrixPCA, phi, 0.));
 
+      auto indexCollision = collision.globalIndex();
+      uint8_t nProngsContributorsPV = 0;
+      if (indexCollision == track0.collisionId() && track0.isPVContributor()) {
+        nProngsContributorsPV += 1;
+      }
+      if (indexCollision == track1.collisionId() && track1.isPVContributor()) {
+        nProngsContributorsPV += 1;
+      }
+
       // fill candidate table rows
-      rowCandidateBase(collision.globalIndex(),
+      rowCandidateBase(indexCollision,
                        primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ(),
                        secondaryVertex[0], secondaryVertex[1], secondaryVertex[2],
                        errorDecayLength, errorDecayLengthXY,
@@ -222,7 +231,7 @@ struct HfCandidateCreator2Prong {
                        pvec1[0], pvec1[1], pvec1[2],
                        impactParameter0.getY(), impactParameter1.getY(),
                        std::sqrt(impactParameter0.getSigmaY2()), std::sqrt(impactParameter1.getSigmaY2()),
-                       rowTrackIndexProng2.prong0Id(), rowTrackIndexProng2.prong1Id(),
+                       rowTrackIndexProng2.prong0Id(), rowTrackIndexProng2.prong1Id(), nProngsContributorsPV,
                        rowTrackIndexProng2.hfflag());
 
       // fill histograms
@@ -336,8 +345,17 @@ struct HfCandidateCreator2Prong {
         topolChi2PerNdfD0 = kfCandD0Topol2PV.GetChi2() / kfCandD0Topol2PV.GetNDF();
       }
 
+      auto indexCollision = collision.globalIndex();
+      uint8_t nProngsContributorsPV = 0;
+      if (indexCollision == track0.collisionId() && track0.isPVContributor()) {
+        nProngsContributorsPV += 1;
+      }
+      if (indexCollision == track1.collisionId() && track1.isPVContributor()) {
+        nProngsContributorsPV += 1;
+      }
+
       // fill candidate table rows
-      rowCandidateBase(collision.globalIndex(),
+      rowCandidateBase(indexCollision,
                        KFPV.GetX(), KFPV.GetY(), KFPV.GetZ(),
                        kfCandD0.GetX(), kfCandD0.GetY(), kfCandD0.GetZ(),
                        errorDecayLength, errorDecayLengthXY,   // TODO: much different from the DCAFitterN one
@@ -346,7 +364,7 @@ struct HfCandidateCreator2Prong {
                        kfNegKaon.GetPx(), kfNegKaon.GetPy(), kfNegKaon.GetPz(),
                        impactParameter0XY, impactParameter1XY,
                        errImpactParameter0XY, errImpactParameter1XY,
-                       rowTrackIndexProng2.prong0Id(), rowTrackIndexProng2.prong1Id(),
+                       rowTrackIndexProng2.prong0Id(), rowTrackIndexProng2.prong1Id(), nProngsContributorsPV,
                        rowTrackIndexProng2.hfflag());
       rowCandidateKF(topolChi2PerNdfD0,
                      massD0, massD0bar);
@@ -361,7 +379,7 @@ struct HfCandidateCreator2Prong {
 
   void processPvRefitWithDCAFitterN(aod::Collisions const& collisions,
                                     soa::Join<aod::Hf2Prongs, aod::HfPvRefit2Prong> const& rowsTrackIndexProng2,
-                                    aod::TracksWCov const& tracks,
+                                    aod::TracksWCovExtra const& tracks,
                                     aod::BCsWithTimestamps const& bcWithTimeStamps)
   {
     runCreator2ProngWithDCAFitterN<true>(collisions, rowsTrackIndexProng2, tracks, bcWithTimeStamps);
@@ -371,7 +389,7 @@ struct HfCandidateCreator2Prong {
 
   void processNoPvRefitWithDCAFitterN(aod::Collisions const& collisions,
                                       aod::Hf2Prongs const& rowsTrackIndexProng2,
-                                      aod::TracksWCov const& tracks,
+                                      aod::TracksWCovExtra const& tracks,
                                       aod::BCsWithTimestamps const& bcWithTimeStamps)
   {
     runCreator2ProngWithDCAFitterN<false>(collisions, rowsTrackIndexProng2, tracks, bcWithTimeStamps);
@@ -381,7 +399,7 @@ struct HfCandidateCreator2Prong {
 
   void processPvRefitWithKFParticle(aod::Collisions const& collisions,
                                     soa::Join<aod::Hf2Prongs, aod::HfPvRefit2Prong> const& rowsTrackIndexProng2,
-                                    soa::Join<aod::TracksWCov, aod::TracksExtra> const& tracks,
+                                    aod::TracksWCovExtra const& tracks,
                                     aod::BCsWithTimestamps const& bcWithTimeStamps)
   {
     runCreator2ProngWithKFParticle<true>(collisions, rowsTrackIndexProng2, tracks, bcWithTimeStamps);
@@ -390,7 +408,7 @@ struct HfCandidateCreator2Prong {
 
   void processNoPvRefitWithKFParticle(aod::Collisions const& collisions,
                                       aod::Hf2Prongs const& rowsTrackIndexProng2,
-                                      soa::Join<aod::TracksWCov, aod::TracksExtra> const& tracks,
+                                      aod::TracksWCovExtra const& tracks,
                                       aod::BCsWithTimestamps const& bcWithTimeStamps)
   {
     runCreator2ProngWithKFParticle<false>(collisions, rowsTrackIndexProng2, tracks, bcWithTimeStamps);

PWGHF/TableProducer/treeCreatorLcToPKPi.cxx

@@ -32,7 +32,6 @@ namespace o2::aod
 namespace full
 {
 DECLARE_SOA_INDEX_COLUMN(Collision, collision);
-DECLARE_SOA_COLUMN(NProngsContributorsPV, nProngsContributorsPV, int8_t);
 DECLARE_SOA_COLUMN(RSecondaryVertex, rSecondaryVertex, float);
 DECLARE_SOA_COLUMN(PtProng0, ptProng0, float);
 DECLARE_SOA_COLUMN(PProng0, pProng0, float);