Added AddTrackParams() method for convenience + some comments
[u/mrichter/AliRoot.git] / ANALYSIS / AliAnalysisTaskKineFilter.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 //
17 //  Analysis task for Kinematic filtering
18 //  Fill AOD tracks from Kinematic stack
19 //
20 //  Code taken from macro CreateAODfromKineTree.C by  Markus Oldenburg
21 //  
22  
23 #include <TChain.h>
24 #include <TFile.h>
25
26 #include "AliAnalysisTaskKineFilter.h"
27 #include "AliAnalysisManager.h"
28 #include "AliAnalysisFilter.h"
29 #include "AliHeader.h"
30 #include "AliStack.h"
31 #include "TParticle.h"
32 #include "AliMCEvent.h"
33 #include "AliAODEvent.h"
34 #include "AliAODHeader.h"
35 #include "AliAODVertex.h"
36 #include "AliAODTrack.h"
37
38 #include "AliLog.h"
39
40 ClassImp(AliAnalysisTaskKineFilter)
41
42 ////////////////////////////////////////////////////////////////////////
43
44 //____________________________________________________________________
45 AliAnalysisTaskKineFilter::AliAnalysisTaskKineFilter():
46     fTrackFilter(0x0)
47 {
48   // Default constructor
49 }
50
51 //____________________________________________________________________
52 AliAnalysisTaskKineFilter::AliAnalysisTaskKineFilter(const char* name):
53     AliAnalysisTaskSE(name),
54     fTrackFilter(0x0)
55 {
56   // Default constructor
57     DefineInput (0, TChain::Class());
58     DefineOutput(0, TTree::Class());
59 }
60
61 //____________________________________________________________________
62 AliAnalysisTaskKineFilter::AliAnalysisTaskKineFilter(const AliAnalysisTaskKineFilter& obj):
63     AliAnalysisTaskSE(obj),
64     fTrackFilter(0)
65 {
66 // Copy constructor
67     fTrackFilter = obj.fTrackFilter;
68 }
69 //____________________________________________________________________
70 AliAnalysisTaskKineFilter::~AliAnalysisTaskKineFilter()
71 {
72   //  if( fTrackFilter ) delete fTrackFilter;
73 }
74
75
76 //____________________________________________________________________
77 AliAnalysisTaskKineFilter& AliAnalysisTaskKineFilter::operator=(const AliAnalysisTaskKineFilter& other)
78 {
79 // Assignment
80     AliAnalysisTaskSE::operator=(other);
81     fTrackFilter = other.fTrackFilter;
82     return *this;
83 }
84
85 //____________________________________________________________________
86 void AliAnalysisTaskKineFilter::UserCreateOutputObjects()
87 {
88 // Create the output container
89     if (OutputTree()) 
90         OutputTree()->GetUserInfo()->Add(fTrackFilter);
91 }
92
93
94 //____________________________________________________________________
95 void AliAnalysisTaskKineFilter::UserExec(Option_t */*option*/)
96 {
97 // Execute analysis for current event
98 //
99
100 // Fill AOD tracks from Kinematic stack
101     
102   // get AliAOD Event 
103   AliAODEvent* aod = AODEvent();
104   if (!aod) {
105       AliWarning("No Output Handler connected, doing nothing !") ;
106       return;
107   }
108   
109 //  aod->CreateStdContent();
110
111   AliStack* stack = MCEvent()->Stack();
112   Int_t nTracks = stack->GetNtrack();
113   Int_t nPrims = stack->GetNprimary();
114   Int_t nPrimsAdd = 0;
115
116   AliAODVertex *primary = NULL; 
117   Int_t nPos = 0;
118   Int_t nNeg = 0;
119   Int_t jVertices = 0;
120   Int_t jTracks = 0; 
121   Float_t p[3];
122   Float_t x[3];
123  
124   // Access to the header
125   AliAODHeader *header = aod->GetHeader();
126
127   Double_t emEnergy[2] = {-999., -999.};
128
129   // fill the header
130   *header = AliAODHeader(MCEvent()->Header()->GetRun(),
131                          0, // bunchX number         
132                          0, // orbit number
133                          0, // period number
134                          nTracks,
135                          nPos,
136                          nNeg,
137                          -999, // mag. field
138                          -999., // muon mag. field
139                          -999., // centrality
140                          -999., // ZDCN1Energy
141                          -999., // ZDCP1Energy
142                          -999., // ZDCN2Energy
143                          -999., // ZDCP2Energy
144                          emEnergy, // emEnergy
145                          0, // TriggerMask
146                          0, // TriggerCluster
147                          0, // EventType
148                          ""); // title
149   
150   // Access to the AOD container of vertices
151   TClonesArray &vertices = *(aod->GetVertices());
152
153   // Access to the AOD container of tracks
154   TClonesArray &tracks = *(aod->GetTracks());
155  
156   aod->ResetStd(nTracks, 1);
157
158   // track loop
159   for (Int_t iTrack = 0; iTrack < nPrims; ++iTrack) {
160                                                               
161     TParticle *part = stack->Particle(iTrack);
162     
163     if (iTrack == 0) {
164         // add primary vertex
165         x[0] = part->Vx(); x[1] = part->Vy(); x[2] = part->Vz();
166         primary = new(vertices[jVertices++])
167             AliAODVertex(x, NULL, -999., NULL, -1, AliAODVertex::kPrimary);  
168     }
169     
170     
171     //
172     // Track selection
173     UInt_t selectInfo = 0;
174     if (fTrackFilter) {
175        selectInfo = fTrackFilter->IsSelected(part);
176        if (!selectInfo) continue;
177     }
178     
179     x[0] = part->Vx(); x[1] = part->Vy(); x[2] = part->Vz();
180     p[0] = part->Px(); p[1] = part->Py(); p[2] = part->Pz();
181
182     // add primary tracks
183     primary->AddDaughter(new(tracks[jTracks++]) AliAODTrack(0, // ID,
184                                                             iTrack, // Label
185                                                             p,
186                                                             kTRUE,
187                                                             x,
188                                                             kFALSE,
189                                                             NULL, 
190                                                             (Short_t)-99,
191                                                             0, // no ITSClusterMap
192                                                             NULL,
193                                                             primary,
194                                                             kFALSE,  // no fit performed
195                                                             kFALSE,  // no fit preformed
196                                                             AliAODTrack::kPrimary,
197                                                             selectInfo));
198
199     AliAODTrack* currTrack = (AliAODTrack*)tracks.Last();
200     SetChargeAndPID(part->GetPdgCode(), currTrack);
201     if (currTrack->Charge() != -99) {
202       if (currTrack->Charge() > 0) {
203         nPos++;
204       } else if (currTrack->Charge() < 0) {
205         nNeg++;
206       }
207     }
208     ++nPrimsAdd;
209     LoopOverSecondaries(part, jTracks, jVertices, nPos, nNeg);
210     
211   } // end of track loop
212     
213   header->SetRefMultiplicityPos(nPos);
214   header->SetRefMultiplicityNeg(nNeg);
215     
216     
217   if( fDebug > 1 ) 
218      AliInfo(Form("primaries: %d secondaries: %d (pos: %d neg: %d), vertices: %d", 
219                    nPrimsAdd, tracks.GetEntriesFast()-nPrimsAdd, nPos, nNeg, vertices.GetEntriesFast() ) );
220   return;
221 }
222
223
224 //____________________________________________________________________
225 Int_t AliAnalysisTaskKineFilter::LoopOverSecondaries(TParticle *mother, Int_t &jTracks, Int_t &jVertices, Int_t &nPos, Int_t &nNeg ) 
226 {
227   
228   if (mother->GetNDaughters() > 0) {
229     
230     AliStack* stack = MCEvent()->Stack();
231     
232     TClonesArray &vertices = *(AODEvent()->GetVertices());
233     TClonesArray &tracks = *(AODEvent()->GetTracks());
234     Float_t p[3];
235     Float_t x[3];  
236     AliAODVertex* secondary = NULL;
237
238     for (Int_t iDaughter = mother->GetFirstDaughter(); iDaughter <= mother->GetLastDaughter(); iDaughter++) {
239       TParticle *part = stack->Particle(iDaughter);
240       // only final particles
241       
242       p[0] = part->Px(); 
243       p[1] = part->Py(); 
244       p[2] = part->Pz();
245       x[0] = part->Vx(); 
246       x[1] = part->Vy(); 
247       x[2] = part->Vz();
248
249       if (iDaughter == mother->GetFirstDaughter()) {
250         // add secondary vertex
251           secondary = new(vertices[jVertices++])
252               AliAODVertex(x, NULL, -999., tracks.Last(), iDaughter, AliAODVertex::kUndef);
253           SetVertexType(part, secondary);
254       }
255
256       UInt_t selectInfo = 0;
257       //
258       // Track selection
259       if (fTrackFilter) {
260         selectInfo = fTrackFilter->IsSelected(part);
261         if (!selectInfo) continue;
262       }
263         
264       // add secondary tracks
265       secondary->AddDaughter(new(tracks[jTracks++]) AliAODTrack(0, // ID
266                                                                 iDaughter, // label
267                                                                 p,
268                                                                 kTRUE,
269                                                                 x,
270                                                                 kFALSE,
271                                                                 NULL, 
272                                                                 (Short_t)-99,
273                                                                 0, // no cluster map available
274                                                                 NULL,
275                                                                 secondary,
276                                                                 kFALSE, // no fit performed
277                                                                 kFALSE, // no fit performed
278                                                                 AliAODTrack::kSecondary,
279                                                                 selectInfo));
280
281       AliAODTrack* currTrack = (AliAODTrack*)tracks.Last();
282       SetChargeAndPID(part->GetPdgCode(), currTrack);
283       if (currTrack->Charge() != -99) {
284         if (currTrack->Charge() > 0) {
285           nPos++;
286         } else if (currTrack->Charge() < 0) {
287           nNeg++;
288         }
289       }
290
291       LoopOverSecondaries(part, jTracks, jVertices, nPos, nNeg);
292     }
293     return 1;
294   } else {
295     return 0;
296   }
297 }
298
299 //____________________________________________________________________
300 void AliAnalysisTaskKineFilter::SetChargeAndPID(Int_t pdgCode, AliAODTrack *track) {
301
302   Float_t PID[10] = {0., 0., 0., 0., 0., 0., 0., 0., 0., 0.};
303
304   switch (pdgCode) {
305
306   case 22: // gamma
307     track->SetCharge(0);
308     PID[AliAODTrack::kUnknown] = 1.;
309     track->SetPID(PID);
310     break;
311
312   case 11: // e- 
313     track->SetCharge(-1);
314     PID[AliAODTrack::kElectron] = 1.;
315     track->SetPID(PID);
316     break;
317     
318   case -11: // e+
319     track->SetCharge(+1);
320     PID[AliAODTrack::kElectron] = 1.;
321     track->SetPID(PID);
322     break;
323     
324   case 13: // mu- 
325     track->SetCharge(-1);
326     PID[AliAODTrack::kMuon] = 1.;
327     track->SetPID(PID);
328     break;
329     
330   case -13: // mu+
331     track->SetCharge(+1);
332     PID[AliAODTrack::kMuon] = 1.;
333     track->SetPID(PID);
334     break;
335     
336   case 111: // pi0
337     track->SetCharge(0);
338     PID[AliAODTrack::kUnknown] = 1.;
339     track->SetPID(PID);
340     break;
341     
342   case 211: // pi+
343     track->SetCharge(+1);
344     PID[AliAODTrack::kPion] = 1.;
345     track->SetPID(PID);
346     break;
347     
348   case -211: // pi-
349     track->SetCharge(-1);
350     PID[AliAODTrack::kPion] = 1.;
351     track->SetPID(PID);
352     break;
353     
354   case 130: // K0L
355     track->SetCharge(0);
356     PID[AliAODTrack::kUnknown] = 1.;
357     track->SetPID(PID);
358     break;
359     
360   case 321: // K+
361     track->SetCharge(+1);
362     PID[AliAODTrack::kKaon] = 1.;
363     track->SetPID(PID);
364     break;
365     
366   case -321: // K- 
367     track->SetCharge(-1);
368     PID[AliAODTrack::kKaon] = 1.;
369     track->SetPID(PID);
370     break;
371     
372   case 2112: // n
373     track->SetCharge(0);
374     PID[AliAODTrack::kUnknown] = 1.;
375     track->SetPID(PID);
376     break;
377     
378   case 2212: // p
379     track->SetCharge(+1);
380     PID[AliAODTrack::kProton] = 1.;
381     track->SetPID(PID);
382     break;
383     
384   case -2212: // anti-p
385     track->SetCharge(-1);
386     PID[AliAODTrack::kProton] = 1.;
387     track->SetPID(PID);
388     break;
389
390   case 310: // K0S
391     track->SetCharge(0);
392     PID[AliAODTrack::kUnknown] = 1.;
393     track->SetPID(PID);
394     break;
395     
396   case 311: // K0
397     track->SetCharge(0);
398     PID[AliAODTrack::kUnknown] = 1.;
399     track->SetPID(PID);
400     break;
401     
402   case -311: // anti-K0
403     track->SetCharge(0);
404     PID[AliAODTrack::kUnknown] = 1.;
405     track->SetPID(PID);
406     break;
407     
408   case 221: // eta
409     track->SetCharge(0);
410     PID[AliAODTrack::kUnknown] = 1.;
411     track->SetPID(PID);
412     break;
413
414   case 3122: // lambda
415     track->SetCharge(0);
416     PID[AliAODTrack::kUnknown] = 1.;
417     track->SetPID(PID);
418     break;
419
420   case 3222: // Sigma+
421     track->SetCharge(+1);
422     PID[AliAODTrack::kUnknown] = 1.;
423     track->SetPID(PID);
424     break;
425
426   case 3212: // Sigma0
427     track->SetCharge(-1);
428     PID[AliAODTrack::kUnknown] = 1.;
429     track->SetPID(PID);
430     break;
431
432   case 3112: // Sigma-
433     track->SetCharge(-1);
434     PID[AliAODTrack::kUnknown] = 1.;
435     track->SetPID(PID);
436     break;
437
438   case 3322: // Xi0
439     track->SetCharge(0);
440     PID[AliAODTrack::kUnknown] = 1.;
441     track->SetPID(PID);
442     break;
443
444   case 3312: // Xi-
445     track->SetCharge(-1);
446     PID[AliAODTrack::kUnknown] = 1.;
447     track->SetPID(PID);
448     break;
449
450   case 3334: // Omega-
451     track->SetCharge(-1);
452     PID[AliAODTrack::kUnknown] = 1.;
453     track->SetPID(PID);
454     break;
455
456   case -2112: // n-bar
457     track->SetCharge(0);
458     PID[AliAODTrack::kUnknown] = 1.;
459     track->SetPID(PID);
460     break;
461
462   case -3122: // anti-Lambda
463     track->SetCharge(0);
464     PID[AliAODTrack::kUnknown] = 1.;
465     track->SetPID(PID);
466     break;
467
468   case -3222: // anti-Sigma-
469     track->SetCharge(-1);
470     PID[AliAODTrack::kUnknown] = 1.;
471     track->SetPID(PID);
472     break;
473
474   case -3212: // anti-Sigma0
475     track->SetCharge(0);
476     PID[AliAODTrack::kUnknown] = 1.;
477     track->SetPID(PID);
478     break;
479
480   case -3112: // anti-Sigma+
481     track->SetCharge(+1);
482     PID[AliAODTrack::kUnknown] = 1.;
483     track->SetPID(PID);
484     break;
485
486   case -3322: // anti-Xi0
487     track->SetCharge(0);
488     PID[AliAODTrack::kUnknown] = 1.;
489     track->SetPID(PID);
490     break;
491
492   case -3312: // anti-Xi+
493     track->SetCharge(+1);
494     break;
495
496   case -3334: // anti-Omega+
497     track->SetCharge(+1);
498     PID[AliAODTrack::kUnknown] = 1.;
499     track->SetPID(PID);
500     break;
501
502   case 411: // D+
503     track->SetCharge(+1);
504     PID[AliAODTrack::kUnknown] = 1.;
505     track->SetPID(PID);
506     break;
507
508   case -411: // D- 
509     track->SetCharge(-1);
510     PID[AliAODTrack::kUnknown] = 1.;
511     track->SetPID(PID);
512     break;
513
514   case 421: // D0
515     track->SetCharge(0);
516     PID[AliAODTrack::kUnknown] = 1.;
517     track->SetPID(PID);
518     break;
519
520   case -421: // anti-D0
521     track->SetCharge(0);
522     PID[AliAODTrack::kUnknown] = 1.;
523     track->SetPID(PID);
524     break;
525
526   default : // unknown
527     track->SetCharge(-99);
528     PID[AliAODTrack::kUnknown] = 1.;
529     track->SetPID(PID);
530  }
531
532   return;
533 }
534
535 //____________________________________________________________________
536 void AliAnalysisTaskKineFilter::SetVertexType(TParticle *part, AliAODVertex *vertex) 
537 {
538   // this whole thing doesn't make much sense. but anyhow...
539   AliStack* stack = MCEvent()->Stack();
540   TParticle *mother = stack->Particle(part->GetFirstMother());
541   Int_t pdgMother = mother->GetPdgCode();
542   Int_t pdgPart = part->GetPdgCode();
543   
544   // kinks
545   if (mother->GetNDaughters() == 2) {
546     Int_t firstPdgCode = stack->Particle(mother->GetFirstDaughter())->GetPdgCode();
547     Int_t lastPdgCode = stack->Particle(mother->GetLastDaughter())->GetPdgCode();
548
549     if (!(pdgMother == 22 || pdgMother == 111 || pdgMother == 130 || 
550           TMath::Abs(pdgMother) == 2112 || pdgMother == 310 || pdgMother == 221 || 
551           TMath::Abs(pdgMother) == 3122 || TMath::Abs(pdgMother) == 3322 || 
552           pdgMother == -3212 || TMath::Abs(pdgMother) == 421 || 
553           TMath::Abs(pdgMother) == 311) // not neutral
554         && (((firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
555               TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
556               firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
557               TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
558               TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // neutral
559              && !(lastPdgCode == 22 || lastPdgCode == 111 || lastPdgCode == 130 || 
560                   TMath::Abs(lastPdgCode) == 2112 || lastPdgCode == 310 || 
561                   lastPdgCode == 221 || TMath::Abs(lastPdgCode) == 3122 || 
562                   TMath::Abs(lastPdgCode) == 3322 || lastPdgCode == -3212 || 
563                   TMath::Abs(lastPdgCode) == 421 || TMath::Abs(pdgMother) == 311)) // not neutral
564             || !((firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
565                   TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
566                   firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
567                   TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
568                   TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not neutral
569                  && (lastPdgCode == 22 || lastPdgCode == 111 || lastPdgCode == 130 || 
570                      TMath::Abs(lastPdgCode) == 2112 || lastPdgCode == 310 || 
571                      lastPdgCode == 221 || TMath::Abs(lastPdgCode) == 3122 || 
572                      TMath::Abs(lastPdgCode) == 3322 || lastPdgCode == -3212 || 
573                      TMath::Abs(lastPdgCode) == 421 || TMath::Abs(pdgMother) == 311)))) { // neutral
574       
575       vertex->SetType(AliAODVertex::kKink);
576   //    jKinks++;
577     }
578   }
579
580   // V0
581   else if (mother->GetNDaughters() == 2) {
582     Int_t firstPdgCode = stack->Particle(mother->GetFirstDaughter())->GetPdgCode();
583     Int_t lastPdgCode = stack->Particle(mother->GetLastDaughter())->GetPdgCode();
584
585     if ((pdgMother == 22 || pdgMother == 111 || pdgMother == 130 || 
586          TMath::Abs(pdgMother) == 2112 || pdgMother == 310 || 
587          pdgMother == 221 || TMath::Abs(pdgMother) == 3122 || 
588          TMath::Abs(pdgMother) == 3322 || pdgMother == -3212 || 
589          TMath::Abs(pdgMother) == 421 || TMath::Abs(pdgMother) == 311) // neutral
590         && !(lastPdgCode == 22 || lastPdgCode == 111 || lastPdgCode == 130 || 
591              TMath::Abs(lastPdgCode) == 2112 || lastPdgCode == 310 || 
592              lastPdgCode == 221 || TMath::Abs(lastPdgCode) == 3122 || 
593              TMath::Abs(lastPdgCode) == 3322 || lastPdgCode == -3212 || 
594              TMath::Abs(lastPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not neutral
595         && !(firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
596              TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
597              firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
598              TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
599              TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311)) { // not neutral
600       
601       vertex->SetType(AliAODVertex::kV0);
602   //    jV0s++;
603     }
604   }
605
606   // Cascade
607   else if (mother->GetNDaughters() == 2) {
608     Int_t firstPdgCode = stack->Particle(mother->GetFirstDaughter())->GetPdgCode();
609     Int_t lastPdgCode = stack->Particle(mother->GetLastDaughter())->GetPdgCode();
610     
611     if ((TMath::Abs(pdgMother) == 3334 || TMath::Abs(pdgMother) == 3312 || TMath::Abs(pdgMother) == 3322) &&
612         (TMath::Abs(pdgPart) == 3122 || TMath::Abs(pdgPart) == 211 || TMath::Abs(pdgPart) == 321)
613         && ((!(firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
614                TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
615                firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
616                TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
617                TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not neutral   
618              && TMath::Abs(lastPdgCode) == 3122) // labmda or anti-lambda
619             || ((!(lastPdgCode == 22 || lastPdgCode == 111 || lastPdgCode == 130 || 
620                    TMath::Abs(lastPdgCode) == 2112 || lastPdgCode == 310 || 
621                    lastPdgCode == 221 || TMath::Abs(lastPdgCode) == 3122 || 
622                    TMath::Abs(lastPdgCode) == 3322 || lastPdgCode == -3212 || 
623                    TMath::Abs(lastPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not neutral
624                  && TMath::Abs(firstPdgCode) == 3122)))) { // lambda or anti-lambda
625       vertex->SetType(AliAODVertex::kCascade);
626   //    jCascades++;
627     }
628   }
629
630   // Multi
631   else if (mother->GetNDaughters() > 2) {
632
633     vertex->SetType(AliAODVertex::kMulti);
634   //  jMultis++;
635   }
636
637   else {
638     vertex->SetType(AliAODVertex::kUndef);
639   }
640 }