If MC handler is connected but MC not available, the InitOK() flag is consistently...
[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(0),
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(kLmPreRead), // was 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(),
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(kLmPreRead), // was 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     if (!fMCEvent) fMCEvent = new AliMCEvent();
148     fMCEvent->ConnectTreeE(fTreeE);
149     fNEvent = fTreeE->GetEntries();
150     //
151     // Tree K
152     fFileK = TFile::Open(Form("%sKinematics%s.root", fPathName->Data(), fkExtension));
153     if (!fFileK) {
154         AliError(Form("AliMCEventHandler:Kinematics.root not found in directory %s ! \n", fPathName->Data()));
155         fInitOk = kFALSE;
156         return kTRUE;
157     }
158     
159     fEventsPerFile = fFileK->GetNkeys() - fFileK->GetNProcessIDs();
160     //
161     // Tree TR
162     if (fReadTR) {
163         fFileTR = TFile::Open(Form("%sTrackRefs%s.root", fPathName->Data(), fkExtension));
164         if (!fFileTR) {
165             AliError(Form("AliMCEventHandler:TrackRefs.root not found in directory %s ! \n", fPathName->Data()));
166             fInitOk = kFALSE;
167             return kTRUE;
168         }
169     }
170     //
171     // Reset the event number
172     fEvent      = -1;
173     fFileNumber =  0;
174     AliInfo(Form("Number of events in this directory %5d \n", fNEvent));
175     fInitOk = kTRUE;
176
177
178     if (fSubsidiaryHandlers) {
179         TIter next(fSubsidiaryHandlers);
180         AliMCEventHandler *handler;
181         while((handler = (AliMCEventHandler*)next())) {
182             handler->Init(opt);
183             handler->SetNumberOfEventsInContainer(fNEvent);
184         }
185     }
186
187     return kTRUE;
188 }
189
190 Bool_t AliMCEventHandler::LoadEvent(Int_t iev)
191 {
192     // Load the event number iev
193     //
194     // Calculate the file number
195   if (!fInitOk) return kFALSE;
196     
197   Int_t inew  = iev / fEventsPerFile;
198   Bool_t firsttree = (fTreeK==0) ? kTRUE : kFALSE;
199 //  Bool_t newtree = firsttree;
200   if (inew != fFileNumber) {
201 //    newtree = kTRUE;
202     fFileNumber = inew;
203     if (!OpenFile(fFileNumber)){
204       return kFALSE;
205     }
206   }
207   // Folder name
208   char folder[20];
209   snprintf(folder, 20, "Event%d", iev);
210   // TreeE
211   fTreeE->GetEntry(iev);
212   // Tree K
213   fFileK->GetObject(folder, fDirK);
214   if (!fDirK) {
215     AliWarning(Form("AliMCEventHandler: Event #%5d - Cannot get kinematics\n", iev));
216     return kFALSE;
217   }
218     
219   fDirK ->GetObject("TreeK", fTreeK);
220   if (!fTreeK) {
221     AliError(Form("AliMCEventHandler: Event #%5d - Cannot get TreeK\n",iev));
222     return kFALSE;
223   }  
224   // Connect TreeK to MCEvent
225   fMCEvent->ConnectTreeK(fTreeK);
226   //Tree TR 
227   if (fFileTR) {
228     // Check which format has been read
229     fFileTR->GetObject(folder, fDirTR);
230     if (!fDirTR) {
231       AliError(Form("AliMCEventHandler: Event #%5d - Cannot get track references\n",iev));
232       return kFALSE;
233     }  
234      
235     fDirTR->GetObject("TreeTR", fTreeTR);
236     //
237     if (!fTreeTR) {
238       AliError(Form("AliMCEventHandler: Event #%5d - Cannot get TreeTR\n",iev));
239       return kFALSE;
240     }  
241     // Connect TR to MCEvent
242     fMCEvent->ConnectTreeTR(fTreeTR);
243   }
244   // Now setup the caches if not yet done
245   if (fCacheSize) {
246     fTreeK->SetCacheSize(fCacheSize);
247     fCacheTK = (TTreeCache*) fFileK->GetCacheRead(fTreeK);
248     TTreeCache::SetLearnEntries(1);
249     fTreeK->AddBranchToCache("*",kTRUE);
250     if (firsttree) Info("LoadEvent","Read cache enabled %lld bytes for TreeK",fCacheSize);
251     if (fTreeTR) {
252       fTreeTR->SetCacheSize(fCacheSize);
253       fCacheTR = (TTreeCache*) fFileTR->GetCacheRead(fTreeTR);
254       TTreeCache::SetLearnEntries(1);
255       fTreeTR->AddBranchToCache("*",kTRUE);
256       if (firsttree) Info("LoadEvent","Read cache enabled %lld bytes for TreeTR",fCacheSize);
257     } 
258 //    } else {
259       // We need to reuse the previous caches and every new event is a new tree
260 //      if (fCacheTK) {
261 //         fCacheTK->ResetCache();
262 //         if (fFileK) fFileK->SetCacheRead(fCacheTK, fTreeK);
263 //         fCacheTK->UpdateBranches(fTreeK);
264 //      }
265 //      if (fCacheTR) {
266 //         fCacheTR->ResetCache();
267 //         if (fFileTR) fFileTR->SetCacheRead(fCacheTR, fTreeTR);
268 //         fCacheTR->UpdateBranches(fTreeTR);
269 //      }   
270   }  
271   return kTRUE;
272 }
273
274 Bool_t AliMCEventHandler::OpenFile(Int_t i)
275 {
276     // Open file i
277     fInitOk = kFALSE;
278     if (i > 0) {
279       fkExtension = Form("%d", i);
280     } else {
281       fkExtension = "";
282     }
283     
284     if (fFileK && fCacheTK) fFileK->SetCacheRead(0, fTreeK);
285     delete fFileK;
286     fFileK = TFile::Open(Form("%sKinematics%s.root", fPathName->Data(), fkExtension));
287     if (!fFileK) {
288       AliError(Form("AliMCEventHandler:Kinematics%s.root not found in directory %s ! \n", fkExtension, fPathName->Data()));
289       delete fMCEvent; fMCEvent=0;
290       return fInitOk;
291     }
292     
293     fInitOk = kTRUE;
294     if (fReadTR) {
295       if (fFileTR && fCacheTR) fFileTR->SetCacheRead(0, fTreeTR);
296       delete fFileTR;
297       fFileTR = TFile::Open(Form("%sTrackRefs%s.root", fPathName->Data(), fkExtension));
298       if (!fFileTR) {
299         AliError(Form("AliMCEventHandler:TrackRefs%s.root not found in directory %s ! \n", fkExtension, fPathName->Data()));
300         return fInitOk;
301       }
302     }
303     return fInitOk;
304 }
305
306 Bool_t AliMCEventHandler::BeginEvent(Long64_t entry)
307
308     // Begin event
309     if (!fInitOk) return kFALSE;
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 (!fInitOk) return;
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     TString dirname = gSystem->DirName(fileName);
463     TString basename = gSystem->BaseName(fileName);
464     Int_t index = basename.Index("#");
465     basename = basename(0, index+1);
466     fileName = dirname;
467     fileName += "/";
468     fileName += basename;
469     if (fileName.BeginsWith("root:")) {
470       fileName.Append("?ZIP=");
471     }
472
473     *fPathName = fileName;
474     AliInfo(Form("Path: -%s-\n", fPathName->Data()));
475     
476     ResetIO();
477     InitIO("");
478
479 // Handle subsidiary handlers
480     if (fSubsidiaryHandlers) {
481         TIter next(fSubsidiaryHandlers);
482         AliMCEventHandler *handler;
483         while((handler = (AliMCEventHandler*) next())) {
484             TString* spath = handler->GetInputPath();
485             if (spath->Contains("merged")) {
486                 if (! fPathName->IsNull()) {
487                     handler->Notify(Form("%s/../.", fPathName->Data()));
488                 } else {
489                     handler->Notify("../");
490                 }
491             }
492         }
493     }
494     
495     return kTRUE;
496 }
497
498 void AliMCEventHandler::ResetIO()
499 {
500 //  Clear header and stack
501     
502     if (fInitOk) fMCEvent->Clean();
503     
504 // Delete Tree E
505     delete fTreeE; fTreeE = 0;
506      
507 // Reset files
508     if (fFileE)  {delete fFileE;  fFileE  = 0;}
509     if (fFileK)  {delete fFileK;  fFileK  = 0;}
510     if (fFileTR) {delete fFileTR; fFileTR = 0;}
511     fkExtension="";
512     fInitOk = kFALSE;
513
514     if (fSubsidiaryHandlers) {
515         TIter next(fSubsidiaryHandlers);
516         AliMCEventHandler *handler;
517         while((handler = (AliMCEventHandler*)next())) {
518             handler->ResetIO();
519         }
520     }
521
522 }
523
524                             
525 Bool_t AliMCEventHandler::FinishEvent()
526 {
527     // Clean-up after each event
528    if (fFileK && fCacheTK) {
529       fTreeK->SetCacheSize(0);
530       fCacheTK = 0;  
531       fFileK->SetCacheRead(0, fTreeK);
532    }   
533    if (fFileTR && fCacheTR) {
534       fTreeTR->SetCacheSize(0);
535       fCacheTR = 0;
536       fFileTR->SetCacheRead(0, fTreeTR);
537    }
538    delete fDirTR;  fDirTR = 0;
539    delete fDirK;   fDirK  = 0;    
540     if (fInitOk) fMCEvent->FinishEvent();
541
542     if (fSubsidiaryHandlers) {
543         TIter next(fSubsidiaryHandlers);
544         AliMCEventHandler *handler;
545         while((handler = (AliMCEventHandler*)next())) {
546             handler->FinishEvent();
547         }
548     }
549
550     return kTRUE;
551 }
552
553 Bool_t AliMCEventHandler::Terminate()
554
555     // Dummy 
556     return kTRUE;
557 }
558
559 Bool_t AliMCEventHandler::TerminateIO()
560
561     // Dummy
562     return kTRUE;
563 }
564     
565
566 void AliMCEventHandler::SetInputPath(const char* fname)
567 {
568     // Set the input path name
569     delete fPathName;
570     fPathName = new TString(fname);
571 }
572
573 void AliMCEventHandler::AddSubsidiaryHandler(AliMCEventHandler* handler)
574 {
575     // Add a subsidiary handler. For example for background events
576
577     if (!fSubsidiaryHandlers) fSubsidiaryHandlers = new TList();
578     fSubsidiaryHandlers->Add(handler);
579 }