[PWGHF] add the v1 for D0 (#8582)

Author: zhangbiao-phy

PWGHF/D2H/Tasks/taskDirectedFlowCharmHadrons.cxx

@@ -13,6 +13,7 @@
 /// \brief Analysis task for charm hadron directed flow
 ///
 /// \author Prottay Das, prottay.das@cern.ch
+/// \author Biao Zhang, biao.zhanng@cern.ch
 
 #include <string>
 #include <vector>
@@ -40,7 +41,9 @@ using namespace o2::framework::expressions;
 using namespace o2::hf_centrality;
 using namespace o2::hf_evsel;
 
-enum DecayChannel { DplusToPiKPi = 0 };
+enum DecayChannel { DplusToPiKPi = 0,
+                    D0ToPiK,
+                    D0ToKPi };
 
 struct HfTaskDirectedFlowCharmHadrons {
   Configurable<int> centEstimator{"centEstimator", 2, "Centrality estimation (FT0A: 1, FT0C: 2, FT0M: 3, FV0A: 4)"};
@@ -63,10 +66,18 @@ struct HfTaskDirectedFlowCharmHadrons {
 
   using CandDplusDataWMl = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi, aod::HfMlDplusToPiKPi>>;
   using CandDplusData = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi>>;
+  using CandD0DataWMl = soa::Filtered<soa::Join<aod::HfCand2Prong, aod::HfSelD0, aod::HfMlD0>>;
+  using CandD0Data = soa::Filtered<soa::Join<aod::HfCand2Prong, aod::HfSelD0>>;
   using CollsWithQvecs = soa::Join<aod::Collisions, aod::EvSels, aod::CentFV0As, aod::CentFT0Ms, aod::CentFT0As, aod::CentFT0Cs, aod::SPCalibrationTables>;
   using TracksWithExtra = soa::Join<aod::Tracks, aod::TracksExtra>;
 
   Filter filterSelectDplusCandidates = aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlag;
+  Filter filterSelectD0Candidates = aod::hf_sel_candidate_d0::isSelD0 >= selectionFlag || aod::hf_sel_candidate_d0::isSelD0bar >= selectionFlag;
+
+  Partition<CandD0Data> selectedD0ToPiK = aod::hf_sel_candidate_d0::isSelD0 >= selectionFlag;
+  Partition<CandD0Data> selectedD0ToKPi = aod::hf_sel_candidate_d0::isSelD0bar >= selectionFlag;
+  Partition<CandD0DataWMl> selectedD0ToPiKWMl = aod::hf_sel_candidate_d0::isSelD0 >= selectionFlag;
+  Partition<CandD0DataWMl> selectedD0ToKPiWMl = aod::hf_sel_candidate_d0::isSelD0bar >= selectionFlag;
 
   SliceCache cache;
   HfHelper hfHelper;
@@ -80,7 +91,7 @@ struct HfTaskDirectedFlowCharmHadrons {
   {
 
     /// check process functions
-    std::array<int, 2> processes = {doprocessDplusStd, doprocessDplusMl};
+    std::array<int, 4> processes = {doprocessDplusStd, doprocessDplusMl, doprocessD0Std, doprocessD0Ml};
     const int nProcesses = std::accumulate(processes.begin(), processes.end(), 0);
     if (nProcesses > 1) {
       LOGP(fatal, "Only one process function should be enabled at a time, please check your configuration");
@@ -197,10 +208,29 @@ struct HfTaskDirectedFlowCharmHadrons {
           for (unsigned int iclass = 0; iclass < classMl->size(); iclass++)
             outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMl->at(iclass)];
         }
+      } else if constexpr (std::is_same_v<T1, CandD0Data> || std::is_same_v<T1, CandD0DataWMl>) {
+        switch (channel) {
+          case DecayChannel::D0ToPiK:
+            massCand = hfHelper.invMassD0ToPiK(candidate);
+            if constexpr (std::is_same_v<T1, CandD0DataWMl>) {
+              for (unsigned int iclass = 0; iclass < classMl->size(); iclass++)
+                outputMl[iclass] = candidate.mlProbD0()[classMl->at(iclass)];
+            }
+            break;
+          case DecayChannel::D0ToKPi:
+            massCand = hfHelper.invMassD0barToKPi(candidate);
+            if constexpr (std::is_same_v<T1, CandD0DataWMl>) {
+              for (unsigned int iclass = 0; iclass < classMl->size(); iclass++)
+                outputMl[iclass] = candidate.mlProbD0bar()[classMl->at(iclass)];
+            }
+            break;
+          default:
+            break;
+        }
       }
 
       auto trackprong0 = candidate.template prong0_as<Trk>();
-      double sign = trackprong0.sign(); // to differentiate between D+ and D-
+      double sign = trackprong0.sign(); // electric charge of the first daughter track to differentiate particle and antiparticle
 
       double ptCand = candidate.pt();
       double etaCand = candidate.eta();
@@ -246,6 +276,28 @@ struct HfTaskDirectedFlowCharmHadrons {
       }
     }
   }
+  // D0 with ML
+  void processD0Ml(CollsWithQvecs::iterator const& collision,
+                   TracksWithExtra const& tracks)
+  {
+    auto candsD0ToPiKWMl = selectedD0ToPiKWMl->sliceByCached(aod::hf_cand::collisionId, collision.globalIndex(), cache);
+    auto candsD0ToKPiWMl = selectedD0ToKPiWMl->sliceByCached(aod::hf_cand::collisionId, collision.globalIndex(), cache);
+    runFlowAnalysis<DecayChannel::D0ToPiK>(collision, candsD0ToPiKWMl, tracks);
+    runFlowAnalysis<DecayChannel::D0ToKPi>(collision, candsD0ToKPiWMl, tracks);
+  }
+  PROCESS_SWITCH(HfTaskDirectedFlowCharmHadrons, processD0Ml, "Process D0 candidates with ML", false);
+
+  // D0 with rectangular cuts
+  void processD0Std(CollsWithQvecs::iterator const& collision,
+                    TracksWithExtra const& tracks)
+  {
+    auto candsD0ToPiK = selectedD0ToPiK->sliceByCached(aod::hf_cand::collisionId, collision.globalIndex(), cache);
+    auto candsD0ToKPi = selectedD0ToKPi->sliceByCached(aod::hf_cand::collisionId, collision.globalIndex(), cache);
+    runFlowAnalysis<DecayChannel::D0ToPiK>(collision, candsD0ToPiK, tracks);
+    runFlowAnalysis<DecayChannel::D0ToKPi>(collision, candsD0ToKPi, tracks);
+  }
+  PROCESS_SWITCH(HfTaskDirectedFlowCharmHadrons, processD0Std, "Process D0 candidates with rectangular cuts", false);
+
   // Dplus with ML
   void processDplusMl(CollsWithQvecs::iterator const& collision,
                       CandDplusDataWMl const& candidatesDplus,