Store the total number of produced particles in the header.
[u/mrichter/AliRoot.git] / EVGEN / AliGenCocktail.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
3  *                                                                        *
4  * Author: The ALICE Off-line Project.                                    *
5  * Contributors are mentioned in the code where appropriate.              *
6  *                                                                        *
7  * Permission to use, copy, modify and distribute this software and its   *
8  * documentation strictly for non-commercial purposes is hereby granted   *
9  * without fee, provided that the above copyright notice appears in all   *
10  * copies and that both the copyright notice and this permission notice   *
11  * appear in the supporting documentation. The authors make no claims     *
12  * about the suitability of this software for any purpose. It is          *
13  * provided "as is" without express or implied warranty.                  *
14  **************************************************************************/
15
16 /* $Id$ */
17
18 // Container class for AliGenerator through recursion.
19 // Container is itself an AliGenerator.
20 // What is stored are not the pointers to the generators directly but to objects of type
21 // AliGenCocktail entry.   
22 // The class provides also iterator functionality.  
23 // Author: andreas.morsch@cern.ch 
24 //
25
26 #include <TList.h>
27 #include <TObjArray.h>
28
29 #include "AliGenCocktail.h"
30 #include "AliGenCocktailEntry.h"
31 #include "AliCollisionGeometry.h"
32 #include "AliRun.h"
33 #include "AliMC.h"
34 #include "AliGenCocktailEventHeader.h"
35
36 ClassImp(AliGenCocktail)
37
38 AliGenCocktail::AliGenCocktail()
39     :AliGenerator(), 
40      fNGenerators(0),
41      fRandom(kFALSE),
42      fUsePerEventRate(kFALSE),
43      fProb(0),
44      fEntries(0),
45      flnk1(0),
46      flnk2(0), 
47      fHeader(0)
48 {
49 // Constructor
50     fName = "Cocktail";
51     fTitle= "Particle Generator using cocktail of generators";
52 }
53
54 AliGenCocktail::~AliGenCocktail()
55 {
56 // Destructor
57     delete fEntries;
58     fEntries = 0;
59     //    delete fHeader; // It is removed in AliRunLoader
60     fHeader = 0;
61 }
62
63 void AliGenCocktail::
64 AddGenerator(AliGenerator *Generator, const char* Name, Float_t RateExp)
65 {
66 //
67 // Add a generator to the list 
68 // First check that list exists
69     if (!fEntries) fEntries = new TList();
70
71 //
72 //  Forward parameters to the new generator
73     if(TestBit(kPtRange) && !(Generator->TestBit(kPtRange)) && !(Generator->TestBit(kMomentumRange))) 
74         Generator->SetPtRange(fPtMin,fPtMax);
75     if(TestBit(kMomentumRange) && !(Generator->TestBit(kPtRange)) && !(Generator->TestBit(kMomentumRange)))
76         Generator->SetMomentumRange(fPMin,fPMax);
77     
78     if (!(Generator->TestBit(kYRange)))    
79         Generator->SetYRange(fYMin,fYMax);
80     if (!(Generator->TestBit(kPhiRange)))   
81         Generator->SetPhiRange(fPhiMin*180/TMath::Pi(),fPhiMax*180/TMath::Pi());
82     if (!(Generator->TestBit(kThetaRange))) 
83         Generator->SetThetaRange(fThetaMin*180/TMath::Pi(),fThetaMax*180/TMath::Pi());
84     if (!(Generator->TestBit(kVertexRange))) {
85         Generator->SetOrigin(fOrigin[0], fOrigin[1], fOrigin[2]);
86         Generator->SetSigma(fOsigma[0], fOsigma[1], fOsigma[2]);
87         Generator->SetVertexSmear(fVertexSmear);
88         Generator->SetVertexSource(kContainer);
89     }
90     Generator->SetTrackingFlag(fTrackIt);
91     Generator->SetContainer(this);
92     
93         
94 //
95 //  Add generator to list   
96     char theName[256];
97     sprintf(theName, "%s_%d",Name, fNGenerators);
98     Generator->SetName(theName);
99
100     AliGenCocktailEntry *entry = 
101         new AliGenCocktailEntry(Generator, Name, RateExp);
102     
103      fEntries->Add(entry);
104      fNGenerators++;
105      flnk1 = 0;
106      flnk2 = 0;
107      fRandom  = kFALSE;
108      fHeader  = 0;
109 }
110
111   void AliGenCocktail::Init()
112 {
113 // Initialisation
114     TIter next(fEntries);
115     AliGenCocktailEntry *entry;
116     //
117     // Loop over generators and initialize
118     while((entry = (AliGenCocktailEntry*)next())) {
119         if (fStack)  entry->Generator()->SetStack(fStack);
120         entry->Generator()->Init();
121     }  
122
123     next.Reset();
124
125     if (fRandom) {
126         fProb.Set(fNGenerators);
127         next.Reset();
128         Float_t sum = 0.;
129         while((entry = (AliGenCocktailEntry*)next())) {
130             sum += entry->Rate();
131         } 
132
133         next.Reset();
134         Int_t i = 0;
135         Float_t psum = 0.;
136         while((entry = (AliGenCocktailEntry*)next())) {
137             psum +=  entry->Rate() / sum;
138             fProb[i++] = psum;
139         }
140     }
141         next.Reset();
142 }
143
144   void AliGenCocktail::FinishRun()
145 {
146 // Initialisation
147     TIter next(fEntries);
148     AliGenCocktailEntry *entry;
149     //
150     // Loop over generators and initialize
151     while((entry = (AliGenCocktailEntry*)next())) {
152         entry->Generator()->FinishRun();
153     }  
154 }
155
156  void AliGenCocktail::Generate()
157 {
158 //
159 // Generate event 
160     TIter next(fEntries);
161     AliGenCocktailEntry *entry = 0;
162     AliGenCocktailEntry *preventry = 0;
163     AliGenerator* gen = 0;
164     if (fHeader) delete fHeader;
165
166     
167     fHeader = new AliGenCocktailEventHeader("Cocktail Header");
168
169     TObjArray *partArray = gAlice->GetMCApp()->Particles();
170
171 //
172 //  Generate the vertex position used by all generators
173 //    
174     if(fVertexSmear == kPerEvent) Vertex();
175
176     TArrayF eventVertex;
177     eventVertex.Set(3);
178     for (Int_t j=0; j < 3; j++) eventVertex[j] = fVertex[j];
179
180     if (!fRandom) {
181         //
182         // Loop over generators and generate events
183         Int_t igen=0;
184         while((entry = (AliGenCocktailEntry*)next())) {
185             if (fUsePerEventRate && (gRandom->Rndm() > entry->Rate())) continue;
186             
187             igen++;
188             if (igen ==1) {
189                 entry->SetFirst(0);
190             } else {
191                 entry->SetFirst((partArray->GetEntriesFast())+1);
192             }
193 //
194 //      Handle case in which current generator needs collision geometry from previous generator
195 //
196             gen = entry->Generator();
197             if (gen->NeedsCollisionGeometry())
198             {
199                 if (preventry && preventry->Generator()->ProvidesCollisionGeometry())
200                 {
201                     gen->SetCollisionGeometry(preventry->Generator()->CollisionGeometry());
202                 } else {
203                     Fatal("Generate()", "No Collision Geometry Provided");
204                 }
205             }
206             entry->Generator()->SetVertex(fVertex.At(0), fVertex.At(1), fVertex.At(2));
207             entry->Generator()->Generate();
208             entry->SetLast(partArray->GetEntriesFast());
209             preventry = entry;
210         }  
211     } else if (fRandom) {
212         //
213         // Select a generator randomly
214         //
215         Int_t i;
216         Float_t p0 =  gRandom->Rndm();
217
218         for (i = 0; i < fNGenerators; i++) {
219             if (p0 < fProb[i]) break;
220         }
221
222         entry = (AliGenCocktailEntry*) fEntries->At(i);
223         entry->SetFirst(0);
224         gen = entry->Generator();
225         entry->Generator()->SetVertex(fVertex.At(0), fVertex.At(1), fVertex.At(2));
226         entry->Generator()->Generate();
227         entry->SetLast(partArray->GetEntriesFast());
228     } 
229     
230     next.Reset();
231
232 // Event Vertex
233     fHeader->SetPrimaryVertex(eventVertex);
234     fHeader->CalcNProduced();
235     gAlice->SetGenEventHeader(fHeader); 
236 }
237
238 void AliGenCocktail::SetVertexSmear(VertexSmear_t smear)
239 {
240 // Set vertex smearing and propagate it to the generators
241
242   AliGenerator::SetVertexSmear(smear);
243   TIter next(fEntries);
244   while (AliGenCocktailEntry* entry = (AliGenCocktailEntry*)next()) {
245     entry->Generator()->SetVertexSmear(smear);
246   }
247 }
248
249 AliGenCocktailEntry *  AliGenCocktail::FirstGenerator()
250 {
251 // Iterator over generators: Initialisation
252     flnk1 = fEntries->FirstLink();
253     if (flnk1) {
254         return (AliGenCocktailEntry*) (flnk1->GetObject());
255     } else {
256         return 0;
257     }
258 }
259
260 AliGenCocktailEntry*  AliGenCocktail::NextGenerator()
261 {
262 // Iterator over generators: Increment
263     flnk1 = flnk1->Next();
264     if (flnk1) {
265         return (AliGenCocktailEntry*) (flnk1->GetObject());
266     } else {
267         return 0;
268     }
269 }
270
271 void AliGenCocktail::
272 FirstGeneratorPair(AliGenCocktailEntry*& e1, AliGenCocktailEntry*& e2)
273 {
274 // Iterator over generator pairs: Initialisation
275     flnk2 = flnk1 = fEntries->FirstLink();
276     if (flnk1) {
277         e2 = e1 = (AliGenCocktailEntry*) (flnk1->GetObject());
278     } else {
279         e2= e1 = 0;
280     }
281 }
282
283 void AliGenCocktail::
284 NextGeneratorPair(AliGenCocktailEntry*& e1, AliGenCocktailEntry*& e2)
285 {
286 // Iterator over generators: Increment
287     flnk2 = flnk2->Next();
288     if (flnk2) {
289         e1 = (AliGenCocktailEntry*) (flnk1->GetObject());
290         e2 = (AliGenCocktailEntry*) (flnk2->GetObject());       
291     } else {
292         flnk2 = flnk1 = flnk1->Next();
293         if (flnk1) {
294             e1 = (AliGenCocktailEntry*) (flnk1->GetObject());
295             e2 = (AliGenCocktailEntry*) (flnk2->GetObject());
296         } else {
297             e1=0;
298             e2=0;
299         }
300     }
301 }
302
303 void AliGenCocktail::AddHeader(AliGenEventHeader* header)
304 {
305 // Add a header to the list 
306     if (fHeader) fHeader->AddHeader(header);
307 }
308
309
310
311