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