1629e8a9aa3cab4faae3253927eea497cdf4c718
[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 */
19
20 ///////////////////////////////////////////////////////////////////////////////
21 //                                                                           //
22 //   The TRD particle identification class                                   //
23 //                                                                           //
24 //   Its main purposes are:                                                  //
25 //      - Creation and bookkeeping of the propability distributions          //
26 //      - Assignment of a e/pi propability to a given track                  //
27 //                                                                           //
28 ///////////////////////////////////////////////////////////////////////////////
29
30 #include <stdlib.h>
31 #include <math.h>
32
33 #include <TROOT.h>
34 #include <TH1.h>
35 #include <TObjArray.h>
36 #include <TTree.h>
37 #include <TFile.h>
38 #include <TParticle.h>
39
40 #include "AliRun.h"
41 #include "AliTRD.h"
42 #include "AliTRDpid.h"
43 #include "AliTRDcluster.h"
44 #include "AliTRDtrack.h"
45 #include "AliTRDtracker.h"
46 #include "AliTRDgeometry.h"
47
48 ClassImp(AliTRDpid)
49
50 //_____________________________________________________________________________
51 AliTRDpid::AliTRDpid():TNamed()
52 {
53   //
54   // AliTRDpid default constructor
55   // 
56
57   fNMom   = 0;
58   fMinMom = 0;
59   fMaxMom = 0;
60   fWidMom = 0;
61
62   fQHist        = NULL;
63   fLQHist       = NULL;
64   fTrackArray   = NULL;
65   fClusterArray = NULL;
66   fGeometry     = NULL;
67
68 }
69
70 //_____________________________________________________________________________
71 AliTRDpid::AliTRDpid(const char* name, const char* title):TNamed(name,title)
72 {
73   //
74   // AliTRDpid constructor
75   //
76
77   fNMom   = 0;
78   fMinMom = 0;
79   fMaxMom = 0;
80   fWidMom = 0;
81
82   fQHist        = NULL;
83   fLQHist       = NULL;
84   fTrackArray   = NULL;
85   fClusterArray = NULL;
86   fGeometry     = NULL;
87
88   Init();
89
90 }
91
92 //_____________________________________________________________________________
93 AliTRDpid::AliTRDpid(const AliTRDpid &p)
94 {
95   //
96   // AliTRDpid copy constructor
97   //
98
99   ((AliTRDpid &) p).Copy(*this);
100
101 }
102
103 //_____________________________________________________________________________
104 AliTRDpid::~AliTRDpid()
105 {
106   //
107   // AliTRDpid destructor
108   //
109
110   if (fClusterArray) {
111     fClusterArray->Delete();
112     delete fClusterArray;
113   }
114
115   if (fTrackArray) {
116     fTrackArray->Delete();
117     delete fTrackArray;
118   }
119
120   if (fQHist) {
121     fQHist->Delete();
122     delete fQHist;
123   }
124
125   if (fLQHist) {
126     fLQHist->Delete();
127     delete fLQHist;
128   }
129
130 }
131
132 //_____________________________________________________________________________
133 AliTRDpid &AliTRDpid::operator=(const AliTRDpid &p)
134 {
135   //
136   // Assignment operator
137   //
138
139   if (this != &p) ((AliTRDpid &) p).Copy(*this);
140   return *this;
141
142 }
143
144 //_____________________________________________________________________________
145 void AliTRDpid::Copy(TObject &p)
146 {
147   //
148   // Copy function
149   //
150
151   fQHist->Copy(*((AliTRDpid &) p).fQHist);
152   fLQHist->Copy(*((AliTRDpid &) p).fLQHist);
153
154   ((AliTRDpid &) p).fTrackArray   = NULL;    
155   ((AliTRDpid &) p).fClusterArray = NULL;    
156   ((AliTRDpid &) p).fGeometry     = NULL;    
157
158   ((AliTRDpid &) p).fNMom   = fNMom;
159   ((AliTRDpid &) p).fMinMom = fMinMom;
160   ((AliTRDpid &) p).fMaxMom = fMaxMom;
161   ((AliTRDpid &) p).fWidMom = fWidMom;
162
163 }
164
165 //_____________________________________________________________________________
166 Bool_t AliTRDpid::Init()
167 {
168   //
169   // Initializes the PID object 
170   //
171
172   fClusterArray = new TObjArray();
173   fTrackArray   = new TObjArray();
174
175   return kTRUE;
176
177 }
178
179 //_____________________________________________________________________________
180 Bool_t AliTRDpid::AssignLQ(TObjArray *tarray)
181 {
182   //
183   // Assigns the e / pi Q-likelihood to all tracks in the array
184   //
185
186   Bool_t status = kTRUE;
187
188   AliTRDtrack *track;
189
190   TIter nextTrack(tarray);  
191   while ((track = (AliTRDtrack *) nextTrack())) {
192     if (!AssignLQ(track)) {
193       status = kFALSE;
194       break;
195     }
196   }
197
198   return status;
199
200 }
201
202 //_____________________________________________________________________________
203 Bool_t AliTRDpid::AssignLQ(AliTRDtrack *t)
204 {
205   //
206   // Assigns the e / pi Q-likelihood to a given track
207   //
208
209   const Int_t kNplane = AliTRDgeometry::Nplan();
210   Float_t charge[kNplane];
211
212   // Calculate the total charge in each plane
213   if (!SumCharge(t,charge)) return kFALSE;
214
215   // Assign the likelihoods 
216   t->SetLikelihoodPion(LQPion(charge));
217   t->SetLikelihoodElectron(LQElectron(charge));
218
219   return kTRUE;  
220
221 }
222
223 //_____________________________________________________________________________
224 Bool_t AliTRDpid::CreateHistograms(const Int_t nmom
225                                  , const Float_t minmom
226                                  , const Float_t maxmom)
227 {
228   //
229   // Creates the likelihood histograms
230   //
231
232   Int_t imom;
233   Int_t ipid;
234   Int_t ipla;
235
236   const Int_t kNpla = AliTRDgeometry::Nplan();
237
238   fNMom   = nmom;
239   fMinMom = minmom;
240   fMaxMom = maxmom;
241   fWidMom = (maxmom - minmom) / ((Float_t) nmom);
242
243   // The L-Q distributions
244   fLQHist = new TObjArray(kNpid * nmom);
245   for (imom = 0; imom < nmom;  imom++) {
246     for (ipid = 0; ipid < kNpid; ipid++) {
247       Int_t index = GetIndexLQ(imom,ipid);
248       Char_t name[10];
249       Char_t title[80];
250       sprintf(name ,"hLQ%03d",index);
251       if (ipid == kElectron) {
252         sprintf(title,"L-Q electron p-bin %03d",imom);
253       }
254       else {
255         sprintf(title,"L-Q pion p-bin %03d",imom);
256       }
257       TH1F *hTmp  = new TH1F(name,title,416,-0.02,1.02);
258       fLQHist->AddAt(hTmp,index);
259     }
260   }
261
262   // The Q-distributions
263   fQHist = new TObjArray(kNpla * kNpid * nmom);
264   for (imom = 0; imom < nmom;  imom++) {
265     for (ipid = 0; ipid < kNpid; ipid++) {
266       for (ipla = 0; ipla < kNpla; ipla++) {
267         Int_t index = GetIndexQ(imom,ipla,ipid);
268         Char_t name[10];
269         Char_t title[80];
270         sprintf(name ,"hQ%03d",index);
271         if (ipid == kElectron) {
272           sprintf(title,"Q electron p-bin %03d plane %d",imom,ipla);
273         }
274         else {
275           sprintf(title,"Q pion p-bin %03d plane %d",imom,ipla);
276         }
277         TH1F *hTmp  = new TH1F(name,title,100,0.0,5000.0);
278         fQHist->AddAt(hTmp,index);
279       }
280     }
281   }
282
283   return kTRUE;
284
285 }
286
287 //_____________________________________________________________________________
288 Bool_t AliTRDpid::FillQspectra()
289 {
290   //
291   // Fills the Q-spectra histograms.
292   //
293   
294   Bool_t status = kTRUE;
295
296   AliTRDtrack *track;
297
298   TIter nextTrack(fTrackArray);
299   while ((track = (AliTRDtrack *) nextTrack())) {
300     if (!FillQspectra(track)) {
301       status = kFALSE;
302       break;
303     }
304   }
305
306   return status;
307
308 }
309
310 //_____________________________________________________________________________
311 Bool_t AliTRDpid::FillQspectra(const AliTRDtrack *t)
312 {
313   //
314   // Fills the Q-spectra histograms with track <t>.
315   //
316
317   const Int_t kNpla = AliTRDgeometry::Nplan();
318
319   if (isnan(t->GetP())) return kTRUE;
320
321   printf("----------------------------------------------------------\n");
322   Float_t charge[kNpla];
323   Float_t mom  = t->GetP();
324   Int_t   ipid = Pid(t);
325
326   if (!SumCharge(t,charge)) return kFALSE;
327
328   printf(" ipid   = %d\n",ipid);
329   printf(" mom    = %f\n",mom);
330   printf(" charge = %8.2f, %8.2f, %8.2f, %8.2f, %8.2f, %8.2f\n"
331         ,charge[0],charge[1],charge[2],charge[3],charge[4],charge[5]);
332
333   for (Int_t ipla = 0; ipla < kNpla; ipla++) {
334     Int_t index = GetIndexQ(mom,ipla,ipid);    
335     if (index > -1) {
336       TH1F *hTmp = (TH1F *) fQHist->UncheckedAt(index);
337       hTmp->Fill(charge[ipla]);
338     }
339   }  
340
341   return kTRUE;
342
343 }
344
345 //_____________________________________________________________________________
346 Bool_t AliTRDpid::Open(const Char_t *name, Int_t event)
347 {
348   //
349   // Opens and reads the kine tree, the geometry, the cluster array.
350   // and the track array from the file <name>
351   //
352
353   Bool_t status = kTRUE;
354
355   status = ReadKine(name,event);
356   status = ReadCluster(name);
357   status = ReadTracks(name);
358
359   return status;
360
361 }
362
363 //_____________________________________________________________________________
364 Bool_t AliTRDpid::Open(const Char_t *namekine
365                      , const Char_t *namecluster
366                      , const Char_t *nametracks
367                      , Int_t event)
368 {
369   //
370   // Opens and reads the kine tree and the geometry from file <namekine>,
371   // the cluster array from file <namecluster>,
372   // and the track array from the file <nametracks>
373   //
374
375   Bool_t status = kTRUE;
376
377   status = ReadKine(namekine,event);
378   status = ReadCluster(namecluster);
379   status = ReadTracks(nametracks);
380
381   return status;
382
383 }
384
385 //_____________________________________________________________________________
386 Bool_t AliTRDpid::ReadKine(const Char_t *name, Int_t event)
387 {
388   //
389   // Opens and reads the kine tree and the geometry from the file <name>.
390   //
391
392   TFile *file = (TFile *) gROOT->GetListOfFiles()->FindObject(name);
393   if (!file) {
394     printf("AliTRDpid::ReadKine -- ");
395     printf("Open file %s\n",name);
396     file = new TFile(name);
397     if (!file) {
398       printf("AliTRDpid::ReadKine -- ");
399       printf("Cannot open file %s\n",name);
400       return kFALSE;
401     }
402   }
403
404   gAlice = (AliRun *) file->Get("gAlice");
405   if (!gAlice) {
406     printf("AliTRDpid::ReadKine -- ");
407     printf("No AliRun object found\n");    
408     return kFALSE;
409   }
410   gAlice->GetEvent(event);
411
412   AliTRD *trd = (AliTRD *) gAlice->GetDetector("TRD");
413   if (!trd) {
414     printf("AliTRDpid::ReadKine -- ");
415     printf("No TRD object found\n");    
416     return kFALSE;
417   }
418
419   fGeometry = trd->GetGeometry();
420   if (!fGeometry) {
421     printf("AliTRDpid::ReadKine -- ");
422     printf("No TRD geometry found\n");
423     return kFALSE;
424   }
425
426   file->Close(); 
427
428   return kTRUE;
429
430 }
431
432 //_____________________________________________________________________________
433 Bool_t AliTRDpid::ReadCluster(const Char_t *name)
434 {
435   //
436   // Opens and reads the cluster array from the file <name>.
437   //
438
439   fClusterArray->Delete();
440
441   printf("AliTRDpid::ReadCluster -- ");
442   printf("Open file %s\n",name);
443
444   AliTRDtracker *tracker = new AliTRDtracker("dummy","dummy");
445   tracker->ReadClusters(fClusterArray,name);
446
447   if (!fClusterArray) {
448     printf("AliTRDpid::ReadCluster -- ");
449     printf("Error reading the cluster array from file %s\n",name);
450     return kFALSE;
451   }
452
453   delete tracker;
454
455   return kTRUE;
456
457 }
458
459 //_____________________________________________________________________________
460 Bool_t AliTRDpid::ReadTracks(const Char_t *name)
461 {
462   //
463   // Opens and reads the track array from the file <name>.
464   //
465
466   fTrackArray->Delete();
467
468   TFile *file = (TFile *) gROOT->GetListOfFiles()->FindObject(name);
469   if (!file) {
470     printf("AliTRDpid::ReadTracks -- ");
471     printf("Open file %s\n",name);
472     file = new TFile(name);
473     if (!file) {
474       printf("AliTRDpid::ReadTracks -- ");
475       printf("Cannot open file %s\n",name);
476       return kFALSE;
477     }
478   }
479
480   TTree   *trackTree   = (TTree *) file->Get("TreeT");
481   TBranch *trackBranch = trackTree->GetBranch("tracks");
482
483   Int_t nEntry = ((Int_t) trackTree->GetEntries());
484   for (Int_t iEntry = 0; iEntry < nEntry; iEntry++) {
485     AliTRDtrack *track = new AliTRDtrack();
486     trackBranch->SetAddress(&track);
487     trackTree->GetEvent(iEntry);
488     fTrackArray->AddLast(track);
489   }
490
491   file->Close();
492
493   return kTRUE;
494
495 }
496
497 //_____________________________________________________________________________
498 Int_t AliTRDpid::GetIndexQ(const Float_t mom, const Int_t ipla, const Int_t ipid)
499 {
500   //
501   // Returns the Q-spectrum histogram index
502   //
503
504   if ((mom < fMinMom) || (mom > fMaxMom))  return -1;
505   Int_t imom = ((Int_t) ((mom - fMinMom) / fWidMom));
506   return GetIndexQ(imom,ipla,ipid);
507
508 }
509
510 //_____________________________________________________________________________
511 Int_t AliTRDpid::GetIndexQ(const Int_t imom, const Int_t ipla, const Int_t ipid)
512 {
513   //
514   // Returns the Q-spectrum histogram index
515   //
516
517   const Int_t kNpla = AliTRDgeometry::Nplan();
518   if ((ipid < 0) || (ipid >= kNpid)) return -1;
519   return imom * (kNpid * kNpla) + ipla * kNpid + ipid; 
520
521 }
522
523 //_____________________________________________________________________________
524 Int_t AliTRDpid::GetIndexLQ(const Float_t mom, const Int_t ipid)
525 {
526   //
527   // Returns the Q-likelihood histogram index
528   //
529
530   if ((mom < fMinMom) || (mom > fMaxMom))  return -1;
531   Int_t imom = ((Int_t) ((mom - fMinMom) / fWidMom));
532   return GetIndexLQ(imom,ipid);
533
534 }
535
536 //_____________________________________________________________________________
537 Int_t AliTRDpid::GetIndexLQ(const Int_t imom, const Int_t ipid)
538 {
539   //
540   // Returns the Q-likelihood histogram index
541   //
542
543   if ((ipid < 0) || (ipid >= kNpid)) return -1;
544   return imom * kNpid + ipid; 
545
546 }
547
548 //_____________________________________________________________________________
549 Int_t AliTRDpid::Pid(const AliTRDtrack *t)
550 {
551   //
552   // Determines the pid of the track <t>
553   // For a given track to be assigned an electron or pion pid it requires
554   // that more than a fraction of 0.7 of all associated cluster are 'pure'
555   // clusters -- meaning to have only contribution from one particle --
556   // of the specific particle type.
557   //
558
559   // PDG codes
560   const Int_t kPdgEl =  11;
561   const Int_t kPdgPi = 211;
562
563   // Minimum fraction of cluster from one particle
564   const Float_t kRatioMin = 0.7;
565
566   AliTRDcluster *cluster;
567   TParticle     *particle;
568
569   Int_t   nClusterEl = 0;
570   Int_t   nClusterPi = 0;
571   Int_t   nClusterNo = 0;
572
573   Float_t ratioEl    = 0;
574   Float_t ratioPi    = 0;
575
576   Int_t   ipid       = -1;
577
578   if (!fClusterArray) {
579     printf("AliTRDpid::Pid -- ");
580     printf("ClusterArray not defined. Initialize the PID object first\n");
581     return -1;  
582   }
583   
584   // Loop through all clusters associated to this track
585   Int_t nCluster = t->GetNclusters();
586   for (Int_t iCluster = 0; iCluster < nCluster; iCluster++) {
587
588     // Get a cluster
589     Int_t index = t->GetClusterIndex(iCluster);
590     if (!(cluster = (AliTRDcluster *) fClusterArray->UncheckedAt(index))) {
591       break;
592     } 
593
594     // Get the first two MC track indices
595     Int_t track0 = cluster->GetTrackIndex(0);
596     Int_t track1 = cluster->GetTrackIndex(1);
597     printf(" track0 = %d, track1 = %d\n",track0,track1);
598
599     // Take only 'pure' cluster 
600     if ((track0 > -1) && (track1 == -1)) {
601
602       particle = gAlice->Particle(track0);
603       switch (TMath::Abs(particle->GetPdgCode())) {
604       case kPdgEl:
605         nClusterEl++;
606         break;
607       case kPdgPi:
608         nClusterPi++;
609         break;
610       default:
611         nClusterNo++;
612         break;
613       };
614
615     }
616
617   }
618   
619   if (nCluster) {
620     ratioEl = ((Float_t) nClusterEl) / ((Float_t) nCluster);
621     ratioPi = ((Float_t) nClusterPi) / ((Float_t) nCluster);
622     if      (ratioEl > kRatioMin) {
623       ipid = kElectron;
624     }
625     else if (ratioPi > kRatioMin) {
626       ipid = kPion;
627     }
628   }
629   printf(" nCluster = %d, nClusterEl = %d, nClusterPi = %d, nClusterNo = %d\n"
630          ,nCluster,nClusterEl,nClusterPi,nClusterNo);
631   printf(" ratioEl = %f, ratioPi = %f\n",ratioEl,ratioPi);
632
633   return ipid;
634
635 }
636
637 //_____________________________________________________________________________
638 Bool_t AliTRDpid::SumCharge(const AliTRDtrack *t, Float_t *charge)
639 {
640   //
641   // Sums up the charge in each plane for track <t>.
642   //
643
644   Bool_t status = kTRUE;
645
646   AliTRDcluster *cluster;
647
648   const Int_t kNplane = AliTRDgeometry::Nplan();
649   for (Int_t iPlane = 0; iPlane < kNplane; iPlane++) {
650     charge[iPlane] = 0.0;
651   }
652
653   if (!fClusterArray) {
654     printf("AliTRDpid::SumCharge -- ");
655     printf("ClusterArray not defined. Initialize the PID object first\n");
656     return kFALSE;  
657   }
658   if (!fGeometry) {
659     printf("AliTRDpid::SumCharge -- ");
660     printf("TRD geometry not defined. Initialize the PID object first\n");
661     return kFALSE;  
662   }
663   
664   // Loop through all clusters associated to this track
665   Int_t nCluster = t->GetNclusters();
666   for (Int_t iCluster = 0; iCluster < nCluster; iCluster++) {
667
668     // Get a cluster
669     Int_t index = t->GetClusterIndex(iCluster);
670     if (!(cluster = (AliTRDcluster *) fClusterArray->UncheckedAt(index))) {
671       status = kFALSE;
672       break;
673     } 
674
675     // Sum up the charge
676     Int_t plane    = fGeometry->GetPlane(cluster->GetDetector());
677     charge[plane] += cluster->GetQ();
678
679   }
680
681   return status;
682
683 }
684
685 //_____________________________________________________________________________
686 Float_t AliTRDpid::LQElectron(const Float_t *charge)
687 {
688   //
689   // Returns the Q-likelihood to be a electron
690   //
691
692   return 1;
693
694 }
695
696 //_____________________________________________________________________________
697 Float_t AliTRDpid::LQPion(const Float_t *charge)
698 {
699   //
700   // Returns the Q-likelihood to be a pion
701   //
702
703   return 0;
704
705 }
706