PWGHF: store charm cand light tables for QA studies

Author: apalasciano

PWGHF/Tasks/taskCharmHadImpactPar.cxx

@@ -24,6 +24,7 @@
 
 using namespace o2;
 using namespace o2::analysis;
+using namespace o2::constants::math;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
 
@@ -33,10 +34,40 @@ enum Channel : uint8_t {
   NChannels
 };
 
+namespace o2::aod
+{
+namespace hf_charm_cand_lite
+{
+DECLARE_SOA_COLUMN(M, m, float);                          //! Invariant mass of candidate (GeV/c2)
+DECLARE_SOA_COLUMN(Pt, pt, float);                        //! Transverse momentum of candidate (GeV/c)
+DECLARE_SOA_COLUMN(Y, y, float);                          //! Rapidity of candidate
+DECLARE_SOA_COLUMN(Eta, eta, float);                      //! Pseudorapidity of candidate
+DECLARE_SOA_COLUMN(Phi, phi, float);                      //! Azimuth angle of candidate
+DECLARE_SOA_COLUMN(MlScoreBkg, mlScoreBkg, float);        //! ML score for background class
+DECLARE_SOA_COLUMN(MlScorePrompt, mlScorePrompt, float);     //! ML Prompt score for prompt class
+DECLARE_SOA_COLUMN(MlScoreNonPrompt, mlScoreNonPrompt, float);  //! ML Non Prompt score for non prompt class
+}
+
+DECLARE_SOA_TABLE(HfCharmCandLites, "AOD", "HFCHARMCANDLITE", //! Table with some B+ properties
+                  hf_charm_cand_lite::M,
+                  hf_charm_cand_lite::Pt,
+                  hf_charm_cand_lite::Y,
+                  hf_charm_cand_lite::Eta,
+                  hf_charm_cand_lite::Phi,
+                  hf_charm_cand_lite::MlScoreBkg,
+                  hf_charm_cand_lite::MlScorePrompt,
+                  hf_charm_cand_lite::MlScoreNonPrompt);
+}
+
 struct HfTaskCharmHadImpactPar {
+  Produces<aod::HfCharmCandLites> hfCharmCandLite;
+
   Configurable<int> selectionFlag{"selectionFlag", 15, "Selection Flag for the considered charm hadron"};
+  Configurable<int> fillLightTreeCandidate{"fillLightTreeCandidate", 0, "Flag to store charm hadron features"};
   ConfigurableAxis axisPt{"axisPt", {0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 8.f, 10.f, 12.f, 16.f, 24.f, 36.f, 50.f}, "axis for pT of charm hadron"};
   ConfigurableAxis axisMass{"axisMass", {250, 1.65f, 2.15f}, "axis for mass of charm hadron"};
+  ConfigurableAxis axisPhi{"axisPhi", {100, 0.f, 2*PI}, "axis for azimuthal angle of charm hadron"};
+  ConfigurableAxis axisEta{"axisEta", {100, -2.f, 2.f}, "axis for pseudorapidity of charm hadron"};
   ConfigurableAxis axisImpPar{"axisImpPar", {2000, -500.f, 500.f}, "axis for impact-parameter of charm hadron"};
   ConfigurableAxis axisMlScore0{"axisMlScore0", {100, 0.f, 1.f}, "axis for ML output score 0"};
   ConfigurableAxis axisMlScore1{"axisMlScore1", {100, 0.f, 1.f}, "axis for ML output score 1"};
@@ -62,8 +93,10 @@ struct HfTaskCharmHadImpactPar {
     }
     if (doprocessDplus || doprocessDzero) {
       registry.add("hMassPtImpPar", ";#it{M} (GeV/#it{c}^{2});#it{p}_{T} (GeV/#it{c});dca XY (#mum);", HistType::kTHnSparseF, {axisMass, axisPt, axisImpPar});
+      registry.add("hMassPtPhiEta", ";#it{M} (GeV/#it{c}^{2});#it{p}_{T} (GeV/#it{c}); phi; eta;", HistType::kTHnSparseF, {axisMass, axisPt, axisImpPar, axisPhi, axisEta});
     } else if (doprocessDplusWithMl || doprocessDzeroWithMl) {
       registry.add("hMassPtImpPar", ";#it{M} (GeV/#it{c}^{2});#it{p}_{T} (GeV/#it{c});dca XY (#mum);ML score 0;ML score 1; ML score 2;", HistType::kTHnSparseF, {axisMass, axisPt, axisImpPar, axisMlScore0, axisMlScore1, axisMlScore2});
+      registry.add("hMassPtPhiEta", ";#it{M} (GeV/#it{c}^{2});#it{p}_{T} (GeV/#it{c}); phi; eta; ML score 0;ML score 1; ML score 2;", HistType::kTHnSparseF, {axisMass, axisPt, axisImpPar, axisPhi, axisEta, axisMlScore0, axisMlScore1, axisMlScore2});
     }
   }
 
