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