[bug #84578] Request to extend AliGenBox for using YRange
[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 #include <TFormula.h>
29
30 #include "AliGenCocktail.h"
31 #include "AliGenCocktailEntry.h"
32 #include "AliCollisionGeometry.h"
33 #include "AliRun.h"
34 #include "AliLog.h"
35 #include "AliMC.h"
36 #include "AliGenCocktailEventHeader.h"
37
38 ClassImp(AliGenCocktail)
39
40 AliGenCocktail::AliGenCocktail()
41     :AliGenerator(), 
42      fNGenerators(0),
43      fTotalRate(0.),
44      fRandom(kFALSE),
45      fUsePerEventRate(kFALSE),
46      fProb(0),
47      fEntries(0),
48      flnk1(0),
49      flnk2(0), 
50      fHeader(0)
51 {
52 // Constructor
53     fName = "Cocktail";
54     fTitle= "Particle Generator using cocktail of generators";
55 }
56
57 AliGenCocktail::~AliGenCocktail()
58 {
59 // Destructor
60     delete fEntries;
61     fEntries = 0;
62     //    delete fHeader; // It is removed in AliRunLoader
63     fHeader = 0;
64 }
65
66 void AliGenCocktail::
67 AddGenerator(AliGenerator *Generator, const char* Name, Float_t RateExp, TFormula* formula)
68 {
69 //
70 // Add a generator to the list 
71 // First check that list exists
72     if (!fEntries) fEntries = new TList();
73     fTotalRate += RateExp;
74 //
75 //  Forward parameters to the new generator
76     if(TestBit(kPtRange) && !(Generator->TestBit(kPtRange)) && !(Generator->TestBit(kMomentumRange))) 
77         Generator->SetPtRange(fPtMin,fPtMax);
78     if(TestBit(kMomentumRange) && !(Generator->TestBit(kPtRange)) && !(Generator->TestBit(kMomentumRange)))
79         Generator->SetMomentumRange(fPMin,fPMax);
80     
81     if (TestBit(kYRange) && !(Generator->TestBit(kYRange)))
82         Generator->SetYRange(fYMin,fYMax);
83     if (TestBit(kPhiRange) && !(Generator->TestBit(kPhiRange)))
84         Generator->SetPhiRange(fPhiMin*180/TMath::Pi(),fPhiMax*180/TMath::Pi());
85     if (TestBit(kThetaRange) && !(Generator->TestBit(kThetaRange)) && !(Generator->TestBit(kEtaRange)))
86         Generator->SetThetaRange(fThetaMin*180/TMath::Pi(),fThetaMax*180/TMath::Pi());
87     if (TestBit(kVertexRange) && !(Generator->TestBit(kVertexRange))) {
88         Generator->SetOrigin(fOrigin[0], fOrigin[1], fOrigin[2]);
89         Generator->SetSigma(fOsigma[0], fOsigma[1], fOsigma[2]);
90         Generator->SetVertexSmear(fVertexSmear);
91         Generator->SetVertexSource(kContainer);
92     }
93     Generator->SetTrackingFlag(fTrackIt);
94     Generator->SetContainer(this);
95
96         
97 //
98 //  Add generator to list   
99     char theName[256];
100     snprintf(theName, 256, "%s_%d",Name, fNGenerators);
101     Generator->SetName(theName);
102
103     AliGenCocktailEntry *entry = 
104         new AliGenCocktailEntry(Generator, Name, RateExp);
105     if (formula) entry->SetFormula(formula);    
106      fEntries->Add(entry);
107      fNGenerators++;
108      flnk1 = 0;
109      flnk2 = 0;
110      fRandom  = kFALSE;
111      fHeader  = 0;
112 }
113
114   void AliGenCocktail::Init()
115 {
116 // Initialisation
117     TIter next(fEntries);
118     AliGenCocktailEntry *entry;
119     //
120     // Loop over generators and initialize
121     while((entry = (AliGenCocktailEntry*)next())) {
122         if (fStack)  entry->Generator()->SetStack(fStack);
123         entry->Generator()->Init();
124     }  
125
126     next.Reset();
127
128     if (fRandom) {
129         fProb.Set(fNGenerators);
130         next.Reset();
131         Float_t sum = 0.;
132         while((entry = (AliGenCocktailEntry*)next())) {
133             sum += entry->Rate();
134         } 
135
136         next.Reset();
137         Int_t i = 0;
138         Float_t psum = 0.;
139         while((entry = (AliGenCocktailEntry*)next())) {
140             psum +=  entry->Rate() / sum;
141             fProb[i++] = psum;
142         }
143     }
144         next.Reset();
145 }
146
147   void AliGenCocktail::FinishRun()
148 {
149 // Initialisation
150     TIter next(fEntries);
151     AliGenCocktailEntry *entry;
152     //
153     // Loop over generators and initialize
154     while((entry = (AliGenCocktailEntry*)next())) {
155         entry->Generator()->FinishRun();
156     }  
157 }
158
159  void AliGenCocktail::Generate()
160 {
161 //
162 // Generate event 
163     TIter next(fEntries);
164     AliGenCocktailEntry *entry = 0;
165     AliGenCocktailEntry *preventry = 0;
166     AliGenCocktailEntry *collentry = 0;
167     AliGenerator* gen = 0;
168     if (fHeader) delete fHeader;
169
170     
171     fHeader = new AliGenCocktailEventHeader("Cocktail Header");
172
173     const TObjArray *partArray = gAlice->GetMCApp()->Particles();
174
175 //
176 //  Generate the vertex position used by all generators
177 //    
178     if(fVertexSmear == kPerEvent) Vertex();
179
180     TArrayF eventVertex;
181     eventVertex.Set(3);
182     for (Int_t j=0; j < 3; j++) eventVertex[j] = fVertex[j];
183
184     if (!fRandom) {
185         //
186         // Loop over generators and generate events
187         Int_t igen   = 0;
188         while((entry = (AliGenCocktailEntry*)next())) {
189           Int_t ntimes = 1;
190           if (fUsePerEventRate && (gRandom->Rndm() > entry->Rate())) continue;
191           
192           igen++;
193           if (igen ==1) {
194             entry->SetFirst(0);
195           } else {
196             entry->SetFirst((partArray->GetEntriesFast())+1);
197           }
198           gen = entry->Generator();
199           if (gen->ProvidesCollisionGeometry()) collentry = entry; 
200           //
201           //      Handle case in which current generator needs collision geometry from previous generator
202           //
203           if (gen->NeedsCollisionGeometry() && (entry->Formula() == 0))
204             {
205               if (preventry && preventry->Generator()->ProvidesCollisionGeometry())
206                 {
207                   gen->SetCollisionGeometry(preventry->Generator()->CollisionGeometry());
208                 } else {
209                 Fatal("Generate()", "No Collision Geometry Provided");
210               }
211             }
212           //
213           //      Number of signals is calculated from Collision Geometry
214           //      and entry with given centrality bin is selected
215           //
216           if (entry->Formula() != 0)
217             {
218               if (!collentry) {
219                 Fatal("Generate()", "No Collision Geometry Provided");
220                 return;
221               }
222               AliCollisionGeometry* coll = (collentry->Generator())->CollisionGeometry();
223               Float_t b  = coll->ImpactParameter();
224               Int_t nsig = Int_t(entry->Formula()->Eval(b));
225               Int_t bin = entry->Bin() - 100;
226               if (bin > 0) {
227                 if (bin != nsig) continue;
228               } else {
229                 if (nsig < 1) nsig = 1;
230                 AliInfo(Form("Signal Events %13.3f %5d %5d\n", b, coll->HardScatters(), nsig));
231                 ntimes = nsig;
232               }
233             }
234           
235           gen->SetVertex(fVertex.At(0), fVertex.At(1), fVertex.At(2));
236           
237           for (Int_t i = 0; i < ntimes; i++) gen->Generate();
238           entry->SetLast(partArray->GetEntriesFast());
239           preventry = entry;
240         }
241     } else if (fRandom) {
242         //
243         // Select a generator randomly
244         //
245         Int_t i;
246         Float_t p0 =  gRandom->Rndm();
247
248         for (i = 0; i < fNGenerators; i++) {
249             if (p0 < fProb[i]) break;
250         }
251
252         entry = (AliGenCocktailEntry*) fEntries->At(i);
253         entry->SetFirst(0);
254         gen = entry->Generator();
255         gen->SetVertex(fVertex.At(0), fVertex.At(1), fVertex.At(2));
256         gen->Generate();
257         entry->SetLast(partArray->GetEntriesFast());
258     } 
259     
260     next.Reset();
261
262 // Event Vertex
263     fHeader->SetPrimaryVertex(eventVertex);
264     fHeader->CalcNProduced();
265     gAlice->SetGenEventHeader(fHeader); 
266 }
267
268 void AliGenCocktail::SetVertexSmear(VertexSmear_t smear)
269 {
270 // Set vertex smearing and propagate it to the generators
271
272   AliGenerator::SetVertexSmear(smear);
273   TIter next(fEntries);
274   while (AliGenCocktailEntry* entry = (AliGenCocktailEntry*)next()) {
275     entry->Generator()->SetVertexSmear(smear);
276   }
277 }
278
279 AliGenCocktailEntry *  AliGenCocktail::FirstGenerator()
280 {
281 // Iterator over generators: Initialisation
282     flnk1 = fEntries->FirstLink();
283     if (flnk1) {
284         return (AliGenCocktailEntry*) (flnk1->GetObject());
285     } else {
286         return 0;
287     }
288 }
289
290 AliGenCocktailEntry*  AliGenCocktail::NextGenerator()
291 {
292 // Iterator over generators: Increment
293     flnk1 = flnk1->Next();
294     if (flnk1) {
295         return (AliGenCocktailEntry*) (flnk1->GetObject());
296     } else {
297         return 0;
298     }
299 }
300
301 void AliGenCocktail::
302 FirstGeneratorPair(AliGenCocktailEntry*& e1, AliGenCocktailEntry*& e2)
303 {
304 // Iterator over generator pairs: Initialisation
305     flnk2 = flnk1 = fEntries->FirstLink();
306     if (flnk1) {
307         e2 = e1 = (AliGenCocktailEntry*) (flnk1->GetObject());
308     } else {
309         e2= e1 = 0;
310     }
311 }
312
313 void AliGenCocktail::
314 NextGeneratorPair(AliGenCocktailEntry*& e1, AliGenCocktailEntry*& e2)
315 {
316 // Iterator over generators: Increment
317     flnk2 = flnk2->Next();
318     if (flnk2) {
319         e1 = (AliGenCocktailEntry*) (flnk1->GetObject());
320         e2 = (AliGenCocktailEntry*) (flnk2->GetObject());       
321     } else {
322         flnk2 = flnk1 = flnk1->Next();
323         if (flnk1) {
324             e1 = (AliGenCocktailEntry*) (flnk1->GetObject());
325             e2 = (AliGenCocktailEntry*) (flnk2->GetObject());
326         } else {
327             e1=0;
328             e2=0;
329         }
330     }
331 }
332
333 void AliGenCocktail::AddHeader(AliGenEventHeader* header)
334 {
335 // Add a header to the list 
336     if (fHeader) fHeader->AddHeader(header);
337 }
338
339
340
341