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