[u/mrichter/AliRoot.git] / TRD / AliTRDsaveTrackableSeeds.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 "AliTrackReference.h"
  
  #include "TFile.h"
  #include "TParticle.h"
  #include "TStopwatch.h"
#endif


void AliTRDsaveTrackableSeeds() {

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

  TFile *geofile =TFile::Open("AliTRDclusters.root");
  AliTRDtracker *Tracker = new AliTRDtracker(geofile);
  Int_t nEvent = 0;
  Tracker->SetEventNumber(nEvent);

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

  Char_t *alifile = "galice.root";

  // 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 tpc seed
  Int_t label = -1, charge = 0;
  Bool_t primary;
  Float_t mass;
  Int_t pdg_code;

  TDirectory *savedir=gDirectory;

  TFile *in=TFile::Open("AliTPCBackTracks.root");
  if (!in->IsOpen()) {
    cerr<<"can't open file AliTPCBackTracks.root  !\n"; return;
  }

  char   tname[100];
  sprintf(tname,"seedsTPCtoTRD_%d",nEvent);
  TTree *seedTree=(TTree*)in->Get(tname);
  if (!seedTree) {
     cerr<<"AliTRDtracker::PropagateBack(): ";
     cerr<<"can't get a tree with seeds from TPC !\n";
  }

  AliTPCtrack *seed=new AliTPCtrack;
  seedTree->SetBranchAddress("tracks",&seed);

  Int_t nSeeds = (Int_t) seedTree->GetEntries();
  for (Int_t is=0; is<nSeeds; is++) {
     seedTree->GetEvent(is);
     Int_t lbl = seed->GetLabel();
     if(TMath::Abs(lbl) > nparticles) { 
       printf("extra high seed label %d \n", lbl);
       continue;
     }
     TParticle *p = gAlice->Particle(TMath::Abs(lbl));

     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();
     }

     AliTRDmcTrack *t = new AliTRDmcTrack(TMath::Abs(lbl), lbl, primary, 
					  mass, charge, pdg_code);
     TRDmcTracks->AddLast(t);
  }
  delete seed;
  delete seedTree;

  savedir->cd();                         

  Int_t i, mctIndex[nparticles];
  for(i=0; i<nparticles; i++) mctIndex[i]=-1;

  Int_t nMCtracks = TRDmcTracks->GetEntriesFast();
  AliTRDmcTrack *mct = 0;
  for(i = 0; i < nMCtracks; i++) {
    mct = (AliTRDmcTrack*) TRDmcTracks->UncheckedAt(i);
    mctIndex[mct->GetTrackIndex()] = i;
  }

  // Loop through the TRD rec points and set XYZ and Pin and Pout

  Double_t xin[6], xout[6];
  Double_t Px, Py, Pz;

  Float_t  pos[3], rot[3];

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

  TTree *trRefTree = gAlice->TreeTR();
  if (!trRefTree) {
    cout << "<AliTRDreadTrackRef> No TR tree found" << endl;
    return;
  }       

  Int_t nBytes   = 0;
  Bool_t fEntrance, fExit;

  Int_t nTrack = (Int_t) trRefTree->GetEntries();
  cout << "<AliTRDreadTrackRef> Found " << nTrack
       << " primary particles with track refs" << endl;

  for (Int_t iTrack = 0; iTrack < nTrack; iTrack++) {

    printf("TrackRefs for track %d out of %d \n",iTrack,nTrack);
    gAlice->ResetTrackReferences();
    nBytes += trRefTree->GetEvent(iTrack);  

    AliTrackReference *tr = 0;
    tr = (AliTrackReference*) fTRD->FirstTrackReference(-1);
    
    while (tr) {

      Int_t   track = tr->GetTrack();
      if(mctIndex[track] >= 0) { 
	mct = (AliTRDmcTrack*) TRDmcTracks->UncheckedAt(mctIndex[track]);
	
	pos[0] = tr->X();
	pos[1] = tr->Y();
	pos[2] = tr->Z();
	
	Double_t phi = TMath::ATan2(pos[1],pos[0]);
	if(phi < 0) phi = 2 * TMath::Pi() + phi;
	
	Int_t sec = ((Int_t) (phi*180/TMath::Pi())) / 20;
	fGeo->Rotate(fGeo->GetDetector(0,0,17-sec), pos, rot);
	
	fEntrance = kFALSE; fExit = kFALSE;
	for(i = 0; i < 6; i++) {
	  if(TMath::Abs(xin[i] - rot[0]) < 1.4) { fEntrance = kTRUE; break; }
	  if(TMath::Abs(xout[i] - rot[0]) < 1.4) { fExit = kTRUE; break; }
	}
      
	Px = tr->Px();
	Py = tr->Py();
	Pz = tr->Pz();
	
	if(fEntrance) {
	  mct->SetXYZin(i, rot[0], rot[1], rot[2]);
	  mct->SetPin(i, Px, Py, Pz);
	  /*
	    printf("entr plane %d: rp_x = %f vs %f = xin \n",i,rot[0],xin[i]); 
	    printf("             : Px, Py, Pz = %f, %f, %f \n",Px,Py,Pz); 
	    printf("             : y, z = %f, %f\n",rot[1],rot[2]); 
	  */
	}
	if(fExit) {
	  mct->SetXYZout(i, rot[0], rot[1], rot[2]);
	  mct->SetPout(i, Px, Py, Pz);
	  /*
	    printf("exit plane %d: rp_x = %f vs %f = xout \n",i,rot[0],xout[i]); 
	    printf("             : Px, Py, Pz = %f, %f, %f \n",Px,Py,Pz); 
	    printf("             : y, z = %f, %f\n",rot[1],rot[2]);
	  */ 
	}	
      }                                   
      tr = (AliTrackReference *) fTRD->NextTrackReference();
    }
  }  

  // Loop through TRD clusters and assign indexes to MC tracks
  
  Char_t *clusterfile = "AliTRDclusters.root";
  TObjArray carray(2000);
  TObjArray *ClustersArray = &carray;  
  Tracker->ReadClusters(ClustersArray,clusterfile);

  printf("done with ReadClusters()\n");

  Int_t nClusters = carray.GetEntriesFast();

  Int_t track, det, det0, det1, ltb, plane, ind0, ind1, ncl;
  Double_t y, y0, y1, q, q0, q1;

  AliTRDcluster *c0 = NULL, *c1 = NULL;
  printf("nClusters = %d \n", nClusters);

  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);
	continue;
      }     
      if(mctIndex[track] < 0) continue;
      mct = (AliTRDmcTrack*) TRDmcTracks->UncheckedAt(mctIndex[track]);

      label = mct->GetTrackIndex();
      if(label != track) { 
	printf("Got label %d while expected %d !\n", label, track);
	continue;
      }

      ncl = mct->GetNumberOfClusters();
      mct->SetNumberOfClusters(ncl+1);

      det=cl->GetDetector();
      plane = fGeo->GetPlane(det);
      ltb=cl->GetLocalTimeBin();
      if((ltb < 0) || (ltb >= kMAX_TB)) continue;

      ind0 = mct->GetClusterIndex(ltb, plane, 0);
      ind1 = mct->GetClusterIndex(ltb, plane, 1);
      if(ind0 < 0) {
	mct->Update(ltb,plane,0,i);	       
      } else {
	c0 = (AliTRDcluster *) ClustersArray->UncheckedAt(ind0);
	det0 = c0->GetDetector();
	y =  cl->GetY();
	y0 = c0->GetY();
	q  = TMath::Abs(cl->GetQ());
	q0 = TMath::Abs(c0->GetQ());
	if((det == det0) && (TMath::Abs(y-y0) < 1.5)) {
	  if(q > q0) mct->Update(ltb,plane,0,i);
	} else {
	  if (ind1 < 0) {
	    mct->Update(ltb,plane,1,i);
	  } 
	  else {
	    c1 = (AliTRDcluster *) ClustersArray->UncheckedAt(ind1);
	    det1 = c1->GetDetector();
	    y1 = c1->GetY();
	    q1 = TMath::Abs(c1->GetQ());
	    if((det == det0) && (TMath::Abs(y-y0) < 1.5)) {
	      if(q > q1) mct->Update(ltb,plane,1,i);
	    }	  
	  }	  
	}
      }
    }
  } 	    
  	
  // Write TRDmcTracks in output file  

  savedir=gDirectory; 
  Char_t *filename   = "AliTRDtrackableSeeds.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);    
    iotrack=pt;
    tracktree->Fill();
    printf("Put track with label %d and %d clusters in the output tree \n",
	   pt->GetTrackIndex(),pt->GetNumberOfClusters());
  }

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

  return;
}
Commit	Line	Data
10a5e8ec	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	#include "AliTrackReference.h"
	15
	16	#include "TFile.h"
	17	#include "TParticle.h"
	18	#include "TStopwatch.h"
	19	#endif
	20
	21
	22	void AliTRDsaveTrackableSeeds() {
	23
	24	TObjArray mctracks(2000);
	25	TObjArray *TRDmcTracks = &mctracks;
	26
	27	TFile *geofile =TFile::Open("AliTRDclusters.root");
	28	AliTRDtracker *Tracker = new AliTRDtracker(geofile);
	29	Int_t nEvent = 0;
	30	Tracker->SetEventNumber(nEvent);
	31
	32	AliTRDgeometry fGeo = (AliTRDgeometry) geofile->Get("TRDgeometry");
	33	//AliTRDparameter fPar = (AliTRDparameter) geofile->Get("TRDparameter");
	34
	35	Char_t *alifile = "galice.root";
	36
	37	// Connect the AliRoot file containing Geometry, Kine, Hits, and Digits
	38	TFile gafl = (TFile) gROOT->GetListOfFiles()->FindObject(alifile);
	39	if (!gafl) {
	40	cout << "Open the ALIROOT-file " << alifile << endl;
	41	gafl = new TFile(alifile);
	42	}
	43	else {
	44	cout << alifile << " is already open" << endl;
	45	}
	46
	47	// Get AliRun object from file or create it if not on file
	48	gAlice = (AliRun*) gafl->Get("gAlice");
	49	if (gAlice)
	50	cout << "AliRun object found on file" << endl;
	51	else
	52	gAlice = new AliRun("gAlice","Alice test program");
	53
	54	AliTRDv1 fTRD = (AliTRDv1) gAlice->GetDetector("TRD");
	55
	56	// Import the Trees for the event nEvent in the file
	57	const Int_t nparticles = gAlice->GetEvent(nEvent);
	58
	59	printf("found %d particles in event %d \n", nparticles, nEvent);
	60
	61	// Create TRDmcTracks for each tpc seed
	62	Int_t label = -1, charge = 0;
	63	Bool_t primary;
	64	Float_t mass;
65	Int_t pdg_code;
66
67	TDirectory *savedir=gDirectory;
68
69	TFile *in=TFile::Open("AliTPCBackTracks.root");
70	if (!in->IsOpen()) {
71	cerr<<"can't open file AliTPCBackTracks.root !\n"; return;
72	}
73
74	char tname[100];
75	sprintf(tname,"seedsTPCtoTRD_%d",nEvent);
76	TTree seedTree=(TTree)in->Get(tname);
77	if (!seedTree) {
78	cerr<<"AliTRDtracker::PropagateBack(): ";
79	cerr<<"can't get a tree with seeds from TPC !\n";
80	}
81
82	AliTPCtrack *seed=new AliTPCtrack;
83	seedTree->SetBranchAddress("tracks",&seed);
84
85	Int_t nSeeds = (Int_t) seedTree->GetEntries();
86	for (Int_t is=0; is<nSeeds; is++) {
87	seedTree->GetEvent(is);
88	Int_t lbl = seed->GetLabel();
89	if(TMath::Abs(lbl) > nparticles) {
90	printf("extra high seed label %d \n", lbl);
91	continue;
92	}
93	TParticle *p = gAlice->Particle(TMath::Abs(lbl));
94
95	primary = kTRUE;
96	if (p->GetFirstMother()>=0) primary = kFALSE;
97
98	pdg_code = (Int_t) p->GetPdgCode();
99
100	if ((pdg_code == 10010020) \|\|
101	(pdg_code == 10010030) \|\|
102	(pdg_code == 50000050) \|\|
103	(pdg_code == 50000051) \|\|
104	(pdg_code == 10020040)) {
105	mass = 0.;
106	charge = 0;
107	}
108	else {
109	TParticlePDG *pdg = p->GetPDG();
110	charge = (Int_t) pdg->Charge();
111	mass=pdg->Mass();
112	}
113
114	AliTRDmcTrack *t = new AliTRDmcTrack(TMath::Abs(lbl), lbl, primary,
115	mass, charge, pdg_code);
116	TRDmcTracks->AddLast(t);
117	}
118	delete seed;
119	delete seedTree;
120
121	savedir->cd();
122
123	Int_t i, mctIndex[nparticles];
124	for(i=0; i<nparticles; i++) mctIndex[i]=-1;
125
126	Int_t nMCtracks = TRDmcTracks->GetEntriesFast();
127	AliTRDmcTrack *mct = 0;
128	for(i = 0; i < nMCtracks; i++) {
129	mct = (AliTRDmcTrack*) TRDmcTracks->UncheckedAt(i);
130	mctIndex[mct->GetTrackIndex()] = i;
131	}
132
133	// Loop through the TRD rec points and set XYZ and Pin and Pout
134
135	Double_t xin[6], xout[6];
136	Double_t Px, Py, Pz;
137
138	Float_t pos[3], rot[3];
139
140	for(Int_t j = 0; j < 6; j++) {
141	xin[j] = Tracker->GetX(0,j,0)-3;
142	xout[j] = Tracker->GetX(0,j,0);
143	}
144
145	TTree *trRefTree = gAlice->TreeTR();
146	if (!trRefTree) {
147	cout << "<AliTRDreadTrackRef> No TR tree found" << endl;
148	return;
149	}
150
151	Int_t nBytes = 0;
152	Bool_t fEntrance, fExit;
153
154	Int_t nTrack = (Int_t) trRefTree->GetEntries();
155	cout << "<AliTRDreadTrackRef> Found " << nTrack
156	<< " primary particles with track refs" << endl;
157
158	for (Int_t iTrack = 0; iTrack < nTrack; iTrack++) {
159
160	printf("TrackRefs for track %d out of %d \n",iTrack,nTrack);
161	gAlice->ResetTrackReferences();
162	nBytes += trRefTree->GetEvent(iTrack);
163
164	AliTrackReference *tr = 0;
165	tr = (AliTrackReference*) fTRD->FirstTrackReference(-1);
166
167	while (tr) {
168
169	Int_t track = tr->GetTrack();
170	if(mctIndex[track] >= 0) {
171	mct = (AliTRDmcTrack*) TRDmcTracks->UncheckedAt(mctIndex[track]);
172
173	pos[0] = tr->X();
174	pos[1] = tr->Y();
175	pos[2] = tr->Z();
176
177	Double_t phi = TMath::ATan2(pos[1],pos[0]);
178	if(phi < 0) phi = 2 * TMath::Pi() + phi;
179
180	Int_t sec = ((Int_t) (phi*180/TMath::Pi())) / 20;
181	fGeo->Rotate(fGeo->GetDetector(0,0,17-sec), pos, rot);
182
183	fEntrance = kFALSE; fExit = kFALSE;
184	for(i = 0; i < 6; i++) {
185	if(TMath::Abs(xin[i] - rot[0]) < 1.4) { fEntrance = kTRUE; break; }
186	if(TMath::Abs(xout[i] - rot[0]) < 1.4) { fExit = kTRUE; break; }
187	}
188
189	Px = tr->Px();
190	Py = tr->Py();
191	Pz = tr->Pz();
192
193	if(fEntrance) {
194	mct->SetXYZin(i, rot[0], rot[1], rot[2]);
195	mct->SetPin(i, Px, Py, Pz);
196	/*
197	printf("entr plane %d: rp_x = %f vs %f = xin \n",i,rot[0],xin[i]);
198	printf(" : Px, Py, Pz = %f, %f, %f \n",Px,Py,Pz);
199	printf(" : y, z = %f, %f\n",rot[1],rot[2]);
200	*/
201	}
202	if(fExit) {
203	mct->SetXYZout(i, rot[0], rot[1], rot[2]);
204	mct->SetPout(i, Px, Py, Pz);
205	/*
206	printf("exit plane %d: rp_x = %f vs %f = xout \n",i,rot[0],xout[i]);
207	printf(" : Px, Py, Pz = %f, %f, %f \n",Px,Py,Pz);
208	printf(" : y, z = %f, %f\n",rot[1],rot[2]);
209	*/
210	}
211	}
212	tr = (AliTrackReference *) fTRD->NextTrackReference();
213	}
214	}
215
216	// Loop through TRD clusters and assign indexes to MC tracks
217
218	Char_t *clusterfile = "AliTRDclusters.root";
219	TObjArray carray(2000);
220	TObjArray *ClustersArray = &carray;
221	Tracker->ReadClusters(ClustersArray,clusterfile);
222
223	printf("done with ReadClusters()\n");
224
225	Int_t nClusters = carray.GetEntriesFast();
226
227	Int_t track, det, det0, det1, ltb, plane, ind0, ind1, ncl;
228	Double_t y, y0, y1, q, q0, q1;
229
230	AliTRDcluster c0 = NULL, c1 = NULL;
231	printf("nClusters = %d \n", nClusters);
232
233	for (Int_t i = 0; i < nClusters; i++) {
234
235	printf("\r assigning cluster %d out of %d", i, nClusters);
236	AliTRDcluster cl = (AliTRDcluster ) ClustersArray->UncheckedAt(i);
237
238	for(Int_t j=0; j<3; j++) {
239	track = cl->GetLabel(j);
240	if(track < 0) continue;
241	if(track >= nparticles) {
242	printf("Track index %d is larger than total number of particles %d\n"
243	,track,nparticles);
244	continue;
245	}
246	if(mctIndex[track] < 0) continue;
247	mct = (AliTRDmcTrack*) TRDmcTracks->UncheckedAt(mctIndex[track]);
248
249	label = mct->GetTrackIndex();
250	if(label != track) {
251	printf("Got label %d while expected %d !\n", label, track);
252	continue;
253	}
254
255	ncl = mct->GetNumberOfClusters();
256	mct->SetNumberOfClusters(ncl+1);
257
258	det=cl->GetDetector();
259	plane = fGeo->GetPlane(det);
260	ltb=cl->GetLocalTimeBin();
261	if((ltb < 0) \|\| (ltb >= kMAX_TB)) continue;
262
263	ind0 = mct->GetClusterIndex(ltb, plane, 0);
264	ind1 = mct->GetClusterIndex(ltb, plane, 1);
265	if(ind0 < 0) {
266	mct->Update(ltb,plane,0,i);
267	} else {
268	c0 = (AliTRDcluster *) ClustersArray->UncheckedAt(ind0);
269	det0 = c0->GetDetector();
270	y = cl->GetY();
271	y0 = c0->GetY();
272	q = TMath::Abs(cl->GetQ());
273	q0 = TMath::Abs(c0->GetQ());
274	if((det == det0) && (TMath::Abs(y-y0) < 1.5)) {
275	if(q > q0) mct->Update(ltb,plane,0,i);
276	} else {
277	if (ind1 < 0) {
278	mct->Update(ltb,plane,1,i);
279	}
280	else {
281	c1 = (AliTRDcluster *) ClustersArray->UncheckedAt(ind1);
282	det1 = c1->GetDetector();
283	y1 = c1->GetY();
284	q1 = TMath::Abs(c1->GetQ());
285	if((det == det0) && (TMath::Abs(y-y0) < 1.5)) {
286	if(q > q1) mct->Update(ltb,plane,1,i);
287	}
288	}
289	}
290	}
291	}
292	}
293
294	// Write TRDmcTracks in output file
295
296	savedir=gDirectory;
297	Char_t *filename = "AliTRDtrackableSeeds.root";
298	TFile *out = new TFile(filename,"RECREATE");
299
300	TTree *tracktree = new TTree("MCtracks","TRD MC tracks");
301
302	AliTRDmcTrack *iotrack=0;
303	tracktree->Branch("MCtracks","AliTRDmcTrack",&iotrack,32000,0);
304
305	Int_t ntracks = TRDmcTracks->GetEntriesFast();
306
307	for (Int_t i=0; i<ntracks; i++) {
308	AliTRDmcTrack pt=(AliTRDmcTrack)TRDmcTracks->UncheckedAt(i);
309	iotrack=pt;
310	tracktree->Fill();
311	printf("Put track with label %d and %d clusters in the output tree \n",
312	pt->GetTrackIndex(),pt->GetNumberOfClusters());
313	}
314
315	tracktree->Write();
316	out->Close();
317	savedir->cd();
318
319	return;
320	}
321