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