bug fixes
[u/mrichter/AliRoot.git] / EVGEN / AliGenExtFile.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 // Event generator that using an instance of type AliGenReader
19 // reads particles from a file and applies cuts.
20 // Example: In your Config.C you can include the following lines
21 //  AliGenExtFile *gener = new AliGenExtFile(-1);
22 //  gener->SetMomentumRange(0,999);
23 //  gener->SetPhiRange(-180.,180.);
24 //  gener->SetThetaRange(0,180);
25 //  gener->SetYRange(-999,999);
26 //  AliGenReaderTreeK * reader = new AliGenReaderTreeK();
27 //  reader->SetFileName("myFileWithTreeK.root");
28 //  gener->SetReader(reader);
29 //  gener->Init();
30
31
32 #include <Riostream.h>
33
34 #include "AliLog.h"
35 #include "AliGenExtFile.h"
36 #include "AliRunLoader.h"
37 #include "AliHeader.h"
38 #include "AliStack.h"
39 #include "AliGenEventHeader.h"
40 #include "AliGenReader.h"
41
42 #include <TParticle.h>
43 #include <TFile.h>
44 #include <TTree.h>
45
46
47 using std::cout;
48 using std::endl;
49 using std::map;
50
51 ClassImp(AliGenExtFile)
52
53 AliGenExtFile::AliGenExtFile()
54     :AliGenMC(),
55      fFileName(0),
56      fReader(0),
57      fStartEvent(0)
58 {
59 //  Constructor
60 //
61 //  Read all particles
62     fNpart  = -1;
63 }
64
65 AliGenExtFile::AliGenExtFile(Int_t npart)
66     :AliGenMC(npart),
67      fFileName(0),
68      fReader(0),
69      fStartEvent(0)
70 {
71 //  Constructor
72     fName   = "ExtFile";
73     fTitle  = "Primaries from ext. File";
74 }
75
76 //____________________________________________________________
77 AliGenExtFile::~AliGenExtFile()
78 {
79 // Destructor
80     delete fReader;
81 }
82
83 //___________________________________________________________
84 void AliGenExtFile::Init()
85 {
86 // Initialize
87     if (fReader) fReader->Init();
88 }
89
90 //___________________________________________________________
91 void AliGenExtFile::Generate()
92 {
93 // Generate particles
94
95   Double_t polar[3]  = {0,0,0};
96   //
97   Double_t origin[3] = {0,0,0};
98   Double_t time = 0.;
99   Double_t p[4];
100   Float_t random[6];
101   Int_t i = 0, j, nt;
102   //
103   //
104   if (fVertexSmear == kPerEvent) Vertex();
105
106   // Fast forward up to start Event
107   for (Int_t ie=0; ie<fStartEvent; ++ie ) {
108     Int_t nTracks = fReader->NextEvent();       
109     if (nTracks == 0) {
110       // printf("\n No more events !!! !\n");
111       Warning("AliGenExtFile::Generate","\nNo more events in external file!!!\nLast event may be empty or incomplete.\n");
112       return;
113     }
114     for (Int_t i=0; i<nTracks; ++i) {
115       if (NULL == fReader->NextParticle())
116         AliFatal("Error while skipping tracks");
117     }
118     cout << "Skipping event " << ie << endl;
119   }  
120   fStartEvent = 0; // do not skip events the second time 
121
122   while(1) {
123     Int_t nTracks = fReader->NextEvent();       
124     if (nTracks == 0) {
125       // printf("\n No more events !!! !\n");
126       Warning("AliGenExtFile::Generate","\nNo more events in external file!!!\nLast event may be empty or incomplete.\n");
127       return;
128     }
129
130     //
131     // Particle selection loop
132     //
133     // The selection criterium for the external file generator is as follows:
134     //
135     // 1) All tracks are subject to the cuts defined by AliGenerator, i.e.
136     //    fThetaMin, fThetaMax, fPhiMin, fPhiMax, fPMin, fPMax, fPtMin, fPtMax,
137     //    fYMin, fYMax.
138     //    If the particle does not satisfy these cuts, it is not put on the
139     //    stack.
140     // 2) If fCutOnChild and some specific child is selected (e.g. if
141     //    fForceDecay==kSemiElectronic) the event is rejected if NOT EVEN ONE
142     //    child falls into the child-cuts.
143     TParticle* iparticle = 0x0;
144     
145     if(fCutOnChild) {
146       // Count the selected children
147       Int_t nSelected = 0;
148       while ((iparticle=fReader->NextParticle()) ) {
149           Int_t kf = CheckPDGCode(iparticle->GetPdgCode());
150           kf = TMath::Abs(kf);
151           if (ChildSelected(kf) && KinematicSelection(iparticle, 1)) {
152               nSelected++;
153           }
154       }
155       if (!nSelected) continue;    // No particle selected:  Go to next event
156       fReader->RewindEvent();
157     }
158
159     //
160     // Stack selection loop
161     //
162     class SelectorLogic { // need to do recursive back tracking, requires a "nested" function
163     private:
164        Int_t idCount;
165        map<Int_t,Int_t> firstMotherMap;
166        map<Int_t,Int_t> secondMotherMap;
167        map<Int_t,Bool_t> selectedIdMap;
168        map<Int_t,Int_t> newIdMap;
169        void selectMothersToo(Int_t particleId) {
170           Int_t mum1 = firstMotherMap[particleId];
171           if (mum1 > -1 && !selectedIdMap[mum1]) {
172              selectedIdMap[mum1] = true;
173              selectMothersToo(mum1);
174           }
175           Int_t mum2 = secondMotherMap[particleId];
176           if (mum2 > -1 && !selectedIdMap[mum2]) {
177              selectedIdMap[mum2] = true;
178              selectMothersToo(mum2);
179           }
180        }
181     public:
182        void init() {
183           idCount = 0;
184        }
185        void setData(Int_t id, Int_t mum1, Int_t mum2, Bool_t selected) {
186           idCount++; // we know that this function is called in succession of ids, so counting is fine to determine max id
187           firstMotherMap[id] = mum1;
188           secondMotherMap[id] = mum2;
189           selectedIdMap[id] = selected;
190        }
191        void reselectCuttedMothersAndRemapIDs() {
192           for (Int_t id = 0; id < idCount; ++id) {
193              if (selectedIdMap[id]) {
194                 selectMothersToo(id);
195              }
196           }
197           Int_t newId = 0;
198           for (Int_t id = 0; id < idCount; id++) {
199              if (selectedIdMap[id]) {
200                 newIdMap[id] = newId; ++newId;
201              } else {
202                 newIdMap[id] = -1;
203              }
204           }
205        }
206        Bool_t isSelected(Int_t id) {
207           return selectedIdMap[id];
208        }
209        Int_t newId(Int_t id) {
210           if (id == -1) return -1;
211           return newIdMap[id];
212        }
213     };
214     SelectorLogic selector;
215     selector.init();
216     for (Int_t i = 0; i < nTracks; i++) {
217        TParticle* jparticle = fReader->NextParticle();
218        selector.setData(i,
219              jparticle->GetFirstMother(),
220              jparticle->GetSecondMother(),
221              KinematicSelection(jparticle,0));
222     }
223     selector.reselectCuttedMothersAndRemapIDs();
224     fReader->RewindEvent();
225
226     //
227     // Stack filling loop
228     //
229     fNprimaries = 0;
230     for (i = 0; i < nTracks; i++) {
231        TParticle* jparticle = fReader->NextParticle();
232        Bool_t selected = selector.isSelected(i);
233        if (!selected) {
234           continue;
235        }
236        Int_t parent = selector.newId(jparticle->GetFirstMother());
237 //       printf("particle %d -> %d, with mother %d -> %d\n", i, selector.newId(i), jparticle->GetFirstMother(), parent);
238
239         p[0] = jparticle->Px();
240         p[1] = jparticle->Py();
241         p[2] = jparticle->Pz();
242         p[3] = jparticle->Energy();
243         
244         Int_t idpart = jparticle->GetPdgCode();
245         if(fVertexSmear==kPerTrack) 
246         {
247             Rndm(random,6);
248             for (j = 0; j < 3; j++) {
249                 origin[j]=fOrigin[j]+
250                     fOsigma[j]*TMath::Cos(2*random[2*j]*TMath::Pi())*
251                     TMath::Sqrt(-2*TMath::Log(random[2*j+1]));
252             }
253             Rndm(random,2);
254             time = fTimeOrigin + fOsigma[2]/TMath::Ccgs()*
255               TMath::Cos(2*random[0]*TMath::Pi())*
256               TMath::Sqrt(-2*TMath::Log(random[1]));
257         } else {
258             origin[0] = fVertex[0] + jparticle->Vx();
259             origin[1] = fVertex[1] + jparticle->Vy();
260             origin[2] = fVertex[2] + jparticle->Vz();
261             time = fTime + jparticle->T();
262         }
263         Int_t doTracking = fTrackIt && selected && (jparticle->TestBit(kTransportBit));
264         
265         PushTrack(doTracking, parent, idpart,
266                   p[0], p[1], p[2], p[3], origin[0], origin[1], origin[2], time,
267                   polar[0], polar[1], polar[2],
268                   kPPrimary, nt, 1., jparticle->GetStatusCode());
269
270         KeepTrack(nt);
271         fNprimaries++;
272     } // track loop
273
274     // Generated event header
275     AliGenEventHeader * header = fReader->GetGenEventHeader();
276     if (!header) header = new AliGenEventHeader();
277     header->SetNProduced(fNprimaries);
278     header->SetPrimaryVertex(fVertex);
279     header->SetInteractionTime(fTime);
280     AddHeader(header);
281     break;
282     
283   } // event loop
284   
285   SetHighWaterMark(nt);
286   CdEventFile();
287 }
288
289 //___________________________________________________________
290 void AliGenExtFile::CdEventFile()
291 {
292 // CD back to the event file
293   AliRunLoader::Instance()->CdGAFile();
294 }
295
296
297
298
299