Using VZERO v6
[u/mrichter/AliRoot.git] / macros / mcminiesd.C
CommitLineData
75c681cb 1// This is an example macro which loops on all ESD events in a file set
2// and stores the selected information in form of mini ESD tree.
3// It also extracts the corresponding MC information in a parallel MC tree.
4//
5// Input: galice.root,Kinematics.root,AliESDs.root (read only)
6// Output: miniesd.root (recreate)
7//
8// Usage: faster if compiled
9// gSystem->SetIncludePath("-I$ROOTSYS/include -I$ALICE_ROOT/include -I$ALICE_ROOT/TPC");
10// gROOT->LoadMacro("mcminiesd.C+");
11// mcminiesd()
12//
13// Author: Peter Hristov
14// e-mail: Peter.Hristov@cern.ch
15
16#if !defined(__CINT__) || defined(__MAKECINT__)
17
18// Root include files
19#include <Riostream.h>
20#include <TFile.h>
21#include <TTree.h>
22#include <TBranch.h>
23#include <TStopwatch.h>
24#include <TObject.h>
25#include <TParticle.h>
26#include <TLorentzVector.h>
27#include <TVector3.h>
28#include <TArrayF.h>
29#include <TDatabasePDG.h>
30#include <TParticlePDG.h>
31
32// AliRoot include files
33#include "AliESD.h"
34#include "AliESDVertex.h"
35#include "AliRunLoader.h"
36#include "AliRun.h"
37#include "AliStack.h"
38#include "AliHeader.h"
39#include "AliGenEventHeader.h"
40#include "AliTPCtrack.h"
c84a5e9e 41#include "AliTracker.h"
75c681cb 42
43#endif
44
45void selectRecTracks(AliESD* esdOut, Int_t * select, Int_t &nkeep0) {
46 // Select also all the reconstructed tracks in case it was not done before
47 Int_t nrec = esdOut->GetNumberOfTracks();
48 for(Int_t irec=0; irec<nrec; irec++) {
49 AliESDtrack * track = esdOut->GetTrack(irec);
50 UInt_t label = TMath::Abs(track->GetLabel());
51 if (label>=10000000) continue; // Underlying event. 10000000 is the
52 // value of fkMASKSTEP in AliRunDigitizer
53 if (!select[label]) {
54 select[label] = 1;
55 nkeep0++;
56 }
57 }
58}
59
60void copyGeneralESDInfo(AliESD* esdIn, AliESD* esdOut) {
61 // Event and run number
62 esdOut->SetEventNumber(esdIn->GetEventNumber());
63 esdOut->SetRunNumber(esdIn->GetRunNumber());
64
65 // Trigger
66 esdOut->SetTrigger(esdIn->GetTrigger());
67
68 // Magnetic field
69 esdOut->SetMagneticField(esdIn->GetMagneticField());
70
71 // Copy ESD vertex
72 const AliESDVertex * vtxIn = esdIn->GetVertex();
73 Double_t pos[3];
74 vtxIn->GetXYZ(pos);
75 Double_t cov[6];
76 vtxIn->GetCovMatrix(cov);
77
78 AliESDVertex * vtxOut = new AliESDVertex(pos,cov,
79 vtxIn->GetChi2(),
80 vtxIn->GetNContributors());
81 Double_t tp[3];
82 vtxIn->GetTruePos(tp);
83 vtxOut->SetTruePos(tp);
84
85 esdOut->SetVertex(vtxOut);
86}
87
88void selectMiniESD(AliESD* esdIn, AliESD* &esdOut) {
89 // This function copies the general ESD information
90 // and selects the reconstructed tracks of interest
91 // The second argument is a reference to pointer since we
92 // want to operate not with the content, but with the pointer itself
93
94 printf("--------------------\n");
95 printf("Selecting data mini ESD\n");
96 TStopwatch timer;
97 timer.Start();
98
99 // Create the new output ESD
100 esdOut = new AliESD();
101
102 // Copy the general information
103 copyGeneralESDInfo(esdIn, esdOut);
104
105 // Select tracks
106 Int_t ntrk = esdIn->GetNumberOfTracks();
107
108 // Loop on tracks
109 for (Int_t itrk=0; itrk<ntrk; itrk++) {
110
111 AliESDtrack * trackIn = esdIn->GetTrack(itrk);
112 UInt_t status=trackIn->GetStatus();
113
114 //select only tracks with TPC or ITS refit
115 if ((status&AliESDtrack::kTPCrefit)==0
116 && (status&AliESDtrack::kITSrefit)==0
117 ) continue;
118
119 //select only tracks with "combined" PID
120 if ((status&AliESDtrack::kESDpid)==0) continue;
121
122 AliESDtrack * trackOut = new AliESDtrack(*trackIn);
123
124 // Reset most of the redundant information
125 trackOut->MakeMiniESDtrack();
126
127 esdOut->AddTrack(trackOut);
128
129 // Do not delete trackIn, it is a second pointer to an
130 // object belonging to the TClonesArray
131 delete trackOut;
132 }
133
134 printf("%d tracks selected\n",esdOut->GetNumberOfTracks());
135 esdOut->Print();
136 timer.Stop();
137 timer.Print();
138 printf("--------------------\n");
139}
140
141void fixMotherDaughter(TClonesArray& particles, Int_t * map) {
142 // Fix mother-daughter relationship
143 // using the map of old to new indexes
144 Int_t nkeep = particles.GetEntriesFast();
145 for (Int_t i=0; i<nkeep; i++) {
146 TParticle * part = (TParticle*)particles[i];
147
148 // mother
149 Int_t mum = part->GetFirstMother();
150 if (mum>-1 && i>map[mum])
151 part->SetFirstMother(map[mum]);
152
153 // old indexes
154 Int_t fd = part->GetFirstDaughter();
155 Int_t ld = part->GetLastDaughter();
156
157 // invalidate daughters
158 part->SetFirstDaughter(-1);
159 part->SetLastDaughter(-1);
160
161 for (Int_t id=TMath::Max(fd,0); id<=ld; id++) {
162 if (map[id]>-1) {
163 // this daughter was selected
164 if(part->GetFirstDaughter() < 0) part->SetFirstDaughter(map[id]);
165 part->SetLastDaughter(map[id]);
166 }
167 }
168 }
169}
170
171void checkConsistency(TClonesArray& particles) {
172 // Check consistency
173 // each mother is before the daughters,
174 // each daughter from the mother's list
175 // points back to its mother
176 Int_t nkeep = particles.GetEntriesFast();
177 for (Int_t i=0; i<nkeep; i++) {
178 TParticle * part = (TParticle*)particles[i];
179
180 Int_t mum = part->GetFirstMother();
181
182 if (mum>-1 && i<mum) {
183 cout << "Problem: mother " << mum << " after daughter " << i << endl;
184 }
185
186 Int_t fd = part->GetFirstDaughter();
187 Int_t ld = part->GetLastDaughter();
188
189 if (fd > ld ) cout << "Problem " << fd << " > " << ld << endl;
190 if (fd > -1 && fd < i )
191 cout << "Problem: mother " << i << " after daughter " << ld << endl;
192
193 for (Int_t id=TMath::Max(fd,0); id<=ld; id++) {
194 TParticle * daughter = (TParticle*)particles[id];
195 if (daughter->GetFirstMother() != i) {
196 cout << "Problem "<< i << " not "
197 << daughter->GetFirstMother() << endl;
198 daughter->Print();
199 }
200 }
201
202 }
203}
204
205
206void kinetree(AliRunLoader* runLoader, AliESD* &esdOut, TClonesArray* &ministack) {
207
208 printf("--------------------\n");
209 printf("Selecting MC mini ESD\n");
210 TStopwatch timer;
211 timer.Start();
212
213 // Particle properties
214 TDatabasePDG * pdgdb = TDatabasePDG::Instance();
215
216 // Particle stack
217 AliStack * stack = runLoader->Stack();
218
219 Int_t npart = stack->GetNtrack();
220
221 // Particle momentum and vertex. Will be extracted from TParticle
222 TLorentzVector momentum;
223 TArrayF vertex(3);
224 runLoader->GetHeader()->GenEventHeader()->PrimaryVertex(vertex);
225 TVector3 dvertex;
226
227 // Counter for selected particles
228 Int_t nkeep0 = 0;
229
230 Int_t * select = new Int_t[npart];
231
232 // Loop on particles: physics selection
233 for (Int_t ipart=0; ipart<npart; ipart++) {
234
235 select[ipart] = 0;
236
237 TParticle * part = stack->Particle(ipart);
238
239 if (!part) {
240 cout << "Empty particle " << ipart << endl;
241 continue;
242 }
243
244 // Particle momentum and vertex of origin
245 part->Momentum(momentum);
246 dvertex.SetXYZ(part->Vx()-vertex[0],
247 part->Vy()-vertex[1],
248 part->Vz()-vertex[2]);
249
250 // Now select only the "interesting" particles
251
252 // Select all particles from the primary vertex:
253 // resonances and products of resonance decays
254 if(dvertex.Mag()<0.0001) {
255 select[ipart] = 1;
256 nkeep0++;
257 continue;
258 }
259
260 // Reject particles born outside ITS
261 if (dvertex.Perp()>100) continue;
262
263 // Reject particles with too low momentum, they don't rich TPC
264 if (part->Pt()<0.1) continue;
265
266 // Reject "final state" neutral particles
267 // This has to be redone after this loop
268 Int_t pdgcode = part->GetPdgCode();
269 TParticlePDG * pdgpart = pdgdb->GetParticle(pdgcode);
270 Int_t charge = 0;
271 if (pdgpart) charge = (Int_t) pdgpart->Charge();
272
273 if (!charge) {
274 if (part->GetFirstDaughter()<0) continue;
275 }
276
277 // Select the rest of particles except the case
278 // particle -> particle + X
279 // for example bremstrahlung and delta electrons
280 Int_t mumid = part->GetFirstMother();
281 TParticle * mother = stack->Particle(mumid);
282 if (mother) {
283 Int_t mumpdg = mother->GetPdgCode();
284 Bool_t skip = kFALSE;
285 for (Int_t id=mother->GetFirstDaughter();
286 id<=mother->GetLastDaughter();
287 id++) {
288 TParticle * daughter=stack->Particle(id);
289 if (!daughter) continue;
290 if (mumpdg == daughter->GetPdgCode()) {
291 skip=kTRUE;
292 break;
293 }
294 }
295 if (skip) continue;
296 }
297
298 // All the criteria are OK, this particle is selected
299 select[ipart] = 1;
300 nkeep0++;
301
302 } // end loop over particles in the current event
303
304 selectRecTracks(esdOut, select, nkeep0);
305
306 // Container for the selected particles
307 TClonesArray * pparticles = new TClonesArray("TParticle",nkeep0);
308 TClonesArray &particles = *pparticles;
309
310 // Map of the indexes in the stack and the new ones in the TClonesArray
311 Int_t * map = new Int_t[npart];
312
313 // Counter for selected particles
314 Int_t nkeep = 0;
315
316 for (Int_t ipart=0; ipart<npart; ipart++) {
317
318 map[ipart] = -99; // Reset the map
319
320 if (select[ipart]) {
321
322 TParticle * part = stack->Particle(ipart);
323 map[ipart] = nkeep;
324 new(particles[nkeep++]) TParticle(*part);
325
326 }
327 }
328
329 if (nkeep0 != nkeep) printf("Strange %d is not %d\n",nkeep0,nkeep);
330
331 delete [] select;
332
333 // Fix mother-daughter relationship
334 fixMotherDaughter(particles,map);
335
336 // Check consistency
337 checkConsistency(particles);
338
339 // Now remove the "final state" neutral particles
340
341 TClonesArray * particles1 = new TClonesArray("TParticle",nkeep);
342 Int_t * map1 = new Int_t[nkeep];
343 Int_t nkeep1 = 0;
344
345 // Loop on particles
346 for (Int_t ipart=0; ipart<nkeep; ipart++) {
347
348 map1[ipart] = -99; // Reset the map
349
350 TParticle * part = (TParticle*)particles[ipart];
351
352 // Reject "final state" neutral particles
353 // This has to be redone after this loop
354 Int_t pdgcode = part->GetPdgCode();
355 TParticlePDG * pdgpart = pdgdb->GetParticle(pdgcode);
356 Int_t charge = 0;
357 if (pdgpart) charge = (Int_t) pdgpart->Charge();
358
359 if (!charge) {
360 if (part->GetFirstDaughter()<0) continue;
361 }
362
363 // Select the particle
364 map1[ipart] = nkeep1;
365 TClonesArray &ar = *particles1;
366 new(ar[nkeep1++]) TParticle(*part);
367 }
368
369 particles.Delete();
370 delete pparticles;
371 cout << nkeep1 << " particles finally selected" << endl;
372
373 fixMotherDaughter(*particles1,map1);
374 checkConsistency(*particles1);
375
376 // Remap the labels of reconstructed tracks
377
378 AliESD * esdNew = new AliESD();
379
380 copyGeneralESDInfo(esdOut,esdNew);
381
75c681cb 382 // Tracks
383 Int_t nrec = esdOut->GetNumberOfTracks();
384 for(Int_t irec=0; irec<nrec; irec++) {
385 AliESDtrack * track = esdOut->GetTrack(irec);
386 UInt_t label = TMath::Abs(track->GetLabel());
387 if (label<10000000) { // Signal event. 10000000 is the
388 // value of fkMASKSTEP in AliRunDigitizer
389 track->SetLabel(TMath::Sign(map1[map[label]],track->GetLabel()));
390 }
391 esdNew->AddTrack(track);
392 }
393
394 delete esdOut;
395 esdOut = esdNew;
396
397 ministack = particles1;
398
399
400 delete [] map;
401 delete [] map1;
402
403
404 timer.Stop();
405 timer.Print();
406 printf("--------------------\n");
407}
408
409
410
411AliTPCtrack * particle2track(TParticle * part) {
412
413 // Converts TParticle to AliTPCtrack
414
415 UInt_t index;
416 Double_t xx[5];
417 Double_t cc[15];
418 Double_t xref;
419 Double_t alpha;
420
421 // Calculate alpha: the rotation angle of the corresponding TPC sector
422 alpha = part->Phi()*180./TMath::Pi();
423 if (alpha<0) alpha+= 360.;
424 if (alpha>360) alpha -= 360.;
425
426 Int_t sector = (Int_t)(alpha/20.);
427 alpha = 10. + 20.*sector;
428 alpha /= 180;
429 alpha *= TMath::Pi();
430
431
432 // Reset the covariance matrix
433 for (Int_t i=0; i<15; i++) cc[i]=0.;
434
435
436 // Index
437 index = part->GetUniqueID();
438
439 // Get the vertex of origin and the momentum
440 TVector3 ver(part->Vx(),part->Vy(),part->Vz());
441 TVector3 mom(part->Px(),part->Py(),part->Pz());
442
443
444 // Rotate to the local (sector) coordinate system
445 ver.RotateZ(-alpha);
446 mom.RotateZ(-alpha);
447
448 // X of the referense plane
449 xref = ver.X();
450
451 // Track parameters
452 // fX = xref X-coordinate of this track (reference plane)
453 // fAlpha = Alpha Rotation angle the local (TPC sector)
454 // fP0 = YL Y-coordinate of a track
455 // fP1 = ZG Z-coordinate of a track
456 // fP2 = C*x0 x0 is center x in rotated frame
457 // fP3 = Tgl tangent of the track momentum dip angle
458 // fP4 = C track curvature
459
460 // Magnetic field
461 TVector3 field(0.,0.,1/AliKalmanTrack::GetConvConst());
462 // radius [cm] of track projection in (x,y)
463 Double_t rho =
464 mom.Pt()*AliKalmanTrack::GetConvConst();
465
466 Double_t charge =
467 TDatabasePDG::Instance()->GetParticle(part->GetPdgCode())->Charge();
468 charge /=3.;
469
470 if (TMath::Abs(charge) < 0.9) charge=1.e-9;
471
472 TVector3 vrho = mom.Cross(field);
473 vrho *= charge;
474 vrho.SetMag(rho);
475
476 vrho += ver;
477 Double_t x0 = vrho.X();
478
479 xx[0] = ver.Y();
480 xx[1] = ver.Z();
481 xx[3] = mom.Pz()/mom.Pt();
482 xx[4] = -charge/rho;
483 xx[2] = xx[4]*x0;
484 // if (TMath::Abs(charge) < 0.9) xx[4] = 0;
485
486 AliTPCtrack * tr = new AliTPCtrack(index, xx, cc, xref, alpha);
487 tr->SetLabel(index);
488 return tr;
489
490}
491
492void mcminiesd() {
493
494 printf("====================\n");
495 printf("Main program\n");
496 TStopwatch timer;
497 timer.Start();
498
499 // Input "data" file, tree, and branch
500 TFile * fIn = TFile::Open("AliESDs.root");
501 TTree * tIn = (TTree*)fIn->Get("esdTree");
502 TBranch * bIn = tIn->GetBranch("ESD");
503
504 AliESD * esdIn = 0; // The input ESD object is put here
505 bIn->SetAddress(&esdIn);
506
507 // Output file, trees, and branches.
508 // Both "data" and MC are stored in the same file
509 TFile * fOut = TFile::Open("miniesd.root","recreate");
510 TTree * tOut = new TTree("esdTree","esdTree");
511 TTree * tMC = new TTree("mcStackTree","mcStackTree");
512
513 AliESD * esdOut = 0; // The newly created ESD object
514 TBranch * bOut = tOut->Branch("ESD","AliESD",&esdOut);
515 bOut->SetCompressionLevel(9);
516
517 // TParticles are stored in TClonesArray
518 // The corresponding branch is created using "dummy" TClonesArray
519 TClonesArray * ministack = new TClonesArray("TParticle");
520 TBranch * bMC = tMC->Branch("Stack",&ministack);
521 ministack->Delete();
522 delete ministack;
523 bMC->SetCompressionLevel(9);
524
525 // Event header
526 AliHeader * headerIn = 0; //
527 TBranch * bHeader = tMC->Branch("Header","AliHeader",&headerIn);
528 bHeader->SetCompressionLevel(9);
529
530 // Run loader
531 AliRunLoader* runLoader = AliRunLoader::Open("galice.root");
532
533 // gAlice
534 runLoader->LoadgAlice();
535 gAlice = runLoader->GetAliRun();
536
c84a5e9e 537 // Magnetic field
538 AliTracker::SetFieldMap(gAlice->Field(),1); // 1 means uniform magnetic field
539
75c681cb 540 // Now load kinematics and event header
541 runLoader->LoadKinematics();
542 runLoader->LoadHeader();
543
544 // Loop on events: check that MC and data contain the same number of events
545 Int_t nevESD = bIn->GetEntries();
546 Int_t nevMC = (Int_t)runLoader->GetNumberOfEvents();
547
548 if (nevESD != nevMC)
549 cout << "Inconsistent number of ESD("<<nevESD<<") and MC("<<nevMC<<") events" << endl;
550
551 // Loop on events
552 Int_t nev=TMath::Min(nevMC,nevESD);
553
554 cout << nev << " events to process" << endl;
555
556 for (Int_t iev=0; iev<nev; iev++) {
557 cout << "Event " << iev << endl;
558 // Get "data" event
559 bIn->GetEntry(iev);
560
561 // "Data" selection
562 selectMiniESD(esdIn,esdOut);
563
564 // Get MC event
565 runLoader->GetEvent(iev);
566
567 // MC selection
568 kinetree(runLoader,esdOut,ministack);
569 bMC->SetAddress(&ministack); // needed because ministack has changed
570
571 // Header
572 headerIn = runLoader->GetHeader();
573
574 // Fill the trees
575 tOut->Fill();
576 tMC->Fill();
577 delete esdOut;
578 esdOut = 0x0;
579 // delete esdIn;
580 // esdIn = 0x0;
581 ministack->Delete();
582 delete ministack;
583
584 }
585
586 fOut = tOut->GetCurrentFile();
587 fOut->Write();
588 fOut->Close();
589
590 fIn->Close();
591
592 timer.Stop();
593 timer.Print();
594 printf("====================\n");
595}