AddFriend in Notify()
[u/mrichter/AliRoot.git] / ANALYSIS / AliAnalysisTaskSE.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 #include <TROOT.h>
19 #include <TSystem.h>
20 #include <TInterpreter.h>
21 #include <TChain.h>
22 #include <TFile.h>
23 #include <TList.h>
24
25 #include "AliAnalysisTaskSE.h"
26 #include "AliAnalysisManager.h"
27 #include "AliAnalysisCuts.h"
28 #include "AliAnalysisDataSlot.h"
29
30 #include "AliESDEvent.h"
31 #include "AliESDfriend.h"
32 #include "AliESD.h"
33 #include "AliAODEvent.h"
34 #include "AliAODHeader.h"
35 #include "AliAODTracklets.h"
36 #include "AliAODCaloCells.h"
37 #include "AliAODMCParticle.h"
38 #include "AliVEvent.h"
39 #include "AliAODHandler.h"
40 #include "AliAODInputHandler.h"
41 #include "AliMCEventHandler.h"
42 #include "AliInputEventHandler.h"
43 #include "AliESDInputHandler.h"
44 #include "AliMCEvent.h"
45 #include "AliStack.h"
46 #include "AliLog.h"
47 #include "AliAODDimuon.h"
48
49
50 ClassImp(AliAnalysisTaskSE)
51
52 ////////////////////////////////////////////////////////////////////////
53 AliAODHeader*    AliAnalysisTaskSE::fgAODHeader         = NULL;
54 TClonesArray*    AliAnalysisTaskSE::fgAODTracks         = NULL;
55 TClonesArray*    AliAnalysisTaskSE::fgAODVertices       = NULL;
56 TClonesArray*    AliAnalysisTaskSE::fgAODV0s            = NULL;
57 TClonesArray*    AliAnalysisTaskSE::fgAODPMDClusters    = NULL;
58 TClonesArray*    AliAnalysisTaskSE::fgAODJets           = NULL;
59 TClonesArray*    AliAnalysisTaskSE::fgAODFMDClusters    = NULL;
60 TClonesArray*    AliAnalysisTaskSE::fgAODCaloClusters   = NULL;
61 TClonesArray*    AliAnalysisTaskSE::fgAODMCParticles    = NULL;
62 AliAODTracklets* AliAnalysisTaskSE::fgAODTracklets      = NULL;
63 AliAODCaloCells* AliAnalysisTaskSE::fgAODEmcalCells     = NULL;
64 AliAODCaloCells* AliAnalysisTaskSE::fgAODPhosCells      = NULL;
65 TClonesArray*    AliAnalysisTaskSE::fgAODDimuons        = NULL;
66
67 AliAnalysisTaskSE::AliAnalysisTaskSE():
68     AliAnalysisTask(),
69     fDebug(0),
70     fEntry(0),
71     fInputEvent(0x0),
72     fESDfriend(0x0),
73     fInputHandler(0x0),
74     fOutputAOD(0x0),
75     fMCEvent(0x0),
76     fTreeA(0x0),
77     fCurrentRunNumber(-1),
78     fHistosQA(0x0),
79     fSelectCollisions(0)
80 {
81   // Default constructor
82 }
83
84 AliAnalysisTaskSE::AliAnalysisTaskSE(const char* name):
85     AliAnalysisTask(name, "AnalysisTaskSE"),
86     fDebug(0),
87     fEntry(0),
88     fInputEvent(0x0),
89     fESDfriend(0x0),
90     fInputHandler(0x0),
91     fOutputAOD(0x0),
92     fMCEvent(0x0),
93     fTreeA(0x0),
94     fCurrentRunNumber(-1),
95     fHistosQA(0x0),
96     fSelectCollisions(0)
97 {
98   // Default constructor
99     DefineInput (0, TChain::Class());
100     DefineOutput(0,  TTree::Class());
101 }
102
103 AliAnalysisTaskSE::AliAnalysisTaskSE(const AliAnalysisTaskSE& obj):
104     AliAnalysisTask(obj),
105     fDebug(0),
106     fEntry(0),
107     fInputEvent(0x0),
108     fESDfriend(0x0),
109     fInputHandler(0x0),
110     fOutputAOD(0x0),
111     fMCEvent(0x0),
112     fTreeA(0x0),
113     fCurrentRunNumber(-1),
114     fHistosQA(0x0),
115     fSelectCollisions(0)
116 {
117 // Copy constructor
118     fDebug            = obj.fDebug;
119     fEntry            = obj.fEntry;
120     fInputEvent       = obj.fInputEvent;
121     fESDfriend        = obj.fESDfriend;
122     fInputHandler     = obj.fInputHandler;
123     fOutputAOD        = obj.fOutputAOD;
124     fMCEvent          = obj.fMCEvent;
125     fTreeA            = obj.fTreeA;    
126     fCurrentRunNumber = obj.fCurrentRunNumber;
127     fHistosQA         = obj.fHistosQA;
128
129 }
130
131
132 AliAnalysisTaskSE& AliAnalysisTaskSE::operator=(const AliAnalysisTaskSE& other)
133 {
134 // Assignment
135     AliAnalysisTask::operator=(other);
136     fDebug            = other.fDebug;
137     fEntry            = other.fEntry;
138     fInputEvent       = other.fInputEvent;
139     fESDfriend        = other.fESDfriend;
140     fInputHandler     = other.fInputHandler;
141     fOutputAOD        = other.fOutputAOD;
142     fMCEvent          = other.fMCEvent;
143     fTreeA            = other.fTreeA;    
144     fCurrentRunNumber = other.fCurrentRunNumber;
145     fHistosQA         = other.fHistosQA;
146     fSelectCollisions = other.fSelectCollisions;
147     return *this;
148 }
149
150
151 void AliAnalysisTaskSE::ConnectInputData(Option_t* /*option*/)
152 {
153 // Connect the input data
154     if (fDebug > 1) printf("AnalysisTaskSE::ConnectInputData() \n");
155 //
156 //  ESD
157 //
158     fInputHandler = (AliInputEventHandler*) 
159          ((AliAnalysisManager::GetAnalysisManager())->GetInputEventHandler());
160 //
161 //  Monte Carlo
162 //
163     AliMCEventHandler*    mcH = 0;
164     mcH = (AliMCEventHandler*) ((AliAnalysisManager::GetAnalysisManager())->GetMCtruthEventHandler());
165     if (mcH) {
166         fMCEvent = mcH->MCEvent();
167     } 
168     
169     if (fInputHandler) {
170         if ((fInputHandler->GetTree())->GetBranch("ESDfriend."))
171             fESDfriend = ((AliESDInputHandler*)fInputHandler)->GetESDfriend();
172
173         fInputEvent = fInputHandler->GetEvent();
174     } else if( fMCEvent ) {
175          AliWarning("No Input Event Handler connected, only MC Truth Event Handler") ; 
176     } else {
177          AliError("No Input Event Handler connected") ; 
178          return ; 
179     }
180 }
181
182 void AliAnalysisTaskSE::CreateOutputObjects()
183 {
184 // Create the output container
185 //
186 //  Default AOD
187     if (fDebug > 1) printf("AnalysisTaskSE::CreateOutPutData() \n");
188
189     AliAODHandler* handler = (AliAODHandler*) 
190          ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
191     
192     Bool_t merging = kFALSE;
193     AliAODInputHandler* aodIH = static_cast<AliAODInputHandler*>((AliAnalysisManager::GetAnalysisManager())->GetInputEventHandler());
194     if (aodIH) {
195         if (aodIH->GetMergeEvents()) merging = kTRUE;
196     }
197     
198
199     // Check if AOD replication has been required
200     
201     if (handler) {
202         fOutputAOD   = handler->GetAOD();
203         fTreeA = handler->GetTree();
204         if (fOutputAOD && !(handler->IsStandard())) {
205             if ((handler->NeedsHeaderReplication()) && !(fgAODHeader)) 
206                 {
207                  if (fDebug > 1) AliInfo("Replicating header");
208                  fgAODHeader = new AliAODHeader;
209                  handler->AddBranch("AliAODHeader", &fgAODHeader);
210                 }
211             if ((handler->NeedsTracksBranchReplication() || merging) && !(fgAODTracks))      
212             {   
213                 if (fDebug > 1) AliInfo("Replicating track branch\n");
214                 fgAODTracks = new TClonesArray("AliAODTrack",500);
215                 fgAODTracks->SetName("tracks");
216                 handler->AddBranch("TClonesArray", &fgAODTracks);
217             }    
218             if ((handler->NeedsVerticesBranchReplication() || merging) && !(fgAODVertices))
219             {
220                 if (fDebug > 1) AliInfo("Replicating vertices branch\n");
221                 fgAODVertices = new TClonesArray("AliAODVertex",500);
222                 fgAODVertices->SetName("vertices");
223                 handler->AddBranch("TClonesArray", &fgAODVertices);
224             }   
225             if ((handler->NeedsV0sBranchReplication()) && !(fgAODV0s))    
226             {   
227                 if (fDebug > 1) AliInfo("Replicating V0s branch\n");
228                 fgAODV0s = new TClonesArray("AliAODv0",500);
229                 fgAODV0s->SetName("v0s");
230                 handler->AddBranch("TClonesArray", &fgAODV0s);
231             }
232             if ((handler->NeedsTrackletsBranchReplication()) && !(fgAODTracklets))        
233             {   
234                 if (fDebug > 1) AliInfo("Replicating Tracklets branch\n");
235                 fgAODTracklets = new AliAODTracklets("tracklets","tracklets");
236                 handler->AddBranch("AliAODTracklets", &fgAODTracklets);
237             }
238             if ((handler->NeedsPMDClustersBranchReplication()) && !(fgAODPMDClusters))    
239             {   
240                 if (fDebug > 1) AliInfo("Replicating PMDClusters branch\n");
241                 fgAODPMDClusters = new TClonesArray("AliAODPmdCluster",500);
242                 fgAODPMDClusters->SetName("pmdClusters");
243                 handler->AddBranch("TClonesArray", &fgAODPMDClusters);
244             }
245             if ((handler->NeedsJetsBranchReplication() || merging) && !(fgAODJets))       
246             {   
247                 if (fDebug > 1) AliInfo("Replicating Jets branch\n");
248                 fgAODJets = new TClonesArray("AliAODJet",500);
249                 fgAODJets->SetName("jets");
250                 handler->AddBranch("TClonesArray", &fgAODJets);
251             }
252             if ((handler->NeedsFMDClustersBranchReplication()) && !(fgAODFMDClusters))    
253             {   
254                 AliInfo("Replicating FMDClusters branch\n");
255                 fgAODFMDClusters = new TClonesArray("AliAODFmdCluster",500);
256                 fgAODFMDClusters->SetName("fmdClusters");
257                 handler->AddBranch("TClonesArray", &fgAODFMDClusters);
258             }
259             if ((handler->NeedsCaloClustersBranchReplication() || merging) && !(fgAODCaloClusters))       
260             {   
261                 if (fDebug > 1) AliInfo("Replicating CaloClusters branch\n");
262                 fgAODCaloClusters = new TClonesArray("AliAODCaloCluster",500);
263                 fgAODCaloClusters->SetName("caloClusters");
264                 handler->AddBranch("TClonesArray", &fgAODCaloClusters);
265
266                 fgAODEmcalCells = new AliAODCaloCells();
267                 fgAODEmcalCells->SetName("emcalCells");
268                 handler->AddBranch("AliAODCaloCells", &fgAODEmcalCells);
269
270                 fgAODPhosCells = new AliAODCaloCells();
271                 fgAODPhosCells->SetName("phosCells");
272                 handler->AddBranch("AliAODCaloCells", &fgAODPhosCells);
273             }
274             if ((handler->NeedsMCParticlesBranchReplication() || merging) && !(fgAODMCParticles))         
275             {   
276                 if (fDebug > 1) AliInfo("Replicating MCParticles branch\n");
277                 fgAODMCParticles = new TClonesArray("AliAODMCParticle",500);
278                 fgAODMCParticles->SetName("mcparticles");
279                 handler->AddBranch("TClonesArray", &fgAODMCParticles);
280             }
281             if ((handler->NeedsDimuonsBranchReplication() || merging) && !(fgAODDimuons))      
282             {   
283                 if (fDebug > 1) AliInfo("Replicating dimuon branch\n");
284                 fgAODDimuons = new TClonesArray("AliAODDimuon",0);
285                 fgAODDimuons->SetName("dimuons");
286                 handler->AddBranch("TClonesArray", &fgAODDimuons);
287             }    
288
289             // cache the pointerd in the AODEvent
290             fOutputAOD->GetStdContent();
291         }
292     } else {
293         AliWarning("No AOD Event Handler connected.") ; 
294     }
295     UserCreateOutputObjects();
296 }
297
298 void AliAnalysisTaskSE::Exec(Option_t* option)
299 {
300 //
301 // Exec analysis of one event
302     if ( fDebug >= 10)
303       printf("Task is active %5d\n", IsActive());
304     
305     if (fDebug > 1) AliInfo("AliAnalysisTaskSE::Exec() \n");
306 //
307     AliAODHandler* handler = (AliAODHandler*) 
308         ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
309     AliAODInputHandler* aodH = dynamic_cast<AliAODInputHandler*>(fInputHandler);
310 //
311 // Was event selected ?
312     Bool_t isSelected = kTRUE;
313     if( fInputHandler && fInputHandler->GetEventSelection() && fSelectCollisions) {
314       isSelected = fInputHandler->IsEventSelected();
315     }
316
317     if (handler) handler->SetFillAOD(isSelected);
318
319     if( fInputHandler ) {
320         fEntry = fInputHandler->GetReadEntry();
321         fESDfriend = ((AliESDInputHandler*)fInputHandler)->GetESDfriend();
322     }
323     
324
325 // Notify the change of run number
326     if (InputEvent() && (InputEvent()->GetRunNumber() != fCurrentRunNumber)) {
327         fCurrentRunNumber = InputEvent()->GetRunNumber();
328         NotifyRun();
329     }    
330            
331     else if( fMCEvent )
332        fEntry = fMCEvent->Header()->GetEvent(); 
333     if ( !((Entry()-1)%100) && fDebug > 0) 
334          AliInfo(Form("%s ----> Processing event # %lld", CurrentFileName(), Entry()));
335
336     
337     
338
339     if (handler && aodH) {
340         fMCEvent = aodH->MCEvent();
341         Bool_t merging = aodH->GetMergeEvents();
342         
343         if (!(handler->IsStandard()) && !(handler->AODIsReplicated())) {
344             if ((handler->NeedsHeaderReplication()) && (fgAODHeader))
345             {
346               // copy the contents by assigment
347               *fgAODHeader =  *(dynamic_cast<AliAODHeader*>(InputEvent()->GetHeader()));
348             }
349             if ((handler->NeedsTracksBranchReplication() || merging) && (fgAODTracks))
350             {
351                 TClonesArray* tracks = (dynamic_cast<AliAODEvent*>(InputEvent()))->GetTracks();
352                 new (fgAODTracks) TClonesArray(*tracks);
353             }
354             if ((handler->NeedsVerticesBranchReplication() || merging) && (fgAODVertices))
355             {
356                 TClonesArray* vertices = (dynamic_cast<AliAODEvent*>(InputEvent()))->GetVertices();
357                 new (fgAODVertices) TClonesArray(*vertices);
358             }
359             if ((handler->NeedsV0sBranchReplication()) && (fgAODV0s))
360             {
361                 TClonesArray* v0s = (dynamic_cast<AliAODEvent*>(InputEvent()))->GetV0s();
362                 new (fgAODV0s) TClonesArray(*v0s);
363             }
364             if ((handler->NeedsTrackletsBranchReplication()) && (fgAODTracklets))
365             {
366               *fgAODTracklets = *(dynamic_cast<AliAODEvent*>(InputEvent()))->GetTracklets();
367             }
368             if ((handler->NeedsPMDClustersBranchReplication()) && (fgAODPMDClusters))
369             {
370                 TClonesArray* pmdClusters = (dynamic_cast<AliAODEvent*>(InputEvent()))->GetPmdClusters();
371                 new (fgAODPMDClusters) TClonesArray(*pmdClusters);
372             }
373             if ((handler->NeedsJetsBranchReplication() || merging) && (fgAODJets))
374             {
375                 TClonesArray* jets = (dynamic_cast<AliAODEvent*>(InputEvent()))->GetJets();
376                 new (fgAODJets) TClonesArray(*jets);
377             }
378             if ((handler->NeedsFMDClustersBranchReplication()) && (fgAODFMDClusters))
379             {
380                 TClonesArray* fmdClusters = (dynamic_cast<AliAODEvent*>(InputEvent()))->GetFmdClusters();
381                 new (fgAODFMDClusters) TClonesArray(*fmdClusters);
382             }
383             if ((handler->NeedsCaloClustersBranchReplication() || merging) && (fgAODCaloClusters))
384             {
385                 TClonesArray* caloClusters = (dynamic_cast<AliAODEvent*>(InputEvent()))->GetCaloClusters();
386                 new (fgAODCaloClusters) TClonesArray(*caloClusters);
387             }
388
389             if ((handler->NeedsMCParticlesBranchReplication() || merging) && (fgAODMCParticles))
390             {
391                 TClonesArray* mcParticles = (TClonesArray*) ((dynamic_cast<AliAODEvent*>(InputEvent()))->FindListObject("mcparticles"));
392                 new (fgAODMCParticles) TClonesArray(*mcParticles);
393             }
394             
395             if ((handler->NeedsDimuonsBranchReplication() || merging) && (fgAODDimuons))
396             {
397                 fgAODDimuons->Clear();
398                 TClonesArray& dimuons = *fgAODDimuons;
399                 TClonesArray& tracksnew = *fgAODTracks;
400                 
401                 Int_t nMuonTrack[10]; 
402                 for(Int_t imuon = 0; imuon < 10; imuon++) nMuonTrack[imuon] = 0;
403                 Int_t nMuons=0;
404                 for(Int_t ii=0; ii < fgAODTracks->GetEntries(); ii++){
405                     AliAODTrack *track = (AliAODTrack*) fgAODTracks->At(ii);
406                     if(track->IsMuonTrack()) {
407                         nMuonTrack[nMuons]= ii;
408                         nMuons++;
409                     }  
410                 }
411                 Int_t jDimuons=0;
412                 if(nMuons >= 2){
413                     for(Int_t i = 0; i < nMuons; i++){
414                         Int_t index0 = nMuonTrack[i];
415                         for(Int_t j = i+1; j < nMuons; j++){
416                             Int_t index1 = nMuonTrack[j];
417                             tracksnew.At(index0)->ResetBit(kIsReferenced);
418                             tracksnew.At(index0)->SetUniqueID(0); 
419                             tracksnew.At(index1)->ResetBit(kIsReferenced);
420                             tracksnew.At(index1)->SetUniqueID(0);
421                             new(dimuons[jDimuons++]) AliAODDimuon(tracksnew.At(index0),tracksnew.At(index1));
422                         }
423                     }    
424                 }
425             }
426             
427             // Additional merging if needed
428             if (merging) {
429                 // mcParticles
430                 TClonesArray* mcparticles = (TClonesArray*) ((aodH->GetEventToMerge())->FindListObject("mcparticles"));
431                 Int_t npart = mcparticles->GetEntries();
432                 Int_t nc = fgAODMCParticles->GetEntries();
433                 Int_t nc0 = nc;
434                 
435                 for (Int_t i = 0; i < npart; i++) {
436                     AliAODMCParticle* particle = (AliAODMCParticle*) mcparticles->At(i);
437                     new((*fgAODMCParticles)[nc++]) AliAODMCParticle(*particle);
438                 }
439                 
440                 // tracks
441                 TClonesArray* tracks = aodH->GetEventToMerge()->GetTracks();
442                 Int_t ntr = tracks->GetEntries();
443                 nc  = fgAODTracks->GetEntries();        
444                 for (Int_t i = 0; i < ntr; i++) {
445                     AliAODTrack*    track = (AliAODTrack*) tracks->At(i);
446                     AliAODTrack* newtrack = new((*fgAODTracks)[nc++]) AliAODTrack(*track);
447                     newtrack->SetLabel(newtrack->GetLabel() + nc0);
448                 }
449
450                 for (Int_t i = 0; i < nc; i++) 
451                 {
452                     AliAODTrack* track = (AliAODTrack*) fgAODTracks->At(i);
453                     track->ResetBit(kIsReferenced);
454                     track->SetUniqueID(0);
455                 }
456                 
457                 
458                 // clusters
459                 TClonesArray* clusters = aodH->GetEventToMerge()->GetCaloClusters();
460                 Int_t ncl  = clusters->GetEntries();
461                 nc         =  fgAODCaloClusters->GetEntries();
462                 for (Int_t i = 0; i < ncl; i++) {
463                     AliAODCaloCluster*    cluster = (AliAODCaloCluster*) clusters->At(i);
464                     new((*fgAODCaloClusters)[nc++]) AliAODCaloCluster(*cluster);
465                 }
466                 // cells
467                 AliAODCaloCells* cellsA = aodH->GetEventToMerge()->GetEMCALCells();
468                 Int_t ncells  = cellsA->GetNumberOfCells();
469                 nc = fgAODEmcalCells->GetNumberOfCells();
470                 
471                 for (Int_t i  = 0; i < ncells; i++) {
472                     Int_t cn  = cellsA->GetCellNumber(i);
473                     Int_t pos = fgAODEmcalCells->GetCellPosition(cn);
474                     if (pos >= 0) {
475                         Double_t amp = cellsA->GetAmplitude(i) + fgAODEmcalCells->GetAmplitude(pos);
476                         fgAODEmcalCells->SetCell(pos, cn, amp);
477                     } else {
478                         Double_t amp = cellsA->GetAmplitude(i);
479                         fgAODEmcalCells->SetCell(nc++, cn, amp);
480                         fgAODEmcalCells->Sort();
481                     }
482                 }
483                 
484                 
485             } // merging
486             
487             handler->SetAODIsReplicated();
488         }
489     }
490
491 // Call the user analysis    
492     if (!fSelectCollisions || isSelected) 
493         UserExec(option);
494     
495 // Added protection in case the derived task is not an AOD producer.
496     AliAnalysisDataSlot *out0 = GetOutputSlot(0);
497     if (out0 && out0->IsConnected()) PostData(0, fTreeA);    
498 }
499
500 const char* AliAnalysisTaskSE::CurrentFileName()
501 {
502 // Returns the current file name    
503     if( fInputHandler )
504       return fInputHandler->GetTree()->GetCurrentFile()->GetName();
505     else if( fMCEvent )
506       return ((AliMCEventHandler*) ((AliAnalysisManager::GetAnalysisManager())->GetMCtruthEventHandler()))->TreeK()->GetCurrentFile()->GetName();
507     else return "";
508 }
509
510 void AliAnalysisTaskSE::AddAODBranch(const char* cname, void* addobj, const char *fname)
511 {
512     // Add a new branch to the aod tree
513     AliAODHandler* handler = (AliAODHandler*) 
514         ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
515     if (handler) {
516         handler->AddBranch(cname, addobj, fname);
517     }
518 }
519
520 Bool_t AliAnalysisTaskSE::IsStandardAOD() const
521 {
522 // Check if the output AOD handler is configured for standard or delta AOD.
523 // Users should first check that AODEvent() returns non-null.
524     AliAODHandler* handler = (AliAODHandler*) 
525          ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
526     if (!handler) {
527        Error("IsStandardAOD", "No AOD handler. Please use AODEvent() to check this first");
528        return kTRUE;
529     }
530     return handler->IsStandard();   
531 }
532
533 Bool_t AliAnalysisTaskSE::Notify()
534 {
535     return (UserNotify());
536 }
537
538