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