c14218ef862eeacc6969058344fadcb8272051b2
[u/mrichter/AliRoot.git] / PWG3 / dielectron / AliAnalysisTaskDielectronFilter.cxx
1 /*************************************************************************
2 * Copyright(c) 1998-2009, 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 ///////////////////////////////////////////////////////////////////////////
17 //                                                                       //
18 //                        Basic Analysis Task                            //
19 //                                                                       //
20 ///////////////////////////////////////////////////////////////////////////
21
22 #include <TChain.h>
23 #include <TH1D.h>
24
25 #include <AliLog.h>
26 #include <AliAODHandler.h>
27 #include <AliAnalysisManager.h>
28 #include <AliVEvent.h>
29 #include <AliInputEventHandler.h>
30 #include <AliESDInputHandler.h>
31 #include <AliAODInputHandler.h>
32
33 #include "AliDielectron.h"
34 #include "AliDielectronMC.h"
35 #include "AliDielectronHistos.h"
36 #include "AliDielectronVarManager.h"
37 #include "AliAnalysisTaskDielectronFilter.h"
38
39 ClassImp(AliAnalysisTaskDielectronFilter)
40
41 //_________________________________________________________________________________
42 AliAnalysisTaskDielectronFilter::AliAnalysisTaskDielectronFilter() :
43 AliAnalysisTaskSE(),
44 fDielectron(0),
45 fSelectPhysics(kTRUE),
46 fTriggerMask(AliVEvent::kMB),
47 fEventStat(0x0)
48 {
49   //
50   // Constructor
51   //
52 }
53
54 //_________________________________________________________________________________
55 AliAnalysisTaskDielectronFilter::AliAnalysisTaskDielectronFilter(const char *name) :
56 AliAnalysisTaskSE(name),
57 fDielectron(0),
58 fSelectPhysics(kTRUE),
59 fTriggerMask(AliVEvent::kMB),
60 fEventStat(0x0)
61 {
62   //
63   // Constructor
64   //
65   DefineInput(0,TChain::Class());
66   DefineOutput(1, THashList::Class());
67   DefineOutput(2, TH1D::Class());
68 }
69
70 //_________________________________________________________________________________
71 void AliAnalysisTaskDielectronFilter::Init()
72 {
73   // Initialization
74   if (fDebug > 1) AliInfo("Init() \n");
75   
76 // require AOD handler
77   AliAODHandler *aodH = (AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
78   if (!aodH) Fatal("Init", "No AOD handler. Halting.");
79   
80 //require dielectron framework 
81   if (!fDielectron) {
82     Error("Init","Dielectron framework class required. Please create and instance with proper cuts and set it via 'SetDielectron' before executing this task!!!");
83     return;
84   }
85   fDielectron->Init();
86   
87   aodH->AddFilteredAOD("AliAOD.Dielectron.root", "DielectronEvents");
88 //   AddAODBranch("AliDielectronCandidates",fDielectron->GetPairArraysPointer(),"deltaAOD.Dielectron.root");
89 }
90
91 //_________________________________________________________________________________
92 void AliAnalysisTaskDielectronFilter::UserCreateOutputObjects()
93 {
94   //
95   // Initilise histograms
96   //
97   if (!fEventStat){
98     fEventStat=new TH1D("hEventStat","Event statistics",5,0,5);
99     fEventStat->GetXaxis()->SetBinLabel(1,"Before Phys. Sel.");
100     fEventStat->GetXaxis()->SetBinLabel(2,"After Phys. Sel.");
101     fEventStat->GetXaxis()->SetBinLabel(3,"After Cand. Sel.");
102   }
103   
104   PostData(2,fEventStat);
105 }
106
107 //_________________________________________________________________________________
108 void AliAnalysisTaskDielectronFilter::UserExec(Option_t *)
109 {
110   //
111   // Main loop. Called for every event
112   //
113   
114   if (!fDielectron) return;
115   
116   AliAnalysisManager *man=AliAnalysisManager::GetAnalysisManager();
117   AliESDInputHandler *esdHandler=0x0;
118   if ( (esdHandler=dynamic_cast<AliESDInputHandler*>(man->GetInputEventHandler())) && esdHandler->GetESDpid() ){
119     AliDielectronVarManager::SetESDpid(esdHandler->GetESDpid());
120   } else {
121     //load esd pid bethe bloch parameters depending on the existance of the MC handler
122     // yes: MC parameters
123     // no:  data parameters
124     
125     //ESD case
126     if (man->GetInputEventHandler()->IsA()==AliESDInputHandler::Class()){
127       if (!AliDielectronVarManager::GetESDpid()){
128         
129         if (AliDielectronMC::Instance()->HasMC()) {
130           AliDielectronVarManager::InitESDpid();
131         } else {
132           AliDielectronVarManager::InitESDpid(1);
133         }
134       }
135     }
136     //AOD case
137     if (man->GetInputEventHandler()->IsA()==AliESDInputHandler::Class()){
138       if (!AliDielectronVarManager::GetAODpidUtil()){
139         if (AliDielectronMC::Instance()->HasMC()) {
140           AliDielectronVarManager::InitAODpidUtil();
141         } else {
142           AliDielectronVarManager::InitAODpidUtil(1);
143         }
144       }
145     }
146   }
147   
148   // Was event selected ?
149   AliInputEventHandler* inputHandler = (AliInputEventHandler*) (man->GetInputEventHandler());
150   UInt_t isSelected = AliVEvent::kAny;
151   if( fSelectPhysics && inputHandler && inputHandler->GetEventSelection() ) {
152     isSelected = inputHandler->IsEventSelected();
153     isSelected&=fTriggerMask;
154   }
155   
156   //Before physics selection
157   fEventStat->Fill(0.);
158   if (isSelected==0) {
159     PostData(2,fEventStat);
160     return;
161   }
162   //after physics selection
163   fEventStat->Fill(1.);
164   
165   //bz for AliKF
166   Double_t bz = InputEvent()->GetMagneticField();
167   AliKFParticle::SetField( bz );
168   
169   fDielectron->Process(InputEvent());
170   
171   if(fDielectron->HasCandidates()){
172     AliAODEvent *aod = ((AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler()))->GetAOD();
173     
174     //replace the references of the legs with the AOD references
175     TObjArray *obj = 0x0;
176     for(Int_t i=0; i < 10; i++ ){
177       obj = (TObjArray*)((*(fDielectron->GetPairArraysPointer()))->UncheckedAt(i));
178       if(!obj) continue;
179       for(int j=0;j<obj->GetEntriesFast();j++){
180         AliAODTrack *leg1 = 0x0;
181         AliAODTrack *leg2 = 0x0;
182         AliDielectronPair *pairObj = (AliDielectronPair*)obj->UncheckedAt(j);
183         Int_t id1 = ((AliVTrack*)pairObj->GetFirstDaughter())->GetID();
184         Int_t id2 = ((AliVTrack*)pairObj->GetSecondDaughter())->GetID();
185         
186         for(Int_t it=0;it<aod->GetNumberOfTracks();it++){
187           if(aod->GetTrack(it)->GetID() == id1) leg1 = aod->GetTrack(it);
188           if(aod->GetTrack(it)->GetID() == id2) leg2 = aod->GetTrack(it);
189         }
190         if(!leg1 || !leg2) continue;
191         
192         if(man->GetInputEventHandler()->IsA()==AliAODInputHandler::Class()){
193           leg1->ResetBit(kIsReferenced);
194           leg1->SetUniqueID(0);
195           leg2->ResetBit(kIsReferenced);
196           leg2->SetUniqueID(0);
197         }
198         pairObj->SetRefFirstDaughter(leg1);
199         pairObj->SetRefSecondDaughter(leg2);
200       }
201     }
202     
203     AliAODExtension *extDielectron = dynamic_cast<AliAODHandler*>       ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler())->GetFilteredAOD("AliAOD.Dielectron.root");
204     extDielectron->SelectEvent();
205     //after candidate selection
206     fEventStat->Fill(2.);
207     
208     //see if dielectron candidate branch exists, if not create is
209     TTree *t=extDielectron->GetTree();
210
211     if(!t->GetListOfBranches()->GetEntries() && man->GetInputEventHandler()->IsA()==AliAODInputHandler::Class())
212       t->Branch(aod->GetList());
213     
214     if (!t->GetBranch("dielectrons")){
215       t->Bronch("dielectrons","TObjArray",fDielectron->GetPairArraysPointer());
216     }
217     
218     if(man->GetInputEventHandler()->IsA()==AliAODInputHandler::Class()) t->Fill();
219   }
220   
221   PostData(1, const_cast<THashList*>(fDielectron->GetHistogramList()));
222   PostData(2,fEventStat);
223 }
224