- Volume name attribute replaced with volume path
[u/mrichter/AliRoot.git] / macros / mcminiesd.C
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"
41 #include "AliTracker.h"
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
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
88 void 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
141 void 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
171 void 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
206 void 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
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
411 AliTPCtrack * 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
492 void 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
537   // Magnetic field
538   AliTracker::SetFieldMap(gAlice->Field(),1); // 1 means uniform magnetic field
539
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 }