changes for Vertex and Tracks classes
[u/mrichter/AliRoot.git] / PWGDQ / dielectron / macrosJPSI / RaaPbPb2010 / AliAnalysisTaskCorrelationTree.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 //    Analysis task for creating a reduced data tree                     //
19 //                                                                       //
20 ///////////////////////////////////////////////////////////////////////////
21
22 #include <iostream>
23 using namespace std;
24
25 #include <TChain.h>
26 #include <TH1D.h>
27 #include <TFile.h>
28
29 #include <AliCFContainer.h>
30 #include <AliInputEventHandler.h>
31 #include <AliESDInputHandler.h>
32 #include <AliAODInputHandler.h>
33 #include <AliAnalysisManager.h>
34 #include <AliVEvent.h>
35 #include <AliESDEvent.h>
36 #include <AliAODEvent.h>
37 #include <AliAODTrack.h>
38 #include <AliTriggerAnalysis.h>
39 #include <AliESDtrackCuts.h>
40 #include <AliVZDC.h>
41 #include <AliESDv0.h>
42 #include <AliESDv0Cuts.h>
43 #include <AliVCluster.h>
44 #include "AliDielectron.h"
45 #include "AliDielectronHistos.h"
46 #include "AliDielectronMC.h"
47 #include "AliDielectronVarManager.h"
48 #include "AliFlowTrackCuts.h"
49 #include "AliFlowBayesianPID.h"
50
51 #include "AliCorrelationReducedEvent.h"
52 #include "AliAnalysisTaskCorrelationTree.h"
53
54 ClassImp(AliAnalysisTaskCorrelationTree)
55
56
57 //_________________________________________________________________________________
58 AliAnalysisTaskCorrelationTree::AliAnalysisTaskCorrelationTree() :
59   AliAnalysisTaskSE(),
60   fListDielectron(),
61   fListHistos(),
62   fSelectPhysics(kFALSE),
63   fTriggerMask(AliVEvent::kAny),
64   fRejectPileup(kFALSE),
65   fEventFilter(0x0),
66   fTrackFilter(0x0),
67   fFlowTrackFilter(0x0),
68   fK0sCuts(0x0),
69   fLambdaCuts(0x0),
70   fK0sPionCuts(0x0),
71   fLambdaProtonCuts(0x0),
72   fLambdaPionCuts(0x0),
73   fK0sMassRange(),
74   fLambdaMassRange(),
75   fV0Histos(0x0),
76   fTreeFile(0x0),
77   fTree(0x0),
78   fFriendTreeFile(0x0),
79   fFriendTree(0x0),
80   fReducedEvent(0x0),
81   fReducedEventFriend(0x0),
82   fFlowTrackCuts(0x0)
83 {
84   //
85   // Constructor
86   //
87 }
88
89 //_________________________________________________________________________________
90 AliAnalysisTaskCorrelationTree::AliAnalysisTaskCorrelationTree(const char *name) :
91   AliAnalysisTaskSE(name),
92   fListDielectron(),
93   fListHistos(),
94   fSelectPhysics(kFALSE),
95   fTriggerMask(AliVEvent::kAny),
96   fRejectPileup(kFALSE),
97   fEventFilter(0x0),
98   fTrackFilter(0x0),
99   fFlowTrackFilter(0x0),
100   fK0sCuts(0x0),
101   fLambdaCuts(0x0),
102   fK0sPionCuts(0x0),
103   fLambdaProtonCuts(0x0),
104   fLambdaPionCuts(0x0),
105   fK0sMassRange(),
106   fLambdaMassRange(),
107   fV0Histos(0x0),
108   fTreeFile(0x0),
109   fTree(0x0),
110   fFriendTreeFile(0x0),
111   fFriendTree(0x0),
112   fReducedEvent(0x0),
113   fReducedEventFriend(0x0),
114   fFlowTrackCuts(0x0)
115 {
116   //
117   // Constructor
118   //
119   fK0sMassRange[0] = 0.4; fK0sMassRange[1] = 0.6;
120   fLambdaMassRange[0] = 1.08; fLambdaMassRange[1] = 1.15;
121   
122   DefineInput(0,TChain::Class());
123   DefineOutput(1, TList::Class());   // QA histograms
124   //DefineOutput(2, TTree::Class());   // reduced information tree
125   //DefineOutput(3, TTree::Class());   // reduced information tree with friends
126   DefineOutput(2, TTree::Class());   // reduced information tree with friends
127   
128   fListHistos.SetName("QAhistograms");
129   fListDielectron.SetOwner();
130   fListHistos.SetOwner(kFALSE);
131 }
132
133
134 //_________________________________________________________________________________
135 void AliAnalysisTaskCorrelationTree::UserCreateOutputObjects()
136 {
137   //
138   // Add all histogram manager histogram lists to the output TList
139   //
140
141   if (!fListHistos.IsEmpty() || fTree || fFriendTree) return; //already initialised
142
143   TIter nextDie(&fListDielectron);
144   AliDielectron *die=0;
145   while ( (die=static_cast<AliDielectron*>(nextDie())) ){
146     die->Init();
147     if (die->GetHistogramList()) fListHistos.Add(const_cast<THashList*>(die->GetHistogramList()));
148   }  
149   if(fV0Histos) fListHistos.Add(const_cast<THashList*>(fV0Histos->GetHistogramList()));
150
151   fFriendTree = new TTree("DstFriendTree","Reduced ESD information");
152   fReducedEventFriend = new AliCorrelationReducedEventFriend();
153   fFriendTree->Branch("Event",&fReducedEventFriend,16000,99);
154   
155   fTreeFile = new TFile("dstTree.root", "RECREATE");
156   fTree = new TTree("DstTree","Reduced ESD information");
157   fReducedEvent = new AliCorrelationReducedEvent();
158   fTree->Branch("Event",&fReducedEvent,16000,99);
159     
160   fFlowTrackCuts = new AliFlowTrackCuts("flow cuts");
161   fFlowTrackCuts->SetPID(AliPID::kPion, AliFlowTrackCuts::kTOFbayesian);
162   
163   PostData(1, &fListHistos);
164   //PostData(2, fTree);
165   //PostData(3, fFriendTree);
166   PostData(2, fFriendTree);
167 }
168
169 //_________________________________________________________________________________
170 void AliAnalysisTaskCorrelationTree::UserExec(Option_t *option)
171 {
172   //
173   // Main loop. Called for every event
174   //  
175   //cout << "Event" << endl;
176   option = option;
177   AliAnalysisManager *man=AliAnalysisManager::GetAnalysisManager();
178   Bool_t isESD=man->GetInputEventHandler()->IsA()==AliESDInputHandler::Class();
179   Bool_t isAOD=man->GetInputEventHandler()->IsA()==AliAODInputHandler::Class();
180   
181   AliInputEventHandler* inputHandler = (AliInputEventHandler*) (man->GetInputEventHandler());
182   if (!inputHandler) return;
183   
184   if ( inputHandler->GetPIDResponse() ){
185     AliDielectronVarManager::SetPIDResponse( inputHandler->GetPIDResponse() );
186   } else {
187     AliFatal("This task needs the PID response attached to the input event handler!");
188   }
189   
190   fFlowTrackCuts->SetEvent(InputEvent());
191   
192   // Was event selected ?
193   UInt_t isSelected = AliVEvent::kAny;
194   if(fSelectPhysics && inputHandler){
195     if((isESD && inputHandler->GetEventSelection()) || isAOD){
196       isSelected = inputHandler->IsEventSelected();
197       isSelected&=fTriggerMask;
198     }
199   }
200
201   if(isSelected==0) {
202     return;
203   }
204
205   // fill event histograms before event filter
206   Double_t values[AliDielectronVarManager::kNMaxValues]={0};
207   AliDielectronVarManager::Fill(InputEvent(),values);
208   
209   TIter nextDie(&fListDielectron);
210   AliDielectron *die=0;
211   while ( (die=static_cast<AliDielectron*>(nextDie())) ){
212     AliDielectronHistos *h=die->GetHistoManager();
213     if (h){
214       if (h->GetHistogramList()->FindObject("Event_noCuts"))
215         h->FillClass("Event_noCuts",AliDielectronVarManager::kNMaxValues,values);
216     }
217   }
218   nextDie.Reset();
219
220   //event filter
221   if (fEventFilter) {
222     if (!fEventFilter->IsSelected(InputEvent())) return;
223   }
224   
225   //pileup
226   if (fRejectPileup){
227     if (InputEvent()->IsPileupFromSPD(3,0.8,3.,2.,5.)) return;
228   }
229
230   //cout << "Event selected" << endl;
231   //bz for AliKF
232   Double_t bz = InputEvent()->GetMagneticField();
233   AliKFParticle::SetField( bz );
234
235   //Process event in all AliDielectron instances
236   fReducedEvent->ClearEvent();
237   fReducedEventFriend->ClearEvent();
238   FillEventInfo();
239   FillV0PairInfo();
240   
241   Short_t idie=0;
242   while((die=static_cast<AliDielectron*>(nextDie()))){
243     die->Process(InputEvent());
244     FillDielectronPairInfo(die, idie);
245     ++idie;
246   }
247   nextDie.Reset();
248   
249   FillTrackInfo();
250   FillFriendEventInfo();              // Q-vector calculation
251   
252   fTree->Fill();
253   fFriendTree->Fill();
254       
255   // if there are candidate pairs, add the information to the reduced tree
256   PostData(1, &fListHistos);
257   //PostData(2, fTree);
258   //PostData(3, fFriendTree);
259   PostData(2, fFriendTree);
260 }
261
262
263 //_________________________________________________________________________________
264 void AliAnalysisTaskCorrelationTree::FillEventInfo() 
265 {
266   //
267   // fill reduced event information
268   //
269   AliVEvent* event = InputEvent();
270   // Was event selected ?
271   AliAnalysisManager *man=AliAnalysisManager::GetAnalysisManager();
272   Bool_t isESD = (event->IsA()==AliESDEvent::Class());
273   Bool_t isAOD = (event->IsA()==AliAODEvent::Class());
274   
275   AliInputEventHandler* inputHandler = (AliInputEventHandler*) (man->GetInputEventHandler());
276   UInt_t isSelected = AliVEvent::kAny;
277   if(inputHandler){
278     if((isESD && inputHandler->GetEventSelection()) || isAOD){
279       isSelected = inputHandler->IsEventSelected();
280       isSelected&=fTriggerMask;
281     }
282   }
283
284   Double_t values[AliDielectronVarManager::kNMaxValues];
285   AliDielectronVarManager::Fill(event, values);
286   
287   fReducedEvent->fRunNo       = event->GetRunNumber();
288   fReducedEvent->fBC          = event->GetBunchCrossNumber();
289   fReducedEvent->fTriggerMask = event->GetTriggerMask();
290   fReducedEvent->fIsPhysicsSelection = (isSelected!=0 ? kTRUE : kFALSE);
291   AliVVertex* eventVtx = 0x0;
292   if(isESD) eventVtx = const_cast<AliESDVertex*>((static_cast<AliESDEvent*>(event))->GetPrimaryVertexTracks());
293   if(isAOD) eventVtx = const_cast<AliAODVertex*>((static_cast<AliAODEvent*>(event))->GetPrimaryVertex());
294   if(eventVtx) {
295     fReducedEvent->fVtx[0] = (isESD ? ((AliESDVertex*)eventVtx)->GetX() : ((AliAODVertex*)eventVtx)->GetX());
296     fReducedEvent->fVtx[1] = (isESD ? ((AliESDVertex*)eventVtx)->GetY() : ((AliAODVertex*)eventVtx)->GetY());
297     fReducedEvent->fVtx[2] = (isESD ? ((AliESDVertex*)eventVtx)->GetZ() : ((AliAODVertex*)eventVtx)->GetZ());
298     fReducedEvent->fNVtxContributors = eventVtx->GetNContributors();
299   }
300   if(isESD) {
301     eventVtx = const_cast<AliESDVertex*>((static_cast<AliESDEvent*>(event))->GetPrimaryVertexTPC());
302     if(eventVtx) {
303       fReducedEvent->fVtxTPC[0] = ((AliESDVertex*)eventVtx)->GetX();
304       fReducedEvent->fVtxTPC[1] = ((AliESDVertex*)eventVtx)->GetY();
305       fReducedEvent->fVtxTPC[2] = ((AliESDVertex*)eventVtx)->GetZ();
306       fReducedEvent->fNVtxTPCContributors = eventVtx->GetNContributors();
307     }
308   }
309   
310   AliCentrality *centrality = event->GetCentrality();
311   if(centrality) {
312     fReducedEvent->fCentrality[0] = centrality->GetCentralityPercentile("V0M");
313     fReducedEvent->fCentrality[1] = centrality->GetCentralityPercentile("CL1");
314     fReducedEvent->fCentrality[2] = centrality->GetCentralityPercentile("TRK");
315     fReducedEvent->fCentrality[3] = centrality->GetCentralityPercentile("ZEMvsZDC");    
316     fReducedEvent->fCentQuality   = centrality->GetQuality();
317   }
318   
319   fReducedEvent->fNtracks[0] = event->GetNumberOfTracks();
320   fReducedEvent->fSPDntracklets = values[AliDielectronVarManager::kNaccTrckltsEsd10Corr];
321
322   AliVVZERO* vzero = event->GetVZEROData();
323   for(Int_t i=0;i<64;++i) 
324     fReducedEvent->fVZEROMult[i] = vzero->GetMultiplicity(i);  
325
326   AliESDZDC* zdc = (isESD ? (static_cast<AliESDEvent*>(event))->GetESDZDC() : 0x0);
327   if(zdc) {
328     for(Int_t i=0; i<4; ++i)  fReducedEvent->fZDCnEnergy[i]   = zdc->GetZN1TowerEnergy()[i];
329     for(Int_t i=4; i<8; ++i)  fReducedEvent->fZDCnEnergy[i]   = zdc->GetZN2TowerEnergy()[i-5];
330   }
331   
332   // EMCAL/PHOS clusters
333   FillCaloClusters();
334   
335   // TODO FMD multiplicities
336   
337 }
338
339
340 //_________________________________________________________________________________
341 void AliAnalysisTaskCorrelationTree::FillCaloClusters() {
342   //
343   // Fill info about the calorimeter clusters
344   //
345   AliVEvent* event = InputEvent();
346   Int_t nclusters = event->GetNumberOfCaloClusters();
347
348   fReducedEvent->fNCaloClusters = 0;
349   for(Int_t iclus=0; iclus<nclusters; ++iclus) {
350     AliVCluster* cluster = event->GetCaloCluster(iclus);
351     
352     TClonesArray& clusters = *(fReducedEvent->fCaloClusters);
353     AliCorrelationReducedCaloCluster *reducedCluster=new(clusters[fReducedEvent->fNCaloClusters]) AliCorrelationReducedCaloCluster();
354     
355     reducedCluster->fType    = (cluster->IsEMCAL() ? AliCorrelationReducedCaloCluster::kEMCAL : AliCorrelationReducedCaloCluster::kPHOS);
356     reducedCluster->fEnergy  = cluster->E();
357     reducedCluster->fTrackDx = cluster->GetTrackDx();
358     reducedCluster->fTrackDz = cluster->GetTrackDz();
359     fReducedEvent->fNCaloClusters += 1;
360   }  // end loop over clusters
361 }
362
363
364 //_________________________________________________________________________________
365 void AliAnalysisTaskCorrelationTree::FillFriendEventInfo() {
366   //
367   // Fill event info into the friend tree
368   //
369   // Add here calculated Q-vector components from all detectors
370   for(Int_t idet=0; idet<AliCorrelationReducedEventFriend::kNdetectors; ++idet) {
371     fReducedEvent->GetQvector(fReducedEventFriend->fQvector[idet], idet);
372     for(Int_t ih=0; ih<fgkNMaxHarmonics; ++ih)
373       fReducedEventFriend->fEventPlaneStatus[idet][ih] = AliCorrelationReducedEventFriend::kRaw;
374   }
375 }
376
377
378 //_________________________________________________________________________________
379 void AliAnalysisTaskCorrelationTree::FillTrackInfo() 
380 {
381   //
382   // fill reduced track information
383   //
384   AliVEvent* event = InputEvent();
385   Bool_t isESD = (event->IsA()==AliESDEvent::Class());
386   Bool_t isAOD = (event->IsA()==AliAODEvent::Class());
387
388   Int_t ntracks=event->GetNumberOfTracks();
389   for(Int_t itrack=0; itrack<ntracks; ++itrack){
390     AliVParticle *particle=event->GetTrack(itrack);
391     //apply track cuts
392     if(fTrackFilter && !fTrackFilter->IsSelected(particle)) continue;
393     
394     TClonesArray& tracks = *(fReducedEvent->fTracks);
395     AliCorrelationReducedTrack *reducedParticle=new(tracks[fReducedEvent->fNtracks[1]]) AliCorrelationReducedTrack();
396         
397     Double_t values[AliDielectronVarManager::kNMaxValues];
398     AliDielectronVarManager::Fill(particle, values);
399     reducedParticle->fStatus        = (ULong_t)values[AliDielectronVarManager::kTrackStatus];
400     reducedParticle->fGlobalPhi     = values[AliDielectronVarManager::kPhi];
401     reducedParticle->fGlobalPt      = values[AliDielectronVarManager::kPt]*values[AliDielectronVarManager::kCharge];
402     reducedParticle->fGlobalEta     = values[AliDielectronVarManager::kEta];    
403     reducedParticle->fMomentumInner = values[AliDielectronVarManager::kPIn];
404     reducedParticle->fDCA[0]        = values[AliDielectronVarManager::kImpactParXY];
405     reducedParticle->fDCA[1]        = values[AliDielectronVarManager::kImpactParZ];
406     
407     reducedParticle->fITSclusterMap = values[AliDielectronVarManager::kITSclusterMap];
408     reducedParticle->fITSsignal     = values[AliDielectronVarManager::kITSsignal];
409     
410     reducedParticle->fTPCNcls      = (UChar_t)values[AliDielectronVarManager::kNclsTPC];
411     reducedParticle->fTPCNclsF     = (UChar_t)values[AliDielectronVarManager::kNFclsTPC];
412     reducedParticle->fTPCNclsIter1 = (UChar_t)values[AliDielectronVarManager::kNclsTPCiter1];
413     reducedParticle->fTPCsignal    = values[AliDielectronVarManager::kTPCsignal];
414     reducedParticle->fTPCnSig[0]   = values[AliDielectronVarManager::kTPCnSigmaEle];
415     reducedParticle->fTPCnSig[1]   = values[AliDielectronVarManager::kTPCnSigmaPio];
416     reducedParticle->fTPCnSig[2]   = values[AliDielectronVarManager::kTPCnSigmaKao];
417     reducedParticle->fTPCnSig[3]   = values[AliDielectronVarManager::kTPCnSigmaPro];
418     
419     reducedParticle->fTOFbeta      = values[AliDielectronVarManager::kTOFbeta];
420     reducedParticle->fTOFnSig[0]   = values[AliDielectronVarManager::kTOFnSigmaEle];
421     reducedParticle->fTOFnSig[1]   = values[AliDielectronVarManager::kTOFnSigmaPio];
422     reducedParticle->fTOFnSig[2]   = values[AliDielectronVarManager::kTOFnSigmaKao];
423     reducedParticle->fTOFnSig[3]   = values[AliDielectronVarManager::kTOFnSigmaPro];
424
425     reducedParticle->fTRDpid[0]    = values[AliDielectronVarManager::kTRDprobEle];
426     reducedParticle->fTRDpid[1]    = values[AliDielectronVarManager::kTRDprobPio];
427     
428     if(fFlowTrackFilter) {
429       // switch on the first bit if this particle should be used for the event plane
430       if(fFlowTrackFilter->IsSelected(particle)) reducedParticle->fFlags |= (1<<0);
431     }
432     
433     if(isESD){
434       AliESDtrack *track=static_cast<AliESDtrack*>(particle);
435       reducedParticle->fTrackId          = (UShort_t)track->GetID();
436       reducedParticle->fTPCCrossedRows   = (UChar_t)track->GetTPCCrossedRows();
437       reducedParticle->fTPCClusterMap    = EncodeTPCClusterMap(track);
438       const AliExternalTrackParam* tpcInner = track->GetTPCInnerParam();
439       reducedParticle->fTPCPhi        = (tpcInner ? tpcInner->Phi() : 0.0);
440       reducedParticle->fTPCPt         = (tpcInner ? tpcInner->Pt() : 0.0);
441       reducedParticle->fTPCEta        = (tpcInner ? tpcInner->Eta() : 0.0);
442       reducedParticle->fTRDntracklets[0] = track->GetTRDntracklets();
443       reducedParticle->fTRDntracklets[1] = track->GetTRDntrackletsPID();
444       fFlowTrackCuts->IsSelected(track);
445       AliFlowBayesianPID* bayesResponse = fFlowTrackCuts->GetBayesianResponse();
446       reducedParticle->fBayesPID[0] = bayesResponse->GetProb()[AliPID::kPion];
447       reducedParticle->fBayesPID[1] = bayesResponse->GetProb()[AliPID::kKaon];
448       reducedParticle->fBayesPID[2] = bayesResponse->GetProb()[AliPID::kProton];
449       if(track->IsEMCAL()) reducedParticle->fCaloClusterId = track->GetEMCALcluster();
450       if(track->IsPHOS()) reducedParticle->fCaloClusterId = track->GetPHOScluster();
451     }
452     if(isAOD) {
453       AliAODTrack *track=static_cast<AliAODTrack*>(particle);
454       reducedParticle->fTrackId = track->GetID(); 
455       if(track->IsEMCAL()) reducedParticle->fCaloClusterId = track->GetEMCALcluster();
456       if(track->IsPHOS()) reducedParticle->fCaloClusterId = track->GetPHOScluster();
457     }
458
459     fReducedEvent->fNtracks[1] += 1;
460   }
461 }
462
463
464 //_________________________________________________________________________________
465 void AliAnalysisTaskCorrelationTree::FillDielectronPairInfo(AliDielectron* die, Short_t iDie) 
466 {
467   //
468   // fill reduced pair information
469   //
470   Bool_t hasMC=AliDielectronMC::Instance()->HasMC();
471
472   for(Int_t iType=0; iType<3; ++iType) {
473     
474     const TObjArray* array = die->GetPairArray(iType);
475     if(!array || array->GetEntriesFast()==0) continue;
476     cout << "dielectron tracks (idie/type/n): " << iDie << " / 0 / " << die->GetTrackArray(0)->GetEntriesFast() << endl;
477     cout << "dielectron tracks (idie/type/n): " << iDie << " / 1 / " << die->GetTrackArray(1)->GetEntriesFast() << endl;
478
479     cout << "diele idie/type/n-candidates: " << iDie << " / " << iType << " / " << array->GetEntriesFast() << endl;
480     for(Int_t iCandidate=0; iCandidate<array->GetEntriesFast(); ++iCandidate) {
481       if(iCandidate%100==0) cout << "iCandidate = " << iCandidate << endl;
482       AliDielectronPair* pair = (AliDielectronPair*)array->At(iCandidate);
483       Double_t values[AliDielectronVarManager::kNMaxValues];
484       AliDielectronVarManager::Fill(pair, values);
485       
486       TClonesArray& tracks = *(fReducedEvent->fCandidates);
487       AliCorrelationReducedPair *reducedParticle= 
488          new (tracks[fReducedEvent->fNV0candidates[1]+fReducedEvent->fNDielectronCandidates]) AliCorrelationReducedPair();
489       // !!! hardcoded flag for dielectron id 
490       reducedParticle->fCandidateId  = (iDie==0 ? AliCorrelationReducedPair::kJpsiToEE : AliCorrelationReducedPair::kPhiToKK);
491       reducedParticle->fPairType     = values[AliDielectronVarManager::kPairType];
492       reducedParticle->fLegIds[0]    = (UShort_t)(static_cast<AliVTrack*>(pair->GetFirstDaughter()))->GetID();
493       reducedParticle->fLegIds[1]    = (UShort_t)(static_cast<AliVTrack*>(pair->GetSecondDaughter()))->GetID();
494       reducedParticle->fMass[0]      = values[AliDielectronVarManager::kM];
495       reducedParticle->fMass[1]      = -999.;
496       reducedParticle->fMass[2]      = -999.;
497       reducedParticle->fPhi          = values[AliDielectronVarManager::kPhi];  // in the [-pi,pi] interval
498       if(reducedParticle->fPhi<0.0) reducedParticle->fPhi = 2.0*TMath::Pi() + reducedParticle->fPhi;  // converted to [0,2pi]
499       reducedParticle->fPt           = values[AliDielectronVarManager::kPt];
500       reducedParticle->fEta          = values[AliDielectronVarManager::kEta];
501       reducedParticle->fLxy          = values[AliDielectronVarManager::kPseudoProperTime];
502       reducedParticle->fOpeningAngle = values[AliDielectronVarManager::kOpeningAngle];     
503      
504       reducedParticle->fMCid         = 0;
505       if(hasMC) {
506         AliDielectronMC::Instance()->ConnectMCEvent();
507         const TObjArray* mcSignals = die->GetMCSignals();
508         for(Int_t iSig=0; iSig<mcSignals->GetEntries(); ++iSig) {
509           if(iSig>31) break;
510           AliDielectronMC *mc=AliDielectronMC::Instance();
511           if(mc->IsMCTruth(pair, (AliDielectronSignalMC*)mcSignals->At(iSig))) {
512             reducedParticle->fMCid = reducedParticle->fMCid | (1<<iSig);
513           }
514         }
515       }   // end if has MC
516       fReducedEvent->fNDielectronCandidates += 1;
517     }    // end loop over candidates
518   }    // end loop over pair type
519 }
520
521
522 //_________________________________________________________________________________
523 void AliAnalysisTaskCorrelationTree::FillV0PairInfo() 
524 {
525   //
526   // fill reduced pair information
527   //
528   AliESDEvent* esd = (AliESDEvent*)InputEvent();
529   const AliESDVertex *primaryVertex = esd->GetPrimaryVertex();
530   AliKFVertex primaryVertexKF(*primaryVertex);
531   
532   fReducedEvent->fNV0candidates[0] = InputEvent()->GetNumberOfV0s();
533   
534   Double_t valuesPos[AliDielectronVarManager::kNMaxValues];
535   Double_t valuesNeg[AliDielectronVarManager::kNMaxValues];
536    
537   cout << "n total V0s: " << InputEvent()->GetNumberOfV0s() << endl;
538   for(Int_t iV0=0; iV0<InputEvent()->GetNumberOfV0s(); ++iV0) {   // loop over V0s
539     if(iV0%1000==0) cout << "iV0 = " << iV0 << endl;
540     AliESDv0 *v0 = esd->GetV0(iV0);
541        
542     AliESDtrack* legPos = esd->GetTrack(v0->GetPindex());
543     AliESDtrack* legNeg = esd->GetTrack(v0->GetNindex());
544  
545     if(legPos->GetSign() == legNeg->GetSign()) {
546       //cout << "V0 rejected: same sign legs" << endl;
547       continue;
548     }
549
550     Bool_t v0ChargesAreCorrect = (legPos->GetSign()==+1 ? kTRUE : kFALSE);
551     legPos = (!v0ChargesAreCorrect ? esd->GetTrack(v0->GetNindex()) : legPos);
552     legNeg = (!v0ChargesAreCorrect ? esd->GetTrack(v0->GetPindex()) : legNeg); 
553     
554     // apply the K0s filter
555     Bool_t goodK0s = kTRUE;
556     if(fK0sPionCuts && (!fK0sPionCuts->IsSelected(legPos) || !fK0sPionCuts->IsSelected(legNeg))) goodK0s = kFALSE;
557     if(fK0sCuts) {
558       TList k0sList;
559       k0sList.Add(v0);
560       k0sList.Add(legPos); k0sList.Add(legNeg);
561       k0sList.Add(const_cast<AliESDVertex*>(primaryVertex));
562       if(!fK0sCuts->IsSelected(&k0sList)) goodK0s = kFALSE;
563     }
564     
565     // apply the lambda filter
566     Bool_t goodLambda = kTRUE;
567     if(fLambdaProtonCuts && !fLambdaProtonCuts->IsSelected(legPos)) goodLambda = kFALSE;
568     if(fLambdaPionCuts && !fLambdaPionCuts->IsSelected(legNeg)) goodLambda = kFALSE;
569     Bool_t goodALambda = kTRUE;
570     if(fLambdaProtonCuts && !fLambdaProtonCuts->IsSelected(legNeg)) goodALambda = kFALSE;
571     if(fLambdaPionCuts && !fLambdaPionCuts->IsSelected(legPos)) goodALambda = kFALSE;
572     if(fLambdaCuts) {
573       TList lambdaList;
574       lambdaList.Add(v0);
575       lambdaList.Add(legPos); lambdaList.Add(legNeg);
576       lambdaList.Add(const_cast<AliESDVertex*>(primaryVertex));
577       if(!fLambdaCuts->IsSelected(&lambdaList)) {goodLambda = kFALSE; goodALambda = kFALSE;}
578     }
579     
580     if(!(goodK0s || goodLambda || goodALambda)) continue;
581     
582     // Fill the V0 information into the tree for 3 hypothesis: K0s, Lambda, Anti-Lambda
583     AliCorrelationReducedPair* k0sReducedPair     = FillV0PairInfo(v0, AliCorrelationReducedPair::kK0sToPiPi,     legPos, legNeg, &primaryVertexKF, v0ChargesAreCorrect);
584     AliCorrelationReducedPair* lambdaReducedPair  = FillV0PairInfo(v0, AliCorrelationReducedPair::kLambda0ToPPi,  legPos, legNeg, &primaryVertexKF, v0ChargesAreCorrect);
585     AliCorrelationReducedPair* alambdaReducedPair = FillV0PairInfo(v0, AliCorrelationReducedPair::kALambda0ToPPi, legPos, legNeg, &primaryVertexKF, v0ChargesAreCorrect);
586
587     if(goodK0s && k0sReducedPair->fMass[0]>fK0sMassRange[0] && k0sReducedPair->fMass[0]<fK0sMassRange[1]) {
588       TClonesArray& tracks = *(fReducedEvent->fCandidates);
589       AliCorrelationReducedPair *goodK0sPair = new (tracks[fReducedEvent->fNV0candidates[1]]) AliCorrelationReducedPair(*k0sReducedPair);
590       goodK0sPair->fMass[0] = k0sReducedPair->fMass[0];
591       goodK0sPair->fMass[1] = lambdaReducedPair->fMass[0];
592       goodK0sPair->fMass[2] = alambdaReducedPair->fMass[0];
593       fReducedEvent->fNV0candidates[1] += 1;
594     } else {goodK0s=kFALSE;}
595     if(goodLambda && lambdaReducedPair->fMass[0]>fLambdaMassRange[0] && lambdaReducedPair->fMass[0]<fLambdaMassRange[1]) {
596       TClonesArray& tracks = *(fReducedEvent->fCandidates);
597       AliCorrelationReducedPair *goodLambdaPair = new (tracks[fReducedEvent->fNV0candidates[1]]) AliCorrelationReducedPair(*lambdaReducedPair);
598       fReducedEvent->fNV0candidates[1] += 1;
599       goodLambdaPair->fMass[0] = k0sReducedPair->fMass[0];
600       goodLambdaPair->fMass[1] = lambdaReducedPair->fMass[0];
601       goodLambdaPair->fMass[2] = alambdaReducedPair->fMass[0];
602     } else {goodLambda=kFALSE;}
603     if(goodALambda && alambdaReducedPair->fMass[0]>fLambdaMassRange[0] && alambdaReducedPair->fMass[0]<fLambdaMassRange[1]) {
604       TClonesArray& tracks = *(fReducedEvent->fCandidates);
605       AliCorrelationReducedPair *goodALambdaPair = new (tracks[fReducedEvent->fNV0candidates[1]]) AliCorrelationReducedPair(*alambdaReducedPair);
606       fReducedEvent->fNV0candidates[1] += 1;  
607       goodALambdaPair->fMass[0] = k0sReducedPair->fMass[0];
608       goodALambdaPair->fMass[1] = lambdaReducedPair->fMass[0];
609       goodALambdaPair->fMass[2] = alambdaReducedPair->fMass[0];
610     } else {goodALambda = kFALSE;}
611     delete k0sReducedPair;
612     delete lambdaReducedPair;
613     delete alambdaReducedPair;
614
615     if(!(goodK0s || goodLambda || goodALambda)) continue;
616     
617     //  Fill histograms and the CF container
618     AliDielectronVarManager::Fill(legPos, valuesPos);
619     AliDielectronVarManager::Fill(legNeg, valuesNeg);
620     
621     if(fV0Histos && fV0Histos->GetHistogramList()->FindObject("V0Track_Pos"))
622       fV0Histos->FillClass("V0Track_Pos", AliDielectronVarManager::kNMaxValues, valuesPos);
623     if(fV0Histos && fV0Histos->GetHistogramList()->FindObject("V0Track_Neg"))
624       fV0Histos->FillClass("V0Track_Neg", AliDielectronVarManager::kNMaxValues, valuesNeg);    
625   }   // end loop over V0s
626 }
627
628
629 //_________________________________________________________________________________
630 AliCorrelationReducedPair* AliAnalysisTaskCorrelationTree::FillV0PairInfo(AliESDv0* v0, Int_t id, 
631                                                     AliESDtrack* legPos, AliESDtrack* legNeg,
632                                                     AliKFVertex* vtxKF, Bool_t chargesAreCorrect) {
633   //
634   // Create a reduced V0 object and fill it
635   //
636   AliCorrelationReducedPair* reducedPair=new AliCorrelationReducedPair();  
637   reducedPair->fCandidateId = id;
638   reducedPair->fPairType    = v0->GetOnFlyStatus();    // on the fly status
639   //reducedPair->fOnTheFly    = v0->GetOnFlyStatus();
640   reducedPair->fLegIds[0]   = legPos->GetID();
641   reducedPair->fLegIds[1]   = legNeg->GetID();
642   if(!reducedPair->fPairType) {    // offline
643     UInt_t pidPos = AliPID::kPion;
644     if(id==AliCorrelationReducedPair::kLambda0ToPPi) pidPos = AliPID::kProton;
645     UInt_t pidNeg = AliPID::kPion;
646     if(id==AliCorrelationReducedPair::kALambda0ToPPi) pidNeg = AliPID::kProton;
647     reducedPair->fMass[0]      = v0->GetEffMass(pidPos, pidNeg);
648     reducedPair->fPhi          = v0->Phi();
649     if(reducedPair->fPhi<0.0) reducedPair->fPhi = 2.0*TMath::Pi() + reducedPair->fPhi;  // converted to [0,2pi]
650     reducedPair->fPt           = v0->Pt();
651     reducedPair->fEta          = v0->Eta();
652     reducedPair->fLxy          = 0.0;           //TODO
653     reducedPair->fOpeningAngle = 0.0;
654   }
655   else {
656     const AliExternalTrackParam *negHelix=v0->GetParamN();
657     const AliExternalTrackParam *posHelix=v0->GetParamP();
658     if(!chargesAreCorrect) {
659       negHelix = v0->GetParamP();
660       posHelix = v0->GetParamN();
661     }
662     AliKFParticle negKF(*(negHelix),(id==AliCorrelationReducedPair::kALambda0ToPPi ? -2212 : -211));
663     AliKFParticle posKF(*(posHelix),(id==AliCorrelationReducedPair::kLambda0ToPPi ? +2212 : +211));
664     AliKFParticle v0Refit;
665     v0Refit += negKF;
666     v0Refit += posKF;
667     Double_t massFit=0.0, massErrFit=0.0;
668     v0Refit.GetMass(massFit,massErrFit);
669     reducedPair->fMass[0] = massFit;
670     reducedPair->fPhi          = v0Refit.GetPhi();
671     if(reducedPair->fPhi<0.0) reducedPair->fPhi = 2.0*TMath::Pi() + reducedPair->fPhi;  // converted to [0,2pi]
672     reducedPair->fPt           = v0Refit.GetPt();
673     reducedPair->fEta          = v0Refit.GetEta();
674     reducedPair->fLxy          = v0Refit.GetPseudoProperDecayTime(*vtxKF, massFit);
675     reducedPair->fOpeningAngle = negKF.GetAngle(posKF);
676   }
677   return reducedPair;
678 }
679
680
681 //_________________________________________________________________________________
682 UChar_t AliAnalysisTaskCorrelationTree::EncodeTPCClusterMap(AliESDtrack* track) {
683   //
684   // Encode the TPC cluster map into an UChar_t
685   // Divide the 159 bits from the bit map into 8 groups of adiacent clusters
686   // For each group enable its corresponding bit if in that group there are more clusters compared to
687   // a threshold.
688   //
689   const UChar_t threshold=5;
690   TBits tpcClusterMap = track->GetTPCClusterMap();
691   UChar_t map=0;
692   UChar_t n=0;
693   UChar_t j=0;
694   for(UChar_t i=0; i<8; ++i) {
695     n=0;
696     for(j=i*20; j<(i+1)*20 && j<159; ++j) n+=tpcClusterMap.TestBitNumber(j);
697     if(n>=threshold) map |= (1<<i);
698   }
699   return map;
700 }
701
702
703 //_________________________________________________________________________________
704 void AliAnalysisTaskCorrelationTree::FinishTaskOutput()
705 {
706   //
707   // Finish Task 
708   //
709   fTreeFile->Write();
710   fTreeFile->Close();
711 }