Add directory structure (needed if new in cvmfs)
[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("aodTree", "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     Double_t emEnergy[2] = {-999., -999.};
128
129     // fill the header
130     *header = AliAODHeader(aliHeader->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     jVertices=0;
153     jKinks=0;
154     jV0s=0;
155     jCascades=0;
156     jMultis=0;
157
158     // Access to the AOD container of tracks
159     TClonesArray &tracks = *(aod->GetTracks());
160     jTracks=0; 
161  
162     aod->ResetStd(nTracks, 1);
163
164     // track loop
165     for (Int_t iTrack = 0; iTrack < nPrims; ++iTrack) {
166                                                                 
167       TParticle *part = stack->Particle(iTrack);
168
169       //if (part->IsPrimary()) { // this will exclude 'funny' primaries, too
170       if (kTRUE) { // this won't
171         p[0] = part->Px(); p[1] = part->Py(); p[2] = part->Pz();
172         x[0] = part->Vx(); x[1] = part->Vy(); x[2] = part->Vz();
173         
174         if (iTrack == 0) {
175           // add primary vertex
176           primary = new(vertices[jVertices++])
177             AliAODVertex(x, NULL, -999., NULL, AliAODVertex::kPrimary);
178         }
179         
180         // add primary tracks
181         primary->AddDaughter(new(tracks[jTracks++]) AliAODTrack(0, // ID,
182                                                                 0, // Label
183                                                                 p,
184                                                                 kTRUE,
185                                                                 x,
186                                                                 kFALSE,
187                                                                 cov, 
188                                                                 (Short_t)-99,
189                                                                 0, // no ITSClusterMap
190                                                                 pid,
191                                                                 primary,
192                                                                 kFALSE,  // no fit performed
193                                                                 kFALSE, // no fit preformed
194                                                                 AliAODTrack::kPrimary));
195         currTrack = (AliAODTrack*)tracks.Last();
196         SetChargeAndPID(part->GetPdgCode(), currTrack);
197         if (currTrack->Charge() != -99) {
198           if (currTrack->Charge() > 0) {
199             nPos++;
200           } else if (currTrack->Charge() < 0) {
201             nNeg++;
202           }         
203         }
204         
205         LoopOverSecondaries(part);
206       } 
207     } // end of track loop
208     
209     header->SetRefMultiplicityPos(nPos);
210     header->SetRefMultiplicityNeg(nNeg);
211     
212     cout << ":: primaries: " << nPrims << " secondaries: " << tracks.GetEntriesFast()-nPrims << 
213       " (pos: " << nPos << ", neg: " << nNeg << "), vertices: " << vertices.GetEntriesFast() << 
214       " (kinks: " << jKinks << ", V0: " << jV0s << ", cascades: " << jCascades << 
215       ", multi: " << jMultis << ")" << endl;
216     /*
217     cout << "  gamma: " << nGamma << " e-: " << nElectron << " e+: " << nPositron << " p: " << nProton << 
218       " pBar: " << nAntiProton << " n: " << nNeutron << " nBar: " << nAntiNeutron << " pi0: " << nPi0 <<
219       " pi+: " << nPiPlus << " pi-: " << nPiMinus << " K0: " << nK0 << " K+: " << nKPlus << 
220       " K-: " << nKMinus << endl;
221     */
222     // fill the tree for this event
223     aodTree->Fill();
224   } // end of event loop
225
226   aodTree->GetUserInfo()->Add(aod);
227   
228   // write the tree to the specified file
229   outFile = aodTree->GetCurrentFile();
230   outFile->cd();
231   aodTree->Write();
232   outFile->Close();
233
234 }
235
236
237 Int_t LoopOverSecondaries(TParticle *mother) {
238   
239   if (mother->GetNDaughters() > 0) {
240
241     TClonesArray &vertices = *(aod->GetVertices());
242     TClonesArray &tracks = *(aod->GetTracks());
243
244     for (Int_t iDaughter = mother->GetFirstDaughter(); iDaughter <= mother->GetLastDaughter(); iDaughter++) {
245       TParticle *part = stack->Particle(iDaughter);
246       p[0] = part->Px(); 
247       p[1] = part->Py(); 
248       p[2] = part->Pz();
249       x[0] = part->Vx(); 
250       x[1] = part->Vy(); 
251       x[2] = part->Vz();
252
253       if (iDaughter == mother->GetFirstDaughter()) {
254         // add secondary vertex
255         secondary = new(vertices[jVertices++])
256           AliAODVertex(x, NULL, -999., tracks.Last(), AliAODVertex::kUndef);
257         
258         SetVertexType(part, secondary);
259       }
260         
261       // add secondary tracks
262       secondary->AddDaughter(new(tracks[jTracks++]) AliAODTrack(0, // ID
263                                                                 0, // label
264                                                                 p,
265                                                                 kTRUE,
266                                                                 x,
267                                                                 kFALSE,
268                                                                 cov, 
269                                                                 (Short_t)-99,
270                                                                 0, // no cluster map available
271                                                                 pid,
272                                                                 secondary,
273                                                                 kFALSE, // no fit performed
274                                                                 kFALSE, // no fit performed
275                                                                 AliAODTrack::kSecondary));
276
277       currTrack = (AliAODTrack*)tracks.Last();
278       SetChargeAndPID(part->GetPdgCode(), currTrack);
279       if (currTrack->Charge() != -99) {
280         if (currTrack->Charge() > 0) {
281           nPos++;
282         } else if (currTrack->Charge() < 0) {
283           nNeg++;
284         }           
285       }
286
287       LoopOverSecondaries(part);
288     }
289     return 1;
290   } else {
291     return 0;
292   }
293 }
294
295
296 void SetChargeAndPID(Int_t pdgCode, AliAODTrack *track) {
297
298   Float_t PID[10] = {0., 0., 0., 0., 0., 0., 0., 0., 0., 0.};
299
300   switch (pdgCode) {
301
302   case 22: // gamma
303     track->SetCharge(0);
304     PID[AliAODTrack::kUnknown] = 1.;
305     track->SetPID(PID);
306     nGamma++;
307     break;
308
309   case 11: // e- 
310     track->SetCharge(-1);
311     PID[AliAODTrack::kElectron] = 1.;
312     track->SetPID(PID);
313     nElectron++;
314     break;
315     
316   case -11: // e+
317     track->SetCharge(+1);
318     PID[AliAODTrack::kElectron] = 1.;
319     track->SetPID(PID);
320     nPositron++;
321     break;
322     
323   case 13: // mu- 
324     track->SetCharge(-1);
325     PID[AliAODTrack::kMuon] = 1.;
326     track->SetPID(PID);
327     nMuon++;
328     break;
329     
330   case -13: // mu+
331     track->SetCharge(+1);
332     PID[AliAODTrack::kMuon] = 1.;
333     track->SetPID(PID);
334     nAntiMuon++;
335     break;
336     
337   case 111: // pi0
338     track->SetCharge(0);
339     PID[AliAODTrack::kUnknown] = 1.;
340     track->SetPID(PID);
341     nPi0++;
342     break;
343     
344   case 211: // pi+
345     track->SetCharge(+1);
346     PID[AliAODTrack::kPion] = 1.;
347     track->SetPID(PID);
348     nPiPlus++;
349     break;
350     
351   case -211: // pi-
352     track->SetCharge(-1);
353     PID[AliAODTrack::kPion] = 1.;
354     track->SetPID(PID);
355     nPiMinus++;
356     break;
357     
358   case 130: // K0L
359     track->SetCharge(0);
360     PID[AliAODTrack::kUnknown] = 1.;
361     track->SetPID(PID);
362     nK0++;
363     break;
364     
365   case 321: // K+
366     track->SetCharge(+1);
367     PID[AliAODTrack::kKaon] = 1.;
368     track->SetPID(PID);
369     nKPlus++;
370     break;
371     
372   case -321: // K- 
373     track->SetCharge(-1);
374     PID[AliAODTrack::kKaon] = 1.;
375     track->SetPID(PID);
376     nKMinus++;
377     break;
378     
379   case 2112: // n
380     track->SetCharge(0);
381     PID[AliAODTrack::kUnknown] = 1.;
382     track->SetPID(PID);
383     nNeutron++;
384     break;
385     
386   case 2212: // p
387     track->SetCharge(+1);
388     PID[AliAODTrack::kProton] = 1.;
389     track->SetPID(PID);
390     nProton++;
391     break;
392     
393   case -2212: // anti-p
394     track->SetCharge(-1);
395     PID[AliAODTrack::kProton] = 1.;
396     track->SetPID(PID);
397     nAntiProton++;
398     break;
399
400   case 310: // K0S
401     track->SetCharge(0);
402     PID[AliAODTrack::kUnknown] = 1.;
403     track->SetPID(PID);
404     nK0++;
405     break;
406     
407   case 311: // K0
408     track->SetCharge(0);
409     PID[AliAODTrack::kUnknown] = 1.;
410     track->SetPID(PID);
411     nK0++;
412     break;
413     
414   case -311: // anti-K0
415     track->SetCharge(0);
416     PID[AliAODTrack::kUnknown] = 1.;
417     track->SetPID(PID);
418     nK0++;
419     break;
420     
421   case 221: // eta
422     track->SetCharge(0);
423     PID[AliAODTrack::kUnknown] = 1.;
424     track->SetPID(PID);
425     break;
426
427   case 3122: // lambda
428     track->SetCharge(0);
429     PID[AliAODTrack::kUnknown] = 1.;
430     track->SetPID(PID);
431     break;
432
433   case 3222: // Sigma+
434     track->SetCharge(+1);
435     PID[AliAODTrack::kUnknown] = 1.;
436     track->SetPID(PID);
437     break;
438
439   case 3212: // Sigma0
440     track->SetCharge(-1);
441     PID[AliAODTrack::kUnknown] = 1.;
442     track->SetPID(PID);
443     break;
444
445   case 3112: // Sigma-
446     track->SetCharge(-1);
447     PID[AliAODTrack::kUnknown] = 1.;
448     track->SetPID(PID);
449     break;
450
451   case 3322: // Xi0
452     track->SetCharge(0);
453     PID[AliAODTrack::kUnknown] = 1.;
454     track->SetPID(PID);
455     break;
456
457   case 3312: // Xi-
458     track->SetCharge(-1);
459     PID[AliAODTrack::kUnknown] = 1.;
460     track->SetPID(PID);
461     break;
462
463   case 3334: // Omega-
464     track->SetCharge(-1);
465     PID[AliAODTrack::kUnknown] = 1.;
466     track->SetPID(PID);
467     break;
468
469   case -2112: // n-bar
470     track->SetCharge(0);
471     PID[AliAODTrack::kUnknown] = 1.;
472     track->SetPID(PID);
473     break;
474
475   case -3122: // anti-Lambda
476     track->SetCharge(0);
477     PID[AliAODTrack::kUnknown] = 1.;
478     track->SetPID(PID);
479     break;
480
481   case -3222: // anti-Sigma-
482     track->SetCharge(-1);
483     PID[AliAODTrack::kUnknown] = 1.;
484     track->SetPID(PID);
485     break;
486
487   case -3212: // anti-Sigma0
488     track->SetCharge(0);
489     PID[AliAODTrack::kUnknown] = 1.;
490     track->SetPID(PID);
491     break;
492
493   case -3112: // anti-Sigma+
494     track->SetCharge(+1);
495     PID[AliAODTrack::kUnknown] = 1.;
496     track->SetPID(PID);
497     break;
498
499   case -3322: // anti-Xi0
500     track->SetCharge(0);
501     PID[AliAODTrack::kUnknown] = 1.;
502     track->SetPID(PID);
503     break;
504
505   case -3312: // anti-Xi+
506     track->SetCharge(+1);
507     break;
508
509   case -3334: // anti-Omega+
510     track->SetCharge(+1);
511     PID[AliAODTrack::kUnknown] = 1.;
512     track->SetPID(PID);
513     break;
514
515   case 411: // D+
516     track->SetCharge(+1);
517     PID[AliAODTrack::kUnknown] = 1.;
518     track->SetPID(PID);
519     break;
520
521   case -411: // D- 
522     track->SetCharge(-1);
523     PID[AliAODTrack::kUnknown] = 1.;
524     track->SetPID(PID);
525     break;
526
527   case 421: // D0
528     track->SetCharge(0);
529     PID[AliAODTrack::kUnknown] = 1.;
530     track->SetPID(PID);
531     break;
532
533   case -421: // anti-D0
534     track->SetCharge(0);
535     PID[AliAODTrack::kUnknown] = 1.;
536     track->SetPID(PID);
537     break;
538
539   default : // unknown
540     track->SetCharge(-99);
541     PID[AliAODTrack::kUnknown] = 1.;
542     track->SetPID(PID);
543  }
544
545   return;
546 }
547
548
549 void SetVertexType(TParticle *part, AliAODVertex *vertex) {
550   // this whole thing doesn't make much sense. but anyhow...
551
552   TParticle *mother = stack->Particle(part->GetFirstMother());
553   Int_t pdgMother = mother->GetPdgCode();
554   Int_t pdgPart = part->GetPdgCode();
555   
556   // kinks
557   if (mother->GetNDaughters() == 2) {
558     Int_t firstPdgCode = stack->Particle(mother->GetFirstDaughter())->GetPdgCode();
559     Int_t lastPdgCode = stack->Particle(mother->GetLastDaughter())->GetPdgCode();
560
561     if (!(pdgMother == 22 || pdgMother == 111 || pdgMother == 130 || 
562           TMath::Abs(pdgMother) == 2112 || pdgMother == 310 || pdgMother == 221 || 
563           TMath::Abs(pdgMother) == 3122 || TMath::Abs(pdgMother) == 3322 || 
564           pdgMother == -3212 || TMath::Abs(pdgMother) == 421 || 
565           TMath::Abs(pdgMother) == 311) // not neutral
566         && (((firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
567               TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
568               firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
569               TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
570               TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // neutral
571              && !(lastPdgCode == 22 || lastPdgCode == 111 || lastPdgCode == 130 || 
572                   TMath::Abs(lastPdgCode) == 2112 || lastPdgCode == 310 || 
573                   lastPdgCode == 221 || TMath::Abs(lastPdgCode) == 3122 || 
574                   TMath::Abs(lastPdgCode) == 3322 || lastPdgCode == -3212 || 
575                   TMath::Abs(lastPdgCode) == 421 || TMath::Abs(pdgMother) == 311)) // not neutral
576             || !((firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
577                   TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
578                   firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
579                   TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
580                   TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not 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)))) { // neutral
586       
587       vertex->SetType(AliAODVertex::kKink);
588       jKinks++;
589     }
590   }
591
592   // V0
593   else if (mother->GetNDaughters() == 2) {
594     Int_t firstPdgCode = stack->Particle(mother->GetFirstDaughter())->GetPdgCode();
595     Int_t lastPdgCode = stack->Particle(mother->GetLastDaughter())->GetPdgCode();
596
597     if ((pdgMother == 22 || pdgMother == 111 || pdgMother == 130 || 
598          TMath::Abs(pdgMother) == 2112 || pdgMother == 310 || 
599          pdgMother == 221 || TMath::Abs(pdgMother) == 3122 || 
600          TMath::Abs(pdgMother) == 3322 || pdgMother == -3212 || 
601          TMath::Abs(pdgMother) == 421 || TMath::Abs(pdgMother) == 311) // neutral
602         && !(lastPdgCode == 22 || lastPdgCode == 111 || lastPdgCode == 130 || 
603              TMath::Abs(lastPdgCode) == 2112 || lastPdgCode == 310 || 
604              lastPdgCode == 221 || TMath::Abs(lastPdgCode) == 3122 || 
605              TMath::Abs(lastPdgCode) == 3322 || lastPdgCode == -3212 || 
606              TMath::Abs(lastPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not neutral
607         && !(firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
608              TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
609              firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
610              TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
611              TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311)) { // not neutral
612       
613       vertex->SetType(AliAODVertex::kV0);
614       jV0s++;
615     }
616   }
617
618   // Cascade
619   else if (mother->GetNDaughters() == 2) {
620     Int_t firstPdgCode = stack->Particle(mother->GetFirstDaughter())->GetPdgCode();
621     Int_t lastPdgCode = stack->Particle(mother->GetLastDaughter())->GetPdgCode();
622     
623     if ((TMath::Abs(pdgMother) == 3334 || TMath::Abs(pdgMother) == 3312 || TMath::Abs(pdgMother) == 3322) &&
624         (TMath::Abs(pdgPart) == 3122 || TMath::Abs(pdgPart) == 211 || TMath::Abs(pdgPart) == 321)
625         && ((!(firstPdgCode == 22 || firstPdgCode == 111 || firstPdgCode == 130 || 
626                TMath::Abs(firstPdgCode) == 2112 || firstPdgCode == 310 || 
627                firstPdgCode == 221 || TMath::Abs(firstPdgCode) == 3122 || 
628                TMath::Abs(firstPdgCode) == 3322 || firstPdgCode == -3212 || 
629                TMath::Abs(firstPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not neutral   
630              && TMath::Abs(lastPdgCode) == 3122) // labmda or anti-lambda
631             || ((!(lastPdgCode == 22 || lastPdgCode == 111 || lastPdgCode == 130 || 
632                    TMath::Abs(lastPdgCode) == 2112 || lastPdgCode == 310 || 
633                    lastPdgCode == 221 || TMath::Abs(lastPdgCode) == 3122 || 
634                    TMath::Abs(lastPdgCode) == 3322 || lastPdgCode == -3212 || 
635                    TMath::Abs(lastPdgCode) == 421 || TMath::Abs(pdgMother) == 311) // not neutral
636                  && TMath::Abs(firstPdgCode) == 3122)))) { // lambda or anti-lambda
637       vertex->SetType(AliAODVertex::kCascade);
638       jCascades++;
639     }
640   }
641
642   // Multi
643   else if (mother->GetNDaughters() > 2) {
644
645     vertex->SetType(AliAODVertex::kMulti);
646     jMultis++;
647   }
648
649   else {
650     vertex->SetType(AliAODVertex::kUndef);
651   }
652 }
653 //  LocalWords:  SetCharge