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 | |
45 | void 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 | |
60 | void 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 | |
90 | void 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 | |
143 | void 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 | |
173 | void 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 | |
208 | void 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 | |
413 | AliTPCtrack * 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 | |
494 | void 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 | } |