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