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