remove temporary task that was used to troubleshoot the gamma flow anaysis
[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   }
84   else
85     vtxOut = new AliESDVertex();
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);
106
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
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
539   // Magnetic field
540   AliTracker::SetFieldMap(gAlice->Field(),1); // 1 means uniform magnetic field
541
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 }