Add detailed geometry and simple simulator
[u/mrichter/AliRoot.git] / TRD / AliTRDpid.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
3  *                                                                        *
4  * Author: The ALICE Off-line Project.                                    *
5  * Contributors are mentioned in the code where appropriate.              *
6  *                                                                        *
7  * Permission to use, copy, modify and distribute this software and its   *
8  * documentation strictly for non-commercial purposes is hereby granted   *
9  * without fee, provided that the above copyright notice appears in all   *
10  * copies and that both the copyright notice and this permission notice   *
11  * appear in the supporting documentation. The authors make no claims     *
12  * about the suitability of this software for any purpose. It is          *
13  * provided "as is" without express or implied warranty.                  *
14  **************************************************************************/
15
16 /*
17 $Log$
18 Revision 1.1  2001/05/07 08:08:05  cblume
19 Update of TRD code
20
21 */
22
23 ///////////////////////////////////////////////////////////////////////////////
24 //                                                                           //
25 //   The TRD particle identification base class                              //
26 //                                                                           //
27 //   Its main purposes are:                                                  //
28 //      - Provide I/O framework for all neccessary files                     //
29 //      - Assignment of a e/pi propability to a given track                  //
30 //                                                                           //
31 ///////////////////////////////////////////////////////////////////////////////
32
33 #include <stdlib.h>
34 #include <math.h>
35
36 #include <TROOT.h>
37 #include <TH1.h>
38 #include <TObjArray.h>
39 #include <TTree.h>
40 #include <TFile.h>
41 #include <TParticle.h>
42
43 #include "AliRun.h"
44 #include "AliTRD.h"
45 #include "AliTRDpid.h"
46 #include "AliTRDcluster.h"
47 #include "AliTRDtrack.h"
48 #include "AliTRDtracker.h"
49 #include "AliTRDgeometry.h"
50
51 ClassImp(AliTRDpid)
52
53 //_____________________________________________________________________________
54 AliTRDpid::AliTRDpid():TNamed()
55 {
56   //
57   // AliTRDpid default constructor
58   // 
59
60   fTrackArray    = NULL;
61   fClusterArray  = NULL;
62   fGeometry      = NULL;
63   fFileKine      = NULL;
64
65   fPIDratioMin   = 0.0;
66   fPIDpurePoints = kFALSE;
67   fPIDindexMin   = 0;
68   fPIDindexMax   = 0;
69
70   fThreePadOnly  = kFALSE;
71
72 }
73
74 //_____________________________________________________________________________
75 AliTRDpid::AliTRDpid(const char* name, const char* title):TNamed(name,title)
76 {
77   //
78   // AliTRDpid constructor
79   //
80
81   fTrackArray   = NULL;
82   fClusterArray = NULL;
83   fGeometry     = NULL;
84   fFileKine     = NULL;
85
86   Init();
87
88 }
89
90 //_____________________________________________________________________________
91 AliTRDpid::AliTRDpid(const AliTRDpid &p)
92 {
93   //
94   // AliTRDpid copy constructor
95   //
96
97   ((AliTRDpid &) p).Copy(*this);
98
99 }
100
101 //_____________________________________________________________________________
102 AliTRDpid::~AliTRDpid()
103 {
104   //
105   // AliTRDpid destructor
106   //
107
108   if (fClusterArray) {
109     fClusterArray->Delete();
110     delete fClusterArray;
111   }
112
113   if (fTrackArray) {
114     fTrackArray->Delete();
115     delete fTrackArray;
116   }
117
118   fFileKine->Close();
119
120 }
121
122 //_____________________________________________________________________________
123 AliTRDpid &AliTRDpid::operator=(const AliTRDpid &p)
124 {
125   //
126   // Assignment operator
127   //
128
129   if (this != &p) ((AliTRDpid &) p).Copy(*this);
130   return *this;
131
132 }
133
134 //_____________________________________________________________________________
135 void AliTRDpid::Copy(TObject &p)
136 {
137   //
138   // Copy function
139   //
140
141   ((AliTRDpid &) p).fTrackArray    = NULL;    
142   ((AliTRDpid &) p).fClusterArray  = NULL;    
143   ((AliTRDpid &) p).fGeometry      = NULL;    
144   ((AliTRDpid &) p).fFileKine      = NULL;
145   ((AliTRDpid &) p).fPIDratioMin   = fPIDratioMin;
146   ((AliTRDpid &) p).fPIDpurePoints = fPIDpurePoints;
147   ((AliTRDpid &) p).fPIDindexMin   = fPIDindexMin;
148   ((AliTRDpid &) p).fPIDindexMax   = fPIDindexMax;
149   ((AliTRDpid &) p).fThreePadOnly  = fThreePadOnly;
150
151 }
152
153 //_____________________________________________________________________________
154 Bool_t AliTRDpid::Init()
155 {
156   //
157   // Initializes the PID object 
158   //
159
160   fClusterArray  = new TObjArray();
161   fTrackArray    = new TObjArray();
162
163   fPIDratioMin   = 0.7;
164   fPIDpurePoints = kTRUE;
165   fPIDindexMin   = -1;
166   fPIDindexMax   = -1;
167
168   fThreePadOnly  = kFALSE;
169
170   return kTRUE;
171
172 }
173
174 //_____________________________________________________________________________
175 Bool_t AliTRDpid::AssignLikelihood()
176 {
177   //
178   // Assigns the e / pi likelihood to all tracks.
179   //
180
181   return AssignLikelihood(fTrackArray);
182
183 }
184
185 //_____________________________________________________________________________
186 Bool_t AliTRDpid::AssignLikelihood(TObjArray *tarray)
187 {
188   //
189   // Assigns the e / pi likelihood to all tracks in the array <tarray>.
190   //
191
192   Bool_t status = kTRUE;
193
194   AliTRDtrack *track;
195
196   TIter nextTrack(tarray);  
197   while ((track = (AliTRDtrack *) nextTrack())) {
198     if (!AssignLikelihood(track)) {
199       status = kFALSE;
200       continue;
201     }
202   }
203
204   return status;
205
206 }
207
208 //_____________________________________________________________________________
209 Bool_t AliTRDpid::FillSpectra()
210 {
211   //
212   // Fills the energy loss distributions with all tracks.
213   //
214
215   return FillSpectra(fTrackArray);
216
217 }
218   
219 //_____________________________________________________________________________
220 Bool_t AliTRDpid::FillSpectra(TObjArray *tarray)
221 {
222   //
223   // Fills the energy loss distributions with all tracks in <tarray>.
224   //
225   
226   Bool_t status = kTRUE;
227
228   AliTRDtrack *track;
229
230   TIter nextTrack(tarray);
231   while ((track = (AliTRDtrack *) nextTrack())) {
232     if (!FillSpectra(track)) {
233       status = kFALSE;
234       continue;
235     }
236   }  
237
238   return kTRUE;
239
240 }
241
242 //_____________________________________________________________________________
243 Bool_t AliTRDpid::Open(const Char_t *name, Int_t event)
244 {
245   //
246   // Opens and reads the kine tree, the geometry, the cluster array
247   // and the track array from the file <name>.
248   //
249
250   Bool_t status = kTRUE;
251
252   status = ReadCluster(name);
253   status = ReadTracks(name);
254   status = ReadKine(name,event);
255
256   return status;
257
258 }
259
260 //_____________________________________________________________________________
261 Bool_t AliTRDpid::Open(const Char_t *namekine
262                      , const Char_t *namecluster
263                      , const Char_t *nametracks
264                      , Int_t event)
265 {
266   //
267   // Opens and reads the kine tree and the geometry from file <namekine>,
268   // the cluster array from file <namecluster>,
269   // and the track array from the file <nametracks>.
270   //
271
272   Bool_t status = kTRUE;
273
274   status = ReadCluster(namecluster);
275   status = ReadTracks(nametracks);
276   status = ReadKine(namekine,event);
277
278   return status;
279
280 }
281
282 //_____________________________________________________________________________
283 Bool_t AliTRDpid::ReadKine(const Char_t *name, Int_t event)
284 {
285   //
286   // Opens and reads the kine tree and the geometry from the file <name>.
287   //
288
289   TFile *fFileKine = (TFile *) gROOT->GetListOfFiles()->FindObject(name);
290   if (!fFileKine) {
291     printf("AliTRDpid::ReadKine -- ");
292     printf("Open file %s\n",name);
293     fFileKine = new TFile(name);
294     if (!fFileKine) {
295       printf("AliTRDpid::ReadKine -- ");
296       printf("Cannot open file %s\n",name);
297       return kFALSE;
298     }
299   }
300
301   gAlice = (AliRun *) fFileKine->Get("gAlice");
302   if (!gAlice) {
303     printf("AliTRDpid::ReadKine -- ");
304     printf("No AliRun object found\n");    
305     return kFALSE;
306   }
307   gAlice->GetEvent(event);
308
309   AliTRD *trd = (AliTRD *) gAlice->GetDetector("TRD");
310   if (!trd) {
311     printf("AliTRDpid::ReadKine -- ");
312     printf("No TRD object found\n");    
313     return kFALSE;
314   }
315
316   fGeometry = trd->GetGeometry();
317   if (!fGeometry) {
318     printf("AliTRDpid::ReadKine -- ");
319     printf("No TRD geometry found\n");
320     return kFALSE;
321   }
322
323   return kTRUE;
324
325 }
326
327 //_____________________________________________________________________________
328 Bool_t AliTRDpid::ReadCluster(const Char_t *name)
329 {
330   //
331   // Opens and reads the cluster array from the file <name>.
332   //
333  
334   TDirectory *savedir = gDirectory;                                                   
335
336   if (fClusterArray) {
337     fClusterArray->Delete();
338   }
339   else {
340     fClusterArray = new TObjArray();
341   }
342
343   printf("AliTRDpid::ReadCluster -- ");
344   printf("Open file %s\n",name);
345
346   AliTRDtracker *tracker = new AliTRDtracker("dummy","dummy");
347   tracker->ReadClusters(fClusterArray,name);
348
349   if (!fClusterArray) {
350     printf("AliTRDpid::ReadCluster -- ");
351     printf("Error reading the cluster array from file %s\n",name);
352     return kFALSE;
353   }
354
355   delete tracker;
356
357   savedir->cd();
358
359   return kTRUE;
360
361 }
362
363 //_____________________________________________________________________________
364 Bool_t AliTRDpid::ReadTracks(const Char_t *name)
365 {
366   //
367   // Opens and reads the track array from the file <name>.
368   //
369  
370   TDirectory *savedir = gDirectory;                                                   
371
372   if (fTrackArray) {
373     fTrackArray->Delete();
374   }
375   else {
376     fTrackArray = new TObjArray();
377   }
378
379   TFile *file = (TFile *) gROOT->GetListOfFiles()->FindObject(name);
380   if (!file) {
381     printf("AliTRDpid::ReadTracks -- ");
382     printf("Open file %s\n",name);
383     file = new TFile(name);
384     if (!file) {
385       printf("AliTRDpid::ReadTracks -- ");
386       printf("Cannot open file %s\n",name);
387       return kFALSE;
388     }
389   }
390
391   TTree   *trackTree   = (TTree *) file->Get("TreeT");
392   TBranch *trackBranch = trackTree->GetBranch("tracks");
393
394   Int_t nEntry = ((Int_t) trackTree->GetEntries());
395   for (Int_t iEntry = 0; iEntry < nEntry; iEntry++) {
396     AliTRDtrack *track = new AliTRDtrack();
397     trackBranch->SetAddress(&track);
398     trackTree->GetEvent(iEntry);
399     fTrackArray->AddLast(track);
400   }
401
402   file->Close();
403
404   savedir->cd();
405
406   return kTRUE;
407
408 }
409
410 //_____________________________________________________________________________
411 Int_t AliTRDpid::MCpid(const AliTRDtrack *t)
412 {
413   //
414   // Determines the pid of the MC track <t>.
415   //
416
417   // PDG codes
418   const Int_t kPdgEl =  11;
419   const Int_t kPdgPi = 211;
420
421   AliTRDcluster *cluster;
422   TParticle     *particle;
423
424   Int_t   nClusterEl   = 0;
425   Int_t   nClusterPi   = 0;
426   Int_t   nClusterPure = 0;
427   Int_t   nClusterMix  = 0;
428
429   Float_t ratioEl      = 0;
430   Float_t ratioPi      = 0;
431
432   Int_t   ipid         = -1;
433
434   if (!fClusterArray) {
435     printf("AliTRDpid::MCpid -- ");
436     printf("ClusterArray not defined. Initialize the PID object first\n");
437     return -1;  
438   }
439   
440   // Loop through all clusters associated to this track
441   Int_t nCluster = t->GetNclusters();
442   for (Int_t iCluster = 0; iCluster < nCluster; iCluster++) {
443
444     // Get a cluster
445     Int_t index = t->GetClusterIndex(iCluster);
446     if (!(cluster = (AliTRDcluster *) fClusterArray->UncheckedAt(index))) {
447       break;
448     } 
449
450     // Get the first two MC track indices
451     Int_t track0 = cluster->GetTrackIndex(0);
452     Int_t track1 = cluster->GetTrackIndex(1);
453
454     // Check on the track index to find the right primaries
455     if ((track0 >  fPIDindexMin) && 
456         (track0 <= fPIDindexMax)) {
457
458       // Check whether it is a pure cluster, i.e. not overlapping
459       Bool_t accept = kTRUE;
460       if ((fPIDpurePoints) && (track1 != -1)) {
461         accept = kFALSE;
462       }
463       if (accept) {
464
465         particle = gAlice->Particle(track0);
466         if (particle->GetFirstMother() == -1) {
467           switch (TMath::Abs(particle->GetPdgCode())) {
468           case kPdgEl:
469             nClusterEl++;
470             break;
471           case kPdgPi:
472             nClusterPi++;
473             break;
474           };
475           nClusterPure++;
476         }
477
478       }
479       else {
480  
481         nClusterMix++;
482
483       }
484
485     }
486
487   }
488   
489   if (nCluster) {
490     ratioEl = ((Float_t) nClusterEl) / ((Float_t) nCluster);
491     ratioPi = ((Float_t) nClusterPi) / ((Float_t) nCluster);
492     if      (ratioEl > fPIDratioMin) {
493       ipid = kElectron;
494     }
495     else if (ratioPi > fPIDratioMin) {
496       ipid = kPion;
497     }
498   }
499 //   printf("AliTRDpid::MCpid -- ");
500 //   printf("nCluster = %d, nClusterEl = %d, nClusterPi = %d\n"
501 //       ,nCluster,nClusterEl,nClusterPi);
502 //   printf("AliTRDpid::MCpid -- ");
503 //   printf("nClusterPure = %d, nClusterMix = %d\n"
504 //       ,nClusterPure,nClusterMix);
505 //   printf("AliTRDpid::MCpid -- ");
506 //   printf("ratioEl = %f, ratioPi = %f, ipid = %d\n",ratioEl,ratioPi,ipid);
507
508   return ipid;
509
510 }
511
512 //_____________________________________________________________________________
513 Int_t AliTRDpid::MCpid(const AliTRDtrack *t, Int_t *pdg
514                                            , Int_t *nFound
515                                            , Int_t *indices)
516 {
517   //
518   // Determines the pid of the MC track <t>.
519   // Returns the number of different MC particles that contribute to this
520   // track. <pdg> contains their PDG code, ordered by their frequency.
521   // <nFound> contains the corresponding number of cluster that contain
522   // contributions to this cluster.
523   //
524
525   const Int_t kNtrack = 3;
526   const Int_t kNpart  = 200;
527
528   AliTRDcluster *cluster;
529   TParticle     *particle;
530
531   Int_t nPart    = 0;
532   Int_t iPart    = 0;
533   Int_t iTrack   = 0;
534   Int_t iCluster = 0;
535
536   if (!fClusterArray) {
537     printf("AliTRDpid::MCpid -- ");
538     printf("ClusterArray not defined. Initialize the PID object first\n");
539     return -1;  
540   }
541   
542   Int_t sort[kNpart];
543   Int_t pdgTmp[kNpart];
544   Int_t nFoundTmp[kNpart];
545   Int_t indicesTmp[kNpart];
546   for (iPart = 0; iPart < kNpart; iPart++) {
547     pdg[iPart]        = 0;
548     nFound[iPart]     = 0;
549     indices[iPart]    = 0;
550     pdgTmp[iPart]     = 0;
551     nFoundTmp[iPart]  = 0;
552     indicesTmp[iPart] = 0;
553     sort[iPart]       = 0;
554   }
555
556   // Loop through all clusters associated to this track
557   Int_t nCluster = t->GetNclusters();
558   for (iCluster = 0; iCluster < nCluster; iCluster++) {
559
560     // Get a cluster
561     Int_t index = t->GetClusterIndex(iCluster);
562     if (!(cluster = (AliTRDcluster *) fClusterArray->UncheckedAt(index))) {
563       break;
564     } 
565
566     // Get the MC track indices
567     for (iTrack = 0; iTrack < kNtrack; iTrack++) {
568
569       Int_t trackIndex = cluster->GetTrackIndex(iTrack);
570       if (trackIndex >= 0) {
571         particle = gAlice->Particle(trackIndex);
572         Int_t  pdgCode = particle->GetPdgCode(); 
573         Bool_t newPart = kTRUE;
574         for (iPart = 0; iPart < nPart; iPart++) {
575          if (trackIndex == indicesTmp[iPart]) {
576             nFoundTmp[iPart]++;
577             newPart = kFALSE;
578             break;
579           }
580         }
581         if ((newPart) && (nPart < kNpart)) {
582           indicesTmp[nPart] = trackIndex;
583           pdgTmp[nPart]     = pdgCode;
584           nFoundTmp[nPart]++;
585           nPart++;
586         }
587       }
588
589     }
590
591   }
592
593   // Sort the arrays by the frequency of the appearance of a MC track
594   TMath::Sort(nPart,nFoundTmp,sort,kTRUE);
595   for (iPart = 0; iPart < nPart; iPart++) {
596     pdg[iPart]     = pdgTmp[sort[iPart]];
597     nFound[iPart]  = nFoundTmp[sort[iPart]];
598     indices[iPart] = indicesTmp[sort[iPart]];
599 //     printf(" iPart = %d, pdg = %d, nFound = %d, indices = %d\n"
600 //        ,iPart,pdg[iPart],nFound[iPart],indices[iPart]);
601   }
602
603   return nPart;
604
605 }
606
607 //_____________________________________________________________________________
608 Bool_t AliTRDpid::SumCharge(const AliTRDtrack *t
609                            , Float_t *charge
610                            , Int_t   *nCluster)
611 {
612   //
613   // Sums up the charge in each plane for track <t>.
614   //
615
616   Bool_t status = kTRUE;
617
618   AliTRDcluster *cluster;
619
620   const Int_t kNplane = AliTRDgeometry::Nplan();
621   for (Int_t iPlane = 0; iPlane < kNplane; iPlane++) {
622     charge[iPlane]   = 0.0;
623     nCluster[iPlane] = 0;
624   }
625
626   if (!fClusterArray) {
627     printf("AliTRDpid::SumCharge -- ");
628     printf("ClusterArray not defined. Initialize the PID object first\n");
629     return kFALSE;  
630   }
631   if (!fGeometry) {
632     printf("AliTRDpid::SumCharge -- ");
633     printf("TRD geometry not defined. Initialize the PID object first\n");
634     return kFALSE;  
635   }
636   
637   // Loop through all clusters associated to this track
638   Int_t nClus = t->GetNclusters();
639   for (Int_t iClus = 0; iClus < nClus; iClus++) {
640
641     // Get a cluster
642     Int_t index = t->GetClusterIndex(iClus);
643     if (!(cluster = (AliTRDcluster *) fClusterArray->UncheckedAt(index))) {
644       status = kFALSE;
645       break;
646     } 
647
648     if (fThreePadOnly) {
649       if ((!cluster->From2pad()) &&
650           (!cluster->From3pad())) continue;
651     }
652
653     // Sum up the charge
654     Int_t plane    = fGeometry->GetPlane(cluster->GetDetector());
655     //charge[plane] += cluster->GetQ();
656     // Corrected charge
657     charge[plane] += t->GetClusterdQdl(iClus);
658     nCluster[plane]++;
659
660   }
661 //   printf("AliTRDpid::SumCharge -- ");
662 //   printf("charge = %d, %d, %d, %d, %d, %d\n"
663 //       ,(Int_t) charge[0]
664 //          ,(Int_t) charge[1]
665 //          ,(Int_t) charge[2]
666 //          ,(Int_t) charge[3]
667 //          ,(Int_t) charge[4]
668 //          ,(Int_t) charge[5]);
669 //   printf("AliTRDpid::SumCharge -- ");
670 //   printf("nCluster = %d, %d, %d, %d, %d, %d\n"
671 //          ,nCluster[0]
672 //          ,nCluster[1]
673 //          ,nCluster[2]
674 //          ,nCluster[3]
675 //          ,nCluster[4]
676 //          ,nCluster[5]);
677
678   return status;
679
680 }