@@ -81,8 +114,10 @@ struct HfTaskCharmHadImpactPar {
           outputMl[iScore] = candidate.mlProbDplusToPiKPi()[iScore];
         }
         registry.fill(HIST("hMassPtImpPar"), invMass, candidate.pt(), candidate.impactParameterXY(), outputMl[0], outputMl[1], outputMl[2]);
+        registry.fill(HIST("hMassPtPhiEta"), invMass, candidate.pt(), candidate.phi(), candidate.eta(), outputMl[0], outputMl[1], outputMl[2]);
       } else {
         registry.fill(HIST("hMassPtImpPar"), invMass, candidate.pt(), candidate.impactParameterXY());
+        registry.fill(HIST("hMassPtPhiEta"), invMass, candidate.pt(), candidate.phi(), candidate.eta());
       }
     } else if constexpr (channel == Channel::DzeroToKPi) {
       if (candidate.isSelD0()) { // D0 -> Kpi
@@ -92,8 +127,10 @@ struct HfTaskCharmHadImpactPar {
             outputMl[iScore] = candidate.mlProbD0()[iScore];
           }
           registry.fill(HIST("hMassPtImpPar"), invMass, candidate.pt(), candidate.impactParameterXY(), outputMl[0], outputMl[1], outputMl[2]);
+          registry.fill(HIST("hMassPtPhiEta"), invMass, candidate.pt(), candidate.phi(), candidate.eta(), outputMl[0], outputMl[1], outputMl[2]);
         } else {
           registry.fill(HIST("hMassPtImpPar"), invMass, candidate.pt(), candidate.impactParameterXY());
+          registry.fill(HIST("hMassPtPhiEta"), invMass, candidate.pt(), candidate.phi(), candidate.eta());
         }
       }
       if (candidate.isSelD0bar()) {
@@ -103,19 +140,72 @@ struct HfTaskCharmHadImpactPar {
             outputMl[iScore] = candidate.mlProbD0bar()[iScore];
           }
           registry.fill(HIST("hMassPtImpPar"), invMass, candidate.pt(), candidate.impactParameterXY(), outputMl[0], outputMl[1], outputMl[2]);
+          registry.fill(HIST("hMassPtPhiEta"), invMass, candidate.pt(), candidate.phi(), candidate.eta(), outputMl[0], outputMl[1], outputMl[2]);
         } else {
           registry.fill(HIST("hMassPtImpPar"), invMass, candidate.pt(), candidate.impactParameterXY());
+          registry.fill(HIST("hMassPtPhiEta"), invMass, candidate.pt(), candidate.phi(), candidate.eta(), outputMl[0], outputMl[1], outputMl[2]);
         }
       }
     }
   }
 
+  // Fill THnSparses for the ML analysis
+  /// \param candidate is a particle candidate
+  template <Channel channel, bool withMl, typename CCands>
+  void fillTree(const CCands& candidate)
+  {
+    std::vector<float> outputMl = {-999., -999., -999.};
+    float invMass{-1.f};
+    float yCand{-999.f};
+    if constexpr (channel == Channel::DplusToKPiPi) { // D+ -> Kpipi
+      invMass = hfHelper.invMassDplusToPiKPi(candidate);
+      yCand = hfHelper.yDplus(candidate);
+      if constexpr (withMl) {
+        for (auto iScore{0u}; iScore < candidate.mlProbDplusToPiKPi().size(); ++iScore) {
+          outputMl[iScore] = candidate.mlProbDplusToPiKPi()[iScore];
+        }
+      }
+    } else if constexpr (channel == Channel::DzeroToKPi) {
+      if (candidate.isSelD0()) { // D0 -> Kpi
+        invMass = hfHelper.invMassD0ToPiK(candidate);
+        yCand = hfHelper.yD0(candidate);
+        if constexpr (withMl) {
+          for (auto iScore{0u}; iScore < candidate.mlProbD0().size(); ++iScore) {
+            outputMl[iScore] = candidate.mlProbD0()[iScore];
+          }
+        }
+      }
+      if (candidate.isSelD0bar()) {
+        invMass = hfHelper.invMassD0barToKPi(candidate);
+        yCand = hfHelper.yD0(candidate);
+        if constexpr (withMl) {
+          for (auto iScore{0u}; iScore < candidate.mlProbD0bar().size(); ++iScore) {
+            outputMl[iScore] = candidate.mlProbD0bar()[iScore];
+          }
+        }
+      }
+    }
+    hfCharmCandLite(
+          // Charm candidate meson features
+          invMass,
+          candidate.pt(),
+          yCand,
+          candidate.eta(),
+          candidate.phi(),
+          outputMl[0],
+          outputMl[1],
+          outputMl[2]);
+  }
+
   /// \param candidates are reconstructed candidates
   template <Channel channel, bool withMl, typename CCands>
   void runAnalysis(const CCands& candidates)
   {
     for (auto const& candidate : candidates) {
       fillSparse<channel, withMl>(candidate);
+      if(fillLightTreeCandidate){
+        fillTree<channel, withMl>(candidate);
+      }
     }
   }