Support for HFE analysis added (Matthias)
[u/mrichter/AliRoot.git] / PWGHF / correlationHF / AliHFCorrelator.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 //             Base class for Heavy Flavour Correlations Analysis
18 //             Single Event and Mixed Event Analysis are implemented
19 //
20 //-----------------------------------------------------------------------
21 //          
22 //
23 //                                                 Author S.Bjelogrlic
24 //                         Utrecht University 
25 //                      sandro.bjelogrlic@cern.ch
26 //
27 //-----------------------------------------------------------------------
28
29 /* $Id$ */
30
31 #include <TParticle.h>
32 #include <TVector3.h>
33 #include <TChain.h>
34 #include "TROOT.h"
35 #include "AliHFCorrelator.h"
36 #include "AliRDHFCutsDStartoKpipi.h"
37 #include "AliHFAssociatedTrackCuts.h"
38 #include "AliEventPoolManager.h"
39 #include "AliReducedParticle.h"
40 #include "AliCentrality.h"
41 #include "AliAODMCParticle.h"
42
43 using std::cout;
44 using std::endl;
45
46 //_____________________________________________________
47 AliHFCorrelator::AliHFCorrelator() :
48 //
49 // default constructor
50 //
51 TNamed(),
52 fPoolMgr(0x0),         
53 fPool(0x0),
54 fhadcuts(0x0),
55 fAODEvent(0x0),
56 fAssociatedTracks(0x0),
57 fmcArray(0x0),
58 fReducedPart(0x0),
59 fD0cand(0x0), 
60 fhypD0(0), 
61 fDCharge(0),
62
63 fmixing(kFALSE),
64 fmontecarlo(kFALSE),
65 fsystem(kFALSE),
66 fUseReco(kTRUE),
67 fselect(kUndefined),
68
69 fUseImpactParameter(0),
70 fPIDmode(0),
71
72 fNofTracks(0),
73 fPoolContent(0),
74
75 fPhiMin(0),
76 fPhiMax(0),
77
78 fPtTrigger(0),
79 fPhiTrigger(0),
80 fEtaTrigger(0),
81
82
83 fDeltaPhi(0),
84 fDeltaEta(0),
85 fk0InvMass(0)
86
87 {
88         // default constructor  
89 }
90
91
92
93 //_____________________________________________________
94 AliHFCorrelator::AliHFCorrelator(const Char_t* name, AliHFAssociatedTrackCuts *cuts, Bool_t ppOrPbPb) :
95 TNamed(name,"title"),
96 fPoolMgr(0x0),         
97 fPool(0x0),
98 fhadcuts(0x0),
99 fAODEvent(0x0),
100 fAssociatedTracks(0x0),
101 fmcArray(0x0),
102 fReducedPart(0x0),
103 fD0cand(0x0), 
104 fhypD0(0),
105 fDCharge(0),
106
107 fmixing(kFALSE),
108 fmontecarlo(kFALSE),
109 fsystem(ppOrPbPb),
110 fUseReco(kTRUE),
111 fselect(kUndefined),
112 fUseImpactParameter(0),
113 fPIDmode(0),
114
115 fNofTracks(0),
116 fPoolContent(0),
117
118 fPhiMin(0),
119 fPhiMax(0),
120
121 fPtTrigger(0),
122 fPhiTrigger(0),
123 fEtaTrigger(0),
124
125
126 fDeltaPhi(0),
127 fDeltaEta(0),
128 fk0InvMass(0)
129 {
130         fhadcuts = cuts; 
131 }
132
133 //_____________________________________________________
134 AliHFCorrelator::~AliHFCorrelator() 
135 {
136 //
137 // destructor
138 //      
139         
140         if(fPoolMgr)  {delete fPoolMgr; fPoolMgr=0;}       
141         if(fPool) {delete fPool; fPool=0;}
142         if(fhadcuts) {delete fhadcuts; fhadcuts=0;}
143         if(fAODEvent) {delete fAODEvent; fAODEvent=0;}
144         if(fAssociatedTracks) {delete fAssociatedTracks; fAssociatedTracks=0;}
145         if(fmcArray) {delete fmcArray; fmcArray=0;}
146         if(fReducedPart) {delete fReducedPart; fReducedPart=0;}
147         if(fD0cand) {delete fD0cand; fD0cand=0;}
148         
149         
150         if(fNofTracks) fNofTracks = 0;
151         
152         if(fPhiMin) fPhiMin = 0;
153         if(fPhiMax) fPhiMax = 0;
154         
155         if(fPtTrigger) fPtTrigger=0;
156         if(fPhiTrigger) fPhiTrigger=0;
157         if(fEtaTrigger) fEtaTrigger=0;
158         
159         if(fDeltaPhi) fDeltaPhi=0;
160         if(fDeltaEta) fDeltaEta=0;
161         
162         if(fk0InvMass) fk0InvMass=0;
163         
164 }
165 //_____________________________________________________
166 Bool_t AliHFCorrelator::DefineEventPool(){
167         // definition of the Pool Manager for Event Mixing
168         
169
170         Int_t MaxNofEvents = fhadcuts->GetMaxNEventsInPool();
171         Int_t MinNofTracks = fhadcuts->GetMinNTracksInPool();
172         Int_t NofCentBins = fhadcuts->GetNCentPoolBins();
173         Double_t * CentBins = fhadcuts->GetCentPoolBins();
174         Int_t NofZVrtxBins = fhadcuts->GetNZvtxPoolBins();
175         Double_t *ZVrtxBins = fhadcuts->GetZvtxPoolBins();
176                 
177                         
178         fPoolMgr = new AliEventPoolManager(MaxNofEvents, MinNofTracks, NofCentBins, CentBins, NofZVrtxBins, ZVrtxBins);
179         if(!fPoolMgr) return kFALSE;
180         return kTRUE;
181 }
182 //_____________________________________________________
183 Bool_t AliHFCorrelator::Initialize(){
184         
185     //  std::cout << "AliHFCorrelator::Initialize"<< std::endl;
186   AliInfo("AliHFCorrelator::Initialize") ;
187   if(!fAODEvent){
188     AliInfo("No AOD event") ;
189     return kFALSE;
190   }
191     //std::cout << "No AOD event" << std::endl;
192         
193         AliCentrality *centralityObj = 0;
194         Int_t multiplicity = -1;
195         Double_t MultipOrCent = -1;
196         
197         // initialize the pool for event mixing
198         if(!fsystem){ // pp
199         multiplicity = fAODEvent->GetNTracks();
200                 MultipOrCent = multiplicity; // convert from Int_t to Double_t
201         }
202         if(fsystem){ // PbPb
203                 
204                 centralityObj = fAODEvent->GetHeader()->GetCentralityP();
205                 MultipOrCent = centralityObj->GetCentralityPercentileUnchecked("V0M");
206                 AliInfo(Form("Centrality is %f", MultipOrCent));
207         }
208         
209         AliAODVertex *vtx = fAODEvent->GetPrimaryVertex();
210         Double_t zvertex = vtx->GetZ(); // zvertex
211         Double_t * CentBins = fhadcuts->GetCentPoolBins();
212         Double_t poolmin=CentBins[0];
213         Double_t poolmax=CentBins[fhadcuts->GetNCentPoolBins()];
214
215         
216                 if(TMath::Abs(zvertex)>=10 || MultipOrCent>poolmax || MultipOrCent < poolmin) {
217                 if(!fsystem)AliInfo(Form("pp Event with Zvertex = %.2f cm and multiplicity = %.0f out of pool bounds, SKIPPING",zvertex,MultipOrCent));
218                 if(fsystem) AliInfo(Form("PbPb Event with Zvertex = %.2f cm and centrality = %.1f  out of pool bounds, SKIPPING",zvertex,MultipOrCent));
219
220                         return kFALSE;
221                 }
222         
223         fPool = fPoolMgr->GetEventPool(MultipOrCent, zvertex);
224         
225         if (!fPool){
226                 AliInfo(Form("No pool found for multiplicity = %f, zVtx = %f cm", MultipOrCent, zvertex));
227             return kFALSE;
228         }
229         //fPool->PrintInfo();
230         return kTRUE;
231 }
232
233 //_____________________________________________________
234 Bool_t AliHFCorrelator::ProcessEventPool(){
235          // analysis on Mixed Events
236         //cout << "AliHFCorrelator::ProcessEventPool"<< endl;
237                 if(!fmixing) return kFALSE;
238                 if(!fPool->IsReady()) return kFALSE;
239                 if(fPool->GetCurrentNEvents()<fhadcuts->GetMinEventsToMix()) return kFALSE;
240         //      fPool->PrintInfo();
241                 fPoolContent = fPool->GetCurrentNEvents();
242                 
243                 return kTRUE;
244         
245 }
246
247 //_____________________________________________________
248 Bool_t AliHFCorrelator::ProcessAssociatedTracks(Int_t EventLoopIndex, const TObjArray* associatedTracks){
249   // TODO: memory leak needs to be fixed, for every call, a new array
250   // is allocated, but the pointer immediately lost. The cleanup is
251   // not straightforward as in the case of event mixing the pointer
252   // will be an external array which must not be deleted.
253         fAssociatedTracks = new TObjArray();
254         
255         if(!fmixing){ // analysis on Single Event
256                 
257                 
258                 
259                 if(fselect==kHadron || fselect ==kKaon) fAssociatedTracks = AcceptAndReduceTracks(fAODEvent);
260                 if(fselect==kKZero) {fAssociatedTracks = AcceptAndReduceKZero(fAODEvent);}      
261                 if(fselect==kElectron && associatedTracks) fAssociatedTracks=new TObjArray(*associatedTracks);
262                 
263         }
264         
265         if(fmixing) { // analysis on Mixed Events
266                 
267                         
268                 fAssociatedTracks = fPool->GetEvent(EventLoopIndex);
269                                 
270                                 
271                         
272
273         } // end if mixing
274         
275         if(!fAssociatedTracks) return kFALSE;
276         
277         fNofTracks = fAssociatedTracks->GetEntriesFast(); 
278                 
279         return kTRUE;
280         
281 }
282 //_____________________________________________________
283 Bool_t AliHFCorrelator::Correlate(Int_t loopindex){
284
285         if(loopindex >= fNofTracks) return kFALSE;
286         if(!fAssociatedTracks) return kFALSE;
287         
288         fReducedPart = (AliReducedParticle*)fAssociatedTracks->At(loopindex);
289         
290
291         fDeltaPhi = SetCorrectPhiRange(fPhiTrigger - fReducedPart->Phi());
292         
293         fDeltaEta = fEtaTrigger - fReducedPart->Eta();
294
295         return kTRUE;
296         
297 }
298                 
299 //_____________________________________________________
300 Bool_t AliHFCorrelator::PoolUpdate(const TObjArray* associatedTracks){
301
302         if(!fmixing) return kFALSE;
303         if(!fPool) return kFALSE;
304         if(fmixing) { // update the pool for Event Mixing
305                 TObjArray* objArr = NULL;
306                 if(fselect==kHadron || fselect==kKaon) objArr = (TObjArray*)AcceptAndReduceTracks(fAODEvent);
307                 else if(fselect==kKZero) objArr = (TObjArray*)AcceptAndReduceKZero(fAODEvent);
308                 else if(fselect==kElectron && associatedTracks){
309                   objArr = new TObjArray(*associatedTracks);
310                 }
311                 else return kFALSE;
312                 if(objArr->GetEntriesFast()>0) fPool->UpdatePool(objArr); // updating the pool only if there are entries in the array
313         }
314                 
315         return kTRUE;
316         
317 }
318                 
319 //_____________________________________________________
320 Double_t AliHFCorrelator::SetCorrectPhiRange(Double_t phi){
321         Double_t pi = TMath::Pi();
322         
323         if(phi<fPhiMin) phi = phi + 2*pi;
324         if(phi>fPhiMax) phi = phi - 2*pi;
325         
326         return phi;
327 }
328
329 //_____________________________________________________
330 TObjArray*  AliHFCorrelator::AcceptAndReduceTracks(AliAODEvent* inputEvent){
331         
332         Int_t nTracks = inputEvent->GetNTracks();
333         AliAODVertex * vtx = inputEvent->GetPrimaryVertex();
334         Double_t Bz = inputEvent->GetMagneticField();
335         
336         
337         TObjArray* tracksClone = new TObjArray;
338         tracksClone->SetOwner(kTRUE);
339         
340         //*******************************************************
341     // use reconstruction
342         if(fUseReco){
343                 for (Int_t iTrack=0; iTrack<nTracks; ++iTrack) {
344                         AliAODTrack* track = inputEvent->GetTrack(iTrack);
345                         if (!track) continue;
346                         if(!fhadcuts->IsHadronSelected(track)) continue; // apply ESD level selections
347                         if(!fhadcuts->Charge(fDCharge,track)) continue; // apply selection on charge, if required
348                 
349                         Double_t pT = track->Pt();
350                 
351                 //compute impact parameter
352                         Double_t d0z0[2],covd0z0[3];
353                         Double_t d0=-999999.;
354                         if(fUseImpactParameter) track->PropagateToDCA(vtx,Bz,100,d0z0,covd0z0);
355                         else d0z0[0] = 1. ; // random number - be careful with the cuts you applied
356                 
357                         if(fUseImpactParameter==1) d0 = TMath::Abs(d0z0[0]); // use impact parameter
358                         if(fUseImpactParameter==2) { // use impact parameter over resolution
359                   if(TMath::Abs(covd0z0[0])>0.00000001) d0 = TMath::Abs(d0z0[0])/TMath::Sqrt(covd0z0[0]); 
360                   else d0 = -1.; // if the resoultion is Zero, rejects the track - to be on the safe side
361                         
362                 }
363                 
364                         
365                 if(!fhadcuts->CheckHadronKinematic(pT,d0)) continue; // apply kinematic cuts
366                         Bool_t rejectsoftpi = kFALSE;
367                         if(fD0cand) rejectsoftpi = fhadcuts->InvMassDstarRejection(fD0cand,track,fhypD0);
368                 
369                         
370                         if(fselect ==kKaon){    
371                                 if(!fhadcuts->CheckKaonCompatibility(track,fmontecarlo,fmcArray,fPIDmode)) continue; // check if it is a Kaon - data and MC
372                         }
373                 
374                 tracksClone->Add(new AliReducedParticle(track->Eta(), track->Phi(), pT,track->GetLabel(),track->GetID(),d0,rejectsoftpi,track->Charge()));
375                 } // end loop on tracks
376         } // end if use reconstruction kTRUE
377         
378         //*******************************************************
379         
380         //use MC truth
381         if(fmontecarlo && !fUseReco){
382                 
383                 for (Int_t iPart=0; iPart<fmcArray->GetEntriesFast(); iPart++) { 
384                         AliAODMCParticle* mcPart = dynamic_cast<AliAODMCParticle*>(fmcArray->At(iPart));
385                         if (!mcPart) {
386                                 AliWarning("MC Particle not found in tree, skipping"); 
387                                 continue;
388                         }
389                         if(!mcPart->Charge()) continue; // consider only charged tracks
390                         
391                         Int_t PDG =TMath::Abs(mcPart->PdgCode()); 
392                         if(!((PDG==321)||(PDG==211)||(PDG==2212)||(PDG==13)||(PDG==11))) continue; // select only if kaon, pion, proton, muon or electron
393                 
394                         Double_t pT = mcPart->Pt();
395             Double_t d0 =1; // set 1 fot the moment - no displacement calculation implemented yet
396                         if(!fhadcuts->CheckHadronKinematic(pT,d0)) continue; // apply kinematic cuts
397                         
398                         tracksClone->Add(new AliReducedParticle(mcPart->Eta(), mcPart->Phi(), pT,mcPart->GetLabel(),-1,d0,kFALSE,mcPart->Charge()));
399                 }
400                 
401         } // end if use  MC truth
402         
403         
404         return tracksClone;
405 }
406
407 //_____________________________________________________
408 TObjArray*  AliHFCorrelator::AcceptAndReduceKZero(AliAODEvent* inputEvent){
409         
410         Int_t nOfVZeros = inputEvent->GetNumberOfV0s();
411         TObjArray* KZeroClone = new TObjArray;
412         AliAODVertex *vertex1 = (AliAODVertex*)inputEvent->GetPrimaryVertex();
413         Int_t v0label = -1;
414         Int_t pdgDgK0toPipi[2] = {211,211};
415         Double_t mPDGK0=0.497614;//TDatabasePDG::Instance()->GetParticle(310)->Mass();
416         const Int_t size = inputEvent->GetNumberOfV0s();
417         Int_t prevnegID[size];
418         Int_t prevposID[size];
419         for(Int_t iv0 =0; iv0< nOfVZeros; iv0++){// loop on all v0 candidates
420                 AliAODv0 *v0 = (static_cast<AliAODEvent*> (inputEvent))->GetV0(iv0);
421                 if(!v0) {
422       AliInfo(Form("is not a v0 at step, %d", iv0)) ;
423         //cout << "is not a v0 at step " << iv0 << endl;
424       continue;
425     }
426                 
427                 if(!fhadcuts->IsKZeroSelected(v0,vertex1)) continue; // check if it is a k0
428             
429                 // checks to avoid double counting
430                 Int_t negID = -999;
431                 Int_t posID = -998;
432                 //Int_t a = 0;
433                 prevnegID[iv0] = -997;
434                 prevposID[iv0]= -996;
435                 negID = v0->GetNegID();
436                 posID = v0->GetPosID();
437                 Bool_t DoubleCounting = kFALSE;
438                 
439                 for(Int_t k=0; k<iv0; k++){
440                         if(((negID==prevnegID[k])&&(posID==prevposID[k]))||((negID==prevposID[k])&&(posID==prevnegID[k]))){
441                                 DoubleCounting = kTRUE;
442                                 //a=k;
443                                 break;
444                         }//end if
445                 } // end for
446                 
447                 if(DoubleCounting) continue;
448                 else{ 
449                         prevposID[iv0] = posID; 
450                         prevnegID[iv0] = negID;
451                 }
452                 
453                 if(fmontecarlo) v0label = v0->MatchToMC(310,fmcArray, 0, pdgDgK0toPipi); //match a K0 short
454                 Double_t k0pt = v0->Pt();
455                 Double_t k0eta = v0->Eta();
456                 Double_t k0Phi = v0->Phi();
457             fk0InvMass = v0->MassK0Short();     
458                 
459                 //if (loopindex == 0) {
460                 //      if(!plotassociation) ((TH2F*)fOutput->FindObject("KZeroSpectra"))->Fill(k0InvMass,k0pt); // spectra for all k0
461                 //      if(plotassociation) ((TH2F*)fOutput->FindObject("KZeroSpectraifHF"))->Fill(k0InvMass,k0pt); // spectra for k0 in association with a D*
462                 //}
463                 // if there are more D* candidates per event, loopindex == 0 makes sure you fill the mass spectra only once!
464                 
465                 if(TMath::Abs(fk0InvMass-mPDGK0)>3*0.004) continue; // select candidates within 3 sigma
466                 KZeroClone->Add(new AliReducedParticle(k0eta,k0Phi,k0pt,v0label));
467                 
468         }
469         
470         return KZeroClone;
471 }