HybridGen: Add missing EventHeader propagation

The HybridGen needs to forward event headers from the underlying generators. In order to access these, we need to make a protected function public. Also adding eventHeader treatment to BoxGenerator.
AliceO2Group · Dec 13, 2024 · 8222c00 · 8222c00
1 parent 3d1e798
commit 8222c00
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 3 deletions.
diff --git a/Generators/include/Generators/BoxGenerator.h b/Generators/include/Generators/BoxGenerator.h
@@ -18,6 +18,7 @@
 #include "TParticle.h"
 #include <vector>
 #include <Generators/BoxGunParam.h>
+#include "SimulationDataFormat/MCEventHeader.h"
 
 namespace o2::eventgen
 {
@@ -92,6 +93,14 @@ class BoxGenerator : public Generator
     return true;
   }
 
+  void updateHeader(o2::dataformats::MCEventHeader* eventHeader) override
+  {
+    using Key = o2::dataformats::MCInfoKeys;
+    if (eventHeader) {
+      eventHeader->putInfo<std::string>(Key::generator, "o2::eventgen::BoxGenerator");
+    }
+  }
+
  private:
   double mPtMin{0.}, mPtMax{0.};       // Transverse momentum range [GeV]
   double mPhiMin{0.}, mPhiMax{360.};   // Azimuth angle range [degree]

diff --git a/Generators/include/Generators/Generator.h b/Generators/include/Generators/Generator.h
@@ -73,6 +73,7 @@ class Generator : public FairGenerator
   /** methods to override **/
   virtual Bool_t generateEvent() = 0;   // generates event (in structure internal to generator)
   virtual Bool_t importParticles() = 0; // fills the mParticles vector (transfer from generator state)
+  virtual void updateHeader(o2::dataformats::MCEventHeader* eventHeader) {};
 
   /** setters **/
   void setMomentumUnit(double val) { mMomentumUnit = val; };
@@ -102,9 +103,6 @@ class Generator : public FairGenerator
   /** operator= **/
   Generator& operator=(const Generator&);
 
-  /** methods that can be overridded **/
-  virtual void updateHeader(o2::dataformats::MCEventHeader* eventHeader){};
-
   /** internal methods **/
   Bool_t addTracks(FairPrimaryGenerator* primGen);
   Bool_t boostEvent();

diff --git a/Generators/include/Generators/GeneratorHybrid.h b/Generators/include/Generators/GeneratorHybrid.h
@@ -61,6 +61,7 @@ class GeneratorHybrid : public Generator
   Bool_t Init() override;
   Bool_t generateEvent() override;
   Bool_t importParticles() override;
+  void updateHeader(o2::dataformats::MCEventHeader* eventHeader) override;
 
   void setNEvents(int n) { mNEvents = n; }
 
@@ -106,6 +107,7 @@ class GeneratorHybrid : public Generator
   bool mIsInitialized = false;
 
   int mNEvents = -1; // the number of events to be done, if known (helps initiating cleanup)
+  o2::dataformats::MCEventHeader mMCEventHeader; // to capture event headers
 
   enum class GenMode {
     kSeq,

diff --git a/Generators/src/GeneratorHybrid.cxx b/Generators/src/GeneratorHybrid.cxx
@@ -334,6 +334,10 @@ bool GeneratorHybrid::importParticles()
   // at this moment the mIndex-th generator is ready to be used
   std::copy(gens[genIndex]->getParticles().begin(), gens[genIndex]->getParticles().end(), std::back_insert_iterator(mParticles));
 
+  // fetch the event Header information from the underlying generator
+  mMCEventHeader.clearInfo();
+  gens[genIndex]->updateHeader(&mMCEventHeader);
+
   mInputTaskQueue.push(genIndex);
   mTasksStarted++;
 
@@ -353,6 +357,17 @@ bool GeneratorHybrid::importParticles()
   return true;
 }
 
+void GeneratorHybrid::updateHeader(o2::dataformats::MCEventHeader* eventHeader)
+{
+  if (eventHeader) {
+    // we forward the original header information if any
+    eventHeader->copyInfoFrom(mMCEventHeader);
+
+    // put additional information about
+    eventHeader->putInfo<std::string>("forwarding-generator", "HybridGen");
+  }
+}
+
 template <typename T>
 std::string GeneratorHybrid::jsonValueToString(const T& value)
 {