[ALICE3] Switch from McPartWithDaus to McParticles with --aod-origin-replace in otf decayer (#16448)

Author: jesgum

ALICE3/Core/Decayer.h

@@ -126,6 +126,7 @@ class Decayer
       particle.setPDG(pdgCodesDaughters[i]);
       particle.setVxVyVz(mVx, mVy, mVz);
       particle.setPxPyPzE(dau.Px(), dau.Py(), dau.Pz(), dau.E());
+      particle.setBitOn(o2::upgrade::DecayerBits::ProducedByDecayer);
       decayProducts.push_back(particle);
     }
 

ALICE3/Core/OTFParticle.h

@@ -21,13 +21,18 @@
 #include <CommonConstants/MathConstants.h>
 
 #include <array>
+#include <bitset>
 #include <cmath>
 #include <cstdint>
 #include <span>
 
 namespace o2::upgrade
 {
 
+enum class DecayerBits { ProducedByDecayer = 0,
+                         IsPrimary,
+                         IsAlive };
+
 class OTFParticle
 {
  public:
@@ -47,20 +52,29 @@ class OTFParticle
     mVy = particle.vy();
     mVz = particle.vz();
     mVt = particle.vt();
+    mFlag = particle.flags();
+    mStatusCode = particle.statusCode();
     mIsFromMcParticles = true;
     if (particle.has_mothers()) {
       mIndicesMother = {particle.mothersIds().front(), particle.mothersIds().back()};
     }
+    if constexpr (requires { particle.decayerBits(); }) {
+      mBits = particle.decayerBits();
+    } else {
+      // If we are here, we created particle in the standard workflow -- without secondaries
+      // Then we should set all particles as physical primaries accordingly
+      setBitOn(DecayerBits::IsPrimary);
+    }
   }
 
   // Setters
-  void setIsAlive(const bool isAlive) { mIsAlive = isAlive; }
   void setIsPrimary(const bool isPrimary) { mIsPrimary = isPrimary; }
   void setCollisionId(const int collisionId) { mCollisionId = collisionId; }
   void setPDG(const int pdg) { mPdgCode = pdg; }
   void setIndicesMother(const int start, const int stop) { mIndicesMother = {start, stop}; }
   void setIndicesDaughter(const int start, const int stop) { mIndicesDaughter = {start, stop}; }
   void setProductionTime(const float vt) { mVt = vt; }
+  void setFlags(uint8_t flag) { mFlag = flag; }
   void setVxVyVz(const float vx, const float vy, const float vz)
   {
     mVx = vx;
@@ -87,16 +101,8 @@ class OTFParticle
     static constexpr float Weight = 1.f;
     return Weight;
   }
-  uint8_t flags() const
-  {
-    static constexpr uint8_t Flags = 1;
-    return Flags; // todo
-  }
-  int statusCode() const
-  {
-    static constexpr int StatusCode = 1;
-    return StatusCode; // todo
-  }
+  uint8_t flags() const { return mFlag; }
+  int statusCode() const { return mStatusCode; }
   float vx() const { return mVx; }
   float vy() const { return mVy; }
   float vz() const { return mVz; }
@@ -112,7 +118,8 @@ class OTFParticle
   float phi() const { return o2::constants::math::PI + std::atan2(-1.0f * py(), -1.0f * px()); }
   float eta() const
   {
-    // As https://github.com/AliceO2Group/AliceO2/blob/dev/Framework/Core/include/Framework/AnalysisDataModel.h#L1943
+    // Conditionally defined to avoid FPEs
+    // As https://github.com/AliceO2Group/AliceO2/blob/dev/Framework/Core/include/Framework/AnalysisDataModel.h#L1959
     static constexpr float Tolerance = 1e-7f;
     if ((p() - mPz) < Tolerance) {
       return (mPz < 0.0f) ? -100.0f : 100.0f;
@@ -122,7 +129,8 @@ class OTFParticle
   }
   float y() const
   {
-    // As https://github.com/AliceO2Group/AliceO2/blob/dev/Framework/Core/include/Framework/AnalysisDataModel.h#L1922
+    // Conditionally defined to avoid FPEs
+    // As https://github.com/AliceO2Group/AliceO2/blob/dev/Framework/Core/include/Framework/AnalysisDataModel.h#L1980
     static constexpr float Tolerance = 1e-7f;
     if ((e() - mPz) < Tolerance) {
       return (mPz < 0.0f) ? -100.0f : 100.0f;
@@ -151,13 +159,27 @@ class OTFParticle
     return (mGlobalIndex != -1);
   }
 
+  // Bits
+  bool checkBit(DecayerBits bit) const { return mBits.test(static_cast<size_t>(bit)); }
+  void setBit(DecayerBits bit, bool value = true) { mBits.set(static_cast<size_t>(bit), value); }
+  void setBitOn(DecayerBits bit) { mBits.set(static_cast<size_t>(bit), true); }
+  void setBitOff(DecayerBits bit) { mBits.set(static_cast<size_t>(bit), false); }
+
+  std::bitset<8> getBits() const { return mBits; }
+  uint8_t getBitsValue() const { return static_cast<uint8_t>(mBits.to_ulong()); }
+  void setBits(std::bitset<8> bits) { mBits = bits; }
+
  private:
   int mPdgCode{}, mGlobalIndex{-1};
   int mCollisionId{};
   float mVx{}, mVy{}, mVz{}, mVt{};
   float mPx{}, mPy{}, mPz{}, mE{};
   bool mIsAlive{}, mIsFromMcParticles{false};
   bool mIsPrimary{};
+
+  int mStatusCode{};
+  uint8_t mFlag{};
+  std::bitset<8> mBits{};
   std::array<int, 2> mIndicesMother{-1, -1}, mIndicesDaughter{-1, -1};
 };
 

ALICE3/DataModel/OTFMCParticle.h

@@ -50,11 +50,15 @@ namespace mcparticle_alice3
 {
 DECLARE_SOA_COLUMN(NHits, nHits, int);     //! number of silicon hits
 DECLARE_SOA_COLUMN(Charge, charge, float); //! particle charge
+DECLARE_SOA_BITMAP_COLUMN(DecayerBits, decayerBits, 8); //! Bit mask for particle produced by the OTF decayer
 } // namespace mcparticle_alice3
 DECLARE_SOA_TABLE(MCParticlesExtraA3, "AOD", "MCParticlesExtraA3",
                   mcparticle_alice3::NHits,
                   mcparticle_alice3::Charge);
 using MCParticleExtraA3 = MCParticlesExtraA3::iterator;
+
+DECLARE_SOA_TABLE(OTFDecayerBits, "AOD", "OTFDecayerBits", mcparticle_alice3::DecayerBits);
+
 } // namespace o2::aod
 
 #endif // ALICE3_DATAMODEL_TRACKSALICE3_H_

ALICE3/DataModel/tracksAlice3.h

@@ -18,7 +18,7 @@
 #include "ALICE3/Core/Decayer.h"
 #include "ALICE3/Core/OTFParticle.h"
 #include "ALICE3/Core/TrackUtilities.h"
-#include "ALICE3/DataModel/OTFMCParticle.h"
+#include "ALICE3/DataModel/tracksAlice3.h"
 
 #include <Framework/AnalysisDataModel.h>
 #include <Framework/AnalysisHelpers.h>
@@ -40,7 +40,6 @@
 #include <array>
 #include <cmath>
 #include <cstdlib>
-#include <map>
 #include <string>
 #include <vector>
 
@@ -67,12 +66,18 @@ static const std::vector<int> pdgCodes{PDG_t::kK0Short,
                                        PDG_t::kOmegaMinus,
                                        PDG_t::kOmegaPlusBar};
 
+namespace o2::aod
+{
+O2ORIGIN("TMP");
+}
+
 struct OnTheFlyDecayer {
-  Produces<aod::McPartWithDaus> tableMcParticlesWithDau;
+  Produces<aod::McCollisions_001> tableMcCollisions;
+  Produces<aod::StoredMcParticles_001> tableMcParticles;
+  Produces<aod::OTFDecayerBits> tableOTFDecayerBits;
 
   o2::upgrade::Decayer decayer;
   Service<o2::framework::O2DatabasePDG> pdgDB;
-  std::map<int, std::vector<o2::upgrade::OTFParticle>> mDecayDaughters;
   HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject};
 
   Configurable<int> seed{"seed", 0, "Set seed for particle decayer"};
@@ -117,18 +122,18 @@ struct OnTheFlyDecayer {
   void decayParticles(const int start, const int stop)
   {
     int ndau = 0;
-    for (int i = start; i < stop; i++) {
+    for (int i = start; i < stop; ++i) {
       o2::upgrade::OTFParticle& particle = allParticles[i];
       if (particle.isFromMcParticles()) {
-        particle.setIsPrimary(true);
-        particle.setIsAlive(true);
+        particle.setBitOn(o2::upgrade::DecayerBits::IsPrimary);
+        particle.setBitOn(o2::upgrade::DecayerBits::IsAlive);
       }
 
       if (!canDecay(particle)) {
         continue;
       }
 
-      particle.setIsAlive(false);
+      particle.setBitOff(o2::upgrade::DecayerBits::IsAlive);
       std::vector<o2::upgrade::OTFParticle> decayStack = decayer.decayParticle(pdgDB, particle);
       const float decayRadius = decayer.getDecayRadius();
       const float trackVelocity = o2::upgrade::computeParticleVelocity(particle.p(), pdgDB->GetParticle(particle.pdgCode())->Mass());
@@ -150,8 +155,8 @@ struct OnTheFlyDecayer {
       for (o2::upgrade::OTFParticle daughter : decayStack) {
         daughter.setIndicesMother(i, i);
         daughter.setCollisionId(mCollisionId);
-        daughter.setIsAlive(true);
-        daughter.setIsPrimary(false);
+        daughter.setBitOn(o2::upgrade::DecayerBits::IsAlive);
+        daughter.setBitOff(o2::upgrade::DecayerBits::IsPrimary);
         daughter.setProductionTime(trackTimeNS);
         allParticles.push_back(daughter);
         ndau++;
@@ -165,42 +170,54 @@ struct OnTheFlyDecayer {
     decayParticles(stop, stop + ndau);
   }
 
-  void process(aod::McCollision const& collision, aod::McParticles const& mcParticles)
+  void process(aod::McCollisions_001From<aod::Hash<"TMP"_h>>::iterator const& collision, aod::McParticles_001From<aod::Hash<"TMP"_h>> const& mcParticles)
   {
-    mCollisionId = collision.globalIndex();
     allParticles.clear();
 
+    // Reproduce collision table to have AOD origin
+    mCollisionId = collision.globalIndex();
+    tableMcCollisions(collision.bcId(),
+                      collision.generatorsID(),
+                      collision.posX(),
+                      collision.posY(),
+                      collision.posZ(),
+                      collision.t(),
+                      collision.weight(),
+                      collision.impactParameter(),
+                      collision.eventPlaneAngle());
+
     // First we copy the particles from the table into a vector that is extendable
-    for (int index{0}; index < static_cast<int>(mcParticles.size()); ++index) {
-      const auto& mcParticle = mcParticles.rawIteratorAt(index);
-      allParticles.push_back(o2::upgrade::OTFParticle{mcParticle});
+    for (const auto& particle : mcParticles) {
+      allParticles.emplace_back(o2::upgrade::OTFParticle{particle});
     }
 
     // Do all decays
     decayParticles(0, allParticles.size());
 
     // Fill output table
-    for (int index{0}; index < static_cast<int>(allParticles.size()); ++index) {
-      const auto& otfParticle = allParticles[index];
-
+    for (const auto& otfParticle : allParticles) {
       if (otfParticle.hasNaN()) {
         histos.fill(HIST("hNaNBookkeeping"), 1);
       } else {
         histos.fill(HIST("hNaNBookkeeping"), 0);
       }
 
-      // todo: status codes
-      tableMcParticlesWithDau(otfParticle.collisionId(), otfParticle.pdgCode(), otfParticle.statusCode(),
-                              otfParticle.flags(), otfParticle.getMotherSpan(), otfParticle.getDaughters().data(), otfParticle.weight(),
-                              otfParticle.px(), otfParticle.py(), otfParticle.pz(), otfParticle.e(),
-                              otfParticle.vx(), otfParticle.vy(), otfParticle.vz(), otfParticle.vt(),
-                              otfParticle.phi(), otfParticle.eta(), otfParticle.pt(), otfParticle.p(), otfParticle.y(),
-                              otfParticle.isAlive(), otfParticle.isPrimary());
+      tableOTFDecayerBits(otfParticle.getBitsValue());
+      tableMcParticles(otfParticle.collisionId(), otfParticle.pdgCode(), otfParticle.statusCode(), otfParticle.flags(),
+                       otfParticle.getMotherSpan(), otfParticle.getDaughters().data(), otfParticle.weight(),
+                       otfParticle.px(), otfParticle.py(), otfParticle.pz(), otfParticle.e(),
+                       otfParticle.vx(), otfParticle.vy(), otfParticle.vz(), otfParticle.vt());
     }
   }
 };
 
+struct OnTheFlyDecayerExtensionSpawner {
+  Spawns<aod::McParticles_001Extension> spawnMcParticlesExtensions;
+  void init(o2::framework::InitContext&) {}
+};
+
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
 {
-  return WorkflowSpec{adaptAnalysisTask<OnTheFlyDecayer>(cfgc)};
+  return WorkflowSpec{adaptAnalysisTask<OnTheFlyDecayer>(cfgc),
+                      adaptAnalysisTask<OnTheFlyDecayerExtensionSpawner>(cfgc)};
 }