Add mini trees to Ds task

Author: fchinu

PWGHF/D2H/Tasks/taskDs.cxx

@@ -93,6 +93,39 @@ enum ResonantChannel : int8_t {
   Kstar0K = 2
 };
 
+namespace o2::aod
+{
+namespace hf_cand_ds_mini
+{
+DECLARE_SOA_COLUMN(Pt, pt, float);                                           //! Transverse momentum of D-meson candidate (GeV/c)
+DECLARE_SOA_COLUMN(M, m, float);                                             //! Invariant mass of D-meson candidates (GeV/c)
+DECLARE_SOA_COLUMN(Centrality, centrality, float);                           //! Centrality of collision
+DECLARE_SOA_COLUMN(ImpactParameter, impactParameter, float);                 //! Impact parameter of D-meson candidate
+DECLARE_SOA_COLUMN(DecayLength, decayLength, float);                         //! Decay length of D-meson candidate
+DECLARE_SOA_COLUMN(DecayLengthXY, decayLengthXY, float);                     //! Transverse decay length of D-meson candidate
+DECLARE_SOA_COLUMN(DecayLengthNormalised, decayLengthNormalised, float);     //! Normalised decay length of D-meson candidate
+DECLARE_SOA_COLUMN(DecayLengthXYNormalised, decayLengthXYNormalised, float); //! Normalised transverse decay length of D-meson candidate
+DECLARE_SOA_COLUMN(FlagMc, flagMc, int);                                     //! MC flag (according to DataType enum)
+} // namespace hf_cand_ds_mini
+
+DECLARE_SOA_TABLE(HfCandDsMinis, "AOD", "HFCANDDSMINI", //! Table with few Ds properties
+                  hf_cand_ds_mini::M,
+                  hf_cand_ds_mini::Pt,
+                  hf_cand_ds_mini::Centrality);
+
+DECLARE_SOA_TABLE(HfCandDsDlMinis, "AOD", "HFCANDDSDLMINI", //! Table with decay length Ds properties
+                  hf_cand_ds_mini::DecayLength,
+                  hf_cand_ds_mini::DecayLengthXY,
+                  hf_cand_ds_mini::DecayLengthNormalised,
+                  hf_cand_ds_mini::DecayLengthXYNormalised);
+
+DECLARE_SOA_TABLE(HfCandDsD0Minis, "AOD", "HFCANDDSD0MINI", //! Table with impact parameter (d0)
+                  hf_cand_ds_mini::ImpactParameter);
+
+DECLARE_SOA_TABLE(HfCandDsMcMinis, "AOD", "HFCANDDSMCMINI", //! Table with MC decay type check
+                  hf_cand_ds_mini::FlagMc);
+} // namespace o2::aod
+
 static std::unordered_map<int8_t, std::unordered_map<int8_t, int8_t>> channelsResonant = {{{Mother::Ds, {{ResonantChannel::PhiPi, hf_decay::hf_cand_3prong::DecayChannelResonant::DsToPhiPi}, {ResonantChannel::Kstar0K, hf_decay::hf_cand_3prong::DecayChannelResonant::DsToKstar0K}}},
                                                                                            {Mother::Dplus, {{ResonantChannel::PhiPi, hf_decay::hf_cand_3prong::DecayChannelResonant::DplusToPhiPi}, {ResonantChannel::Kstar0K, hf_decay::hf_cand_3prong::DecayChannelResonant::DplusToKstar0K}}}}};
 
