783d8f5a26a64e95470b94177350c3a01397893b
[u/mrichter/AliRoot.git] / PWG2 / RESONANCES / AliRsnPID.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 //                      Class AliRsnPID
18 //                     -------------------
19 //           Simple collection of reconstructed tracks
20 //           selected from an ESD event
21 //           to be used for analysis.
22 //           .........................................
23 // 
24 // author: A. Pulvirenti             (email: alberto.pulvirenti@ct.infn.it)
25 //-------------------------------------------------------------------------
26
27 #include <TMath.h>
28 #include <TString.h>
29 #include <TClonesArray.h>
30 #include <TDatabasePDG.h>
31
32 #include "AliLog.h"
33 #include "AliRsnParticle.h"
34 #include "AliRsnDaughter.h"
35 #include "AliRsnEvent.h"
36
37 #include "AliRsnPID.h"
38
39 ClassImp(AliRsnPID)
40
41 const char* AliRsnPID::fgkParticleName[AliRsnPID::kSpecies + 1] = 
42 {
43   "electron",
44   "muon",
45   "pion",
46   "kaon",
47   "proton",
48   "unknown"
49 };
50
51 const Int_t AliRsnPID::fgkParticlePDG[AliRsnPID::kSpecies + 1] = 
52 {
53     11,
54     13,
55     211,
56     321,
57     2212,
58     0
59 };
60
61 //_____________________________________________________________________________
62 AliRsnPID::AliRsnPID() :
63   fMethod(kRealistic),
64   fMaxPt(10.0),
65   fMinProb(0.0)
66 {
67 //=========================================================
68 // Default constructor 
69 // Sets a default setup:
70 //  - realistic PID
71 //  - no lower limit for probability
72 //  - upper limit of 10 GeV for Pt
73 //  - suitable values for prior probabilities
74 //=========================================================
75     
76     fPrior[kElectron] = 0.20;
77         fPrior[kMuon    ] = 0.20;
78         fPrior[kPion    ] = 0.83;
79         fPrior[kKaon    ] = 0.07;
80         fPrior[kProton  ] = 0.06;
81 }
82
83 //_____________________________________________________________________________
84 AliRsnPID::AliRsnPID(const AliRsnPID &event) :
85   TObject((TObject)event),
86   fMethod(event.fMethod),
87   fMaxPt(event.fMaxPt),
88   fMinProb(event.fMinProb)
89 {
90 //=========================================================
91 // Copy constructor.
92 // Implemented to manage the array safely.
93 //=========================================================
94     
95     Int_t i;
96     for (i = 0; i < kSpecies; i++) fPrior[i] = event.fPrior[i];
97 }
98
99 //_____________________________________________________________________________
100 AliRsnPID& AliRsnPID::operator=(const AliRsnPID &event)
101 {
102 //=========================================================
103 // Assignment operator.
104 // Implemented to manage the array safely.
105 //=========================================================
106     
107     fMethod = event.fMethod;
108     fMaxPt = event.fMaxPt;
109     fMinProb = event.fMinProb;
110     
111     Int_t i;
112     for (i = 0; i < kSpecies; i++) fPrior[i] = event.fPrior[i];
113     
114     // return the newly created object
115     return (*this);
116 }
117
118 //_____________________________________________________________________________
119 AliRsnPID::EType AliRsnPID::InternalType(Int_t pdg)
120 //=========================================================
121 // Return the internal enum value corresponding to the PDG
122 // code passed as argument, if possible.
123 // Otherwise, returns 'kUnknown' by default.
124 //=========================================================
125 {
126     EType value;
127     Int_t absPDG = TMath::Abs(pdg);
128     
129         switch (absPDG) {
130         case 11:
131             value = kElectron;
132             break;
133         case 13:
134             value = kMuon;
135             break;
136         case 211:
137             value = kPion;
138             break;
139         case 321:
140             value = kKaon;
141             break;
142         case 2212:
143             value = kProton;
144             break;
145         default:
146             value = kUnknown;
147     }
148     return value;
149 }
150
151
152 //_____________________________________________________________________________
153 Int_t AliRsnPID::PDGCode(EType type)
154 {
155 //=========================================================
156 // Returns the PDG code of the particle type
157 // specified as argument (w.r. to the internal enum)
158 //=========================================================
159     
160     if (type >= kElectron && type <= kUnknown) {
161         return fgkParticlePDG[type];
162     }
163     else {
164         return 0;
165     }
166 }
167
168 //_____________________________________________________________________________
169 const char* AliRsnPID::ParticleName(EType type)
170 {
171 //=========================================================
172 // Returns the name of the particle type
173 // specified as argument (w.r. to the internal enum)
174 //=========================================================
175
176     if (type >= kElectron && type <= kUnknown) {
177         return fgkParticleName[type];
178     }
179     else {
180         return "unknown";
181     }
182 }
183
184 //_____________________________________________________________________________
185 Double_t AliRsnPID::ParticleMass(EType type)
186 {
187 //=========================================================
188 // Returns the mass corresponding to the particle type
189 // specified as argument (w.r. to the internal enum)
190 //=========================================================
191     TDatabasePDG *db = TDatabasePDG::Instance();
192     Int_t pdg = PDGCode(type);
193     return db->GetParticle(pdg)->Mass();
194 }
195
196 //_____________________________________________________________________________
197 Bool_t AliRsnPID::IdentifyRealistic(AliRsnDaughter *daughter)
198 {
199 //=========================================================
200 // Uses the Bayesian combination of prior probabilities
201 // with the PID weights of the passed object to compute
202 // the overall PID probabilities for each particle type.
203 //
204 // Once this computation is done, the argument is assigned
205 // the PID corresponding to the largest probability,
206 // and its data members are updated accordingly.
207 // If the track Pt is larger than the cut defined (fMaxPt)
208 // or the probability is smaller than the cut defined (fMinProb),
209 // the track is considered unidentified.
210 //
211 // If the identification procedure encounters errors, 
212 // the return value will be "FALSE", otherwise it is "TRUE".
213 //=========================================================
214
215     // retrieve PID weights from argument
216     Int_t i;
217     Double_t *prob   = new Double_t[kSpecies];
218     Double_t *weight = new Double_t[kSpecies];
219     for (i = 0; i < kSpecies; i++) weight[i] = (daughter->PID())[i];
220         
221     // step 1 - compute the normalization factor
222     Double_t sum = 0.0;
223     for (i = 0; i < kSpecies; i++) {
224         prob[i] = fPrior[i] * weight[i];
225         sum += prob[i];
226     }
227     if (sum <= 0.0) {
228         AliError(Form("Sum of weights = %f < 0.0", sum));
229         Unidentify(daughter);
230         return kFALSE;
231     }
232         
233         // step 2 - normalize PID weights
234     for (i = 0; i < kSpecies; i++) {
235         prob[i] /= sum;
236     }
237         
238         // step 3 - finds the maximum probability
239     Int_t imax = 0;
240     for (i = 1; i < kSpecies; i++) {
241         if (prob[i] > prob[imax]) imax = i;
242     }
243     
244     // step 4 - update daughter data members
245     //          and check pt & prob
246     Double_t pt      = daughter->Pt();
247     EType    type    = (EType)imax;
248     Double_t maxProb = prob[imax];
249     if (pt <= fMaxPt && maxProb >= fMinProb) {
250         daughter->SetPIDType(type);
251         daughter->SetPIDProb(maxProb);
252         daughter->SetM(ParticleMass(type));
253     }
254     else {
255         Unidentify(daughter);
256     }
257     
258     return kTRUE;
259 }
260
261 //_____________________________________________________________________________
262 Bool_t AliRsnPID::IdentifyPerfect(AliRsnDaughter *daughter)
263 {
264 //=========================================================
265 // Uses the true PDG code to make a perfect identification.
266 // If the true PDG code does not correspond to any 
267 // of the expected PID types, gives a warning, sets the 
268 // PID to 'unknown' and returns kFALSE.
269 // Otherwise, returns kTRUE.
270 //=========================================================
271
272     // retrieve true PDG code from argument
273     AliRsnParticle *particle = daughter->GetParticle();
274     if (!particle) {
275         AliWarning("Particle object not initialized: impossible to do perfect PID");
276         Unidentify(daughter);
277         return kFALSE;
278     }
279     Int_t pdgCode = particle->PDG();
280     EType type = InternalType(pdgCode);
281     if (type == kUnknown) {
282         //AliWarning(Form("Unrecognized PDG code = %d -- set PID to Unknown", pdgCode));
283         Unidentify(daughter);
284         return kFALSE;
285     }
286     else {
287         daughter->SetPIDType(type);
288         daughter->SetPIDProb(1.0);
289         daughter->SetM(ParticleMass(type));
290         return kTRUE;
291     }
292 }
293
294 //_____________________________________________________________________________
295 Bool_t AliRsnPID::Unidentify(AliRsnDaughter *daughter)
296 {
297 //=========================================================
298 // Sets the PID to 'unknown' to every track.
299 //=========================================================
300
301     daughter->SetPIDType(kUnknown);
302     return kTRUE;
303 }
304
305 //_____________________________________________________________________________
306 Bool_t AliRsnPID::Identify(AliRsnDaughter *daughter)
307 {
308 //=========================================================
309 // Recalls one of the above methods, according to the one
310 // defined in the related data member.
311 // If the method is not recognized, returns kFALSE and 
312 // gives an alert. Otherwise, returns kTRUE.
313 //=========================================================
314
315     switch (fMethod) {
316         case kNone:
317             Unidentify(daughter);
318             return kTRUE;
319         case kRealistic:
320             IdentifyRealistic(daughter);
321             return kTRUE;
322         case kPerfect:
323             IdentifyPerfect(daughter);
324             return kTRUE;
325         default:
326             AliError(Form("PID method '%d' unrecognized. Nothing done.", fMethod));
327             return kFALSE;
328     }
329 }
330
331 //_____________________________________________________________________________
332 Bool_t AliRsnPID::Identify(AliRsnEvent *event)
333 {
334 //=========================================================
335 // Performs identification for all tracks in a given event.
336 // Returns the logical AND of all PID operations.
337 //=========================================================
338
339     Bool_t check = kTRUE;
340     AliRsnDaughter *daughter = 0;
341     TObjArrayIter iter(event->GetAllTracks());
342     while ( (daughter = (AliRsnDaughter*)iter.Next()) ) {
343         check = check && Identify(daughter);
344     }
345     event->FillPIDArrays();
346
347     return check;
348 }
349
350
351 //_____________________________________________________________________________
352 void AliRsnPID::SetPriorProbability(EType type, Double_t p)
353 {
354 //=========================================================
355 // Sets the prior probability for Realistic PID, for a
356 // given particle species.
357 //=========================================================
358     
359     if (type >= kElectron && type < kSpecies) {
360         fPrior[type] = p;
361     }
362 }