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