Clean-up. Call to ConnectTree to store correctly the stack
[u/mrichter/AliRoot.git] / test / genkine / gen / fastGen.C
1 // Example: generation of kinematics tree with selected properties.
2 // Below we select events containing the decays D* -> D0 pi, D0 -> K- pi+
3 // inside the barrel part of the ALICE detector (45 < theta < 135)
4
5 #if !defined(__CINT__) || defined(__MAKECINT__)
6 #include <Riostream.h>
7 #include <TH1F.h>
8 #include <TStopwatch.h>
9 #include <TDatime.h>
10 #include <TRandom.h>
11 #include <TDatabasePDG.h>
12 #include <TParticle.h>
13 #include <TArrayI.h>
14
15 #include "AliGenerator.h"
16 #include "AliPDG.h"
17 #include "AliRunLoader.h"
18 #include "AliRun.h"
19 #include "AliStack.h"
20 #include "AliHeader.h"
21 #include "PYTHIA6/AliGenPythia.h"
22 #include "PYTHIA6/AliPythia.h"
23 #endif
24
25 Float_t EtaToTheta(Float_t arg);
26 void GetFinalDecayProducts(Int_t ind, AliStack & stack , TArrayI & ar);
27
28 void fastGen(Int_t nev = 1, char* filename = "galice.root")
29 {
30   AliPDG::AddParticlesToPdgDataBase();
31   TDatabasePDG::Instance();
32  
33
34
35   // Run loader
36   AliRunLoader* rl = AliRunLoader::Open("galice.root","FASTRUN","recreate");
37   
38   rl->SetCompressionLevel(2);
39   rl->SetNumberOfEventsPerFile(nev);
40   rl->LoadKinematics("RECREATE");
41   rl->MakeTree("E");
42   gAlice->SetRunLoader(rl);
43   
44   //  Create stack
45   rl->MakeStack();
46   AliStack* stack = rl->Stack();
47   
48   //  Header
49   AliHeader* header = rl->GetHeader();
50   
51   //  Create and Initialize Generator
52  
53   // Example of charm generation taken from Config_PythiaHeavyFlavours.C
54   AliGenPythia *gener = new AliGenPythia(-1);
55   gener->SetEnergyCMS(14000.);
56   gener->SetMomentumRange(0,999999);
57   gener->SetPhiRange(0., 360.);
58   gener->SetThetaRange(0.,180.);
59   //  gener->SetProcess(kPyCharmppMNR); // Correct Pt distribution, wrong mult
60   gener->SetProcess(kPyMb); // Correct multiplicity, wrong Pt
61   gener->SetStrucFunc(kCTEQ4L);
62   gener->SetPtHard(2.1,-1.0);
63   gener->SetFeedDownHigherFamily(kFALSE);
64   gener->SetStack(stack);
65   gener->Init();
66
67   // Go to galice.root
68   rl->CdGAFile();
69
70   // Forbid some decays. Do it after gener->Init(0, because
71   // the initialization of the generator includes reading of the decay table.
72
73   AliPythia * py= AliPythia::Instance();
74   py->SetMDME(731,1,0); //forbid D*+->D+ + pi0
75   py->SetMDME(732,1,0);//forbid D*+->D+ + gamma
76
77   // Forbid all D0 decays except D0->K- pi+
78   for(Int_t d=741; d<=756; d++){ 
79     py->SetMDME(d,1,0);
80   }
81   // decay 757 is D0->K- pi+
82   for(Int_t d=758; d<=801; d++){ 
83     py->SetMDME(d,1,0);
84   }
85
86   //
87   //                        Event Loop
88   //
89   
90   TStopwatch timer;
91   timer.Start();
92   for (Int_t iev = 0; iev < nev; iev++) {
93     
94     cout <<"Event number "<< iev << endl;
95     
96     //  Initialize event
97     header->Reset(0,iev);
98     rl->SetEventNumber(iev);
99     stack->Reset();
100     rl->MakeTree("K");
101     
102     //  Generate event
103     Int_t nprim = 0;
104     Int_t ntrial = 0;
105     Int_t ndstar = 0;
106    
107    
108
109     //-------------------------------------------------------------------------------------
110
111     while(!ndstar) {
112       // Selection of events with D*
113       stack->Reset();
114       stack->ConnectTree(rl->TreeK());
115       gener->Generate();
116       ntrial++;
117       nprim = stack->GetNprimary();
118       
119       for(Int_t ipart =0; ipart < nprim; ipart++){
120         TParticle * part = stack->Particle(ipart);
121         if(part)    {
122           
123           if (TMath::Abs(part->GetPdgCode())== 413) {
124
125             TArrayI daughtersId;
126
127             GetFinalDecayProducts(ipart,*stack,daughtersId);
128
129             Bool_t kineOK = kTRUE;
130
131             Double_t thetaMin = TMath::Pi()/4;
132             Double_t thetaMax = 3*TMath::Pi()/4;
133
134             for (Int_t id=1; id<=daughtersId[0]; id++) {
135               TParticle * daughter = stack->Particle(daughtersId[id]);
136               if (!daughter) {
137                 kineOK = kFALSE;
138                 break;
139               }
140
141               Double_t theta = daughter->Theta();
142               if (theta<thetaMin || theta>thetaMax) {
143                 kineOK = kFALSE;
144                 break;
145               }
146             }
147
148             if (!kineOK) continue;
149
150             part->Print();
151             ndstar++;     
152             
153           }
154         }
155       }   
156     }   
157       
158     cout << "Number of particles " << nprim << endl;
159     cout << "Number of trials " << ntrial << endl;
160     
161     //  Finish event
162     header->SetNprimary(stack->GetNprimary());
163     header->SetNtrack(stack->GetNtrack());  
164     
165     //      I/O
166     stack->FinishEvent();
167     header->SetStack(stack);
168     rl->TreeE()->Fill();
169     rl->WriteKinematics("OVERWRITE");
170     
171   } // event loop
172   timer.Stop();
173   timer.Print();
174   
175   //                         Termination
176   //  Generator
177   gener->FinishRun();
178   //  Write file
179   rl->WriteHeader("OVERWRITE");
180   gener->Write();
181   rl->Write();
182 }
183
184
185
186 Float_t EtaToTheta(Float_t arg){
187   return (180./TMath::Pi())*2.*atan(exp(-arg));
188 }
189
190
191 void GetFinalDecayProducts(Int_t ind, AliStack & stack , TArrayI & ar){
192
193   // Recursive algorithm to get the final decay products of a particle
194   //
195   // ind is the index of the particle in the AliStack
196   // stack is the particle stack from the generator
197   // ar contains the indexes of the final decay products
198   // ar[0] is the number of final decay products
199
200   if (ind<0 || ind>stack.GetNtrack()) {
201     cerr << "Invalid index of the particle " << ind << endl;
202     return;
203   } 
204   if (ar.GetSize()==0) {
205     ar.Set(10);
206     ar[0] = 0;
207   }
208
209   TParticle * part = stack.Particle(ind);
210
211   Int_t iFirstDaughter = part->GetFirstDaughter();
212   if( iFirstDaughter<0) {
213     // This particle is a final decay product, add its index to the array
214     ar[0]++;
215     if (ar.GetSize() <= ar[0]) ar.Set(ar.GetSize()+10); // resize if needed
216     ar[ar[0]] = ind;
217     return;
218   } 
219
220   Int_t iLastDaughter = part->GetLastDaughter();
221
222   for (Int_t id=iFirstDaughter; id<=iLastDaughter;id++) {
223     // Now search for final decay products of the daughters
224     GetFinalDecayProducts(id,stack,ar);
225   }
226 }