AliAODEvent::GetHeader() returns AliVHeader
[u/mrichter/AliRoot.git] / PWGDQ / 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 #include <AliLog.h>
25 #include <AliAODHandler.h>
26 #include <AliAODInputHandler.h>
27 #include <AliAnalysisManager.h>
28 #include <AliVEvent.h>
29 #include <AliTriggerAnalysis.h>
30 #include <AliInputEventHandler.h>
31 #include <AliAODInputHandler.h>
32 #include <AliESDInputHandler.h>
33
34 #include "AliDielectron.h"
35 #include "AliDielectronMC.h"
36 #include "AliDielectronHistos.h"
37 #include "AliDielectronVarManager.h"
38 #include "AliAnalysisTaskDielectronFilter.h"
39 #include "AliAODCaloCluster.h"
40
41 ClassImp(AliAnalysisTaskDielectronFilter)
42
43 //_________________________________________________________________________________
44 AliAnalysisTaskDielectronFilter::AliAnalysisTaskDielectronFilter() :
45 AliAnalysisTaskSE(),
46 fDielectron(0),
47 fSelectPhysics(kTRUE),
48 fTriggerMask(AliVEvent::kMB),
49 fExcludeTriggerMask(0),
50 fTriggerOnV0AND(kFALSE),
51 fRejectPileup(kFALSE),
52 fEventStat(0x0),
53 fTriggerLogic(kAny),
54 fTriggerAnalysis(0x0),
55 fStoreLikeSign(kFALSE),
56 fStoreRotatedPairs(kFALSE),
57 fStoreEventsWithSingleTracks(kFALSE),
58 fCreateNanoAOD(kFALSE),
59 fStoreHeader(kFALSE),
60 fEventFilter(0x0)
61 {
62   //
63   // Constructor
64   //
65 }
66
67 //_________________________________________________________________________________
68 AliAnalysisTaskDielectronFilter::AliAnalysisTaskDielectronFilter(const char *name) :
69 AliAnalysisTaskSE(name),
70 fDielectron(0),
71 fSelectPhysics(kTRUE),
72 fTriggerMask(AliVEvent::kMB),
73 fExcludeTriggerMask(0),
74 fTriggerOnV0AND(kFALSE),
75 fRejectPileup(kFALSE),
76 fEventStat(0x0),
77 fTriggerLogic(kAny),
78 fTriggerAnalysis(0x0),
79 fStoreLikeSign(kFALSE),
80 fStoreRotatedPairs(kFALSE),
81 fStoreEventsWithSingleTracks(kFALSE),
82 fCreateNanoAOD(kFALSE),
83 fStoreHeader(kFALSE),
84 fEventFilter(0x0)
85 {
86   //
87   // Constructor
88   //
89   DefineInput(0,TChain::Class());
90   DefineOutput(1, THashList::Class());
91   DefineOutput(2, TH1D::Class());
92 }
93
94 //_________________________________________________________________________________
95 void AliAnalysisTaskDielectronFilter::Init()
96 {
97   // Initialization
98   if (fDebug > 1) AliInfo("Init() \n");
99   
100 // require AOD handler
101   AliAODHandler *aodH = (AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
102   if (!aodH) AliFatal("No AOD handler. Halting.");
103     
104   aodH->AddFilteredAOD("AliAOD.Dielectron.root", "DielectronEvents");
105 //   AddAODBranch("AliDielectronCandidates",fDielectron->GetPairArraysPointer(),"deltaAOD.Dielectron.root");
106 }
107
108 //_________________________________________________________________________________
109 void AliAnalysisTaskDielectronFilter::UserCreateOutputObjects()
110 {
111   //
112   // Initilise histograms
113   //
114
115   //require dielectron framework
116   if (!fDielectron) {
117     AliFatal("Dielectron framework class required. Please create and instance with proper cuts and set it via 'SetDielectron' before executing this task!!!");
118     return;
119   }
120   if(fStoreRotatedPairs) fDielectron->SetStoreRotatedPairs(kTRUE);
121   fDielectron->SetDontClearArrays(); 
122   fDielectron->Init();
123
124   Int_t nbins=kNbinsEvent+2;
125   if (!fEventStat){
126     fEventStat=new TH1D("hEventStat","Event statistics",nbins,0,nbins);
127     fEventStat->GetXaxis()->SetBinLabel(1,"Before Phys. Sel.");
128     fEventStat->GetXaxis()->SetBinLabel(2,"After Phys. Sel.");
129
130     //default names
131     fEventStat->GetXaxis()->SetBinLabel(3,"Bin3 not used");
132     fEventStat->GetXaxis()->SetBinLabel(4,"Bin4 not used");
133     fEventStat->GetXaxis()->SetBinLabel(5,"Bin5 not used");
134
135     if(fTriggerOnV0AND) fEventStat->GetXaxis()->SetBinLabel(3,"V0and triggers");
136     if (fEventFilter) fEventStat->GetXaxis()->SetBinLabel(4,"After Event Filter");
137     if (fRejectPileup) fEventStat->GetXaxis()->SetBinLabel(5,"After Pileup rejection");
138
139     fEventStat->GetXaxis()->SetBinLabel((kNbinsEvent+1),Form("#splitline{1 candidate}{%s}",fDielectron->GetName()));
140     fEventStat->GetXaxis()->SetBinLabel((kNbinsEvent+2),Form("#splitline{With >1 candidate}{%s}",fDielectron->GetName()));
141    }
142
143 Bool_t isAOD=AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler()->IsA()==AliAODInputHandler::Class(); 
144 if(fCreateNanoAOD && isAOD){
145   AliAODHandler *aodH = (AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
146   AliAODExtension *extDielectron = aodH->GetFilteredAOD("AliAOD.Dielectron.root");
147    TClonesArray *nanoAODTracks = new TClonesArray("AliAODTrack",500);
148    nanoAODTracks->SetName("tracks");
149    extDielectron->AddBranch("TClonesArray", &nanoAODTracks);
150    TClonesArray *nanoAODVertices = new TClonesArray("AliAODVertex",500);
151    nanoAODVertices->SetName("vertices");
152    extDielectron->AddBranch("TClonesArray", &nanoAODVertices);
153    TClonesArray *nanoAODCaloCluster = new TClonesArray("AliAODCaloCluster",500);
154    nanoAODCaloCluster->SetName("caloClusters");
155    extDielectron->AddBranch("TClonesArray", &nanoAODCaloCluster);  
156    extDielectron->GetAOD()->GetStdContent();
157    }else if(fCreateNanoAOD && !isAOD){AliWarning("Filtered-Nano AODs creation works only on AODs ");  } 
158   
159   PostData(1, const_cast<THashList*>(fDielectron->GetHistogramList()));
160   PostData(2,fEventStat);
161 }
162
163 //_________________________________________________________________________________
164 void AliAnalysisTaskDielectronFilter::UserExec(Option_t *)
165 {
166   //
167   // Main loop. Called for every event
168   //
169   
170   if (!fDielectron) return;
171   
172   AliAnalysisManager *man=AliAnalysisManager::GetAnalysisManager();
173   Bool_t isESD=man->GetInputEventHandler()->IsA()==AliESDInputHandler::Class();
174   Bool_t isAOD=man->GetInputEventHandler()->IsA()==AliAODInputHandler::Class();
175
176   
177   AliInputEventHandler* inputHandler = (AliInputEventHandler*) (man->GetInputEventHandler());
178   if (!inputHandler) return;
179
180   if ( inputHandler->GetPIDResponse() ){
181     AliDielectronVarManager::SetPIDResponse( inputHandler->GetPIDResponse() );
182   } else {
183     AliFatal("This task needs the PID response attached to the input event handler!");
184   }
185   
186   // Was event selected ?
187   ULong64_t isSelected = AliVEvent::kAny;
188   Bool_t isRejected = kFALSE;
189   if( fSelectPhysics && inputHandler){
190     if((isESD && inputHandler->GetEventSelection()) || isAOD){
191       isSelected = inputHandler->IsEventSelected();
192       if (fExcludeTriggerMask && (isSelected&fExcludeTriggerMask)) isRejected=kTRUE;
193       if (fTriggerLogic==kAny) isSelected&=fTriggerMask;
194       else if (fTriggerLogic==kExact) isSelected=((isSelected&fTriggerMask)==fTriggerMask);
195     }
196    }
197  
198   //before physics selection
199   fEventStat->Fill(kAllEvents);
200   if (isSelected==0||isRejected) {
201     PostData(2,fEventStat);
202     return;
203   }
204   //after physics selection
205   fEventStat->Fill(kSelectedEvents);
206
207   //V0and
208   if(fTriggerOnV0AND){
209   if(isESD){if (!fTriggerAnalysis->IsOfflineTriggerFired(static_cast<AliESDEvent*>(InputEvent()), AliTriggerAnalysis::kV0AND))
210             return;}
211   if(isAOD){if(!((static_cast<AliAODEvent*>(InputEvent()))->GetVZEROData()->GetV0ADecision() == AliVVZERO::kV0BB &&
212             (static_cast<AliAODEvent*>(InputEvent()))->GetVZEROData()->GetV0CDecision() == AliVVZERO::kV0BB) )
213             return;}
214    }
215
216    fEventStat->Fill(kV0andEvents);
217
218   //Fill Event histograms before the event filter
219   AliDielectronHistos *h=fDielectron->GetHistoManager();
220
221   Double_t values[AliDielectronVarManager::kNMaxValues]={0};
222   Double_t valuesMC[AliDielectronVarManager::kNMaxValues]={0};
223   if(h)  AliDielectronVarManager::SetFillMap(h->GetUsedVars());
224   else   AliDielectronVarManager::SetFillMap(0x0);
225   AliDielectronVarManager::SetEvent(InputEvent());
226   AliDielectronVarManager::Fill(InputEvent(),values);
227   AliDielectronVarManager::Fill(InputEvent(),valuesMC);
228
229   Bool_t hasMC=AliDielectronMC::Instance()->HasMC();
230   if (hasMC) {
231     if (AliDielectronMC::Instance()->ConnectMCEvent())
232       AliDielectronVarManager::Fill(AliDielectronMC::Instance()->GetMCEvent(),valuesMC);
233   }
234
235   if (h){
236     if (h->GetHistogramList()->FindObject("Event_noCuts"))
237       h->FillClass("Event_noCuts",AliDielectronVarManager::kNMaxValues,values);
238     if (hasMC && h->GetHistogramList()->FindObject("MCEvent_noCuts"))
239       h->FillClass("Event_noCuts",AliDielectronVarManager::kNMaxValues,valuesMC);
240   }
241
242   //event filter
243   if (fEventFilter) {
244   if (!fEventFilter->IsSelected(InputEvent())) return;
245   }
246   fEventStat->Fill(kFilteredEvents);
247
248   //pileup
249   if (fRejectPileup){
250   if (InputEvent()->IsPileupFromSPD(3,0.8,3.,2.,5.)) return;
251   }
252   fEventStat->Fill(kPileupEvents);
253
254   //bz for AliKF
255   Double_t bz = InputEvent()->GetMagneticField();
256   AliKFParticle::SetField( bz );
257
258   AliDielectronPID::SetCorrVal((Double_t)InputEvent()->GetRunNumber());
259   
260   fDielectron->Process(InputEvent());
261   
262   Bool_t hasCand = kFALSE;
263   if(fStoreLikeSign) hasCand = (fDielectron->HasCandidates() || fDielectron->HasCandidatesLikeSign());
264   else hasCand = (fDielectron->HasCandidates());
265
266   if(fStoreRotatedPairs) hasCand = (hasCand || fDielectron->HasCandidatesTR());
267  
268   if(fStoreEventsWithSingleTracks) hasCand = (hasCand || fDielectron->GetTrackArray(0) || fDielectron->GetTrackArray(1));
269
270   
271   AliAODHandler *aodH=(AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
272   AliAODExtension *extDielectron = aodH->GetFilteredAOD("AliAOD.Dielectron.root");
273   if(hasCand){
274     AliAODEvent *aod = aodH->GetAOD();
275     
276     // reset bit for all tracks
277     if(isAOD){
278     for(Int_t it=0;it<aod->GetNumberOfTracks();it++){
279         aod->GetTrack(it)->ResetBit(kIsReferenced);  aod->GetTrack(it)->SetUniqueID(0);
280         }
281     }
282
283     //replace the references of the legs with the AOD references
284     TObjArray *obj = 0x0;
285     for(Int_t i=0; i < 11; i++ ){
286       obj = (TObjArray*)((*(fDielectron->GetPairArraysPointer()))->UncheckedAt(i));
287       if(!obj) continue;
288       for(int j=0;j<obj->GetEntriesFast();j++){
289         AliDielectronPair *pairObj = (AliDielectronPair*)obj->UncheckedAt(j);
290         Int_t id1 = ((AliVTrack*)pairObj->GetFirstDaughterP())->GetID();
291         Int_t id2 = ((AliVTrack*)pairObj->GetSecondDaughterP())->GetID();
292         
293         for(Int_t it=0;it<aod->GetNumberOfTracks();it++){
294           if(aod->GetTrack(it)->GetID() == id1) pairObj->SetRefFirstDaughter(aod->GetTrack(it)); 
295           if(aod->GetTrack(it)->GetID() == id2) pairObj->SetRefSecondDaughter(aod->GetTrack(it));
296         }
297       }
298     }
299     
300     extDielectron->SelectEvent();
301     Int_t ncandidates=fDielectron->GetPairArray(1)->GetEntriesFast();
302     if (ncandidates==1) fEventStat->Fill((kNbinsEvent));
303     else if (ncandidates>1) fEventStat->Fill((kNbinsEvent+1));
304  
305     //see if dielectron candidate branch exists, if not create is
306     TTree *t=extDielectron->GetTree();
307
308     if(!t->GetListOfBranches()->GetEntries() && isAOD)
309       t->Branch(aod->GetList());
310     
311     if (!t->GetBranch("dielectrons"))
312       t->Bronch("dielectrons","TObjArray",fDielectron->GetPairArraysPointer());
313       
314       // store positive and negative tracks
315       if(fCreateNanoAOD && isAOD){
316       Int_t nTracks = (fDielectron->GetTrackArray(0))->GetEntries() + (fDielectron->GetTrackArray(1))->GetEntries();
317        AliAODEvent *nanoEv = extDielectron->GetAOD();
318        nanoEv->GetTracks()->Clear();      
319        nanoEv->GetVertices()->Clear();      
320        nanoEv->GetCaloClusters()->Clear();      
321        
322        AliAODVertex* tmp = ((static_cast<AliAODEvent*>(InputEvent()))->GetPrimaryVertex())->CloneWithoutRefs();
323        nanoEv->AddVertex(tmp);
324        AliAODVertex* tmpSpd = ((static_cast<AliAODEvent*>(InputEvent()))->GetPrimaryVertexSPD())->CloneWithoutRefs();
325        nanoEv->AddVertex(tmpSpd);
326        nanoEv->GetVertex(0)->SetNContributors((static_cast<AliAODEvent*>(InputEvent()))->GetPrimaryVertex()->GetNContributors());
327        nanoEv->GetVertex(1)->SetNContributors((static_cast<AliAODEvent*>(InputEvent()))->GetPrimaryVertexSPD()->GetNContributors());
328        // set event plane 
329        AliAODHeader * header = dynamic_cast<AliAODHeader*>(nanoEv->GetHeader());
330        if(!header) AliFatal("Not a standard AOD");
331
332
333        header->SetEventplane(((AliVAODHeader*)static_cast<AliAODEvent*>(InputEvent())->GetHeader())->GetEventplaneP());
334        header->ResetEventplanePointer(); 
335        // set multiplicity
336        header->SetRefMultiplicity((Int_t)values[AliDielectronVarManager::kNTrk]);
337        header->SetRefMultiplicityPos((Int_t)values[AliDielectronVarManager::kNacc]);
338        //nanoEv->GetHeader()->SetRefMultiplicityNeg(values[AliDielectronVarManager::kMatchEffITSTPC]);
339
340          for(int kj=0; kj<(fDielectron->GetTrackArray(0))->GetEntries(); kj++){
341          Int_t posit = nanoEv->AddTrack((AliAODTrack*)fDielectron->GetTrackArray(0)->At(kj));
342          Int_t posVtx = nanoEv->AddVertex(((AliAODTrack*)fDielectron->GetTrackArray(0)->At(kj))->GetProdVertex()); 
343          nanoEv->GetVertex(posVtx)->ResetBit(kIsReferenced);  
344          nanoEv->GetVertex(posVtx)->SetUniqueID(0);
345          nanoEv->GetVertex(posVtx)->RemoveDaughters();
346          nanoEv->GetTrack(posit)->ResetBit(kIsReferenced);  
347          nanoEv->GetTrack(posit)->SetUniqueID(0); 
348          // calo cluster
349          Int_t caloIndex = ((AliAODTrack*)fDielectron->GetTrackArray(0)->At(kj))->GetEMCALcluster();
350          if(caloIndex > 0 && (static_cast<AliAODEvent*>(InputEvent()))->GetCaloCluster(caloIndex)){
351          Int_t posCaloCls = nanoEv->AddCaloCluster(static_cast<AliAODEvent*>(InputEvent())->GetCaloCluster(caloIndex));
352          nanoEv->GetTrack(posit)->SetEMCALcluster(posCaloCls);
353          AliAODCaloCluster *clCls = nanoEv->GetCaloCluster(posCaloCls);
354          for(int u=0; u<clCls->GetNTracksMatched(); u++) clCls->RemoveTrackMatched(clCls->GetTrackMatched(u)); 
355          nanoEv->GetCaloCluster(posCaloCls)->AddTrackMatched((AliAODTrack*)nanoEv->GetTrack(posit)); 
356          }
357          // set references for vtx
358          nanoEv->GetTrack(posit)->SetProdVertex(nanoEv->GetVertex(posVtx));
359          }
360
361          for(int kj=0; kj<(fDielectron->GetTrackArray(1))->GetEntries(); kj++){
362          Int_t negat = nanoEv->AddTrack((AliAODTrack*)fDielectron->GetTrackArray(1)->At(kj));
363          Int_t negVtx = nanoEv->AddVertex(((AliAODTrack*)fDielectron->GetTrackArray(1)->At(kj))->GetProdVertex());    
364          nanoEv->GetVertex(negVtx)->ResetBit(kIsReferenced);  
365          nanoEv->GetVertex(negVtx)->SetUniqueID(0);
366          nanoEv->GetVertex(negVtx)->RemoveDaughters();
367          nanoEv->GetTrack(negat)->ResetBit(kIsReferenced);  
368          nanoEv->GetTrack(negat)->SetUniqueID(0);
369          // calo cluster
370          Int_t caloIndex = ((AliAODTrack*)fDielectron->GetTrackArray(1)->At(kj))->GetEMCALcluster(); 
371          if(caloIndex > 0 && (static_cast<AliAODEvent*>(InputEvent()))->GetCaloCluster(caloIndex)){
372          Int_t negCaloCls = nanoEv->AddCaloCluster(static_cast<AliAODEvent*>(InputEvent())->GetCaloCluster(caloIndex));
373          nanoEv->GetTrack(negat)->SetEMCALcluster(negCaloCls);
374          AliAODCaloCluster *clCls = nanoEv->GetCaloCluster(negCaloCls); 
375          for(int u=0; u<clCls->GetNTracksMatched(); u++) clCls->RemoveTrackMatched(clCls->GetTrackMatched(u));
376          nanoEv->GetCaloCluster(negCaloCls)->AddTrackMatched((AliAODTrack*)nanoEv->GetTrack(negat));
377          }
378          nanoEv->GetTrack(negat)->SetProdVertex(nanoEv->GetVertex(negVtx)); 
379          }  
380         delete tmp; delete tmpSpd; 
381         nanoEv->GetTracks()->Expand(nTracks); 
382         nanoEv->GetVertices()->Expand(nTracks+2); 
383         nanoEv->GetCaloClusters()->Expand(nanoEv->GetNumberOfCaloClusters()); 
384    
385     }
386     if(isAOD) t->Fill();
387   }
388  
389   if(fCreateNanoAOD && isAOD && (!hasCand) &&  fStoreHeader)  
390    {
391      // set event plane 
392      AliAODHeader * header = dynamic_cast<AliAODHeader*>(extDielectron->GetAOD()->GetHeader());
393      if(!header) AliFatal("Not a standard AOD");     
394      header->SetEventplane(((AliAODHeader*)(static_cast<AliAODEvent*>(InputEvent()))->GetHeader())->GetEventplaneP());
395      header->ResetEventplanePointer();
396      extDielectron->GetTree()->Fill(); // fill header for all events without tracks
397    }
398  
399   PostData(1, const_cast<THashList*>(fDielectron->GetHistogramList()));
400   PostData(2,fEventStat);
401   return;
402 }
403
404