Feat: add first part of MC processing

Author: ariedel-cern

PWGCF/Femto/Core/cascadeBuilder.h

@@ -379,8 +379,8 @@ class CascadeBuilder
     LOG(info) << "Initialization done...";
   }
 
-  template <modes::System system, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8>
-  void fillCascades(T1 const& col, T2& collisionBuilder, T3& collisionProducts, T4& trackProducts, T5& cascadeProducts, T6 const& fullCascades, T7 const& fullTracks, T8& trackBuilder)
+  template <modes::System system, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9>
+  void fillCascades(T1 const& col, T2& collisionBuilder, T3& collisionProducts, T4& trackProducts, T5& cascadeProducts, T6 const& cascades, T7 const& tracks, T8 const& tracksWithItsPid, T9& trackBuilder)
   {
     if (!mFillAnyTable) {
       return;
@@ -389,18 +389,21 @@ class CascadeBuilder
     int64_t bachelorIndex = 0;
     int64_t posDaughterIndex = 0;
     int64_t negDaughterIndex = 0;
-    for (const auto& cascade : fullCascades) {
+    for (const auto& cascade : cascades) {
       if (!mCascadeSelection.checkFilters(cascade)) {
         continue;
       }
-      mCascadeSelection.applySelections(cascade, fullTracks, col);
+      mCascadeSelection.applySelections(cascade, tracks, col);
       if (!mCascadeSelection.passesAllRequiredSelections()) {
         continue;
       }
 
-      auto bachelor = cascade.template bachelor_as<T7>();
-      auto posDaughter = cascade.template posTrack_as<T7>();
-      auto negDaughter = cascade.template negTrack_as<T7>();
+      // cleaner, but without ITS pid: auto bachelor = cascade.template bachelor_as<T7>();
+      auto bachelor = tracksWithItsPid.iteratorAt(cascade.bachelorId() - tracksWithItsPid.offset());
+      // cleaner, but without ITS pid: auto posDaughter = cascade.template posTrack_as<T7>();
+      auto posDaughter = tracksWithItsPid.iteratorAt(cascade.posTrackId() - tracksWithItsPid.offset());
+      // cleaner, but without ITS pid: auto negDaughter = cascade.template negTrack_as<T7>();
+      auto negDaughter = tracksWithItsPid.iteratorAt(cascade.negTrackId() - tracksWithItsPid.offset());
 
       collisionBuilder.template fillCollision<system>(collisionProducts, col);
 

PWGCF/Femto/Core/cascadeHistManager.h

@@ -164,24 +164,24 @@ template <const char* cascadePrefix,
           const char* bachelorPrefix,
           const char* posDauPrefix,
           const char* negDauPrefix,
-          modes::Mode mode,
           modes::Cascade cascade>
 class CascadeHistManager
 {
  public:
   CascadeHistManager() = default;
   ~CascadeHistManager() = default;
 
+  template <modes::Mode mode>
   void init(o2::framework::HistogramRegistry* registry,
             std::map<CascadeHist, std::vector<o2::framework::AxisSpec>> const& cascadeSpecs,
             std::map<trackhistmanager::TrackHist, std::vector<o2::framework::AxisSpec>> const& BachelorSpecs,
             std::map<trackhistmanager::TrackHist, std::vector<o2::framework::AxisSpec>> const& PosDauSpecs,
             std::map<trackhistmanager::TrackHist, std::vector<o2::framework::AxisSpec>> const& NegDauSpecs)
   {
     mHistogramRegistry = registry;
-    mBachelorManager.init(registry, BachelorSpecs);
-    mPosDauManager.init(registry, PosDauSpecs);
-    mNegDauManager.init(registry, NegDauSpecs);
+    mBachelorManager.template init<mode>(registry, BachelorSpecs);
+    mPosDauManager.template init<mode>(registry, PosDauSpecs);
+    mNegDauManager.template init<mode>(registry, NegDauSpecs);
     if constexpr (modes::isFlagSet(mode, modes::Mode::kAnalysis)) {
       initAnalysis(cascadeSpecs);
     }
@@ -199,7 +199,7 @@ class CascadeHistManager
     mPlot2d = CascadeConfBinningQa.plot2d.value;
   }
 
-  template <typename T1, typename T2, typename T3, typename T4>
+  template <modes::Mode mode, typename T1, typename T2, typename T3, typename T4>
   void init(o2::framework::HistogramRegistry* registry,
             std::map<CascadeHist, std::vector<o2::framework::AxisSpec>> const& cascadeSpecs,
             T1 const& CascadeConfBinningQa,
@@ -211,22 +211,22 @@ class CascadeHistManager
             T4 const& NegDauConfBinningQa)
   {
     enableOptionalHistograms(CascadeConfBinningQa, BachelorConfBinningQa, PosDauConfBinningQa, NegDauConfBinningQa);
-    init(registry, cascadeSpecs, BachelorSpecs, PosDauSpecs, NegDauSpecs);
+    this->template init<mode>(registry, cascadeSpecs, BachelorSpecs, PosDauSpecs, NegDauSpecs);
   }
 
-  template <typename T1, typename T2>
+  template <modes::Mode mode, typename T1, typename T2>
   void fill(T1 const& cascadeCandidate, T2 const& tracks)
   {
     // this used to work, still under investigation
     // auto bachelor = cascadeCandidate.template bachelor_as<T2>();
     // auto posDaughter = cascadeCandidate.template posDau_as<T2>();
     // auto negDaughter = cascadeCandidate.template negDau_as<T2>();
     auto posDaughter = tracks.rawIteratorAt(cascadeCandidate.posDauId() - tracks.offset());
-    mPosDauManager.fill(posDaughter, tracks);
+    mPosDauManager.template fill<mode>(posDaughter, tracks);
     auto negDaughter = tracks.rawIteratorAt(cascadeCandidate.negDauId() - tracks.offset());
-    mNegDauManager.fill(negDaughter, tracks);
+    mNegDauManager.template fill<mode>(negDaughter, tracks);
     auto bachelor = tracks.rawIteratorAt(cascadeCandidate.bachelorId() - tracks.offset());
-    mBachelorManager.fill(bachelor, tracks);
+    mBachelorManager.template fill<mode>(bachelor, tracks);
 
     if constexpr (modes::isFlagSet(mode, modes::Mode::kAnalysis)) {
       fillAnalysis(cascadeCandidate);
@@ -310,9 +310,9 @@ class CascadeHistManager
 
   o2::framework::HistogramRegistry* mHistogramRegistry = nullptr;
   bool mPlot2d = true;
-  trackhistmanager::TrackHistManager<bachelorPrefix, mode> mBachelorManager;
-  trackhistmanager::TrackHistManager<posDauPrefix, mode> mPosDauManager;
-  trackhistmanager::TrackHistManager<negDauPrefix, mode> mNegDauManager;
+  trackhistmanager::TrackHistManager<bachelorPrefix> mBachelorManager;
+  trackhistmanager::TrackHistManager<posDauPrefix> mPosDauManager;
+  trackhistmanager::TrackHistManager<negDauPrefix> mNegDauManager;
 };
 }; // namespace cascadehistmanager
 }; // namespace o2::analysis::femto

PWGCF/Femto/Core/collisionBuilder.h

@@ -59,6 +59,7 @@ struct ConfCollisionFilters : o2::framework::ConfigurableGroup {
   o2::framework::Configurable<float> sphericityMax{"sphericityMax", 1.f, "Maximum sphericity"};
   o2::framework::Configurable<int> magFieldMin{"magFieldMin", -5, "Minimum magnetic field strength (kG)"};
   o2::framework::Configurable<int> magFieldMax{"magFieldMax", 5, "Maximum magnetic field strength (kG)"};
+  o2::framework::Configurable<int> subGeneratorId{"subGeneratorId", 0, "MC ONLY: If positive, keep 0 = MB, <0 triggered on something"};
 };
 
 struct ConfCollisionBits : o2::framework::ConfigurableGroup {
@@ -390,6 +391,7 @@ class CollisionBuilder
       mRctFlagsChecker.init(confRct.label.value, confRct.useZdc.value, confRct.treatLimitedAcceptanceAsBad.value);
     }
     mGrpPath = confCcdb.grpPath.value;
+    mSubGeneratorId = confFilter.subGeneratorId.value;
 
     mCollisionSelection.configure(registry, confFilter, confBits);
     mCollisionSelection.printSelections(colSelsName);
@@ -439,6 +441,28 @@ class CollisionBuilder
            mCollisionSelection.passesAllRequiredSelections();
   }
 
+  template <typename T1, typename T2>
+  bool checkCollision(T1 const& col, T2 const& /*mcCols*/)
+  {
+    // check sub generator id of associated generated collision
+    if (mSubGeneratorId >= 0) {
+      if (col.has_mcCollision()) {
+        auto mcCol = col.template mcCollision_as<T2>();
+        if (mcCol.getSubGeneratorId() != mSubGeneratorId) {
+          return false;
+        }
+      }
+    }
+
+    // check RCT flags first
+    if (mUseRctFlags && !mRctFlagsChecker(col)) {
+      return false;
+    }
+    // make other checks
+    return mCollisionSelection.checkFilters(col) &&
+           mCollisionSelection.passesAllRequiredSelections();
+  }
+
   template <modes::System system, typename T1, typename T2>
   void fillCollision(T1& collisionProducts, T2 const& col)
   {
@@ -518,6 +542,7 @@ class CollisionBuilder
   aod::rctsel::RCTFlagsChecker mRctFlagsChecker;
   bool mUseRctFlags = false;
   std::string mTriggerNames = std::string("");
+  int mSubGeneratorId = -1;
   bool mFillAnyTable = false;
   bool mProducedCollisions = false;
   bool mProducedCollisionMasks = false;

PWGCF/Femto/Core/collisionHistManager.h

@@ -51,60 +51,89 @@ enum ColHist {
   kCentVsMult,
   kCentVsSphericity,
   kMultVsSphericity,
+  // mc
+  kMcCent,
+  kMcMult,
   kColHistLast
 };
 
 constexpr std::string_view ColAnalysisDir = "Collisions/Analysis/";
 constexpr std::string_view ColQaDir = "Collisions/QA/";
+constexpr std::string_view ColMcDir = "Collisions/MC/";
 
 constexpr std::array<histmanager::HistInfo<ColHist>, kColHistLast> HistTable = {
   {
     {kPosZ, o2::framework::kTH1F, "hPosZ", "Vertex Z; V_{Z} (cm); Entries"},
     {kMult, o2::framework::kTH1F, "hMult", "Multiplicity; Multiplicity; Entries"},
     {kCent, o2::framework::kTH1F, "hCent", "Centrality; Centrality (%); Entries"},
     {kMagField, o2::framework::kTH1F, "hMagField", "Magnetic Field; B (kG); Entries"},
+    // qa
     {kPosX, o2::framework::kTH1F, "hPosX", "Vertex X; V_{X} (cm); Entries"},
-    {kPosY, o2::framework::kTH1F, "hPosY", "Vertex Z; V_{Y} (cm); Entries"},
+    {kPosY, o2::framework::kTH1F, "hPosY", "Vertex Y; V_{Y} (cm); Entries"},
     {kPos, o2::framework::kTH1F, "hPos", "Primary vertex; V_{pos} (cm); Entries"},
     {kSphericity, o2::framework::kTH1F, "hSphericity", "Sphericity; Sphericity; Entries"},
     {kOccupancy, o2::framework::kTH1F, "hOccupancy", "Occupancy; Occupancy; Entries"},
+    // 2d
     {kPoszVsMult, o2::framework::kTH2F, "hPoszVsMult", "Vertex Z vs Multiplicity; V_{Z} (cm); Multiplicity"},
     {kPoszVsCent, o2::framework::kTH2F, "hPoszVsCent", "Vertex Z vs Centrality; V_{Z} (cm); Centrality (%)"},
     {kCentVsMult, o2::framework::kTH2F, "hCentVsMult", "Centrality vs Multiplicity; Centrality (%); Multiplicity"},
     {kMultVsSphericity, o2::framework::kTH2F, "hMultVsSphericity", "Multiplicity vs Sphericity; Multiplicity; Sphericity"},
     {kCentVsSphericity, o2::framework::kTH2F, "hCentVsSphericity", "Centrality vs Sphericity; Centrality (%); Sphericity"},
+    // mc
+    {kMcCent, o2::framework::kTH1F, "hMcCent", "Monte Carlo Centrality; Centrality (%); Entries"},
+    {kMcMult, o2::framework::kTH1F, "hMcMult", "Monte Carlo Multiplicity; Multiplicity; Entries"},
   }};
 
+#define COL_HIST_ANALYSIS_MAP(conf) \
+  {kPosZ, {conf.vtxZ}},             \
+    {kMult, {conf.mult}},           \
+    {kCent, {conf.cent}},           \
+    {kMagField, {conf.magField}},
+
+#define COL_HIST_QA_MAP(confAnalysis, confQa)                    \
+  {kPosX, {confQa.vtxXY}},                                       \
+    {kPosY, {confQa.vtxXY}},                                     \
+    {kPos, {confQa.vtx}},                                        \
+    {kSphericity, {confQa.sphericity}},                          \
+    {kOccupancy, {confQa.occupancy}},                            \
+    {kPoszVsMult, {confAnalysis.vtxZ, confAnalysis.mult}},       \
+    {kPoszVsCent, {confAnalysis.vtxZ, confAnalysis.cent}},       \
+    {kCentVsMult, {confAnalysis.cent, confAnalysis.mult}},       \
+    {kMultVsSphericity, {confAnalysis.mult, confQa.sphericity}}, \
+    {kCentVsSphericity, {confBinningAnalysis.cent, confQa.sphericity}},
+
+#define COL_HIST_MC_QA_MAP(conf) \
+  {kMcMult, {conf.mult}},        \
+    {kMcCent, {conf.cent}},
+
 template <typename T>
 auto makeColHistSpecMap(const T& confBinningAnalysis)
 {
   return std::map<ColHist, std::vector<framework::AxisSpec>>{
-    {kPosZ, {confBinningAnalysis.vtxZ}},
-    {kMult, {confBinningAnalysis.mult}},
-    {kCent, {confBinningAnalysis.cent}},
-    {kMagField, {confBinningAnalysis.magField}}};
+    COL_HIST_ANALYSIS_MAP(confBinningAnalysis)};
 }
 
 template <typename T1, typename T2>
 auto makeColQaHistSpecMap(const T1& confBinningAnalysis, const T2& confBinningQa)
 {
   return std::map<ColHist, std::vector<framework::AxisSpec>>{
-    {kPosZ, {confBinningAnalysis.vtxZ}},
-    {kMult, {confBinningAnalysis.mult}},
-    {kCent, {confBinningAnalysis.cent}},
-    {kMagField, {confBinningAnalysis.magField}},
-    {kPosX, {confBinningQa.vtxXY}},
-    {kPosY, {confBinningQa.vtxXY}},
-    {kPos, {confBinningQa.vtx}},
-    {kSphericity, {confBinningQa.sphericity}},
-    {kOccupancy, {confBinningQa.occupancy}},
-    {kPoszVsMult, {confBinningAnalysis.vtxZ, confBinningAnalysis.mult}},
-    {kPoszVsCent, {confBinningAnalysis.vtxZ, confBinningAnalysis.cent}},
-    {kCentVsMult, {confBinningAnalysis.cent, confBinningAnalysis.mult}},
-    {kMultVsSphericity, {confBinningAnalysis.mult, confBinningQa.sphericity}},
-    {kCentVsSphericity, {confBinningAnalysis.cent, confBinningQa.sphericity}}};
+    COL_HIST_ANALYSIS_MAP(confBinningAnalysis)
+      COL_HIST_QA_MAP(confBinningAnalysis, confBinningQa)};
+}
+
+template <typename T1, typename T2>
+auto makeColMcQaHistSpecMap(const T1& confBinningAnalysis, const T2& confBinningQa)
+{
+  return std::map<ColHist, std::vector<framework::AxisSpec>>{
+    COL_HIST_ANALYSIS_MAP(confBinningAnalysis)
+      COL_HIST_QA_MAP(confBinningAnalysis, confBinningQa)
+        COL_HIST_MC_QA_MAP(confBinningAnalysis)};
 }
 
+#undef COL_HIST_ANALYSIS_MAP
+#undef COL_HIST_QA_MAP
+#undef COL_HIST_MC_QA_MAP
+
 struct ConfCollisionBinning : o2::framework::ConfigurableGroup {
   std::string prefix = std::string("CollisionBinning");
   o2::framework::ConfigurableAxis vtxZ{"vtxZ", {200, -10, 10}, "Vertex Z binning"};
@@ -122,14 +151,13 @@ struct ConfCollisionQaBinning : o2::framework::ConfigurableGroup {
   o2::framework::ConfigurableAxis occupancy{"occupancy", {500, 0.f, 5000.f}, "Spericity Binning"};
 };
 
-template <modes::Mode mode>
 class CollisionHistManager
 {
  public:
   CollisionHistManager() = default;
   ~CollisionHistManager() = default;
-  /// Initializes histograms for the task
-  /// \param registry Histogram registry to be passed
+
+  template <modes::Mode mode>
   void init(o2::framework::HistogramRegistry* registry, std::map<ColHist, std::vector<o2::framework::AxisSpec>> const& Specs)
   {
     mHistogramRegistry = registry;
@@ -139,6 +167,9 @@ class CollisionHistManager
     if constexpr (isFlagSet(mode, modes::Mode::kQa)) {
       initQa(Specs);
     }
+    if constexpr (isFlagSet(mode, modes::Mode::kMc)) {
+      initMc(Specs);
+    }
   }
 
   template <typename T>
@@ -147,14 +178,14 @@ class CollisionHistManager
     mPlot2d = ConfBinningQa.plot2d.value;
   }
 
-  template <typename T>
+  template <modes::Mode mode, typename T>
   void init(o2::framework::HistogramRegistry* registry, std::map<ColHist, std::vector<o2::framework::AxisSpec>> const& Specs, T const& ConfBinningQa)
   {
     enableOptionalHistograms(ConfBinningQa);
-    init(registry, Specs);
+    this->template init<mode>(registry, Specs);
   }
 
-  template <typename T>
+  template <modes::Mode mode, typename T>
   void fill(T const& col)
   {
     if constexpr (isFlagSet(mode, modes::Mode::kAnalysis)) {
@@ -165,6 +196,20 @@ class CollisionHistManager
     }
   }
 
+  template <modes::Mode mode, typename T1, typename T2>
+  void fill(T1 const& col, T2 const& mcCols)
+  {
+    if constexpr (isFlagSet(mode, modes::Mode::kAnalysis)) {
+      fillAnalysis(col);
+    }
+    if constexpr (isFlagSet(mode, modes::Mode::kQa)) {
+      fillQa(col);
+    }
+    if constexpr (isFlagSet(mode, modes::Mode::kMc)) {
+      fillMc(col, mcCols);
+    }
+  }
+
  private:
   void initAnalysis(std::map<ColHist, std::vector<o2::framework::AxisSpec>> const& Specs)
   {
@@ -192,6 +237,13 @@ class CollisionHistManager
     }
   }
 
+  void initMc(std::map<ColHist, std::vector<o2::framework::AxisSpec>> const& Specs)
+  {
+    std::string mcDir = std::string(ColMcDir);
+    mHistogramRegistry->add(mcDir + getHistNameV2(kMcMult, HistTable), getHistDesc(kMcMult, HistTable), getHistType(kMcMult, HistTable), {Specs.at(kMcMult)});
+    mHistogramRegistry->add(mcDir + getHistNameV2(kMcCent, HistTable), getHistDesc(kMcCent, HistTable), getHistType(kMcCent, HistTable), {Specs.at(kMcCent)});
+  }
+
   template <typename T>
   void fillAnalysis(T const& col)
   {
@@ -218,6 +270,14 @@ class CollisionHistManager
     }
   }
 
+  template <typename T1, typename T2>
+  void fillMc(T1 const& col, T2 const& /*mcCols*/)
+  {
+    auto genCol = col.template fMcCol_as<T2>();
+    mHistogramRegistry->fill(HIST(ColMcDir) + HIST(getHistName(kMcMult, HistTable)), genCol.multMc());
+    mHistogramRegistry->fill(HIST(ColMcDir) + HIST(getHistName(kMcCent, HistTable)), genCol.centMc());
+  }
+
   o2::framework::HistogramRegistry* mHistogramRegistry = nullptr;
   bool mPlot2d = true;
 }; // namespace femtounitedcolhistmanager

PWGCF/Femto/Core/femtoUtils.h

@@ -207,6 +207,12 @@ using HasMass = decltype(std::declval<T&>().mass());
 template <typename T>
 using HasSign = decltype(std::declval<T&>().sign());
 
+template <typename T>
+inline int signum(T x)
+{
+  return (T(0) < x) - (x < T(0)); // works for floats too
+}
+
 }; // namespace utils
 }; // namespace o2::analysis::femto
 #endif // PWGCF_FEMTO_CORE_FEMTOUTILS_H_