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