[PWGCF] Add MC processing to derived-to-derived femto producer (#17852)

Author: ariedel-cern

PWGCF/Femto/Core/mcBuilder.h

@@ -748,6 +748,182 @@ class McBuilder
   std::unordered_map<int64_t, int64_t> mMcMotherMap;
   std::unordered_map<int64_t, int64_t> mMcPartonicMotherMap;
 };
+
+struct McBuilderDerivedToDerivedProducts : o2::framework::ProducesGroup {
+  o2::framework::Produces<o2::aod::StoredFMcCols> producedMcCollisions;
+  o2::framework::Produces<o2::aod::StoredFMcParticles> producedMcParticles;
+  o2::framework::Produces<o2::aod::StoredFMcMothers> producedMothers;
+  o2::framework::Produces<o2::aod::StoredFMcPartMoths> producedPartonicMothers;
+  o2::framework::Produces<o2::aod::StoredFMcMotherLabels> producedMcMotherLabels;
+
+  o2::framework::Produces<o2::aod::StoredFColLabels> producedCollisionLabels;
+  o2::framework::Produces<o2::aod::StoredFTrackLabels> producedTrackLabels;
+};
+
+struct ConfMcTablesDerivedToDerived : o2::framework::ConfigurableGroup {
+  std::string prefix = std::string("McTables");
+  o2::framework::Configurable<bool> requireMcLabel{"requireMcLabel", false, "Reject particles without an associated MC particle instead of writing a -1 label"}; // dummy configuration for now
+};
+
+/// Copies MC information from one femto derived file into another.
+/// The ancestry is already resolved upstream (FMcMothers / FMcPartMoths /
+/// FMcMotherLabels), so this only remaps indices: each source row is copied on
+/// first use and cached for the rest of the dataframe.
+class McBuilderDerivedToDerived
+{
+ public:
+  McBuilderDerivedToDerived() = default;
+  ~McBuilderDerivedToDerived() = default;
+
+  template <typename T>
+  void init(T& config)
+  {
+    LOG(info) << "Initialize derived-to-derived monte carlo builder...";
+    mRequireMcLabel = config.requireMcLabel.value;
+    LOG(info) << "Initialization done...";
+  }
+
+  /// Write one FColLabels row. Call exactly once per produced collision,
+  /// directly after collisionBuilder.processCollision().
+  template <typename T1, typename T2, typename T3>
+  void fillCollisionWithLabel(T1 const& col, T2 const& /*mcCols*/, T3& mcProducts)
+  {
+    if (!col.has_fMcCol()) {
+      mcProducts.producedCollisionLabels(-1);
+      return;
+    }
+    auto mcCol = col.template fMcCol_as<T2>();
+    mcProducts.producedCollisionLabels(this->getOrCreateMcCollisionRow(mcCol, mcProducts));
+  }
+
+  /// True if this track has no MC particle and we are configured to drop such tracks.
+  template <typename T>
+  bool rejectParticle(T const& particle) const
+  {
+    return mRequireMcLabel && !particle.has_fMcParticle();
+  }
+
+  /// Write one FTrackLabels row. Call exactly once per produced track row.
+  /// mcParticles must be joined with FMcMotherLabels so the mother indices resolve.
+  template <typename T1, typename T2, typename T3, typename T4, typename T5, typename T6>
+  void fillTrackWithLabel(T1 const& track, T2 const& mcCols, T3 const& /*mcParticles*/, T4 const& mcMothers, T5 const& mcPartonicMothers, T6& mcProducts)
+  {
+    if (!track.has_fMcParticle()) {
+      mcProducts.producedTrackLabels(-1);
+      return;
+    }
+    auto mcParticle = track.template fMcParticle_as<T3>();
+    mcProducts.producedTrackLabels(
+      this->getOrCreateMcParticleRow(mcParticle, mcCols, mcMothers, mcPartonicMothers, mcProducts));
+  }
+
+  template <typename T1, typename T2>
+  void reset(T1 const& mcCols, T2 const& mcParticles)
+  {
+    mCollisionMap.clear();
+    mCollisionMap.reserve(mcCols.size());
+    mMcParticleMap.clear();
+    mMcParticleMap.reserve(mcParticles.size());
+    mMcMotherMap.clear();
+    mMcMotherMap.reserve(mcParticles.size());
+    mMcPartonicMotherMap.clear();
+    mMcPartonicMotherMap.reserve(mcParticles.size());
+  }
+
+ private:
+  template <typename T1, typename T2>
+  int64_t getOrCreateMcCollisionRow(T1 const& mcCol, T2& mcProducts)
+  {
+    const int64_t gid = mcCol.globalIndex();
+    auto it = mCollisionMap.find(gid);
+    if (it != mCollisionMap.end()) {
+      return it->second;
+    }
+    mcProducts.producedMcCollisions(mcCol.posZ(),
+                                    mcCol.mult(),
+                                    mcCol.cent());
+    const int64_t row = mcProducts.producedMcCollisions.lastIndex();
+    mCollisionMap.emplace(gid, row);
+    return row;
+  }
+
+  /// Find-or-create the FMcParticles row. On first creation the mother and
+  /// partonic mother rows are copied and the single matching FMcMotherLabels row
+  /// is written, so the two tables stay in lockstep exactly as in McBuilder.
+  template <typename T1, typename T2, typename T3, typename T4, typename T5>
+  int64_t getOrCreateMcParticleRow(T1 const& mcParticle, T2 const& /*mcCols*/, T3 const& /*mcMothers*/, T4 const& /*mcPartonicMothers*/, T5& mcProducts)
+  {
+    const int64_t gid = mcParticle.globalIndex();
+    auto it = mMcParticleMap.find(gid);
+    if (it != mMcParticleMap.end()) {
+      return it->second;
+    }
+
+    // NOTE: the MC collision of the particle, not of the reconstructed collision.
+    // These differ for wrongly associated particles, which is the point of keeping it.
+    int64_t mcColId = -1;
+    if (mcParticle.has_fMcCol()) {
+      auto mcCol = mcParticle.template fMcCol_as<T2>();
+      mcColId = this->getOrCreateMcCollisionRow(mcCol, mcProducts);
+    }
+
+    mcProducts.producedMcParticles(mcColId,
+                                   mcParticle.origin(),
+                                   mcParticle.pdgCode(),
+                                   mcParticle.signedPt(),
+                                   mcParticle.eta(),
+                                   mcParticle.phi());
+    const int64_t row = mcProducts.producedMcParticles.lastIndex();
+    mMcParticleMap.emplace(gid, row);
+
+    // --- mother ---
+    int64_t mcMotherRow = -1;
+    if (mcParticle.has_fMcMother()) {
+      auto mother = mcParticle.template fMcMother_as<T3>();
+      const int64_t motherGid = mother.globalIndex();
+      auto itM = mMcMotherMap.find(motherGid);
+      if (itM != mMcMotherMap.end()) {
+        mcMotherRow = itM->second;
+      } else {
+        mcProducts.producedMothers(mother.origin(),
+                                   mother.pdgCode(),
+                                   mother.signedPt(),
+                                   mother.eta(),
+                                   mother.phi());
+        mcMotherRow = mcProducts.producedMothers.lastIndex();
+        mMcMotherMap.emplace(motherGid, mcMotherRow);
+      }
+    }
+
+    // --- partonic mother ---
+    int64_t mcPartonicMotherRow = -1;
+    if (mcParticle.has_fMcPartMoth()) {
+      auto partonicMother = mcParticle.template fMcPartMoth_as<T4>();
+      const int64_t partonicGid = partonicMother.globalIndex();
+      auto itPM = mMcPartonicMotherMap.find(partonicGid);
+      if (itPM != mMcPartonicMotherMap.end()) {
+        mcPartonicMotherRow = itPM->second;
+      } else {
+        mcProducts.producedPartonicMothers(partonicMother.pdgCode());
+        mcPartonicMotherRow = mcProducts.producedPartonicMothers.lastIndex();
+        mMcPartonicMotherMap.emplace(partonicGid, mcPartonicMotherRow);
+      }
+    }
+
+    // exactly one FMcMotherLabels row per FMcParticles row, written here and only here
+    mcProducts.producedMcMotherLabels(mcMotherRow, mcPartonicMotherRow);
+
+    return row;
+  }
+
+  bool mRequireMcLabel = false;
+
+  std::unordered_map<int64_t, int64_t> mCollisionMap;
+  std::unordered_map<int64_t, int64_t> mMcParticleMap;
+  std::unordered_map<int64_t, int64_t> mMcMotherMap;
+  std::unordered_map<int64_t, int64_t> mMcPartonicMotherMap;
+};
+
 } // namespace o2::analysis::femto::mcbuilder
 
 #endif // PWGCF_FEMTO_CORE_MCBUILDER_H_

