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