PrimaryGenerator that can propagate mother/daughter indices (#4361)

preghenella · web-flow · commit cc195872d7b8 · 2020-09-24T14:16:33.000+02:00
* PrimaryGenerator that can propagate mother/daughter indices

* Pass mother2 as secondparent when pushing track to stack
diff --git a/Generators/include/Generators/PrimaryGenerator.h b/Generators/include/Generators/PrimaryGenerator.h
@@ -55,18 +55,14 @@ class PrimaryGenerator : public FairPrimaryGenerator
       **/
   Bool_t GenerateEvent(FairGenericStack* pStack) override;
 
-  /** Public method AddTrack
-      Adding a track to the MC stack. To be called within the ReadEvent
-      methods of the registered generators.
-      *@param pdgid Particle ID (PDG code)
-      *@param px,py,pz Momentum coordinates [GeV]
-      *@param vx,vy,vz Track origin relative to event vertex
-      **/
+  /** ALICE-o2 AddTrack with mother/daughter indices **/
   void AddTrack(Int_t pdgid, Double_t px, Double_t py, Double_t pz,
                 Double_t vx, Double_t vy, Double_t vz,
-                Int_t parent = -1, Bool_t wanttracking = true,
+                Int_t mother1 = -1, Int_t mother2 = -1,
+                Int_t daughter1 = -1, Int_t daughter2 = -1,
+                Bool_t wanttracking = true,
                 Double_t e = -9e9, Double_t tof = 0.,
-                Double_t weight = 0., TMCProcess proc = kPPrimary) override;
+                Double_t weight = 0., TMCProcess proc = kPPrimary);
 
   /** initialize the generator **/
   Bool_t Init() override;
diff --git a/Generators/src/Generator.cxx b/Generators/src/Generator.cxx
@@ -12,6 +12,7 @@
 
 #include "Generators/Generator.h"
 #include "Generators/Trigger.h"
+#include "Generators/PrimaryGenerator.h"
 #include "FairPrimaryGenerator.h"
 #include "FairLogger.h"
 #include <cmath>
@@ -97,20 +98,29 @@ Bool_t
 {
   /** add tracks **/
 
+  auto o2primGen = dynamic_cast<PrimaryGenerator*>(primGen);
+  if (!o2primGen) {
+    LOG(FATAL) << "PrimaryGenerator is not a o2::eventgen::PrimaryGenerator";
+    return kFALSE;
+  }
+
   /** loop over particles **/
   for (const auto& particle : mParticles) {
-    primGen->AddTrack(particle.GetPdgCode(),
-                      particle.Px() * mMomentumUnit,
-                      particle.Py() * mMomentumUnit,
-                      particle.Pz() * mMomentumUnit,
-                      particle.Vx() * mPositionUnit,
-                      particle.Vy() * mPositionUnit,
-                      particle.Vz() * mPositionUnit,
-                      particle.GetMother(0),
-                      particle.GetStatusCode() == 1,
-                      particle.Energy() * mEnergyUnit,
-                      particle.T() * mTimeUnit,
-                      particle.GetWeight());
+    o2primGen->AddTrack(particle.GetPdgCode(),
+                        particle.Px() * mMomentumUnit,
+                        particle.Py() * mMomentumUnit,
+                        particle.Pz() * mMomentumUnit,
+                        particle.Vx() * mPositionUnit,
+                        particle.Vy() * mPositionUnit,
+                        particle.Vz() * mPositionUnit,
+                        particle.GetMother(0),
+                        particle.GetMother(1),
+                        particle.GetDaughter(0),
+                        particle.GetDaughter(1),
+                        particle.GetStatusCode() == 1,
+                        particle.Energy() * mEnergyUnit,
+                        particle.T() * mTimeUnit,
+                        particle.GetWeight());
   }
 
   /** success **/
diff --git a/Generators/src/PrimaryGenerator.cxx b/Generators/src/PrimaryGenerator.cxx
@@ -116,7 +116,9 @@ Bool_t PrimaryGenerator::GenerateEvent(FairGenericStack* pStack)
 
 void PrimaryGenerator::AddTrack(Int_t pdgid, Double_t px, Double_t py, Double_t pz,
                                 Double_t vx, Double_t vy, Double_t vz,
-                                Int_t parent, Bool_t wanttracking,
+                                Int_t mother1, Int_t mother2,
+                                Int_t daughter1, Int_t daughter2,
+                                Bool_t wanttracking,
                                 Double_t e, Double_t tof,
                                 Double_t weight, TMCProcess proc)
 {
@@ -146,9 +148,19 @@ void PrimaryGenerator::AddTrack(Int_t pdgid, Double_t px, Double_t py, Double_t
   Int_t ntr = 0;    // Track number; to be filled by the stack
   Int_t status = 0; // Generation status
 
-  if (parent != -1) {
-    parent += fMCIndexOffset;
-  } // correct for tracks which are in list before generator is called
+  // correct for tracks which are in list before generator is called
+  if (mother1 != -1) {
+    mother1 += fMCIndexOffset;
+  }
+  if (mother2 != -1) {
+    mother2 += fMCIndexOffset;
+  }
+  if (daughter1 != -1) {
+    daughter1 += fMCIndexOffset;
+  }
+  if (daughter2 != -1) {
+    daughter2 += fMCIndexOffset;
+  }
 
   /** if it is a K0/antiK0 to be tracked, convert it into K0s/K0L.
       
@@ -170,9 +182,9 @@ void PrimaryGenerator::AddTrack(Int_t pdgid, Double_t px, Double_t py, Double_t
   }
 
   /** add track to stack **/
-  fStack->PushTrack(doTracking, parent, pdgid, px, py, pz,
+  fStack->PushTrack(doTracking, mother1, pdgid, px, py, pz,
                     e, vx, vy, vz, tof, polx, poly, polz, proc, ntr,
-                    weight, status, parent);
+                    weight, status, mother2);
 
   fNTracks++;
 }