PWGCF/Femto/Core/trackBuilder.h

@@ -134,7 +134,7 @@ struct ConfTrackSelection : public o2::framework::ConfigurableGroup {
   std::string prefix = Prefix; // Unique prefix based on the template argument
   // configuration parameters
   o2::framework::Configurable<int> pdgCodeAbs{"pdgCodeAbs", 2212, "Absolute value of PDG code. Set sign of charge to -1 for antiparticle."};
-  o2::framework::Configurable<int> chargeAbs{"chargeAbs", 1, "Absolute value of charge (e.g. 1 for most tracks, 2 for He3). Set sign of charge to -1 for antiparticle"};
+  o2::framework::Configurable<int> chargeAbs{"chargeAbs", 1, "Absolute value of charge (e.g. 1 for most tracks, 2 for He3)"};
   o2::framework::Configurable<int> chargeSign{"chargeSign", 1, "Track charge sign: +1 for positive, -1 for negative, 0 for both"};
   // filters for kinematics
   o2::framework::Configurable<float> ptMin{"ptMin", 0.0f, "Minimum pT (GeV/c)"};
@@ -920,6 +920,41 @@ class TrackBuilderDerivedToDerived
     return idx;
   }
 
+  template <typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9, typename T10, typename T11, typename T12>
+  void processTracksWithMc(T1& col, T2& /*trackTable*/, T3& partitionTrack1, T4& partitionTrack2, T5& cache,
+                           T6& newTrackTable, T7& newCollisionTable,
+                           T8& mcBuilder, T9 const& mcCols, T10 const& mcParticles, T11 const& mcMothers, T12 const& mcPartonicMothers, auto& mcProducts)
+  {
+    if (mLimitTrack1 > 0) {
+      auto trackSlice1 = partitionTrack1->sliceByCached(o2::aod::femtobase::stored::fColId, col.globalIndex(), cache);
+      for (auto const& track : trackSlice1) {
+        this->fillTrackWithMcLabel(track, newTrackTable, newCollisionTable, mcBuilder, mcCols, mcParticles, mcMothers, mcPartonicMothers, mcProducts);
+      }
+    }
+    if (mLimitTrack2 > 0) {
+      auto trackSlice2 = partitionTrack2->sliceByCached(o2::aod::femtobase::stored::fColId, col.globalIndex(), cache);
+      for (auto const& track : trackSlice2) {
+        this->fillTrackWithMcLabel(track, newTrackTable, newCollisionTable, mcBuilder, mcCols, mcParticles, mcMothers, mcPartonicMothers, mcProducts);
+      }
+    }
+  }
+
+  /// Same as fillTrack, but writes the matching FTrackLabels row. The indexMap
+  /// lookup happens first, so a track selected by both partitions produces
+  /// exactly one track row and exactly one label row.
+  template <typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9>
+  int64_t fillTrackWithMcLabel(T1 const& track, T2& trackProducts, T3& collisionProducts,
+                               T4& mcBuilder, T5 const& mcCols, T6 const& mcParticles, T7 const& mcMothers, T8 const& mcPartonicMothers, T9& mcProducts)
+  {
+    auto index = utils::getIndex(track.globalIndex(), indexMap);
+    if (index) {
+      return index.value();
+    }
+    const int64_t idx = this->fillTrack(track, trackProducts, collisionProducts);
+    mcBuilder.fillTrackWithLabel(track, mcCols, mcParticles, mcMothers, mcPartonicMothers, mcProducts);
+    return idx;
+  }
+
   template <typename T1, typename T2, typename T3>
   int64_t getDaughterIndex(const T1& daughter, T2& trackProducts, T3& collisionProducts)
   {

PWGCF/Femto/DataModel/FemtoTables.h

@@ -1295,11 +1295,13 @@ namespace femtomccollisions
 
 DECLARE_SOA_TABLE_STAGED_VERSIONED(FMcCols_001, "FMCCOL", 1, //! femto mc collisions
                                    o2::soa::Index<>,
-                                   femtocollisions::PosZ, //! Multiplicity of the event as given by the generator in |eta|<0.8
+                                   femtocollisions::PosZ,
                                    femtocollisions::Mult,
                                    femtocollisions::Cent);
 using FMcCols = FMcCols_001;
 using FMcCol = FMcCols_001::iterator;
+using StoredFMcCols = StoredFMcCols_001;
+using StoredFMcCol = StoredFMcCols_001::iterator;
 
 namespace femtomcparticle
 {
@@ -1326,6 +1328,8 @@ DECLARE_SOA_TABLE_STAGED_VERSIONED(FMcParticles_001, "FMCPARTICLE", 1, //! femto
                                    femtobase::dynamic::Theta<femtobase::stored::Eta>);
 using FMcParticles = FMcParticles_001;
 using FMcParticle = FMcParticles::iterator;
+using StoredFMcParticles = StoredFMcParticles_001;
+using StoredFMcParticle = StoredFMcParticles::iterator;
 
 DECLARE_SOA_TABLE_STAGED_VERSIONED(FMcMothers_001, "FMCMOTHER", 1, //! first direct mother of the monte carlo particle
                                    o2::soa::Index<>,               // no collision index needed since the mother is retrieved from the daughter mc particle
@@ -1344,12 +1348,16 @@ DECLARE_SOA_TABLE_STAGED_VERSIONED(FMcMothers_001, "FMCMOTHER", 1, //! first dir
 
 using FMcMothers = FMcMothers_001;
 using FMcMother = FMcMothers::iterator;
+using StoredFMcMothers = StoredFMcMothers_001;
+using StoredFMcMother = StoredFMcMothers::iterator;
 
 DECLARE_SOA_TABLE_STAGED_VERSIONED(FMcPartMoths_001, "FMCPARTMOTH", 1, //! first partonic mother of the monte carlo particle after hadronization
                                    o2::soa::Index<>,
                                    femtomcparticle::PdgCode);
 using FMcPartMoths = FMcPartMoths_001;
 using FMcPartMoth = FMcPartMoths::iterator;
+using StoredFMcPartMoths = StoredFMcPartMoths_001;
+using StoredFMcPartMoth = StoredFMcPartMoths::iterator;
 
 namespace femtolabels
 {
@@ -1360,30 +1368,87 @@ DECLARE_SOA_INDEX_COLUMN(FMcMother, fMcMother);     //!
 DECLARE_SOA_INDEX_COLUMN(FMcPartMoth, fMcPartMoth); //!
 } // namespace femtolabels
 
-DECLARE_SOA_TABLE(FColLabels, "AOD", "FCOLMCLABEL", femtolabels::FMcColId);
-
-DECLARE_SOA_TABLE(FTrackLabels, "AOD", "FTRACKLABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FLambdaLabels, "AOD", "FLAMBDALABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FK0shortLabels, "AOD", "FK0SHORTLABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FD0Labels, "AOD", "FD0LABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FLcLabels, "AOD", "FLCLABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FSigmaLabels, "AOD", "FSIGMALABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FSigmaPlusLabels, "AOD", "FSIGMAPLUSLABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FXiLabels, "AOD", "FXILABEL", femtolabels::FMcParticleId);
-
-DECLARE_SOA_TABLE(FOmegaLabels, "AOD", "FOMEGALABEL", femtolabels::FMcParticleId);
+// Labels are staged like the tables they point into: the derived-to-derived
+// producer writes the Stored* variant, so MC information survives a second
+// derived-data pass.
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FColLabels_001, "FCOLMCLABEL", 1, //! label from femto collision to femto mc collision
+                                   femtolabels::FMcColId);
+using FColLabels = FColLabels_001;
+using FColLabel = FColLabels::iterator;
+using StoredFColLabels = StoredFColLabels_001;
+using StoredFColLabel = StoredFColLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FTrackLabels_001, "FTRACKLABEL", 1, //! label from femto track to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FTrackLabels = FTrackLabels_001;
+using FTrackLabel = FTrackLabels::iterator;
+using StoredFTrackLabels = StoredFTrackLabels_001;
+using StoredFTrackLabel = StoredFTrackLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FLambdaLabels_001, "FLAMBDALABEL", 1, //! label from femto lambda to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FLambdaLabels = FLambdaLabels_001;
+using FLambdaLabel = FLambdaLabels::iterator;
+using StoredFLambdaLabels = StoredFLambdaLabels_001;
+using StoredFLambdaLabel = StoredFLambdaLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FK0shortLabels_001, "FK0SHORTLABEL", 1, //! label from femto k0short to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FK0shortLabels = FK0shortLabels_001;
+using FK0shortLabel = FK0shortLabels::iterator;
+using StoredFK0shortLabels = StoredFK0shortLabels_001;
+using StoredFK0shortLabel = StoredFK0shortLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FD0Labels_001, "FD0LABEL", 1, //! label from femto d0 to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FD0Labels = FD0Labels_001;
+using FD0Label = FD0Labels::iterator;
+using StoredFD0Labels = StoredFD0Labels_001;
+using StoredFD0Label = StoredFD0Labels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FLcLabels_001, "FLCLABEL", 1, //! label from femto lc to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FLcLabels = FLcLabels_001;
+using FLcLabel = FLcLabels::iterator;
+using StoredFLcLabels = StoredFLcLabels_001;
+using StoredFLcLabel = StoredFLcLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FSigmaLabels_001, "FSIGMALABEL", 1, //! label from femto sigma to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FSigmaLabels = FSigmaLabels_001;
+using FSigmaLabel = FSigmaLabels::iterator;
+using StoredFSigmaLabels = StoredFSigmaLabels_001;
+using StoredFSigmaLabel = StoredFSigmaLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FSigmaPlusLabels_001, "FSIGMAPLUSLABEL", 1, //! label from femto sigma plus to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FSigmaPlusLabels = FSigmaPlusLabels_001;
+using FSigmaPlusLabel = FSigmaPlusLabels::iterator;
+using StoredFSigmaPlusLabels = StoredFSigmaPlusLabels_001;
+using StoredFSigmaPlusLabel = StoredFSigmaPlusLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FXiLabels_001, "FXILABEL", 1, //! label from femto xi to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FXiLabels = FXiLabels_001;
+using FXiLabel = FXiLabels::iterator;
+using StoredFXiLabels = StoredFXiLabels_001;
+using StoredFXiLabel = StoredFXiLabels::iterator;
+
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FOmegaLabels_001, "FOMEGALABEL", 1, //! label from femto omega to femto mc particle
+                                   femtolabels::FMcParticleId);
+using FOmegaLabels = FOmegaLabels_001;
+using FOmegaLabel = FOmegaLabels::iterator;
+using StoredFOmegaLabels = StoredFOmegaLabels_001;
+using StoredFOmegaLabel = StoredFOmegaLabels::iterator;
 
 // for mc only processing, we also need Labels pointing from mc particles to mothers and partonic mothers
-DECLARE_SOA_TABLE(FMcMotherLabels, "AOD", "FMCMOTHERLABEL",
-                  femtolabels::FMcMotherId,
-                  femtolabels::FMcPartMothId);
+DECLARE_SOA_TABLE_STAGED_VERSIONED(FMcMotherLabels_001, "FMCMOTHERLABEL", 1, //! labels from femto mc particle to its mother and partonic mother
+                                   femtolabels::FMcMotherId,
+                                   femtolabels::FMcPartMothId);
+using FMcMotherLabels = FMcMotherLabels_001;
+using FMcMotherLabel = FMcMotherLabels::iterator;
+using StoredFMcMotherLabels = StoredFMcMotherLabels_001;
+using StoredFMcMotherLabel = StoredFMcMotherLabels::iterator;
 
 } // namespace o2::aod
 #endif // PWGCF_FEMTO_DATAMODEL_FEMTOTABLES_H_