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