PWGHF: Handle ambiguous tracks in B0 candidate creator and get bz from CCDB (#2280)

* Change structure: loop over aod::Collisions, over D candidates attached to this collision and over tracks attached to this collision

* Repropagate 3prong daughters to this collision if needed + Apply fixes

* PR comments

* Modify initialization of dca for D daughters

* Get bz from CCDB

---------

Co-authored-by: Alexandre Bigot <abigot@sbgat402.in2p3.fr>

Author: AlexBigO

PWGHF/TableProducer/candidateCreatorB0.cxx

@@ -19,10 +19,12 @@
 #include "Framework/AnalysisTask.h"
 #include "DCAFitter/DCAFitterN.h"
 #include "Common/Core/trackUtilities.h"
+#include "Common/DataModel/CollisionAssociation.h"
 #include "ReconstructionDataFormats/DCA.h"
 #include "ReconstructionDataFormats/V0.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
+#include "PWGHF/Utils/utilsBfieldCCDB.h"
 
 using namespace o2;
 using namespace o2::aod;
@@ -51,7 +53,7 @@ struct HfCandidateCreatorB0 {
   Produces<aod::HfCandB0Config> rowCandidateConfig;
 
   // vertexing
-  Configurable<double> bz{"bz", 5., "magnetic field"};
+  // Configurable<double> bz{"bz", 5., "magnetic field"};
   Configurable<bool> propagateToPCA{"propagateToPCA", true, "create tracks version propagated to PCA"};
   Configurable<bool> useAbsDCA{"useAbsDCA", false, "Minimise abs. distance rather than chi2"};
   Configurable<bool> useWeightedFinalPCA{"useWeightedFinalPCA", false, "Recalculate vertex position using track covariances, effective only if useAbsDCA is true"};
@@ -68,16 +70,32 @@ struct HfCandidateCreatorB0 {
   Configurable<int> selectionFlagD{"selectionFlagD", 1, "Selection Flag for D"};
   // FIXME: store B0 creator configurable (until https://alice.its.cern.ch/jira/browse/O2-3582 solved)
   bool isHfCandB0ConfigFilled = false;
+  // magnetic field setting from CCDB
+  Configurable<bool> isRun2{"isRun2", false, "enable Run 2 or Run 3 GRP objects for magnetic field"};
+  Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
+  Configurable<std::string> ccdbPathLut{"ccdbPathLut", "GLO/Param/MatLUT", "Path for LUT parametrization"};
+  Configurable<std::string> ccdbPathGeo{"ccdbPathGeo", "GLO/Config/GeometryAligned", "Path of the geometry file"};
+  Configurable<std::string> ccdbPathGrp{"ccdbPathGrp", "GLO/GRP/GRP", "Path of the grp file (Run 2)"};
+  Configurable<std::string> ccdbPathGrpMag{"ccdbPathGrpMag", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object (Run 3)"};
+
+  Service<o2::ccdb::BasicCCDBManager> ccdb;
+  o2::base::MatLayerCylSet* lut;
+  o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT;
+  int runNumber;
 
   double massPi = RecoDecay::getMassPDG(kPiPlus);
   double massD = RecoDecay::getMassPDG(pdg::Code::kDMinus);
   double massB0 = RecoDecay::getMassPDG(pdg::Code::kB0);
   double massDPi{0.};
+  double bz{0.};
 
   using TracksWithSel = soa::Join<aod::BigTracksExtended, aod::TrackSelection>;
 
-  Filter filterSelectTracks = (!usePionIsGlobalTrackWoDCA) || ((usePionIsGlobalTrackWoDCA) && requireGlobalTrackWoDCAInFilter());
   Filter filterSelectCandidates = (aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagD);
+  using CandsDFiltered = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi>>;
+  Preslice<CandsDFiltered> candsDPerCollision = aod::track_association::collisionId;
+
+  Preslice<aod::TrackAssoc> trackIndicesPerCollision = aod::track_association::collisionId;
 
   OutputObj<TH1F> hMassDToPiKPi{TH1F("hMassDToPiKPi", "D^{#minus} candidates;inv. mass (p^{#minus} K^{#plus} #pi^{#minus}) (GeV/#it{c}^{2});entries", 500, 0., 5.)};
   OutputObj<TH1F> hPtD{TH1F("hPtD", "D^{#minus} candidates;D^{#minus} candidate #it{p}_{T} (GeV/#it{c});entries", 100, 0., 10.)};
@@ -87,6 +105,18 @@ struct HfCandidateCreatorB0 {
   OutputObj<TH1F> hCovPVXX{TH1F("hCovPVXX", "2-prong candidates;XX element of cov. matrix of prim. vtx. position (cm^{2});entries", 100, 0., 1.e-4)};
   OutputObj<TH1F> hCovSVXX{TH1F("hCovSVXX", "2-prong candidates;XX element of cov. matrix of sec. vtx. position (cm^{2});entries", 100, 0., 0.2)};
 
+  void init(InitContext const&)
+  {
+    ccdb->setURL(ccdbUrl);
+    ccdb->setCaching(true);
+    ccdb->setLocalObjectValidityChecking();
+    lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->get<o2::base::MatLayerCylSet>(ccdbPathLut));
+    if (!o2::base::GeometryManager::isGeometryLoaded()) {
+      ccdb->get<TGeoManager>(ccdbPathGeo);
+    }
+    runNumber = 0;
+  }
+
   /// Single-track cuts for pions on dcaXY
   /// \param track is a track
   /// \return true if track passes all cuts
@@ -107,10 +137,11 @@ struct HfCandidateCreatorB0 {
     return true;
   }
 
-  void process(aod::Collision const& collision,
-               soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi>> const& candsD,
+  void process(aod::Collisions const& collisions,
+               CandsDFiltered const& candsD,
+               aod::TrackAssoc const& trackIndices,
                TracksWithSel const&,
-               soa::Filtered<TracksWithSel> const& tracksPion)
+               aod::BCsWithTimestamps const&)
   {
     // FIXME: store B0 creator configurable (until https://alice.its.cern.ch/jira/browse/O2-3582 solved)
     if (!isHfCandB0ConfigFilled) {
@@ -120,7 +151,7 @@ struct HfCandidateCreatorB0 {
     }
     // Initialise fitter for B vertex (2-prong vertex filter)
     o2::vertexing::DCAFitterN<2> df2;
-    df2.setBz(bz);
+    // df2.setBz(bz);
     df2.setPropagateToPCA(propagateToPCA);
     df2.setMaxR(maxR);
     df2.setMaxDZIni(maxDZIni);
@@ -131,7 +162,7 @@ struct HfCandidateCreatorB0 {
 
     // Initial fitter to redo D-vertex to get extrapolated daughter tracks (3-prong vertex filter)
     o2::vertexing::DCAFitterN<3> df3;
-    df3.setBz(bz);
+    // df3.setBz(bz);
     df3.setPropagateToPCA(propagateToPCA);
     df3.setMaxR(maxR);
     df3.setMaxDZIni(maxDZIni);
@@ -140,119 +171,181 @@ struct HfCandidateCreatorB0 {
     df3.setUseAbsDCA(useAbsDCA);
     df3.setWeightedFinalPCA(useWeightedFinalPCA);
 
-    // loop over D candidates
-    for (const auto& candD : candsD) {
-      hMassDToPiKPi->Fill(invMassDplusToPiKPi(candD), candD.pt());
-      hPtD->Fill(candD.pt());
-      hCPAD->Fill(candD.cpa());
-
-      // track0 <-> pi, track1 <-> K, track2 <-> pi
-      auto track0 = candD.prong0_as<TracksWithSel>();
-      auto track1 = candD.prong1_as<TracksWithSel>();
-      auto track2 = candD.prong2_as<TracksWithSel>();
-      auto trackParCov0 = getTrackParCov(track0);
-      auto trackParCov1 = getTrackParCov(track1);
-      auto trackParCov2 = getTrackParCov(track2);
-
-      // reconstruct 3-prong secondary vertex (D±)
-      if (df3.process(trackParCov0, trackParCov1, trackParCov2) == 0) {
-        continue;
-      }
+    static int ncol = 0;
 
-      const auto& secondaryVertex = df3.getPCACandidate();
-      trackParCov0.propagateTo(secondaryVertex[0], bz);
-      trackParCov1.propagateTo(secondaryVertex[0], bz);
-      trackParCov2.propagateTo(secondaryVertex[0], bz);
-
-      // D∓ → π∓ K± π∓
-      array<float, 3> pVecpiK = {track0.px() + track1.px(), track0.py() + track1.py(), track0.pz() + track1.pz()};
-      array<float, 3> pVecD = {pVecpiK[0] + track2.px(), pVecpiK[1] + track2.py(), pVecpiK[2] + track2.pz()};
-      auto trackParCovPiK = o2::dataformats::V0(df3.getPCACandidatePos(), pVecpiK, df3.calcPCACovMatrixFlat(),
-                                                trackParCov0, trackParCov1, {0, 0}, {0, 0});
-      auto trackParCovD = o2::dataformats::V0(df3.getPCACandidatePos(), pVecD, df3.calcPCACovMatrixFlat(),
-                                              trackParCovPiK, trackParCov2, {0, 0}, {0, 0});
-
-      int index0D = track0.globalIndex();
-      int index1D = track1.globalIndex();
-      int index2D = track2.globalIndex();
-
-      // loop over pions
-      for (const auto& trackPion : tracksPion) {
-        // minimum pT selection
-        if (trackPion.pt() < ptPionMin || !isSelectedTrackDCA(trackPion)) {
-          continue;
+    for (const auto& collision : collisions) {
+      auto primaryVertex = getPrimaryVertex(collision);
+      auto covMatrixPV = primaryVertex.getCov();
+
+      if (ncol % 10000 == 0) {
+        LOG(debug) << ncol << " collisions parsed";
+      }
+      ncol++;
+
+      /// Set the magnetic field from ccdb.
+      /// The static instance of the propagator was already modified in the HFTrackIndexSkimCreator,
+      /// but this is not true when running on Run2 data/MC already converted into AO2Ds.
+      auto bc = collision.bc_as<aod::BCsWithTimestamps>();
+      if (runNumber != bc.runNumber()) {
+        LOG(info) << ">>>>>>>>>>>> Current run number: " << runNumber;
+        initCCDB(bc, runNumber, ccdb, isRun2 ? ccdbPathGrp : ccdbPathGrpMag, lut, isRun2);
+        bz = o2::base::Propagator::Instance()->getNominalBz();
+        LOG(info) << ">>>>>>>>>>>> Magnetic field: " << bz;
+      }
+      df2.setBz(bz);
+      df3.setBz(bz);
+
+      auto thisCollId = collision.globalIndex();
+      auto candsDThisColl = candsD.sliceBy(candsDPerCollision, thisCollId);
+
+      for (const auto& candD : candsDThisColl) { // start loop over filtered D candidates indices as associated to this collision in candidateCreator3Prong.cxx
+        hMassDToPiKPi->Fill(invMassDplusToPiKPi(candD), candD.pt());
+        hPtD->Fill(candD.pt());
+        hCPAD->Fill(candD.cpa());
+
+        // track0 <-> pi, track1 <-> K, track2 <-> pi
+        auto track0 = candD.prong0_as<TracksWithSel>();
+        auto track1 = candD.prong1_as<TracksWithSel>();
+        auto track2 = candD.prong2_as<TracksWithSel>();
+        auto trackParCov0 = getTrackParCov(track0);
+        auto trackParCov1 = getTrackParCov(track1);
+        auto trackParCov2 = getTrackParCov(track2);
+
+        std::array<float, 3> pVec0 = {track0.px(), track0.py(), track0.pz()};
+        std::array<float, 3> pVec1 = {track1.px(), track1.py(), track1.pz()};
+        std::array<float, 3> pVec2 = {track2.px(), track2.py(), track2.pz()};
+
+        auto dca0 = o2::dataformats::DCA(track0.dcaXY(), track0.dcaZ(), track0.cYY(), track0.cZY(), track0.cZZ());
+        auto dca1 = o2::dataformats::DCA(track1.dcaXY(), track1.dcaZ(), track1.cYY(), track1.cZY(), track1.cZZ());
+        auto dca2 = o2::dataformats::DCA(track2.dcaXY(), track2.dcaZ(), track2.cYY(), track2.cZY(), track2.cZZ());
+
+        // repropagate tracks to this collision if needed
+        if (track0.collisionId() != thisCollId) {
+          trackParCov0.propagateToDCA(primaryVertex, bz, &dca0);
         }
-        // reject pions that are D daughters
-        if (trackPion.globalIndex() == index0D || trackPion.globalIndex() == index1D || trackPion.globalIndex() == index2D) {
-          continue;
+
+        if (track1.collisionId() != thisCollId) {
+          trackParCov1.propagateToDCA(primaryVertex, bz, &dca1);
         }
-        // reject pi and D with same sign
-        if (trackPion.sign() * track0.sign() > 0) {
-          continue;
+
+        if (track2.collisionId() != thisCollId) {
+          trackParCov2.propagateToDCA(primaryVertex, bz, &dca2);
         }
 
-        hPtPion->Fill(trackPion.pt());
-        array<float, 3> pVecPion = {trackPion.px(), trackPion.py(), trackPion.pz()};
-        auto trackParCovPi = getTrackParCov(trackPion);
         // ---------------------------------
-        // reconstruct the 2-prong B0 vertex
-        if (df2.process(trackParCovD, trackParCovPi) == 0) {
+        // reconstruct 3-prong secondary vertex (D±)
+        if (df3.process(trackParCov0, trackParCov1, trackParCov2) == 0) {
           continue;
         }
 
-        // calculate invariant mass
-        massDPi = RecoDecay::m(array{pVecD, pVecPion}, array{massD, massPi});
-        if (std::abs(massDPi - massB0) > invMassWindowB0) {
-          continue;
-        }
-        hMassB0ToDPi->Fill(massDPi);
-
-        // calculate relevant properties
-        const auto& secondaryVertexB0 = df2.getPCACandidate();
-        auto chi2PCA = df2.getChi2AtPCACandidate();
-        auto covMatrixPCA = df2.calcPCACovMatrixFlat();
-
-        // must be called before getTrack query
-        df2.propagateTracksToVertex();
-        // track.getPxPyPzGlo(pVec) modifies pVec of track
-        df2.getTrack(0).getPxPyPzGlo(pVecD);
-        df2.getTrack(1).getPxPyPzGlo(pVecPion);
-
-        auto primaryVertex = getPrimaryVertex(collision);
-        auto covMatrixPV = primaryVertex.getCov();
-        o2::dataformats::DCA impactParameter0;
-        o2::dataformats::DCA impactParameter1;
-        trackParCovD.propagateToDCA(primaryVertex, bz, &impactParameter0);
-        trackParCovPi.propagateToDCA(primaryVertex, bz, &impactParameter1);
-
-        hCovSVXX->Fill(covMatrixPCA[0]);
-        hCovPVXX->Fill(covMatrixPV[0]);
-
-        // get uncertainty of the decay length
-        double phi, theta;
-        // getPointDirection modifies phi and theta
-        getPointDirection(array{collision.posX(), collision.posY(), collision.posZ()}, secondaryVertexB0, phi, theta);
-        auto errorDecayLength = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, theta) + getRotatedCovMatrixXX(covMatrixPCA, phi, theta));
-        auto errorDecayLengthXY = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, 0.) + getRotatedCovMatrixXX(covMatrixPCA, phi, 0.));
-
-        int hfFlag = BIT(hf_cand_b0::DecayType::B0ToDPi);
-
-        // fill the candidate table for the B0 here:
-        rowCandidateBase(collision.globalIndex(),
-                         collision.posX(), collision.posY(), collision.posZ(),
-                         secondaryVertexB0[0], secondaryVertexB0[1], secondaryVertexB0[2],
-                         errorDecayLength, errorDecayLengthXY,
-                         chi2PCA,
-                         pVecD[0], pVecD[1], pVecD[2],
-                         pVecPion[0], pVecPion[1], pVecPion[2],
-                         impactParameter0.getY(), impactParameter1.getY(),
-                         std::sqrt(impactParameter0.getSigmaY2()), std::sqrt(impactParameter1.getSigmaY2()),
-                         candD.globalIndex(), trackPion.globalIndex(),
-                         hfFlag);
-      } // pi loop
-    }   // D loop
-  }     // process
-};      // struct
+        const auto& secondaryVertexD = df3.getPCACandidate();
+        // propagate the 3 prongs to the secondary vertex
+        trackParCov0.propagateTo(secondaryVertexD[0], bz);
+        trackParCov1.propagateTo(secondaryVertexD[0], bz);
+        trackParCov2.propagateTo(secondaryVertexD[0], bz);
+
+        // update pVec of tracks
+        df3.getTrack(0).getPxPyPzGlo(pVec0);
+        df3.getTrack(1).getPxPyPzGlo(pVec1);
+        df3.getTrack(2).getPxPyPzGlo(pVec2);
+
+        // D∓ → π∓ K± π∓
+        array<float, 3> pVecPiK = RecoDecay::pVec(pVec0, pVec1);
+        array<float, 3> pVecD = RecoDecay::pVec(pVec0, pVec1, pVec2);
+        auto trackParCovPiK = o2::dataformats::V0(df3.getPCACandidatePos(), pVecPiK, df3.calcPCACovMatrixFlat(),
+                                                  trackParCov0, trackParCov1, {0, 0}, {0, 0});
+        auto trackParCovD = o2::dataformats::V0(df3.getPCACandidatePos(), pVecD, df3.calcPCACovMatrixFlat(),
+                                                trackParCovPiK, trackParCov2, {0, 0}, {0, 0});
+
+        int indexTrack0 = track0.globalIndex();
+        int indexTrack1 = track1.globalIndex();
+        int indexTrack2 = track2.globalIndex();
+
+        auto trackIdsThisCollision = trackIndices.sliceBy(trackIndicesPerCollision, thisCollId);
+
+        for (const auto& trackId : trackIdsThisCollision) { // start loop over track indices associated to this collision
+          auto trackPion = trackId.track_as<TracksWithSel>();
+
+          // check isGlobalTrackWoDCA status for pions if wanted
+          if (usePionIsGlobalTrackWoDCA && !trackPion.isGlobalTrackWoDCA()) {
+            continue;
+          }
+
+          // minimum pT selection
+          if (trackPion.pt() < ptPionMin || !isSelectedTrackDCA(trackPion)) {
+            continue;
+          }
+          // reject pions that are D daughters
+          if (trackPion.globalIndex() == indexTrack0 || trackPion.globalIndex() == indexTrack1 || trackPion.globalIndex() == indexTrack2) {
+            continue;
+          }
+          // reject pi and D with same sign
+          if (trackPion.sign() * track0.sign() > 0) {
+            continue;
+          }
+
+          hPtPion->Fill(trackPion.pt());
+          array<float, 3> pVecPion = {trackPion.px(), trackPion.py(), trackPion.pz()};
+          auto trackParCovPi = getTrackParCov(trackPion);
+
+          // ---------------------------------
+          // reconstruct the 2-prong B0 vertex
+          if (df2.process(trackParCovD, trackParCovPi) == 0) {
+            continue;
+          }
+
+          // calculate relevant properties
+          const auto& secondaryVertexB0 = df2.getPCACandidate();
+          auto chi2PCA = df2.getChi2AtPCACandidate();
+          auto covMatrixPCA = df2.calcPCACovMatrixFlat();
+          hCovSVXX->Fill(covMatrixPCA[0]);
+          hCovPVXX->Fill(covMatrixPV[0]);
+
+          // propagate D and Pi to the B0 vertex
+          df2.propagateTracksToVertex();
+          // track.getPxPyPzGlo(pVec) modifies pVec of track
+          df2.getTrack(0).getPxPyPzGlo(pVecD);    // momentum of D at the B0 vertex
+          df2.getTrack(1).getPxPyPzGlo(pVecPion); // momentum of Pi at the B0 vertex
+
+          // calculate invariant mass and apply selection
+          massDPi = RecoDecay::m(array{pVecD, pVecPion}, array{massD, massPi});
+          if (std::abs(massDPi - massB0) > invMassWindowB0) {
+            continue;
+          }
+          hMassB0ToDPi->Fill(massDPi);
+
+          // compute impact parameters of D and Pi
+          o2::dataformats::DCA dcaD;
+          o2::dataformats::DCA dcaPion;
+          trackParCovD.propagateToDCA(primaryVertex, bz, &dcaD);
+          trackParCovPi.propagateToDCA(primaryVertex, bz, &dcaPion);
+
+          // get uncertainty of the decay length
+          double phi, theta;
+          // getPointDirection modifies phi and theta
+          getPointDirection(array{collision.posX(), collision.posY(), collision.posZ()}, secondaryVertexB0, phi, theta);
+          auto errorDecayLength = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, theta) + getRotatedCovMatrixXX(covMatrixPCA, phi, theta));
+          auto errorDecayLengthXY = std::sqrt(getRotatedCovMatrixXX(covMatrixPV, phi, 0.) + getRotatedCovMatrixXX(covMatrixPCA, phi, 0.));
+
+          int hfFlag = BIT(hf_cand_b0::DecayType::B0ToDPi);
+
+          // fill the candidate table for the B0 here:
+          rowCandidateBase(thisCollId,
+                           collision.posX(), collision.posY(), collision.posZ(),
+                           secondaryVertexB0[0], secondaryVertexB0[1], secondaryVertexB0[2],
+                           errorDecayLength, errorDecayLengthXY,
+                           chi2PCA,
+                           pVecD[0], pVecD[1], pVecD[2],
+                           pVecPion[0], pVecPion[1], pVecPion[2],
+                           dcaD.getY(), dcaPion.getY(),
+                           std::sqrt(dcaD.getSigmaY2()), std::sqrt(dcaPion.getSigmaY2()),
+                           candD.globalIndex(), trackPion.globalIndex(),
+                           hfFlag);
+        } // pi loop
+      }   // D loop
+    }     // collision loop
+  }       // process
+};        // struct
 
 /// Extends the base table with expression columns and performs MC matching.
 struct HfCandidateCreatorB0Expressions {