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