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
  42 #endif
  43
  44 void selectRecTracks(AliESD* esdOut, Int_t * select, Int_t &nkeep0) {
  45   // Select also all the reconstructed tracks in case it was not done before
  46   Int_t nrec = esdOut->GetNumberOfTracks();
  47   for(Int_t irec=0; irec<nrec; irec++) {
  48     AliESDtrack * track = esdOut->GetTrack(irec);
  49     UInt_t label = TMath::Abs(track->GetLabel());
  50     if (label>=10000000) continue; // Underlying event. 10000000 is the
  51                                    // value of fkMASKSTEP in AliRunDigitizer
  52     if (!select[label]) {
  53       select[label] = 1;
  54       nkeep0++;
  55     }
  56   }
  57 }
  58
  59 void copyGeneralESDInfo(AliESD* esdIn, AliESD* esdOut) {
  60   // Event and run number
  61   esdOut->SetEventNumber(esdIn->GetEventNumber());
  62   esdOut->SetRunNumber(esdIn->GetRunNumber());
  63
  64   // Trigger
  65   esdOut->SetTrigger(esdIn->GetTrigger());
  66
  67   // Magnetic field
  68   esdOut->SetMagneticField(esdIn->GetMagneticField());
  69
  70   // Copy ESD vertex
  71   const AliESDVertex * vtxIn = esdIn->GetVertex();
  72   Double_t pos[3];
  73   vtxIn->GetXYZ(pos);
  74   Double_t cov[6];
  75   vtxIn->GetCovMatrix(cov);
  76
  77   AliESDVertex * vtxOut = new AliESDVertex(pos,cov,
  78                                            vtxIn->GetChi2(),
  79                                            vtxIn->GetNContributors());
  80   Double_t tp[3];
  81   vtxIn->GetTruePos(tp);
  82   vtxOut->SetTruePos(tp);
  83
  84   esdOut->SetVertex(vtxOut);
  85 }
  86
  87 void selectMiniESD(AliESD* esdIn, AliESD* &esdOut) {
  88   // This function copies the general ESD information
  89   // and selects the reconstructed tracks of interest
  90   // The second argument is a reference to pointer since we
  91   // want to operate not with the content, but with the pointer itself
  92
  93   printf("--------------------\n");
  94   printf("Selecting data mini ESD\n");
  95   TStopwatch timer;
  96   timer.Start();
  97
  98   // Create the new output ESD
  99   esdOut = new AliESD();
 100
 101   // Copy the general information
 102   copyGeneralESDInfo(esdIn, esdOut);
 103
 104   // Select tracks
 105   Int_t ntrk = esdIn->GetNumberOfTracks();
 106
 107   // Loop on tracks
 108   for (Int_t itrk=0; itrk<ntrk; itrk++) {
 109
 110     AliESDtrack * trackIn = esdIn->GetTrack(itrk);
 111     UInt_t status=trackIn->GetStatus();
 112
 113     //select only tracks with TPC or ITS refit
 114     if ((status&AliESDtrack::kTPCrefit)==0
 115         && (status&AliESDtrack::kITSrefit)==0
 116         ) continue;
 117
 118     //select only tracks with "combined" PID
 119     if ((status&AliESDtrack::kESDpid)==0) continue;
 120
 121     AliESDtrack * trackOut = new AliESDtrack(*trackIn);
 122
 123     // Reset most of the redundant information
 124     trackOut->MakeMiniESDtrack();
 125
 126     esdOut->AddTrack(trackOut);
 127
 128     // Do not delete trackIn, it is a second pointer to an
 129     // object belonging to the TClonesArray
 130     delete trackOut;
 131   }
 132
 133   printf("%d tracks selected\n",esdOut->GetNumberOfTracks());
 134   esdOut->Print();
 135   timer.Stop();
 136   timer.Print();
 137   printf("--------------------\n");
 138 }
 139
 140 void fixMotherDaughter(TClonesArray& particles, Int_t * map) {
 141   // Fix mother-daughter relationship
 142   // using the map of old to new indexes
 143   Int_t nkeep = particles.GetEntriesFast();
 144   for (Int_t i=0; i<nkeep; i++) {
 145     TParticle * part = (TParticle*)particles[i];
 146
 147     // mother
 148     Int_t mum = part->GetFirstMother();
 149     if (mum>-1 && i>map[mum])
 150       part->SetFirstMother(map[mum]);
 151
 152     // old indexes
 153     Int_t fd = part->GetFirstDaughter();
 154     Int_t ld = part->GetLastDaughter();
 155
 156     // invalidate daughters
 157     part->SetFirstDaughter(-1);
 158     part->SetLastDaughter(-1);
 159
 160     for (Int_t id=TMath::Max(fd,0); id<=ld; id++) {
 161       if (map[id]>-1) {
 162         // this daughter was selected
 163         if(part->GetFirstDaughter() < 0) part->SetFirstDaughter(map[id]);
 164         part->SetLastDaughter(map[id]);
 165       }
 166     }
 167   }
 168 }
 169
 170 void checkConsistency(TClonesArray& particles) {
 171   // Check consistency
 172   // each mother is before the daughters,
 173   // each daughter from the mother's list
 174   // points back to its mother
 175   Int_t nkeep = particles.GetEntriesFast();
 176   for (Int_t i=0; i<nkeep; i++) {
 177     TParticle * part = (TParticle*)particles[i];
 178
 179     Int_t mum = part->GetFirstMother();
 180
 181     if (mum>-1 && i<mum) {
 182       cout << "Problem: mother " << mum << " after daughter " << i << endl;
 183     }
 184
 185     Int_t fd = part->GetFirstDaughter();
 186     Int_t ld = part->GetLastDaughter();
 187
 188     if (fd > ld ) cout << "Problem " << fd << " > " << ld << endl;
 189     if (fd > -1 && fd < i )
 190       cout << "Problem: mother " << i << " after daughter " << ld << endl;
 191
 192     for (Int_t id=TMath::Max(fd,0); id<=ld; id++) {
 193       TParticle * daughter = (TParticle*)particles[id];
 194       if (daughter->GetFirstMother() != i) {
 195         cout << "Problem "<< i << " not "
 196              << daughter->GetFirstMother()  << endl;
 197         daughter->Print();
 198       }
 199     }
 200
 201   }
 202 }
 203
 204
 205 void kinetree(AliRunLoader* runLoader, AliESD* &esdOut, TClonesArray* &ministack) {
 206
 207   printf("--------------------\n");
 208   printf("Selecting MC mini ESD\n");
 209   TStopwatch timer;
 210   timer.Start();
 211
 212   // Particle properties
 213   TDatabasePDG * pdgdb = TDatabasePDG::Instance();
 214
 215   // Particle stack
 216   AliStack * stack = runLoader->Stack();
 217
 218   Int_t npart = stack->GetNtrack();
 219
 220   // Particle momentum and vertex. Will be extracted from TParticle
 221   TLorentzVector momentum;
 222   TArrayF vertex(3);
 223   runLoader->GetHeader()->GenEventHeader()->PrimaryVertex(vertex);
 224   TVector3 dvertex;
 225
 226   // Counter for selected particles
 227   Int_t nkeep0 = 0;
 228
 229   Int_t * select = new Int_t[npart];
 230
 231   // Loop on particles: physics selection
 232   for (Int_t ipart=0; ipart<npart; ipart++) {
 233
 234     select[ipart] = 0;
 235
 236     TParticle * part = stack->Particle(ipart);
 237
 238     if (!part) {
 239       cout << "Empty particle " << ipart << endl;
 240       continue;
 241     }
 242
 243     // Particle momentum and vertex of origin
 244     part->Momentum(momentum);
 245     dvertex.SetXYZ(part->Vx()-vertex[0],
 246                    part->Vy()-vertex[1],
 247                    part->Vz()-vertex[2]);
 248
 249     // Now select only the "interesting" particles
 250
 251     // Select all particles from the primary vertex:
 252     // resonances and products of resonance decays
 253     if(dvertex.Mag()<0.0001) {
 254       select[ipart] = 1;
 255       nkeep0++;
 256       continue;
 257     }
 258
 259     // Reject particles born outside ITS
 260     if (dvertex.Perp()>100) continue;
 261
 262     // Reject particles with too low momentum, they don't rich TPC
 263     if (part->Pt()<0.1) continue;
 264
 265     // Reject "final state" neutral particles
 266     // This has to be redone after this loop
 267     Int_t pdgcode = part->GetPdgCode();
 268     TParticlePDG * pdgpart = pdgdb->GetParticle(pdgcode);
 269     Int_t charge  = 0;
 270     if (pdgpart) charge = (Int_t) pdgpart->Charge();
 271
 272     if (!charge) {
 273       if (part->GetFirstDaughter()<0) continue;
 274     }
 275
 276     // Select the rest of particles except the case
 277     // particle -> particle + X
 278     // for example bremstrahlung and delta electrons
 279     Int_t mumid = part->GetFirstMother();
 280     TParticle * mother = stack->Particle(mumid);
 281     if (mother) {
 282       Int_t mumpdg = mother->GetPdgCode();
 283       Bool_t skip = kFALSE;
 284       for (Int_t id=mother->GetFirstDaughter();
 285            id<=mother->GetLastDaughter();
 286            id++) {
 287         TParticle * daughter=stack->Particle(id);
 288         if (!daughter) continue;
 289         if (mumpdg == daughter->GetPdgCode()) {
 290           skip=kTRUE;
 291           break;
 292         }
 293       }
 294       if (skip) continue;
 295     }
 296
 297     // All the criteria are OK, this particle is selected
 298     select[ipart] = 1;
 299     nkeep0++;
 300
 301   } // end loop over particles in the current event
 302
 303   selectRecTracks(esdOut, select, nkeep0);
 304
 305   // Container for the selected particles
 306   TClonesArray * pparticles = new TClonesArray("TParticle",nkeep0);
 307   TClonesArray &particles = *pparticles;
 308
 309   // Map of the indexes in the stack and the new ones in the TClonesArray
 310   Int_t * map = new Int_t[npart];
 311
 312   // Counter for selected particles
 313   Int_t nkeep = 0;
 314
 315   for (Int_t ipart=0; ipart<npart; ipart++) {
 316
 317     map[ipart] = -99; // Reset the map
 318
 319     if (select[ipart]) {
 320
 321       TParticle * part = stack->Particle(ipart);
 322       map[ipart] = nkeep;
 323       new(particles[nkeep++]) TParticle(*part);
 324
 325     }
 326   }
 327
 328   if (nkeep0 != nkeep) printf("Strange %d is not %d\n",nkeep0,nkeep);
 329
 330   delete [] select;
 331
 332   // Fix mother-daughter relationship
 333   fixMotherDaughter(particles,map);
 334
 335   // Check consistency
 336   checkConsistency(particles);
 337
 338   // Now remove the "final state" neutral particles
 339
 340   TClonesArray * particles1 = new TClonesArray("TParticle",nkeep);
 341   Int_t * map1 = new Int_t[nkeep];
 342   Int_t nkeep1 = 0;
 343
 344   // Loop on particles
 345   for (Int_t ipart=0; ipart<nkeep; ipart++) {
 346
 347     map1[ipart] = -99; // Reset the map
 348
 349     TParticle * part = (TParticle*)particles[ipart];
 350
 351     // Reject "final state" neutral particles
 352     // This has to be redone after this loop
 353     Int_t pdgcode = part->GetPdgCode();
 354     TParticlePDG * pdgpart = pdgdb->GetParticle(pdgcode);
 355     Int_t charge  =  0;
 356     if (pdgpart) charge = (Int_t) pdgpart->Charge();
 357
 358     if (!charge) {
 359       if (part->GetFirstDaughter()<0) continue;
 360     }
 361
 362     // Select the particle
 363     map1[ipart] = nkeep1;
 364     TClonesArray &ar = *particles1;
 365     new(ar[nkeep1++]) TParticle(*part);
 366   }
 367
 368   particles.Delete();
 369   delete pparticles;
 370   cout << nkeep1 << " particles finally selected" << endl;
 371
 372   fixMotherDaughter(*particles1,map1);
 373   checkConsistency(*particles1);
 374
 375   // Remap the labels of reconstructed tracks
 376
 377   AliESD * esdNew = new AliESD();
 378
 379   copyGeneralESDInfo(esdOut,esdNew);
 380
 381   // Needed by the TPC track
 382   AliKalmanTrack::SetConvConst(1000/0.299792458/esdOut->GetMagneticField());
 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   // Now load kinematics and event header
 540   runLoader->LoadKinematics();
 541   runLoader->LoadHeader();
 542
 543   // Loop on events: check that MC and data contain the same number of events
 544   Int_t nevESD = bIn->GetEntries();
 545   Int_t nevMC = (Int_t)runLoader->GetNumberOfEvents();
 546
 547   if (nevESD != nevMC)
 548     cout << "Inconsistent number of ESD("<<nevESD<<") and MC("<<nevMC<<") events" << endl;
 549
 550   // Loop on events
 551   Int_t nev=TMath::Min(nevMC,nevESD);
 552
 553   cout << nev << " events to process" << endl;
 554
 555   for (Int_t iev=0; iev<nev; iev++) {
 556     cout << "Event " << iev << endl;
 557     // Get "data" event
 558     bIn->GetEntry(iev);
 559
 560     // "Data" selection
 561     selectMiniESD(esdIn,esdOut);
 562
 563     // Get MC event
 564     runLoader->GetEvent(iev);
 565
 566     // MC selection
 567     kinetree(runLoader,esdOut,ministack);
 568     bMC->SetAddress(&ministack); // needed because ministack has changed
 569
 570     // Header
 571     headerIn = runLoader->GetHeader();
 572
 573     // Fill the trees
 574     tOut->Fill();
 575     tMC->Fill();
 576     delete esdOut;
 577     esdOut = 0x0;
 578     //    delete esdIn;
 579     //    esdIn = 0x0;
 580     ministack->Delete();
 581     delete ministack;
 582
 583   }
 584
 585   fOut = tOut->GetCurrentFile();
 586   fOut->Write();
 587   fOut->Close();
 588
 589   fIn->Close();
 590
 591   timer.Stop();
 592   timer.Print();
 593   printf("====================\n");
 594 }