File path is obtained in AliAnalysisManager and is sent to AliMCEventHeandler as...
[u/mrichter/AliRoot.git] / STEER / AliMCEventHandler.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 //                          Class AliMCEventHandler
19 // This class gives access to MC truth during the analysis.
20 // Monte Carlo truth is containe in the kinematics tree (produced particles) and 
21 // the tree of reference hits.
22 //      
23 // Origin: Andreas Morsch, CERN, andreas.morsch@cern.ch 
24 //---------------------------------------------------------------------------------
25
26
27
28 #include "AliMCEventHandler.h"
29 #include "AliTrackReference.h"
30 #include "AliHeader.h"
31 #include "AliStack.h"
32
33 #include <TTree.h>
34 #include <TFile.h>
35 #include <TParticle.h>
36 #include <TClonesArray.h>
37 #include <TDirectoryFile.h>
38 #include <TArrow.h>
39 #include <TMarker.h>
40 #include <TH2F.h>
41
42
43 ClassImp(AliMCEventHandler)
44
45 AliMCEventHandler::AliMCEventHandler() :
46     AliVirtualEventHandler(),
47     fFileE(0),
48     fFileK(0),
49     fFileTR(0),
50     fTmpFileTR(0),
51     fTreeE(0),
52     fTreeK(0),
53     fTreeTR(0),
54     fTmpTreeTR(0),
55     fStack(0),
56     fHeader(0),
57     fTrackReferences(0),
58     fNEvent(-1),
59     fEvent(-1),
60     fNprimaries(-1),
61     fNparticles(-1),
62     fPathName("./"),
63     fExtension(""),
64     fFileNumber(0),
65     fEventsPerFile(0)
66 {
67     // Default constructor
68 }
69
70 AliMCEventHandler::AliMCEventHandler(const char* name, const char* title) :
71     AliVirtualEventHandler(name, title),
72     fFileE(0),
73     fFileK(0),
74     fFileTR(0),
75     fTmpFileTR(0),
76     fTreeE(0),
77     fTreeK(0),
78     fTreeTR(0),
79     fTmpTreeTR(0),
80     fStack(0),
81     fHeader(new AliHeader()),
82     fTrackReferences(new TClonesArray("AliTrackReference", 200)),
83     fNEvent(-1),
84     fEvent(-1),
85     fNprimaries(-1),
86     fNparticles(-1),
87     fPathName("./"),
88     fExtension(""),
89     fFileNumber(0),
90     fEventsPerFile(0)
91 {
92     // Constructor
93 }
94 AliMCEventHandler::~AliMCEventHandler()
95
96     // Destructor
97     delete fFileE;
98     delete fFileK;
99     delete fFileTR;
100 }
101
102 Bool_t AliMCEventHandler::InitIO(Option_t* /*opt*/) 
103
104     // Initialize input
105     //
106     fFileE = new TFile(Form("%sgalice.root", fPathName));
107     if (!fFileE) printf("AliMCEventHandler:galice.root not found in directory %s ! \n", fPathName);
108
109     fFileE->GetObject("TE", fTreeE);
110     fTreeE->SetBranchAddress("Header", &fHeader);
111     fNEvent = fTreeE->GetEntries();
112     //
113     // Tree K
114     fFileK = new TFile(Form("%sKinematics%s.root", fPathName, fExtension));
115     if (!fFileK) printf("AliMCEventHandler:Kinematics.root not found in directory %s ! \n", fPathName);
116     fEventsPerFile = fFileK->GetNkeys() - fFileK->GetNProcessIDs();
117     //
118     // Tree TR
119     fFileTR = new TFile(Form("%sTrackRefs%s.root", fPathName, fExtension));
120     if (!fFileTR) printf("AliMCEventHandler:TrackRefs.root not found in directory %s ! \n", fPathName);
121     //
122     // Reset the event number
123     fEvent = -1;
124     fFileNumber = 0;
125     
126     printf("AliMCEventHandler:Number of events in this directory %5d \n", fNEvent);
127     return kTRUE;
128     
129 }
130
131 Bool_t AliMCEventHandler::GetEvent(Int_t iev)
132 {
133     // Load the event number iev
134     Int_t inew  = iev/fEventsPerFile;
135     if (inew != fFileNumber) {
136         fFileNumber = inew;
137         if (!OpenFile(fFileNumber)){
138             return kFALSE;
139         }
140     }
141     
142     char folder[20];
143     sprintf(folder, "Event%d", iev);
144     // TreeE
145     fTreeE->GetEntry(iev);
146     fStack = fHeader->Stack();
147     // Tree K
148     TDirectoryFile* dirK  = 0;
149     fFileK->GetObject(folder, dirK);
150     if (!dirK) {
151         printf("AliMCEventHandler: Event #%5d not found\n", iev);
152         return kFALSE;
153     }
154     dirK->GetObject("TreeK", fTreeK);
155     fStack->ConnectTree(fTreeK);
156     fStack->GetEvent();
157     //Tree TR 
158     TDirectoryFile* dirTR = 0;
159     fFileTR->GetObject(folder, dirTR);
160     dirTR->GetObject("TreeTR", fTreeTR);
161     if (fTreeTR->GetBranch("AliRun")) {
162         // This is an old format with one branch per detector not in synch with TreeK
163         ReorderAndExpandTreeTR();
164     } else {
165         // New format 
166         fTreeTR->SetBranchAddress("TrackReferences", &fTrackReferences);
167     }
168         
169     //
170     fNparticles = fStack->GetNtrack();
171     fNprimaries = fStack->GetNprimary();
172     printf("AliMCEventHandler: Event#: %5d Number of particles %5d \n", fEvent, fNparticles);
173     
174     return kTRUE;
175
176 }
177
178 Bool_t AliMCEventHandler::OpenFile(Int_t i)
179 {
180     // Open file i
181     Bool_t ok = kTRUE;
182     if (i > 0) {
183         fExtension = Form("%d", i);
184     } else {
185         fExtension = "";
186     }
187     
188     
189     delete fFileK;
190     fFileK = new TFile(Form("%sKinematics%s.root", fPathName, fExtension));
191     if (!fFileK) {
192         printf("AliMCEventHandler:Kinematics%s.root not found in directory %s ! \n", fExtension, fPathName);
193         ok = kFALSE;
194     }
195     
196     delete fFileTR;
197     fFileTR = new TFile(Form("%sTrackRefs%s.root", fPathName, fExtension));
198     if (!fFileTR) {
199         printf("AliMCEventHandler:TrackRefs%s.root not found in directory %s ! \n", fExtension, fPathName);
200         ok = kFALSE;
201     }
202     
203     return ok;
204 }
205
206 Bool_t AliMCEventHandler::BeginEvent()
207
208     // Read the next event
209     fEvent++;
210     if (fEvent >= fNEvent) {
211         printf("AliMCEventHandler: Event number out of range %5d\n", fEvent);
212         return kFALSE;
213     }
214     return GetEvent(fEvent);
215 }
216
217 Int_t AliMCEventHandler::GetParticleAndTR(Int_t i, TParticle*& particle, TClonesArray*& trefs)
218 {
219     // Retrieve entry i
220     if (i > -1 && i < fNparticles) {
221         fTreeTR->GetEntry(fStack->TreeKEntry(i));
222     } else {
223         printf("AliMCEventHandler::GetEntry: Index out of range");
224     }
225     particle = fStack->Particle(i);
226     trefs    = fTrackReferences;
227     return trefs->GetEntries();
228 }
229
230 void AliMCEventHandler::DrawCheck(Int_t i, Bool_t search)
231 {
232     // Retrieve entry i and draw momentum vector and hits
233     if (i > -1 && i < fNparticles) {
234         fTreeTR->GetEntry(fStack->TreeKEntry(i));
235     } else {
236         printf("AliMCEventHandler::GetEntry: Index out of range");
237     }
238
239     Int_t nh = fTrackReferences->GetEntries();
240     
241     
242     if (search) {
243         while(nh == 0 && i < fNparticles - 1) {
244             i++;
245             fTreeTR->GetEntry(fStack->TreeKEntry(i));
246             nh =  fTrackReferences->GetEntries();
247         }
248         printf("Found Hits at %5d\n", i);
249     }
250     TParticle* particle = fStack->Particle(i);
251
252     TH2F*    h = new TH2F("", "", 100, -500, 500, 100, -500, 500);
253     Float_t x0 = particle->Vx();
254     Float_t y0 = particle->Vy();
255
256     Float_t x1 = particle->Vx() + particle->Px() * 50.;
257     Float_t y1 = particle->Vy() + particle->Py() * 50.;
258     
259     TArrow*  a = new TArrow(x0, y0, x1, y1, 0.01);
260     h->Draw();
261     a->SetLineColor(2);
262     
263     a->Draw();
264     
265     for (Int_t ih = 0; ih < nh; ih++) {
266         AliTrackReference* ref = (AliTrackReference*) fTrackReferences->At(ih);
267         TMarker* m = new TMarker(ref->X(), ref->Y(), 20);
268         m->Draw();
269         m->SetMarkerSize(0.4);
270         
271     }
272 }
273
274 Bool_t AliMCEventHandler::Notify(const char *path)
275 {
276     // Notify about directory change
277     // The directory is taken from the 'path' argument
278     // Reconnect trees
279
280     printf("AliMCEventHandler::Notify() file: %s\n", path);
281     fPathName = Form("%s",  path);
282     ResetIO();
283     InitIO("");
284     return kTRUE;
285 }
286     
287 void AliMCEventHandler::ResetIO()
288 {
289     // Reset files
290     if (fFileE)  delete fFileE;
291     if (fFileK)  delete fFileK;
292     if (fFileTR) delete fFileTR;
293 }
294
295                             
296 Bool_t AliMCEventHandler::FinishEvent()
297 {
298     // Dummy 
299     return kTRUE;
300 }
301
302 Bool_t AliMCEventHandler::Terminate()
303
304     // Dummy 
305     return kTRUE;
306 }
307
308 Bool_t AliMCEventHandler::TerminateIO()
309
310     // Dummy
311     return kTRUE;
312 }
313     
314 void AliMCEventHandler::ReorderAndExpandTreeTR()
315 {
316 //
317 //  Reorder and expand the track reference tree in order to match the kinematics tree.
318 //  Copy the information from different branches into one
319 //
320 //  TreeTR
321     if (fTmpTreeTR) delete fTmpTreeTR;
322     if (fTmpFileTR) {
323         fTmpFileTR->Close();
324         delete fTmpFileTR;
325     }
326
327     fTmpFileTR = new TFile("TrackRefsTmp.root", "recreate");
328     fTmpTreeTR = new TTree("TreeTR", "Track References");
329     if (!fTrackReferences)  fTrackReferences = new TClonesArray("AliTrackReference", 100);
330     fTmpTreeTR->Branch("TrackReferences", "TClonesArray", &fTrackReferences, 4000);
331 //
332     TClonesArray* trefs[7];
333     for (Int_t i = 0; i < 7; i++) trefs[i] = 0;
334     if (fTreeTR){
335         // make branch for central track references
336         if (fTreeTR->GetBranch("AliRun")) fTreeTR->SetBranchAddress("AliRun", &trefs[0]);
337         if (fTreeTR->GetBranch("ITS"))    fTreeTR->SetBranchAddress("ITS",    &trefs[1]);
338         if (fTreeTR->GetBranch("TPC"))    fTreeTR->SetBranchAddress("TPC",    &trefs[2]);
339         if (fTreeTR->GetBranch("TRD"))    fTreeTR->SetBranchAddress("TRD",    &trefs[3]);
340         if (fTreeTR->GetBranch("TOF"))    fTreeTR->SetBranchAddress("TOF",    &trefs[4]);
341         if (fTreeTR->GetBranch("FRAME"))  fTreeTR->SetBranchAddress("FRAME",  &trefs[5]);
342         if (fTreeTR->GetBranch("MUON"))   fTreeTR->SetBranchAddress("MUON",   &trefs[6]);
343     }
344
345     Int_t np = fStack->GetNprimary();
346     Int_t nt = fTreeTR->GetEntries();
347     //
348     // Loop over tracks and find the secondaries with the help of the kine tree
349     Int_t ifills = 0;
350     Int_t it     = 0;
351     Int_t itlast = 0;
352     
353     for (Int_t ip = np - 1; ip > -1; ip--) {
354         TParticle *part = fStack->Particle(ip);
355 //      printf("Particle %5d %5d %5d %5d %5d %5d \n", 
356 //             ip, part->GetPdgCode(), part->GetFirstMother(), part->GetFirstDaughter(), 
357 //             part->GetLastDaughter(), part->TestBit(kTransportBit));
358
359         // Determine range of secondaries produced by this primary during transport     
360         Int_t dau1  = part->GetFirstDaughter();
361         if (dau1 < np) continue;  // This particle has no secondaries produced during transport
362         Int_t dau2  = -1;
363         if (dau1 > -1) {
364             Int_t inext = ip - 1;
365             while (dau2 < 0) {
366                 if (inext >= 0) {
367                     part = fStack->Particle(inext);
368                     dau2 =  part->GetFirstDaughter();
369                     if (dau2 == -1 || dau2 < np) {
370                         dau2 = -1;
371                     } else {
372                         dau2--;
373                     }
374                 } else {
375                     dau2 = fStack->GetNtrack() - 1;
376                 }
377                 inext--;
378             } // find upper bound
379         }  // dau2 < 0
380         
381
382 //      printf("Check (1) %5d %5d %5d %5d %5d \n", ip, np, it, dau1, dau2);
383 //
384 // Loop over reference hits and find secondary label
385 // First the tricky part: find the entry in treeTR than contains the hits or
386 // make sure that no hits exist.
387 //
388         Bool_t hasHits   = kFALSE;
389         Bool_t isOutside = kFALSE;
390
391         it = itlast;
392         while (!hasHits && !isOutside && it < nt) {
393             fTreeTR->GetEntry(it++);
394             for (Int_t ib = 0; ib < 7; ib++) {
395                 if (!trefs[ib]) continue;
396                 Int_t nh = trefs[ib]->GetEntries();
397                 for (Int_t ih = 0; ih < nh; ih++) {
398                     AliTrackReference* tr = (AliTrackReference*) trefs[ib]->At(ih);
399                     Int_t label = tr->Label();
400                     if (label >= dau1 && label <= dau2) {
401                         hasHits = kTRUE;
402                         itlast = it - 1;
403                         break;
404                     }
405                     if (label > dau2 || label < ip) {
406                         isOutside = kTRUE;
407                         itlast = it - 1;
408                         break;
409                     }
410                 } // hits
411                 if (hasHits || isOutside) break;
412             } // branches
413         } // entries
414
415         if (!hasHits) {
416             // Write empty entries
417             for (Int_t id = dau1; (id <= dau2); id++) {
418                 fTmpTreeTR->Fill();
419                 ifills++;
420             } 
421         } else {
422             // Collect all hits
423             fTreeTR->GetEntry(itlast);
424             for (Int_t id = dau1; (id <= dau2) && (dau1 > -1); id++) {
425                 for (Int_t ib = 0; ib < 7; ib++) {
426                     if (!trefs[ib]) continue;
427                     Int_t nh = trefs[ib]->GetEntries();
428                     for (Int_t ih = 0; ih < nh; ih++) {
429                         AliTrackReference* tr = (AliTrackReference*) trefs[ib]->At(ih);
430                         Int_t label = tr->Label();
431                         // Skip primaries
432                         if (label == ip) continue;
433                         if (label > dau2 || label < dau1) 
434                             printf("AliMCEventHandler::Track Reference Label out of range !: %5d %5d %5d %5d \n", 
435                                    itlast, label, dau1, dau2);
436                         if (label == id) {
437                             // secondary found
438                             tr->SetDetectorId(ib-1);
439                             Int_t nref =  fTrackReferences->GetEntriesFast();
440                             TClonesArray &lref = *fTrackReferences;
441                             new(lref[nref]) AliTrackReference(*tr);
442                         }
443                     } // hits
444                 } // branches
445                 fTmpTreeTR->Fill();
446                 fTrackReferences->Clear();
447                 ifills++;
448             } // daughters
449         } // has hits
450     } // tracks
451     //
452     // Now loop again and write the primaries
453     //
454     it = nt - 1;
455     for (Int_t ip = 0; ip < np; ip++) {
456         Int_t labmax = -1;
457         while (labmax < ip && it > -1) {
458             fTreeTR->GetEntry(it--);
459             for (Int_t ib = 0; ib < 7; ib++) {
460                 if (!trefs[ib]) continue;
461                 Int_t nh = trefs[ib]->GetEntries();
462                 // 
463                 // Loop over reference hits and find primary labels
464                 for (Int_t ih = 0; ih < nh; ih++) {
465                     AliTrackReference* tr = (AliTrackReference*)  trefs[ib]->At(ih);
466                     Int_t label = tr->Label();
467                     if (label < np && label > labmax) {
468                         labmax = label;
469                     }
470                     
471                     if (label == ip) {
472                         tr->SetDetectorId(ib-1);
473                         Int_t nref = fTrackReferences->GetEntriesFast();
474                         TClonesArray &lref = *fTrackReferences;
475                         new(lref[nref]) AliTrackReference(*tr);
476                     }
477                 } // hits
478             } // branches
479         } // entries
480         it++;
481         fTmpTreeTR->Fill();
482         fTrackReferences->Clear();
483         ifills++;
484     } // tracks
485     // Check
486     if (ifills != fStack->GetNtrack()) 
487         printf("AliMCEventHandler:Number of entries in TreeTR (%5d) unequal to TreeK (%5d) \n", 
488                ifills, fStack->GetNtrack());
489     fTreeTR = fTmpTreeTR;
490 }