@@ -104,6 +137,10 @@ concept HasDsMlInfo = requires(T candidate) {
 
 /// Ds± analysis task
 struct HfTaskDs {
+  Produces<aod::HfCandDsMinis> hfCandDsMinis;
+  Produces<aod::HfCandDsDlMinis> hfCandDsDlMinis;
+  Produces<aod::HfCandDsD0Minis> hfCandDsD0Minis;
+  Produces<aod::HfCandDsMcMinis> hfCandDsMcMinis;
 
   Configurable<int> decayChannel{"decayChannel", 1, "Switch between resonant decay channels: 1 for Ds/Dplus->PhiPi->KKpi, 2 for Ds/Dplus->K0*K->KKPi"};
   Configurable<bool> fillDplusMc{"fillDplusMc", true, "Switch to fill Dplus MC information"};
@@ -119,6 +156,11 @@ struct HfTaskDs {
   Configurable<bool> storeOccupancy{"storeOccupancy", false, "Flag to store occupancy information"};
   Configurable<int> occEstimator{"occEstimator", 0, "Occupancy estimation (None: 0, ITS: 1, FT0C: 2)"};
   Configurable<bool> fillMcBkgHistos{"fillMcBkgHistos", false, "Flag to fill and store histograms for MC background"};
+  struct : ConfigurableGroup {
+    Configurable<bool> produceMiniTrees{"produceMiniTrees", false, "Flag to produce mini trees"};
+    Configurable<bool> extendWithDecayLength{"extendWithDecayLength", false, "Flag to extend trees with decay length information"};
+    Configurable<bool> extendWithImpactParameter{"extendWithImpactParameter", false, "Flag to extend trees with impact parameter information"};
+  } miniTrees;
 
   struct : ConfigurableGroup {
     Configurable<std::string> ccdburl{"ccdburl", "http://alice-ccdb.cern.ch", "The CCDB endpoint url address"};
@@ -542,6 +584,24 @@ struct HfTaskDs {
     }
   }
 
+  template <bool IsMc, typename Coll, typename Cand>
+  void fillMiniTrees(const Cand& candidate, DataType dataType, FinalState finalState)
+  {
+    auto mass = finalState == FinalState::KKPi ? HfHelper::invMassDsToKKPi(candidate) : HfHelper::invMassDsToPiKK(candidate);
+    auto pt = candidate.pt();
+
+    hfCandDsMinis(mass, pt, evaluateCentralityCand<Coll>(candidate));
+    if (miniTrees.extendWithDecayLength) {
+      hfCandDsDlMinis(candidate.decayLength(), candidate.decayLengthXY(), candidate.decayLengthNormalised(), candidate.decayLengthXYNormalised());
+    }
+    if (miniTrees.extendWithImpactParameter) {
+      hfCandDsD0Minis(candidate.impactParameterXY());
+    }
+    if constexpr (IsMc) {
+      hfCandDsMcMinis(dataType);
+    }
+  }
+
   /// Fill histograms of quantities for the KKPi daugther-mass hypothesis
   /// \param candidate is candidate
   /// \param dataType is data class, as defined in DataType enum
@@ -602,7 +662,9 @@ struct HfTaskDs {
         }
         fillHisto(candidate, dataType);
         fillHistoKKPi<true, Coll>(candidate, dataType);
-
+        if (miniTrees.produceMiniTrees) {
+          fillMiniTrees<true, Coll>(candidate, dataType, FinalState::KKPi);
+        }
         if (TESTBIT(candidate.isSelDsToKKPi(), aod::SelectionStep::RecoSkims)) {
           std::get<TH2Ptr>(histosPtr[dataType]["hPtVsYRecoSkim"])->Fill(pt, yCand);
         }
@@ -620,6 +682,9 @@ struct HfTaskDs {
         }
         fillHisto(candidate, dataType);
         fillHistoPiKK<true, Coll>(candidate, dataType);
+        if (miniTrees.produceMiniTrees) {
+          fillMiniTrees<true, Coll>(candidate, dataType, FinalState::PiKK);
+        }
 
         if (TESTBIT(candidate.isSelDsToPiKK(), aod::SelectionStep::RecoSkims)) {
           std::get<TH2Ptr>(histosPtr[dataType]["hPtVsYRecoSkim"])->Fill(pt, yCand);
@@ -643,13 +708,19 @@ struct HfTaskDs {
       }
       fillHisto(candidate, DataType::Data);
       fillHistoKKPi<false, Coll>(candidate, DataType::Data);
+      if (miniTrees.produceMiniTrees) {
+        fillMiniTrees<true, Coll>(candidate, DataType::Data, FinalState::KKPi);
+      }
     }
     if (candidate.isSelDsToPiKK() >= selectionFlagDs) { // PiKK
       if (yCandRecoMax >= 0. && std::abs(candidate.y(HfHelper::invMassDsToPiKK(candidate))) > yCandRecoMax) {
         return;
       }
       fillHisto(candidate, DataType::Data);
       fillHistoPiKK<false, Coll>(candidate, DataType::Data);
+      if (miniTrees.produceMiniTrees) {
+        fillMiniTrees<true, Coll>(candidate, DataType::Data, FinalState::PiKK);
+      }
     }
   }