[u/mrichter/AliRoot.git] / TRD / AliTRDsaveMCtracks.C

#ifndef __CINT__
  #include <iostream.h>
  #include "AliTRDtracker.h"

  #include "AliTRDcluster.h" 
  #include "AliTRDhit.h" 
  #include "AliTRDv1.h"
  #include "AliTRDgeometry.h"    
  #include "AliTRDparameter.h"    
  #include "alles.h"  
  #include "AliTRDmcTrack.h"

  #include "AliTRDtrack.h"
  
  #include "TFile.h"
  #include "TParticle.h"
  #include "TStopwatch.h"
#endif

Int_t Find(Int_t prtcl, TObjArray *mctarray) {

// Returns index of the mct which corresponds to particle <prtcl> 

  Int_t b=0, e=mctarray->GetEntriesFast(), m=(b+e)/2;
  AliTRDmcTrack *mct = 0; 
  while ((b<e)&&(e!=0)) {
    m=(b+e)/2;
    mct = (AliTRDmcTrack*) mctarray->UncheckedAt(m);
    if (prtcl > mct->GetTrackIndex()) b=m+1;
    else e=m;
  }

  if(mct->GetTrackIndex() == prtcl) return m;

  if((m+1) < mctarray->GetEntriesFast()) {
    mct = (AliTRDmcTrack*) mctarray->UncheckedAt(m+1);
    if(mct->GetTrackIndex() == prtcl) return m+1;
  }

  return -1;
}                    

