e64cbf6e3dfe9307b0f161e966b4dfd24855ae35
[u/mrichter/AliRoot.git] / STEER / STEERBase / 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 "AliMCEvent.h"
30 #include "AliMCParticle.h"
31 #include "AliPDG.h"
32 #include "AliTrackReference.h"
33 #include "AliHeader.h"
34 #include "AliStack.h"
35 #include "AliLog.h"
36
37 #include <TTree.h>
38 #include <TSystem.h>
39 #include <TTreeCache.h>
40 #include <TFile.h>
41 #include <TList.h>
42 #include <TParticle.h>
43 #include <TString.h>
44 #include <TClonesArray.h>
45 #include <TDirectoryFile.h>
46
47 ClassImp(AliMCEventHandler)
48
49 AliMCEventHandler::AliMCEventHandler() :
50     AliInputEventHandler(),
51     fMCEvent(new AliMCEvent()),
52     fFileE(0),
53     fFileK(0),
54     fFileTR(0),
55     fTreeE(0),
56     fTreeK(0),
57     fTreeTR(0),
58     fDirK(0),
59     fDirTR(0),
60     fParticleSelected(0),
61     fLabelMap(0),
62     fNEvent(-1),
63     fEvent(-1),
64     fPathName(new TString("./")),
65     fkExtension(""),
66     fFileNumber(0),
67     fEventsPerFile(0),
68     fReadTR(kTRUE),
69     fInitOk(kFALSE),
70     fSubsidiaryHandlers(0),
71     fEventsInContainer(0),
72     fPreReadMode(kNoPreRead),
73     fCacheSize(0),
74     fCacheTK(0),
75     fCacheTR(0)
76 {
77   //
78   // Default constructor
79   //
80   // Be sure to add all particles to the PDG database
81   AliPDG::AddParticlesToPdgDataBase();
82 }
83
84 AliMCEventHandler::AliMCEventHandler(const char* name, const char* title) :
85     AliInputEventHandler(name, title),
86     fMCEvent(new AliMCEvent()),
87     fFileE(0),
88     fFileK(0),
89     fFileTR(0),
90     fTreeE(0),
91     fTreeK(0),
92     fTreeTR(0),
93     fDirK(0),
94     fDirTR(0),
95     fParticleSelected(0),
96     fLabelMap(0),
97     fNEvent(-1),
98     fEvent(-1),
99     fPathName(new TString("./")),
100     fkExtension(""),
101     fFileNumber(0),
102     fEventsPerFile(0),
103     fReadTR(kTRUE),
104     fInitOk(kFALSE),
105     fSubsidiaryHandlers(0),
106     fEventsInContainer(0),
107     fPreReadMode(kNoPreRead),
108     fCacheSize(0),
109     fCacheTK(0),
110     fCacheTR(0)
111 {
112   //
113   // Constructor
114   //
115   // Be sure to add all particles to the PDG database
116   AliPDG::AddParticlesToPdgDataBase();
117 }
118 AliMCEventHandler::~AliMCEventHandler()
119
120     // Destructor
121   delete fPathName;
122     delete fMCEvent;
123     delete fFileE;
124     delete fFileK;
125     delete fFileTR;
126     delete fCacheTK;
127     delete fCacheTR;
128 }
129
130 Bool_t AliMCEventHandler::Init(Option_t* opt)
131
132     // Initialize input
133     //
134     if (!(strcmp(opt, "proof")) || !(strcmp(opt, "local"))) return kTRUE;
135     //
136     fFileE = TFile::Open(Form("%sgalice.root", fPathName->Data()));
137     if (!fFileE) {
138         AliError(Form("AliMCEventHandler:galice.root not found in directory %s ! \n", fPathName->Data()));
139         fInitOk = kFALSE;
140         return kFALSE;
141     }
142     
143     //
144     // Tree E
145     fFileE->GetObject("TE", fTreeE);
146     // Connect Tree E to the MCEvent
147     fMCEvent->ConnectTreeE(fTreeE);
148     fNEvent = fTreeE->GetEntries();
149     //
150     // Tree K
151     fFileK = TFile::Open(Form("%sKinematics%s.root", fPathName->Data(), fkExtension));
152     if (!fFileK) {
153         AliError(Form("AliMCEventHandler:Kinematics.root not found in directory %s ! \n", fPathName->Data()));
154         fInitOk = kFALSE;
155         return kTRUE;
156     }
157     
158     fEventsPerFile = fFileK->GetNkeys() - fFileK->GetNProcessIDs();
159     //
160     // Tree TR
161     if (fReadTR) {
162         fFileTR = TFile::Open(Form("%sTrackRefs%s.root", fPathName->Data(), fkExtension));
163         if (!fFileTR) {
164             AliError(Form("AliMCEventHandler:TrackRefs.root not found in directory %s ! \n", fPathName->Data()));
165             fInitOk = kFALSE;
166             return kTRUE;
167         }
168     }
169     //
170     // Reset the event number
171     fEvent      = -1;
172     fFileNumber =  0;
173     AliInfo(Form("Number of events in this directory %5d \n", fNEvent));
174     fInitOk = kTRUE;
175
176
177     if (fSubsidiaryHandlers) {
178         TIter next(fSubsidiaryHandlers);
179         AliMCEventHandler *handler;
180         while((handler = (AliMCEventHandler*)next())) {
181             handler->Init(opt);
182             handler->SetNumberOfEventsInContainer(fNEvent);
183         }
184     }
185
186     return kTRUE;
187 }
188
189 Bool_t AliMCEventHandler::LoadEvent(Int_t iev)
190 {
191     // Load the event number iev
192     //
193     // Calculate the file number
194   if (!fInitOk) return kFALSE;
195     
196   Int_t inew  = iev / fEventsPerFile;
197   Bool_t firsttree = (fTreeK==0) ? kTRUE : kFALSE;
198 //  Bool_t newtree = firsttree;
199   if (inew != fFileNumber) {
200 //    newtree = kTRUE;
201     fFileNumber = inew;
202     if (!OpenFile(fFileNumber)){
203       return kFALSE;
204     }
205   }
206   // Folder name
207   char folder[20];
208   snprintf(folder, 20, "Event%d", iev);
209   // TreeE
210   fTreeE->GetEntry(iev);
211   // Tree K
212   fFileK->GetObject(folder, fDirK);
213   if (!fDirK) {
214     AliWarning(Form("AliMCEventHandler: Event #%5d - Cannot get kinematics\n", iev));
215     return kFALSE;
216   }
217     
218   fDirK ->GetObject("TreeK", fTreeK);
219   if (!fTreeK) {
220     AliError(Form("AliMCEventHandler: Event #%5d - Cannot get TreeK\n",iev));
221     return kFALSE;
222   }  
223   // Connect TreeK to MCEvent
224   fMCEvent->ConnectTreeK(fTreeK);
225   //Tree TR 
226   if (fFileTR) {
227     // Check which format has been read
228     fFileTR->GetObject(folder, fDirTR);
229     if (!fDirTR) {
230       AliError(Form("AliMCEventHandler: Event #%5d - Cannot get track references\n",iev));
231       return kFALSE;
232     }  
233      
234     fDirTR->GetObject("TreeTR", fTreeTR);
235     //
236     if (!fTreeTR) {
237       AliError(Form("AliMCEventHandler: Event #%5d - Cannot get TreeTR\n",iev));
238       return kFALSE;
239     }  
240     // Connect TR to MCEvent
241     fMCEvent->ConnectTreeTR(fTreeTR);
242   }
243   // Now setup the caches if not yet done
244   if (fCacheSize) {
245     fTreeK->SetCacheSize(fCacheSize);
246     fCacheTK = (TTreeCache*) fFileK->GetCacheRead(fTreeK);
247     TTreeCache::SetLearnEntries(1);
248     fTreeK->AddBranchToCache("*",kTRUE);
249     if (firsttree) Info("LoadEvent","Read cache enabled %lld bytes for TreeK",fCacheSize);
250     if (fTreeTR) {
251       fTreeTR->SetCacheSize(fCacheSize);
252       fCacheTR = (TTreeCache*) fFileTR->GetCacheRead(fTreeTR);
253       TTreeCache::SetLearnEntries(1);
254       fTreeTR->AddBranchToCache("*",kTRUE);
255       if (firsttree) Info("LoadEvent","Read cache enabled %lld bytes for TreeTR",fCacheSize);
256     } 
257 //    } else {
258       // We need to reuse the previous caches and every new event is a new tree
259 //      if (fCacheTK) {
260 //         fCacheTK->ResetCache();
261 //         if (fFileK) fFileK->SetCacheRead(fCacheTK, fTreeK);
262 //         fCacheTK->UpdateBranches(fTreeK);
263 //      }
264 //      if (fCacheTR) {
265 //         fCacheTR->ResetCache();
266 //         if (fFileTR) fFileTR->SetCacheRead(fCacheTR, fTreeTR);
267 //         fCacheTR->UpdateBranches(fTreeTR);
268 //      }   
269   }  
270   return kTRUE;
271 }
272
273 Bool_t AliMCEventHandler::OpenFile(Int_t i)
274 {
275     // Open file i
276     if (i > 0) {
277         fkExtension = Form("%d", i);
278     } else {
279         fkExtension = "";
280     }
281     
282     if (fFileK && fCacheTK) fFileK->SetCacheRead(0, fTreeK);
283     delete fFileK;
284     fFileK = TFile::Open(Form("%sKinematics%s.root", fPathName->Data(), fkExtension));
285     if (!fFileK) {
286         AliError(Form("AliMCEventHandler:Kinematics%s.root not found in directory %s ! \n", fkExtension, fPathName->Data()));
287         fInitOk = kFALSE;
288         return kFALSE;
289     }
290     
291     if (fReadTR) {
292       if (fFileTR && fCacheTR) fFileTR->SetCacheRead(0, fTreeTR);
293       delete fFileTR;
294       fFileTR = TFile::Open(Form("%sTrackRefs%s.root", fPathName->Data(), fkExtension));
295       if (!fFileTR) {
296         AliWarning(Form("AliMCEventHandler:TrackRefs%s.root not found in directory %s ! \n", fkExtension, fPathName->Data()));
297         fInitOk = kFALSE;
298         return kFALSE;
299       }
300     }
301     
302     fInitOk = kTRUE;
303
304     return kTRUE;
305 }
306
307 Bool_t AliMCEventHandler::BeginEvent(Long64_t entry)
308
309     // Begin event
310     fParticleSelected.Delete();
311     fLabelMap.Delete();
312     // Read the next event
313
314     if (fEventsInContainer != 0) {
315         entry = (Long64_t) ( entry * Float_t(fNEvent) / Float_t (fEventsInContainer));
316     }
317
318
319     if (entry == -1) {
320         fEvent++;
321         entry = fEvent;
322     } else {
323         fEvent = entry;
324     }
325
326     if (entry >= fNEvent) {
327         AliWarning(Form("AliMCEventHandler: Event number out of range %5lld %5d\n", entry, fNEvent));
328         return kFALSE;
329     }
330     
331     Bool_t result = LoadEvent(entry);
332
333     if (fSubsidiaryHandlers) {
334         TIter next(fSubsidiaryHandlers);
335         AliMCEventHandler *handler;
336         while((handler = (AliMCEventHandler*)next())) {
337             handler->BeginEvent(entry);
338         }
339         next.Reset();
340         while((handler = (AliMCEventHandler*)next())) {
341             fMCEvent->AddSubsidiaryEvent(handler->MCEvent());
342         }
343         fMCEvent->InitEvent();
344     }
345     
346     if (fPreReadMode == kLmPreRead) {
347         fMCEvent->PreReadAll();
348     }
349
350     return result;
351     
352 }
353
354 void AliMCEventHandler::SelectParticle(Int_t i){
355   // taking the absolute values here, need to take care 
356   // of negative daughter and mother
357   // IDs when setting!
358     if (TMath::Abs(i) >= AliMCEvent::BgLabelOffset()) i =  fMCEvent->BgLabelToIndex(TMath::Abs(i));
359     if(!IsParticleSelected(TMath::Abs(i)))fParticleSelected.Add(TMath::Abs(i),1);
360 }
361
362 Bool_t AliMCEventHandler::IsParticleSelected(Int_t i)  {
363   // taking the absolute values here, need to take 
364   // care with negative daughter and mother
365   // IDs when setting!
366   return (fParticleSelected.GetValue(TMath::Abs(i))==1);
367 }
368
369
370 void AliMCEventHandler::CreateLabelMap(){
371
372   //
373   // this should be called once all selections where done 
374   //
375
376   fLabelMap.Delete();
377   if(!fMCEvent){
378     fParticleSelected.Delete();
379     return;
380   }
381
382   VerifySelectedParticles();
383
384   Int_t iNew = 0;
385   for(int i = 0;i < fMCEvent->GetNumberOfTracks();++i){
386     if(IsParticleSelected(i)){
387       fLabelMap.Add(i,iNew);
388       iNew++;
389     }
390   }
391 }
392
393 Int_t AliMCEventHandler::GetNewLabel(Int_t i) {
394   // Gets the label from the new created Map
395   // Call CreatLabelMap before
396   // otherwise only 0 returned
397   return fLabelMap.GetValue(TMath::Abs(i));
398 }
399
400 void  AliMCEventHandler::VerifySelectedParticles(){
401
402   //  
403   // Make sure that each particle has at least it's predecessors
404   // selected so that we have the complete ancestry tree
405   // Private, should be only called by CreateLabelMap
406
407   if(!fMCEvent){
408       fParticleSelected.Delete();
409       return;
410   }
411
412   Int_t nprim = fMCEvent->GetNumberOfPrimaries();
413
414   for(int i = 0;i < fMCEvent->GetNumberOfTracks(); ++i){
415       if(i < nprim){
416           SelectParticle(i);// take all primaries
417           continue;
418       }
419
420       if(!IsParticleSelected(i))continue;
421
422       AliMCParticle* mcpart = (AliMCParticle*) fMCEvent->GetTrack(i);
423       Int_t imo = mcpart->GetMother();
424       while((imo >= nprim)&&!IsParticleSelected(imo)){
425           // Mother not yet selected
426           SelectParticle(imo);
427           mcpart = (AliMCParticle*) fMCEvent->GetTrack(imo);
428           imo = mcpart->GetMother();
429       }
430     // after last step we may have an unselected primary
431     // mother
432     if(imo>=0){
433       if(!IsParticleSelected(imo))
434         SelectParticle(imo);
435     } 
436   }// loop over all tracks
437 }
438
439 Int_t AliMCEventHandler::GetParticleAndTR(Int_t i, TParticle*& particle, TClonesArray*& trefs)
440 {
441     // Retrieve entry i
442     if (!fInitOk) {
443         return 0;
444     } else {
445         return (fMCEvent->GetParticleAndTR(i, particle, trefs));
446     }
447 }
448
449 void AliMCEventHandler::DrawCheck(Int_t i, Int_t search)
450 {
451     // Retrieve entry i and draw momentum vector and hits
452     fMCEvent->DrawCheck(i, search);
453 }
454
455 Bool_t AliMCEventHandler::Notify(const char *path)
456 {
457   // Notify about directory change
458   // The directory is taken from the 'path' argument
459   // Reconnect trees
460     TString fileName(path);
461     TString dirname = gSystem->DirName(fileName);
462     TString basename = gSystem->BaseName(fileName);
463     Int_t index = basename.Index("#");
464     basename = basename(0, index+1);
465     fileName = dirname;
466     fileName += "/";
467     fileName += basename;
468     if (fileName.BeginsWith("root:")) {
469       fileName.Append("?ZIP=");
470     }
471
472     *fPathName = fileName;
473     AliInfo(Form("Path: -%s-\n", fPathName->Data()));
474     
475     ResetIO();
476     InitIO("");
477
478 // Handle subsidiary handlers
479     if (fSubsidiaryHandlers) {
480         TIter next(fSubsidiaryHandlers);
481         AliMCEventHandler *handler;
482         while((handler = (AliMCEventHandler*) next())) {
483             TString* spath = handler->GetInputPath();
484             if (spath->Contains("merged")) {
485                 if (! fPathName->IsNull()) {
486                     handler->Notify(Form("%s/../.", fPathName->Data()));
487                 } else {
488                     handler->Notify("../");
489                 }
490             }
491         }
492     }
493     
494     return kTRUE;
495 }
496
497 void AliMCEventHandler::ResetIO()
498 {
499 //  Clear header and stack
500     
501     if (fInitOk) fMCEvent->Clean();
502     
503 // Delete Tree E
504     delete fTreeE; fTreeE = 0;
505      
506 // Reset files
507     if (fFileE)  {delete fFileE;  fFileE  = 0;}
508     if (fFileK)  {delete fFileK;  fFileK  = 0;}
509     if (fFileTR) {delete fFileTR; fFileTR = 0;}
510     fkExtension="";
511     fInitOk = kFALSE;
512
513     if (fSubsidiaryHandlers) {
514         TIter next(fSubsidiaryHandlers);
515         AliMCEventHandler *handler;
516         while((handler = (AliMCEventHandler*)next())) {
517             handler->ResetIO();
518         }
519     }
520
521 }
522
523                             
524 Bool_t AliMCEventHandler::FinishEvent()
525 {
526     // Clean-up after each event
527    if (fFileK && fCacheTK) {
528       fTreeK->SetCacheSize(0);
529       fCacheTK = 0;  
530       fFileK->SetCacheRead(0, fTreeK);
531    }   
532    if (fFileTR && fCacheTR) {
533       fTreeTR->SetCacheSize(0);
534       fCacheTR = 0;
535       fFileTR->SetCacheRead(0, fTreeTR);
536    }
537    delete fDirTR;  fDirTR = 0;
538    delete fDirK;   fDirK  = 0;    
539     if (fInitOk) fMCEvent->FinishEvent();
540
541     if (fSubsidiaryHandlers) {
542         TIter next(fSubsidiaryHandlers);
543         AliMCEventHandler *handler;
544         while((handler = (AliMCEventHandler*)next())) {
545             handler->FinishEvent();
546         }
547     }
548
549     return kTRUE;
550 }
551
552 Bool_t AliMCEventHandler::Terminate()
553
554     // Dummy 
555     return kTRUE;
556 }
557
558 Bool_t AliMCEventHandler::TerminateIO()
559
560     // Dummy
561     return kTRUE;
562 }
563     
564
565 void AliMCEventHandler::SetInputPath(const char* fname)
566 {
567     // Set the input path name
568     delete fPathName;
569     fPathName = new TString(fname);
570 }
571
572 void AliMCEventHandler::AddSubsidiaryHandler(AliMCEventHandler* handler)
573 {
574     // Add a subsidiary handler. For example for background events
575
576     if (!fSubsidiaryHandlers) fSubsidiaryHandlers = new TList();
577     fSubsidiaryHandlers->Add(handler);
578 }