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 }