[PWGCF] Update femto framework (#14330)

Author: ariedel-cern

PWGCF/Femto/Core/cascadeBuilder.h

@@ -380,7 +380,7 @@ class CascadeBuilder
   }
 
   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& fullCascades, T7 const& fullTracks, T8& trackBuilder, T9& indexMap)
+  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,24 +389,27 @@ 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);
 
-      bachelorIndex = trackBuilder.template getDaughterIndex<modes::Track::kCascadeBachelor>(bachelor, trackProducts, collisionProducts, indexMap);
-      posDaughterIndex = trackBuilder.template getDaughterIndex<modes::Track::kV0Daughter>(posDaughter, trackProducts, collisionProducts, indexMap);
-      negDaughterIndex = trackBuilder.template getDaughterIndex<modes::Track::kV0Daughter>(negDaughter, trackProducts, collisionProducts, indexMap);
+      bachelorIndex = trackBuilder.template getDaughterIndex<modes::Track::kCascadeBachelor>(bachelor, trackProducts, collisionProducts);
+      posDaughterIndex = trackBuilder.template getDaughterIndex<modes::Track::kV0Daughter>(posDaughter, trackProducts, collisionProducts);
+      negDaughterIndex = trackBuilder.template getDaughterIndex<modes::Track::kV0Daughter>(negDaughter, trackProducts, collisionProducts);
 
       fillCascade(collisionProducts, cascadeProducts, cascade, col, bachelorIndex, posDaughterIndex, negDaughterIndex);
     }

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,14 +441,36 @@ 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)
   {
     if (!mFillAnyTable) {
       return;
     }
 
-    if (mCollisionAleadyFilled) {
+    if (mCollisionAlreadyFilled) {
       return;
     }
 
@@ -488,17 +512,28 @@ class CollisionBuilder
       }
     }
 
-    mCollisionAleadyFilled = true;
+    mCollisionAlreadyFilled = true;
+  }
+
+  template <modes::System system, typename T1, typename T2, typename T3, typename T4, typename T5>
+  void fillMcCollision(T1& collisionProducts, T2 const& col, T3 const& mcCols, T4& mcProducts, T5& mcBuilder)
+  {
+    if (mCollisionAlreadyFilled) {
+      return;
+    }
+    this->template fillCollision<system>(collisionProducts, col);
+    mcBuilder.template fillMcCollisionWithLabel<system>(mcProducts, col, mcCols);
   }
 
-  void reset()
+  void
+    reset()
   {
-    mCollisionAleadyFilled = false;
+    mCollisionAlreadyFilled = false;
   }
 
  private:
   CollisionSelection<HistName> mCollisionSelection;
-  bool mCollisionAleadyFilled = false;
+  bool mCollisionAlreadyFilled = false;
   Zorro mZorro;
   bool mUseTrigger = false;
   int mRunNumber = -1;
@@ -507,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;