TRD/AliTRDsaveMCtracks.C

   1 #ifndef __CINT__
   2   #include <iostream.h>
   3   #include "AliTRDtracker.h"
   4
   5   #include "AliTRDcluster.h"
   6   #include "AliTRDhit.h"
   7   #include "AliTRDv1.h"
   8   #include "AliTRDgeometry.h"
   9   #include "AliTRDparameter.h"
  10   #include "alles.h"
  11   #include "AliTRDmcTrack.h"
  12
  13   #include "AliTRDtrack.h"
  14
  15   #include "TFile.h"
  16   #include "TParticle.h"
  17   #include "TStopwatch.h"
  18 #endif
  19
  20 Int_t Find(Int_t prtcl, TObjArray *mctarray) {
  21
  22 // Returns index of the mct which corresponds to particle <prtcl>
  23
  24   Int_t b=0, e=mctarray->GetEntriesFast(), m=(b+e)/2;
  25   AliTRDmcTrack *mct = 0;
  26   while ((b<e)&&(e!=0)) {
  27     m=(b+e)/2;
  28     mct = (AliTRDmcTrack*) mctarray->UncheckedAt(m);
  29     if (prtcl > mct->GetTrackIndex()) b=m+1;
  30     else e=m;
  31   }
  32
  33   if(mct->GetTrackIndex() == prtcl) return m;
  34
  35   if((m+1) < mctarray->GetEntriesFast()) {
  36     mct = (AliTRDmcTrack*) mctarray->UncheckedAt(m+1);
  37     if(mct->GetTrackIndex() == prtcl) return m+1;
  38   }
  39
  40   return -1;
  41 }
  42
  43 void AliTRDsaveMCtracks() {
  44
  45   TObjArray mctracks(2000);
  46   TObjArray *TRDmcTracks = &mctracks;
  47
  48   Char_t *alifile = "galice.root";
  49   Int_t   nEvent  = 0;
  50   Float_t low_pt_cut = 0.05;
  51   Float_t low_p_cut = 0.02;
  52   Int_t min_no_of_clusters = 10;
  53
  54   // Connect the AliRoot file containing Geometry, Kine, Hits, and Digits
  55   TFile *gafl = (TFile*) gROOT->GetListOfFiles()->FindObject(alifile);
  56   if (!gafl) {
  57     cout << "Open the ALIROOT-file " << alifile << endl;
  58     gafl = new TFile(alifile);
  59   }
  60   else {
  61     cout << alifile << " is already open" << endl;
  62   }
  63
  64   // Get AliRun object from file or create it if not on file
  65   gAlice = (AliRun*) gafl->Get("gAlice");
  66   if (gAlice)
  67     cout << "AliRun object found on file" << endl;
  68   else
  69     gAlice = new AliRun("gAlice","Alice test program");
  70
  71   AliTRDv1       *fTRD     = (AliTRDv1*) gAlice->GetDetector("TRD");
  72
  73   // Import the Trees for the event nEvent in the file
  74   const Int_t nparticles = gAlice->GetEvent(nEvent);
  75
  76   printf("found %d particles in event %d \n", nparticles, nEvent);
  77
  78   // Create TRDmcTracks for each particle
  79   Int_t label = -1, charge = 0;
  80   Bool_t primary;
  81   Float_t mass;
  82   Int_t pdg_code;
  83
  84   printf("\n");
  85
  86   for (Int_t ii=0; ii<nparticles; ii++) {
  87
  88     printf("\r particle %d out of %d",ii,nparticles);
  89
  90     TParticle *p = gAlice->Particle(ii);
  91     if(p->P() < low_p_cut) continue;
  92
  93     primary = kTRUE;
  94     if (p->GetFirstMother()>=0) primary = kFALSE;
  95
  96     pdg_code = (Int_t) p->GetPdgCode();
  97
  98     if ((pdg_code == 10010020) ||
  99         (pdg_code == 10010030) ||
 100         (pdg_code == 50000050) ||
 101         (pdg_code == 50000051) ||
 102         (pdg_code == 10020040)) {
 103
 104       mass = 0.;
 105       charge = 0;
 106     }
 107     else {
 108       TParticlePDG *pdg = p->GetPDG();
 109       charge = (Int_t) pdg->Charge();
 110       mass=pdg->Mass();
 111     }
 112     if(charge == 0) continue;
 113
 114     AliTRDmcTrack *t = new AliTRDmcTrack(ii, primary, mass, charge, pdg_code);
 115     TRDmcTracks->AddLast(t);
 116   }
 117   printf("\r\n");
 118
 119   // Loop through TRD clusters and assign indexes to MC tracks
 120
 121   TFile *geofile =TFile::Open("AliTRDclusters.root");
 122   AliTRDtracker *Tracker = new AliTRDtracker(geofile);
 123   Tracker->SetEventNumber(nEvent);
 124
 125   AliTRDgeometry *fGeo   = (AliTRDgeometry*) geofile->Get("TRDgeometry");
 126   AliTRDparameter *fPar   = (AliTRDparameter*) geofile->Get("TRDparameter");
 127
 128
 129   Char_t *clusterfile = "AliTRDclusters.root";
 130   TObjArray carray(2000);
 131   TObjArray *ClustersArray = &carray;
 132   Tracker->ReadClusters(ClustersArray,clusterfile);
 133
 134   Int_t nClusters = carray.GetEntriesFast();
 135
 136   Int_t track, index;
 137   AliTRDmcTrack *tpr = NULL;
 138
 139   for (Int_t i = 0; i < nClusters; i++) {
 140
 141     printf("\r assigning cluster %d out of %d", i, nClusters);
 142     AliTRDcluster *cl = (AliTRDcluster *) ClustersArray->UncheckedAt(i);
 143
 144     for(Int_t j=0; j<3; j++) {
 145       track = cl->GetLabel(j);
 146       if(track < 0) continue;
 147       if(track >= nparticles)
 148         printf("Track index %d is larger than total number of particles %d\n"
 149                ,track,nparticles);
 150       else {
 151         index = Find(track, TRDmcTracks);
 152         tpr = 0;
 153         if(index > 0) tpr = (AliTRDmcTrack *) TRDmcTracks->UncheckedAt(index);
 154
 155         if(tpr) {
 156           label = tpr->GetTrackIndex();
 157           if(label != track) printf("Got label %d while expected %d !\n",
 158                                     label, track);
 159           TParticle *p = gAlice->Particle(track);
 160
 161           if (p->Pt() > low_pt_cut) tpr->Update(i);
 162         }
 163       }
 164     }
 165   }
 166
 167   // Loop through the TRD hits and set XYZ and Pin and Pout
 168
 169   TTree *hitTree = gAlice->TreeH();
 170   gAlice->ResetHits();
 171   Int_t   preamp, amp, det, plane = 0, nBytes = 0;
 172   Int_t   nTrack = (Int_t) hitTree->GetEntries();
 173   Float_t  pos[3], rot[3];
 174   Float_t  prepos[3], prerot[3];
 175
 176   Bool_t  entrance = kFALSE;
 177   Bool_t  exit = kFALSE;
 178
 179
 180   Double_t xin[6], xout[6];
 181   for(Int_t j = 0; j < 6; j++) {
 182     xin[j] = Tracker->GetX(0,j,14);
 183     xout[j] = Tracker->GetX(0,j,0);
 184   }
 185
 186   AliTRDhit *hit;
 187
 188   Int_t current_track = -1;
 189
 190   printf("\n");
 191   while(nTrack--) {
 192
 193     gAlice->ResetHits();
 194     nBytes += hitTree->GetEvent(nTrack);
 195     entrance = kTRUE;
 196     preamp = 1111;
 197     amp = 1111;
 198     for(Int_t j = 0; j < 3; j++) { pos[j] = 1111.; rot[j] = 1111.;}
 199
 200     for(hit = (AliTRDhit *) fTRD->FirstHit(-1);
 201         hit;
 202         hit = (AliTRDhit *) fTRD->NextHit()) {
 203
 204       preamp = amp;
 205       for(Int_t j = 0; j < 3; j++) { prepos[j] = pos[j]; prerot[j] = rot[j];}
 206
 207       amp = hit->GetCharge();
 208       if(amp != 0) continue;
 209       track   = hit->GetTrack();
 210
 211       TParticle *p = gAlice->Particle(track);
 212       if (p->Pt() <= low_pt_cut) continue;
 213
 214       if(track != current_track) {
 215         current_track = track;
 216         //      printf("\n\n");
 217       }
 218
 219       det = hit->GetDetector();
 220       plane    = fGeo->GetPlane(det);
 221       pos[0] = hit->X();
 222       pos[1] = hit->Y();
 223       pos[2] = hit->Z();
 224       fGeo->Rotate(det,pos,rot);
 225
 226       if(TMath::Abs(rot[0] - xin[plane]) < 1.2) entrance = kTRUE;
 227       else  entrance = kFALSE;
 228       if(TMath::Abs(rot[0] - xout[plane]) < 1.2) exit = kTRUE;
 229       else  exit = kFALSE;
 230
 231       if(entrance && (preamp == 0) && (prerot[0] < 200)) {
 232         index = Find(track, TRDmcTracks);
 233         if(index > 0) {
 234           tpr = 0;
 235           tpr = (AliTRDmcTrack *) TRDmcTracks->UncheckedAt(index);
 236           if(tpr) {
 237             label = tpr->GetTrackIndex();
 238             if(label != track) printf("Got label %d while expected %d !\n",
 239                                       label, track);
 240             else {
 241               /*
 242               printf("\n momentum at plane %d entrance, Pxyz: %f, %f, %f\n",
 243                      plane, prepos[0], prepos[1], prepos[2]);
 244               printf(" XYZ at plane %d entrance: %f, %f, %f\n",
 245                      plane, rot[0], rot[1], rot[2]);
 246               printf("expected x range: %f .. %f\n",xin[plane]-1.2,xin[plane]+1.2);
 247               */
 248               tpr->SetPin(plane, prepos[0], prepos[1], prepos[2]);
 249               tpr->SetXYZin(plane, rot[0], rot[1], rot[2]);
 250             }
 251           }
 252         }
 253       }
 254
 255
 256       if(exit && (preamp == 0) && (prerot[0] < 200)) {
 257         index = Find(track, TRDmcTracks);
 258         if(index > 0) {
 259           tpr = 0;
 260           tpr = (AliTRDmcTrack *) TRDmcTracks->UncheckedAt(index);
 261           if(tpr) {
 262             label = tpr->GetTrackIndex();
 263             if(label != track) printf("Got label %d while expected %d !\n",
 264                                       label, track);
 265             else {
 266               /*
 267               printf("\n momentum at plane %d exit, Pxyz: %f, %f, %f\n",
 268                      plane, prepos[0], prepos[1], prepos[2]);
 269               printf(" XYZ at plane %d exit: %f, %f, %f\n",
 270                      plane, rot[0], rot[1], rot[2]);
 271               printf("expected x range: %f .. %f\n",xout[plane]-1.2,xout[plane]+1.2);
 272               */
 273               tpr->SetPout(plane, prepos[0], prepos[1], prepos[2]);
 274               tpr->SetXYZout(plane, rot[0], rot[1], rot[2]);
 275             }
 276           }
 277         }
 278       }
 279     }
 280   }
 281
 282
 283   // Write TRDmcTracks in output file
 284
 285   TDirectory *savedir=gDirectory;
 286   Char_t *filename   = "AliTRDmcTracks.root";
 287   TFile *out = new TFile(filename,"RECREATE");
 288
 289   TTree *tracktree = new TTree("MCtracks","TRD MC tracks");
 290
 291   AliTRDmcTrack *iotrack=0;
 292   tracktree->Branch("MCtracks","AliTRDmcTrack",&iotrack,32000,0);
 293
 294   Int_t ntracks = TRDmcTracks->GetEntriesFast();
 295
 296   for (Int_t i=0; i<ntracks; i++) {
 297     AliTRDmcTrack *pt=(AliTRDmcTrack*)TRDmcTracks->UncheckedAt(i);
 298
 299     Int_t n = pt->GetNumberOfClusters();
 300     if(n > min_no_of_clusters) {
 301       iotrack=pt;
 302       tracktree->Fill();
 303       printf("Put track with label %d and %d clusters in the output tree \n",
 304              pt->GetTrackIndex(),n);
 305     }
 306   }
 307
 308   tracktree->Write();
 309   out->Close();
 310   savedir->cd();
 311
 312   return;
 313 }
 314