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