Improve Vertex handling in MCEventHeader

Author: jackal1-66

Generators/include/Generators/GeneratorFromFile.h

@@ -144,7 +144,11 @@ class GeneratorFromEventPool : public o2::eventgen::Generator
 
   void updateHeader(o2::dataformats::MCEventHeader* eventHeader) override
   {
+    // Copy current vertex position from the event header
+    const double xyz[3] = {eventHeader->GetX(), eventHeader->GetY(), eventHeader->GetZ()};
     mO2KineGenerator->updateHeader(eventHeader);
+    // Event pool uses vertex position from current simulation, only extKinO2 takes the one from the file instead
+    eventHeader->SetVertex(xyz[0], xyz[1], xyz[2]);
   }
 
   // determine the collection of available files

Generators/include/Generators/GeneratorHybrid.h

@@ -120,6 +120,7 @@ class GeneratorHybrid : public Generator
   bool mIsInitialized = false;
 
   o2::dataformats::MCEventHeader mMCEventHeader; // to capture event headers
+  int mHeaderGeneratorIndex = -1;                // index of the generator that updated the header in current event
 
   enum class GenMode {
     kSeq,

Generators/src/GeneratorHybrid.cxx

@@ -417,7 +417,10 @@ bool GeneratorHybrid::importParticles()
 
   // Clear particles and event header
   mParticles.clear();
-  mMCEventHeader.clearInfo();
+  // event header of underlying generator must be fully reset
+  // this is important when using extKinO2 or event pools where the full header information is forwarded from the generator
+  // otherwise some events might have mixed header information from different generators
+  mMCEventHeader.Reset();
   if (mCocktailMode) {
     // in cocktail mode we need to merge the particles from the different generators
     bool baseGen = true; // first generator of the cocktail is used as reference to update the event header information
@@ -445,6 +448,7 @@ bool GeneratorHybrid::importParticles()
       if (baseGen) {
         gens[subIndex]->updateHeader(&mMCEventHeader);
         baseGen = false;
+        mHeaderGeneratorIndex = subIndex; // store index of generator updating the header
       }
       mInputTaskQueue.push(subIndex);
       mTasksStarted++;
@@ -467,6 +471,7 @@ bool GeneratorHybrid::importParticles()
 
     // fetch the event Header information from the underlying generator
     gens[genIndex]->updateHeader(&mMCEventHeader);
+    mHeaderGeneratorIndex = genIndex; // store index of generator updating the header
     mInputTaskQueue.push(genIndex);
     mTasksStarted++;
   }
@@ -484,6 +489,11 @@ bool GeneratorHybrid::importParticles()
 void GeneratorHybrid::updateHeader(o2::dataformats::MCEventHeader* eventHeader)
 {
   if (eventHeader) {
+    // Overwrite current vertex position only if extkinO2 is not used as underlying generator
+    if (mGens[mHeaderGeneratorIndex] != "extkinO2") {
+      mMCEventHeader.SetVertex(eventHeader->GetX(), eventHeader->GetY(), eventHeader->GetZ());
+    }
+    mHeaderGeneratorIndex = -1; // reset header generator index for next event
     // Forward the base class fields from FairMCEventHeader
     static_cast<FairMCEventHeader&>(*eventHeader) = static_cast<FairMCEventHeader&>(mMCEventHeader);
     // Copy the key-value store info