void AliTRDsaveMCtracks() {

  TObjArray mctracks(2000);
  TObjArray *TRDmcTracks = &mctracks; 

  Char_t *alifile = "galice.root";
  Int_t   nEvent  = 0; 
  Float_t low_pt_cut = 0.05;
  Float_t low_p_cut = 0.02;
  Int_t min_no_of_clusters = 10;

  // Connect the AliRoot file containing Geometry, Kine, Hits, and Digits
  TFile *gafl = (TFile*) gROOT->GetListOfFiles()->FindObject(alifile);
  if (!gafl) {
    cout << "Open the ALIROOT-file " << alifile << endl;
    gafl = new TFile(alifile);
  }
  else {
    cout << alifile << " is already open" << endl;
  }

  // Get AliRun object from file or create it if not on file
  gAlice = (AliRun*) gafl->Get("gAlice");
  if (gAlice)
    cout << "AliRun object found on file" << endl;
  else
    gAlice = new AliRun("gAlice","Alice test program");

  AliTRDv1       *fTRD     = (AliTRDv1*) gAlice->GetDetector("TRD");   
  
  // Import the Trees for the event nEvent in the file
  const Int_t nparticles = gAlice->GetEvent(nEvent);

  printf("found %d particles in event %d \n", nparticles, nEvent);
    
  // Create TRDmcTracks for each particle
  Int_t label = -1, charge = 0;
  Bool_t primary;
  Float_t mass;
  Int_t pdg_code;

  printf("\n");

  for (Int_t ii=0; ii<nparticles; ii++) {

    printf("\r particle %d out of %d",ii,nparticles);

    TParticle *p = gAlice->Particle(ii);
    if(p->P() < low_p_cut) continue;

    primary = kTRUE;
    if (p->GetFirstMother()>=0) primary = kFALSE;
    
    pdg_code = (Int_t) p->GetPdgCode();

    if ((pdg_code == 10010020) ||
 	(pdg_code == 10010030) || 
 	(pdg_code == 50000050) || 
 	(pdg_code == 50000051) || 
 	(pdg_code == 10020040)) {

      mass = 0.;
      charge = 0;
    }
    else {
      TParticlePDG *pdg = p->GetPDG();
      charge = (Int_t) pdg->Charge(); 
      mass=pdg->Mass();
    }
    if(charge == 0) continue;

    AliTRDmcTrack *t = new AliTRDmcTrack(ii, primary, mass, charge, pdg_code);
    TRDmcTracks->AddLast(t);
  }
  printf("\r\n");

  // Loop through TRD clusters and assign indexes to MC tracks
  
  TFile *geofile =TFile::Open("AliTRDclusters.root");   
  AliTRDtracker *Tracker = new AliTRDtracker(geofile);
  Tracker->SetEventNumber(nEvent);

  AliTRDgeometry *fGeo   = (AliTRDgeometry*) geofile->Get("TRDgeometry"); 
  AliTRDparameter *fPar   = (AliTRDparameter*) geofile->Get("TRDparameter"); 


  Char_t *clusterfile = "AliTRDclusters.root";
  TObjArray carray(2000);
  TObjArray *ClustersArray = &carray;  
  Tracker->ReadClusters(ClustersArray,clusterfile);

  Int_t nClusters = carray.GetEntriesFast();

  Int_t track, index;
  AliTRDmcTrack *tpr = NULL;

  for (Int_t i = 0; i < nClusters; i++) {

    printf("\r assigning cluster %d out of %d", i, nClusters);
    AliTRDcluster *cl = (AliTRDcluster *) ClustersArray->UncheckedAt(i);

    for(Int_t j=0; j<3; j++) {
      track = cl->GetLabel(j);
      if(track < 0) continue;
      if(track >= nparticles) 
	printf("Track index %d is larger than total number of particles %d\n"
               ,track,nparticles);
      else {
	index = Find(track, TRDmcTracks);
	tpr = 0;
	if(index > 0) tpr = (AliTRDmcTrack *) TRDmcTracks->UncheckedAt(index);
	
	if(tpr) {
	  label = tpr->GetTrackIndex();
	  if(label != track) printf("Got label %d while expected %d !\n",
				    label, track);
	  TParticle *p = gAlice->Particle(track);

	  if (p->Pt() > low_pt_cut) tpr->Update(i);	  
	}	  
      }
    }
  }
  
  // Loop through the TRD hits and set XYZ and Pin and Pout

  TTree *hitTree = gAlice->TreeH();
  gAlice->ResetHits();
  Int_t   preamp, amp, det, plane = 0, nBytes = 0;
  Int_t   nTrack = (Int_t) hitTree->GetEntries();
  Float_t  pos[3], rot[3];
  Float_t  prepos[3], prerot[3];

  Bool_t  entrance = kFALSE;
  Bool_t  exit = kFALSE;


  Double_t xin[6], xout[6];
  for(Int_t j = 0; j < 6; j++) {
    xin[j] = Tracker->GetX(0,j,14);
    xout[j] = Tracker->GetX(0,j,0);
  }

  AliTRDhit *hit;

  Int_t current_track = -1;

  printf("\n");
  while(nTrack--) {

    gAlice->ResetHits();
    nBytes += hitTree->GetEvent(nTrack);
    entrance = kTRUE;
    preamp = 1111;
    amp = 1111;
    for(Int_t j = 0; j < 3; j++) { pos[j] = 1111.; rot[j] = 1111.;}

    for(hit = (AliTRDhit *) fTRD->FirstHit(-1); 
	hit; 
	hit = (AliTRDhit *) fTRD->NextHit()) {

      preamp = amp;
      for(Int_t j = 0; j < 3; j++) { prepos[j] = pos[j]; prerot[j] = rot[j];}

      amp = hit->GetCharge();
      if(amp != 0) continue; 
      track   = hit->GetTrack(); 

      TParticle *p = gAlice->Particle(track);
      if (p->Pt() <= low_pt_cut) continue; 

      if(track != current_track) {
	current_track = track;
	//	printf("\n\n");
      }

      det = hit->GetDetector();
      plane    = fGeo->GetPlane(det); 
      pos[0] = hit->X();
      pos[1] = hit->Y();
      pos[2] = hit->Z();
      fGeo->Rotate(det,pos,rot);

      if(TMath::Abs(rot[0] - xin[plane]) < 1.2) entrance = kTRUE;
      else  entrance = kFALSE;
      if(TMath::Abs(rot[0] - xout[plane]) < 1.2) exit = kTRUE;
      else  exit = kFALSE;
      
      if(entrance && (preamp == 0) && (prerot[0] < 200)) {
	index = Find(track, TRDmcTracks);
	if(index > 0) {
	  tpr = 0;
	  tpr = (AliTRDmcTrack *) TRDmcTracks->UncheckedAt(index);
	  if(tpr) {
	    label = tpr->GetTrackIndex();
	    if(label != track) printf("Got label %d while expected %d !\n",
				      label, track);	    
	    else {
	      /*
	      printf("\n momentum at plane %d entrance, Pxyz: %f, %f, %f\n",
		     plane, prepos[0], prepos[1], prepos[2]);
	      printf(" XYZ at plane %d entrance: %f, %f, %f\n",
		     plane, rot[0], rot[1], rot[2]);
	      printf("expected x range: %f .. %f\n",xin[plane]-1.2,xin[plane]+1.2);
	      */
	      tpr->SetPin(plane, prepos[0], prepos[1], prepos[2]);
	      tpr->SetXYZin(plane, rot[0], rot[1], rot[2]);
	    }
	  }
	}
      }


      if(exit && (preamp == 0) && (prerot[0] < 200)) {
	index = Find(track, TRDmcTracks);
	if(index > 0) {
	  tpr = 0;
	  tpr = (AliTRDmcTrack *) TRDmcTracks->UncheckedAt(index);
	  if(tpr) {
	    label = tpr->GetTrackIndex();
	    if(label != track) printf("Got label %d while expected %d !\n",
				      label, track);	    
	    else {
	      /*
	      printf("\n momentum at plane %d exit, Pxyz: %f, %f, %f\n",
		     plane, prepos[0], prepos[1], prepos[2]);
	      printf(" XYZ at plane %d exit: %f, %f, %f\n",
		     plane, rot[0], rot[1], rot[2]);
	      printf("expected x range: %f .. %f\n",xout[plane]-1.2,xout[plane]+1.2);
	      */
	      tpr->SetPout(plane, prepos[0], prepos[1], prepos[2]);
	      tpr->SetXYZout(plane, rot[0], rot[1], rot[2]);
	    }
	  }
	}
      }
    }
  }
  
	
  // Write TRDmcTracks in output file  

  TDirectory *savedir=gDirectory; 
  Char_t *filename   = "AliTRDmcTracks.root";
  TFile *out = new TFile(filename,"RECREATE");

  TTree *tracktree = new TTree("MCtracks","TRD MC tracks");

  AliTRDmcTrack *iotrack=0;
  tracktree->Branch("MCtracks","AliTRDmcTrack",&iotrack,32000,0);
  
  Int_t ntracks = TRDmcTracks->GetEntriesFast();
  
  for (Int_t i=0; i<ntracks; i++) {
    AliTRDmcTrack *pt=(AliTRDmcTrack*)TRDmcTracks->UncheckedAt(i);
    
    Int_t n = pt->GetNumberOfClusters();
    if(n > min_no_of_clusters) { 
      iotrack=pt;
      tracktree->Fill();
      printf("Put track with label %d and %d clusters in the output tree \n",
	     pt->GetTrackIndex(),n);
    }
  }

  tracktree->Write();
  out->Close();     
  savedir->cd(); 

  return;
}
Commit	Line	Data
59dfcadd	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