Major update of the HFE package (comments inside the code
[u/mrichter/AliRoot.git] / PWG3 / hfe / AliHFEpriVtx.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 // QA class of primary vertex study for Heavy Flavor electrons
17 // this has functionality to reject electrons from primary vertex
18 // and check primary vertex characteristics
19 //  
20 // Authors:
21 //   MinJung Kweon <minjung@physi.uni-heidelberg.de>
22 //
23
24
25 #include <TH2F.h>
26 #include <TParticle.h>
27
28 #include "AliESDEvent.h"
29 #include <AliESDtrack.h>
30 #include <AliMCEvent.h>
31
32 #include <AliLog.h>
33 #include <AliKFParticle.h>
34 #include <AliKFVertex.h>
35
36 #include "AliHFEpriVtx.h"
37
38
39 ClassImp(AliHFEpriVtx)
40
41 //_______________________________________________________________________________________________
42 AliHFEpriVtx::AliHFEpriVtx():
43         fESD1(0x0)
44         ,fMCEvent(0x0)
45         ,fNtrackswoPid(0)
46         ,fHNtrackswoPid(0x0)
47         ,fNESDprimVtxContributor(0x0)
48         ,fNESDprimVtxIndices(0x0)
49         ,fDiffDCAvsPt(0x0)
50         ,fDiffDCAvsNt(0x0)
51         ,fNsectrk2prim(0)
52         ,fPVxRe(-999.)
53         ,fPVyRe(-999.)
54         ,fPVzRe(-999.)
55
56         //
57         // Default constructor
58         //
59
60         Init();
61
62 }
63
64 //_______________________________________________________________________________________________
65 AliHFEpriVtx::AliHFEpriVtx(const AliHFEpriVtx &p):
66          TObject(p)
67         ,fESD1(0x0)
68         ,fMCEvent(0x0)
69         ,fNtrackswoPid(p.fNtrackswoPid)
70         ,fHNtrackswoPid(0x0)
71         ,fNESDprimVtxContributor(0x0)
72         ,fNESDprimVtxIndices(0x0)
73         ,fDiffDCAvsPt(0x0)
74         ,fDiffDCAvsNt(0x0)
75         ,fNsectrk2prim(p.fNsectrk2prim)
76         ,fPVxRe(p.fPVxRe)
77         ,fPVyRe(p.fPVyRe)
78         ,fPVzRe(p.fPVzRe)
79 {
80         //
81         // Copy constructor
82         //
83 }
84
85 //_______________________________________________________________________________________________
86 AliHFEpriVtx&
87 AliHFEpriVtx::operator=(const AliHFEpriVtx &)
88 {
89         //
90         // Assignment operator
91         //
92
93         AliInfo("Not yet implemented.");
94         return *this;
95 }
96
97 //_______________________________________________________________________________________________
98 AliHFEpriVtx::~AliHFEpriVtx()
99 {
100         //
101         // Destructor
102         //
103
104         AliInfo("Analysis Done.");
105 }
106
107 //__________________________________________
108 void AliHFEpriVtx::Init()
109 {
110         //
111         // initialize counters
112         //
113
114         fNtrackswoPid = 0;
115         for (int i=0; i<10; i++){       
116                 fPrimVtx[i].fNtrackCount = 0 ;
117                 fPrimVtx[i].fNprimVtxContributorCount = 0 ;
118         }
119 }
120
121 //_______________________________________________________________________________________________
122 void AliHFEpriVtx::CreateHistograms(TString hnopt)
123
124         //
125         // create histograms
126         //
127
128         fkSourceLabel[kAll]="all";
129         fkSourceLabel[kDirectCharm]="directCharm";
130         fkSourceLabel[kDirectBeauty]="directBeauty";
131         fkSourceLabel[kBeautyCharm]="beauty2charm";
132         fkSourceLabel[kGamma]="gamma";
133         fkSourceLabel[kPi0]="pi0";
134         fkSourceLabel[kElse]="others";
135         fkSourceLabel[kBeautyGamma]="beauty22gamma";
136         fkSourceLabel[kBeautyPi0]="beauty22pi0";
137         fkSourceLabel[kBeautyElse]="beauty22others";
138
139
140         TString hname;
141         for (Int_t isource = 0; isource < 10; isource++ ){
142
143            hname=hnopt+"ntracks_"+fkSourceLabel[isource];
144            fPrimVtx[isource].fNtracks = new TH1F(hname,hname,50,0,50);
145            hname=hnopt+"nPrimVtxContributor_"+fkSourceLabel[isource];
146            fPrimVtx[isource].fNprimVtxContributor = new TH1F(hname,hname,100,0,100);
147            hname=hnopt+"PtElec_"+fkSourceLabel[isource];
148            fPrimVtx[isource].fPtElec = new TH1F(hname,hname,250,0,50);
149            hname=hnopt+"PtElecContributor_"+fkSourceLabel[isource];
150            fPrimVtx[isource].fPtElecContributor = new TH1F(hname,hname,250,0,50);
151
152         }
153
154         hname=hnopt+"ntrackswopid";
155         fHNtrackswoPid = new TH1F(hname,hname,50,0,50);
156         hname=hnopt+"nESDprimVtxContributor";
157         fNESDprimVtxContributor = new TH1I(hname,hname,100,0,100);
158         hname=hnopt+"nESDprimVtxIndices";
159         fNESDprimVtxIndices= new TH1I(hname,hname,100,0,100);
160         hname=hnopt+"diffDCAvsPt";
161         fDiffDCAvsPt = new TH2F(hname,hname,250,0,50,500,0,1);
162         hname=hnopt+"diffDCAvsNt";
163         fDiffDCAvsNt = new TH2F(hname,hname,100,0,100,500,0,1);
164
165 }
166
167 //_______________________________________________________________________________________________
168 void AliHFEpriVtx::CountNtracks(Int_t sourcePart, Int_t recpid, Double_t recprob)
169 {
170         //
171         // count number of tracks passed certain cuts
172         //
173
174         fNtrackswoPid++;
175
176         if (!recpid && recprob>0.5)     
177         fPrimVtx[kAll].fNtrackCount++;
178         if(sourcePart<0) return;
179         fPrimVtx[sourcePart].fNtrackCount++;
180
181 }
182
183 //_______________________________________________________________________________________________
184 void AliHFEpriVtx::FillNtracks()
185 {
186         //
187         // count number of tracks passed certain cuts
188         //
189
190         fHNtrackswoPid->Fill(fNtrackswoPid);
191         for (int i=0; i<10; i++){
192           fPrimVtx[i].fNtracks->Fill(fPrimVtx[i].fNtrackCount);
193         }
194
195 }
196
197 //_______________________________________________________________________________________________
198 Int_t AliHFEpriVtx::GetMCPID(AliESDtrack const *track)
199 {
200         //
201         // get MC pid
202         //
203
204         AliMCParticle *mctrack = NULL;
205         if(!(mctrack = dynamic_cast<AliMCParticle *>(fMCEvent->GetTrack(TMath::Abs(track->GetLabel()))))) return 0; 
206         TParticle *mcpart = mctrack->Particle();
207
208         if ( !mcpart ) return 0;
209         Int_t pdgCode = mcpart->GetPdgCode();
210
211         return pdgCode;
212 }
213
214 //_______________________________________________________________________________________________
215 void AliHFEpriVtx::GetNPriVxtContributor() 
216 {
217         //
218         // count number of primary vertex contributor
219         //
220
221         const AliESDVertex *primvtx = fESD1->GetPrimaryVertex();
222         fNESDprimVtxContributor->Fill(primvtx->GetNContributors());
223         fNESDprimVtxIndices->Fill(primvtx->GetNIndices());
224 }
225
226 //_______________________________________________________________________________________________
227 void AliHFEpriVtx::CountPriVxtElecContributor(AliESDtrack *ESDelectron, Int_t sourcePart, Int_t recpid, Double_t recprob) 
228 {
229         //
230         // count number of electrons contributing to the primary vertex
231         //
232
233
234         if (recpid || recprob<0.5) return;
235
236         // get track id of our selected electron
237         Int_t elecTrkID = ESDelectron->GetID();
238
239         AliMCParticle *mctrack = NULL;
240         if(!(mctrack = dynamic_cast<AliMCParticle *>(fMCEvent->GetTrack(TMath::Abs(ESDelectron->GetLabel()))))) return; 
241         TParticle *mcpart = mctrack->Particle();
242
243
244         if(!mcpart){
245           AliDebug(1, "no mc particle, return\n");
246           return;
247         }
248
249         AliKFParticle::SetField(fESD1->GetMagneticField());
250         AliKFParticle kfElectron(*ESDelectron,11);
251  
252         // prepare kfprimary vertex
253         AliKFVertex kfESDprimary;
254
255         // Reconstructed Primary Vertex (with ESD tracks)
256         const AliESDVertex *primvtx = fESD1->GetPrimaryVertex();
257 //         Int_t nt = primvtx->GetNContributors();
258         Int_t n=primvtx->GetNIndices();
259
260         if (n>0 && primvtx->GetStatus()){
261
262                 kfESDprimary = AliKFVertex(*primvtx);
263                 Double_t dcaBFrmElec = kfElectron.GetDistanceFromVertexXY(kfESDprimary);
264
265                 UShort_t *priIndex = primvtx->GetIndices();
266
267                 fPrimVtx[kAll].fPtElec->Fill(mcpart->Pt());
268                 if(sourcePart>=0) fPrimVtx[sourcePart].fPtElec->Fill(mcpart->Pt());
269
270                 for (Int_t i=0;i<n;i++){
271
272                         Int_t idx = Int_t(priIndex[i]);
273                         if (idx == elecTrkID){
274                                 fPrimVtx[kAll].fNprimVtxContributorCount++;
275                                 fPrimVtx[kAll].fPtElecContributor->Fill(mcpart->Pt());
276                                 if(sourcePart<0) continue;
277                                 fPrimVtx[sourcePart].fNprimVtxContributorCount++;
278                                 fPrimVtx[sourcePart].fPtElecContributor->Fill(mcpart->Pt());
279
280                                 kfESDprimary -= kfElectron;
281                                 Double_t dcaAFrmElec = kfElectron.GetDistanceFromVertexXY(kfESDprimary);
282                                 fDiffDCAvsPt->Fill(mcpart->Pt(),dcaBFrmElec-dcaAFrmElec);
283                                 fDiffDCAvsNt->Fill(n,dcaBFrmElec-dcaAFrmElec);
284                         }
285                 } 
286         }  
287
288 }
289
290 //_______________________________________________________________________________________________
291 void AliHFEpriVtx::FillNprimVtxContributor() const
292 {
293         //
294         // Fill histogram with number of electrons contributing to the primary vertex
295         //
296
297         for (int i=0; i<10; i++){
298           fPrimVtx[i].fNprimVtxContributor->Fill(fPrimVtx[i].fNprimVtxContributorCount);
299         }
300
301 }
302
303 //_______________________________________________________________________________________________
304 Double_t AliHFEpriVtx::GetDistanceFromRecalVertexXY(AliESDtrack * const ESDelectron) 
305 {
306         //
307         // return recalculated DCA after removing input track from the primary vertex
308         //
309
310         // get track id of our selected electron
311         Int_t elecTrkID = ESDelectron->GetID();
312
313         AliKFParticle::SetField(fESD1->GetMagneticField());
314         AliKFParticle kfElectron(*ESDelectron,11);
315  
316         // prepare kfprimary vertex
317         AliKFVertex kfESDprimary;
318
319         // Reconstructed Primary Vertex (with ESD tracks)
320         const AliESDVertex *primvtx = fESD1->GetPrimaryVertex();
321         Int_t n=primvtx->GetNIndices();
322
323         if (n>0 && primvtx->GetStatus()){
324
325                 kfESDprimary = AliKFVertex(*primvtx);
326                 UShort_t *priIndex = primvtx->GetIndices();
327
328                 for (Int_t i=0;i<n;i++){
329
330                         Int_t idx = Int_t(priIndex[i]);
331                         if (idx == elecTrkID){
332                                 kfESDprimary -= kfElectron;
333                                 Double_t dcaAFrmElec = kfElectron.GetDistanceFromVertexXY(kfESDprimary);
334
335                                 return dcaAFrmElec;
336                         }
337                 } 
338         }  
339               return -1;
340
341 }
342
343 void AliHFEpriVtx::RecalcPrimvtx(Int_t nkftrk, const Int_t * const trkid, const AliKFParticle * const kftrk)
344 {
345         //
346         // recalculate primary vertex after removing the input track
347         //
348
349         const AliESDVertex *primvtx = fESD1->GetPrimaryVertex();
350
351         AliKFVertex kfESDprimary;
352         Int_t n = primvtx->GetNIndices();
353         fNsectrk2prim = 0;
354         fPVxRe = -999.;
355         fPVyRe = -999.;
356         fPVyRe = -999.;
357
358         if (n>0 && primvtx->GetStatus()){
359           kfESDprimary = AliKFVertex(*primvtx);
360           UShort_t *priIndex = primvtx->GetIndices();
361           for (Int_t j=0; j<nkftrk; j++){
362             for (Int_t i=0;i<n;i++){
363               Int_t idx = Int_t(priIndex[i]);
364               if (idx == trkid[j]){
365                 kfESDprimary -= kftrk[j];
366                 fNsectrk2prim++;
367               }
368             }
369           }
370         }
371
372         fPVxRe = kfESDprimary.GetX();
373         fPVyRe = kfESDprimary.GetY();
374         fPVzRe = kfESDprimary.GetZ();
375
376 }
377
378
379 //_______________________________________________________________________________________________
380 void AliHFEpriVtx::RecalcPrimvtx(AliESDtrack * const ESDelectron)
381 {
382         //
383         // recalculate primary vertex after removing the input track
384         //
385
386         // get track id of our selected electron
387         Int_t elecTrkID = ESDelectron->GetID();
388
389         AliKFParticle::SetField(fESD1->GetMagneticField());
390         AliKFParticle kfElectron(*ESDelectron,11);
391
392         const AliESDVertex *primvtx = fESD1->GetPrimaryVertex();
393
394         AliKFVertex kfESDprimary;
395         Int_t n = primvtx->GetNIndices();
396         fPVxRe = -999.;
397         fPVyRe = -999.;
398         fPVyRe = -999.;
399         
400         if (n>0 && primvtx->GetStatus()){
401           kfESDprimary = AliKFVertex(*primvtx);
402           UShort_t *priIndex = primvtx->GetIndices();
403           for (Int_t i=0;i<n;i++){
404             Int_t idx = Int_t(priIndex[i]);
405             if (idx == elecTrkID){
406               kfESDprimary -= kfElectron;
407             }
408           }
409         }     
410         
411         fPVxRe = kfESDprimary.GetX();
412         fPVyRe = kfESDprimary.GetY();
413         fPVzRe = kfESDprimary.GetZ();
414
415 }