MCEventHandler derives from InputEventHandler
[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     fExtension(""),
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     fExtension(""),
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(), fExtension));
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(), fExtension));
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         fExtension = Form("%d", i);
240     } else {
241         fExtension = "";
242     }
243     
244     
245     delete fFileK;
246     fFileK = TFile::Open(Form("%sKinematics%s.root", fPathName->Data(), fExtension));
247     if (!fFileK) {
248         AliError(Form("AliMCEventHandler:Kinematics%s.root not found in directory %s ! \n", fExtension, 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(), fExtension));
256         if (!fFileTR) {
257             AliWarning(Form("AliMCEventHandler:TrackRefs%s.root not found in directory %s ! \n", fExtension, 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("AliAOD.root")){
429         fileName.ReplaceAll("AliAOD.root", "");
430     }
431     else if(fileName.Contains("galice.root")){
432         // for running with galice and kinematics alone...
433         fileName.ReplaceAll("galice.root", "");
434     }
435     else if (fileName.BeginsWith("root:")) {
436       fileName.Append("?ZIP=");
437     }
438
439     *fPathName = fileName;
440     AliInfo(Form("Path: -%s-\n", fPathName->Data()));
441     
442     ResetIO();
443     InitIO("");
444
445 // Handle subsidiary handlers
446     if (fSubsidiaryHandlers) {
447         TIter next(fSubsidiaryHandlers);
448         AliMCEventHandler *handler;
449         while((handler = (AliMCEventHandler*) next())) {
450             TString* spath = handler->GetInputPath();
451             if (spath->Contains("merged")) {
452                 if (! fPathName->IsNull()) {
453                     handler->Notify(Form("%s/../.", fPathName->Data()));
454                 } else {
455                     handler->Notify("../");
456                 }
457             }
458         }
459     }
460     
461     return kTRUE;
462 }
463
464 void AliMCEventHandler::ResetIO()
465 {
466 //  Clear header and stack
467     
468     if (fInitOk) fMCEvent->Clean();
469     
470 // Delete Tree E
471     delete fTreeE; fTreeE = 0;
472      
473 // Reset files
474     if (fFileE)  {delete fFileE;  fFileE  = 0;}
475     if (fFileK)  {delete fFileK;  fFileK  = 0;}
476     if (fFileTR) {delete fFileTR; fFileTR = 0;}
477     fExtension="";
478     fInitOk = kFALSE;
479
480     if (fSubsidiaryHandlers) {
481         TIter next(fSubsidiaryHandlers);
482         AliMCEventHandler *handler;
483         while((handler = (AliMCEventHandler*)next())) {
484             handler->ResetIO();
485         }
486     }
487
488 }
489
490                             
491 Bool_t AliMCEventHandler::FinishEvent()
492 {
493     // Clean-up after each event
494     delete fDirTR;  fDirTR = 0;
495     delete fDirK;   fDirK  = 0;    
496     if (fInitOk) fMCEvent->FinishEvent();
497
498     if (fSubsidiaryHandlers) {
499         TIter next(fSubsidiaryHandlers);
500         AliMCEventHandler *handler;
501         while((handler = (AliMCEventHandler*)next())) {
502             handler->FinishEvent();
503         }
504     }
505
506     return kTRUE;
507 }
508
509 Bool_t AliMCEventHandler::Terminate()
510
511     // Dummy 
512     return kTRUE;
513 }
514
515 Bool_t AliMCEventHandler::TerminateIO()
516
517     // Dummy
518     return kTRUE;
519 }
520     
521
522 void AliMCEventHandler::SetInputPath(const char* fname)
523 {
524     // Set the input path name
525     delete fPathName;
526     fPathName = new TString(fname);
527 }
528
529 void AliMCEventHandler::AddSubsidiaryHandler(AliMCEventHandler* handler)
530 {
531     // Add a subsidiary handler. For example for background events
532
533     if (!fSubsidiaryHandlers) fSubsidiaryHandlers = new TList();
534     fSubsidiaryHandlers->Add(handler);
535 }