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 }