1+ #include "Generators/DecayerPythia8.h"
2+ #include "Generators/GeneratorPythia8.h"
3+ #include "SimulationDataFormat/MCGenProperties.h"
4+ #include "SimulationDataFormat/ParticleStatus.h"
5+ #include "TLorentzVector.h"
6+
7+ namespace o2 {
8+ namespace eventgen {
9+
10+ class DecayerPythia8ForceDecays : public DecayerPythia8 {
11+ public :
12+ DecayerPythia8ForceDecays (){
13+ mPythia .readString ("Random:setSeed = on" );
14+ char * alien_proc_id = getenv ("ALIEN_PROC_ID" );
15+ int seed ;
16+ if (alien_proc_id != NULL ) {
17+ seed = atoi (alien_proc_id );
18+ LOG (info ) << "Seed for DecayerPythia8 set to ALIEN_PROC_ID: " << seed ;
19+ } else {
20+ LOG (info ) << "Unable to retrieve ALIEN_PROC_ID" ;
21+ LOG (info ) << "Setting seed for DecayerPyhtia8 to 0 (random)" ;
22+ seed = 0 ;
23+ }
24+ mPythia .readString ("Random:seed = " + std ::to_string (seed ));
25+ }
26+ ~DecayerPythia8ForceDecays () = default ;
27+
28+
29+ void calculateWeights (std ::vector < int > & pdgs ) {
30+ TLorentzVector mom = TLorentzVector (0. , 0. , 0. , 9999999. );
31+ for (int pdg : pdgs ) {
32+ Decay (pdg , & mom ); // do one fake decay to initalize everything correctly
33+ auto particleData = mPythia .particleData .particleDataEntryPtr (pdg );
34+ double weight = 0. ;
35+ for (int c = 0 ; c < particleData -> sizeChannels (); c ++ ) {
36+ weight += particleData -> channel (c ).currentBR ();
37+ }
38+ LOG (info ) << "PDG = " << pdg
39+ << ": sum of branching ratios of active decay channels = "
40+ << weight ;
41+ mWeights [pdg ] = weight ;
42+ mPythia .particleData .mayDecay (pdg , false);
43+ }
44+ }
45+
46+ void forceDecays (std ::vector < TParticle > & mParticles , int mother_pos ) {
47+ TParticle * p = & mParticles [mother_pos ];
48+ int pdg = p -> GetPdgCode ();
49+ TLorentzVector mom = TLorentzVector (p -> Px (), p -> Py (), p -> Pz (), p -> Energy ());
50+ Decay (pdg , & mom );
51+ TClonesArray daughters = TClonesArray ("TParticle" );
52+ int nParticles = ImportParticles (& daughters );
53+ int mcGenStatus = o2 ::mcgenstatus ::getGenStatusCode (p -> GetStatusCode ());
54+ p -> SetStatusCode (o2 ::mcgenstatus ::MCGenStatusEncoding (2 , - mcGenStatus ).fullEncoding );
55+ p -> SetBit (ParticleStatus ::kToBeDone , false);
56+ double mother_weight = p -> GetWeight ();
57+ TParticle * mother = static_cast < TParticle * > (daughters [0 ]);
58+ int mParticles_size = mParticles .size ();
59+ p -> SetFirstDaughter (mother -> GetFirstDaughter () + mParticles_size - 1 );
60+ p -> SetLastDaughter (mother -> GetLastDaughter () + mParticles_size - 1 );
61+ for (int j = 1 ; j < nParticles ;
62+ j ++ ) { // start loop at 1 to not include mother
63+ TParticle * d = static_cast < TParticle * > (daughters [j ]);
64+ double decay_weight = mWeights [abs (pdg )];
65+ if (decay_weight == 0 ) {
66+ LOG (error ) << "Decaying particle (PDG = " << pdg
67+ << ") with decay weight = 0. Did you set the pdg codes for "
68+ "calculating weights correctly?" ;
69+ }
70+ d -> SetWeight (decay_weight * mother_weight );
71+ if (d -> GetStatusCode () == 1 ) {
72+ d -> SetStatusCode (
73+ o2 ::mcgenstatus ::MCGenStatusEncoding (1 , 91 ).fullEncoding );
74+ d -> SetBit (ParticleStatus ::kToBeDone , true);
75+ } else {
76+ d -> SetStatusCode (
77+ o2 ::mcgenstatus ::MCGenStatusEncoding (2 , -91 ).fullEncoding );
78+ d -> SetBit (ParticleStatus ::kToBeDone , false);
79+ }
80+ int firstmother = d -> GetFirstMother ();
81+ int firstdaughter = d -> GetFirstDaughter ();
82+ int lastdaughter = d -> GetLastDaughter ();
83+ if (firstmother == 0 ) {
84+ d -> SetFirstMother (mother_pos );
85+ d -> SetLastMother (mother_pos );
86+ } else {
87+ d -> SetFirstMother (firstmother + mParticles_size - 1 );
88+ d -> SetLastMother (firstmother + mParticles_size - 1 );
89+ }
90+ if (firstdaughter == 0 ) {
91+ d -> SetFirstDaughter (-1 );
92+ } else {
93+ d -> SetFirstDaughter (firstdaughter + mParticles_size - 1 );
94+ }
95+ if (lastdaughter == 0 ) {
96+ d -> SetLastDaughter (-1 );
97+ } else {
98+ d -> SetLastDaughter (lastdaughter + mParticles_size - 1 );
99+ }
100+ mParticles .push_back (* d );
101+ }
102+ }
103+
104+ private :
105+ std ::map < int , double > mWeights ;
106+ };
107+
108+ class GeneratorPythia8ForcedDecays : public GeneratorPythia8 {
109+
110+ public :
111+ GeneratorPythia8ForcedDecays (){
112+ mPythia .readString ("Random:setSeed = on" );
113+ char * alien_proc_id = getenv ("ALIEN_PROC_ID" );
114+ int seed ;
115+ if (alien_proc_id != NULL ) {
116+ seed = atoi (alien_proc_id );
117+ LOG (info ) << "Seed for GeneratorPythia8 set to ALIEN_PROC_ID: " << seed ;
118+ } else {
119+ LOG (info ) << "Unable to retrieve ALIEN_PROC_ID" ;
120+ LOG (info ) << "Setting seed for GeneratorPyhtia8 to 0 (random)" ;
121+ seed = 0 ;
122+ }
123+ mPythia .readString ("Random:seed = " + std ::to_string (seed ));
124+ }
125+ ~GeneratorPythia8ForcedDecays () = default ;
126+
127+ // overriden methods
128+ bool Init () override { return GeneratorPythia8 ::Init () && InitDecayer (); };
129+ bool importParticles () override {
130+ return GeneratorPythia8 ::importParticles () && makeForcedDecays ();
131+ };
132+
133+ void setPDGs (TString pdgs ) {
134+ TObjArray * obj = pdgs .Tokenize (";" );
135+ for (int i = 0 ; i < obj -> GetEntriesFast (); i ++ ) {
136+ std ::string spdg = obj -> At (i )-> GetName ();
137+ int pdg = std ::stoi (spdg );
138+ mPdgCodes .push_back (pdg );
139+ LOG (info ) << "Force decay of PDG = " << pdg ;
140+ }
141+ }
142+
143+ protected :
144+ bool InitDecayer () {
145+ mDecayer = new DecayerPythia8ForceDecays ();
146+ mDecayer -> Init ();
147+ mDecayer -> calculateWeights (mPdgCodes );
148+ for (int pdg : mPdgCodes ) {
149+ mPythia .particleData .mayDecay (pdg , false );
150+ }
151+ return true;
152+ }
153+
154+ bool makeForcedDecays () {
155+ int mParticles_size = mParticles .size ();
156+ for (int i = 0 ; i < mParticles_size ; i ++ ) {
157+ int pdg = mParticles [i ].GetPdgCode ();
158+ if (std ::find (mPdgCodes .begin (), mPdgCodes .end (), abs (pdg )) !=
159+ mPdgCodes .end ()) {
160+ mDecayer -> forceDecays (mParticles , i );
161+ mParticles_size = mParticles .size ();
162+ }
163+ }
164+ return true;
165+ }
166+
167+ private :
168+ DecayerPythia8ForceDecays * mDecayer ;
169+ std ::vector < int > mPdgCodes ;
170+ };
171+
172+ } // namespace eventgen
173+ } // namespace o2
174+
175+ FairGenerator * GeneratePythia8ForcedDecays (TString pdgs ) {
176+ auto gen = new o2 ::eventgen ::GeneratorPythia8ForcedDecays ();
177+ gen -> setPDGs (pdgs );
178+ return gen ;
179+ }
0 commit comments