Improved functionality of AliRsnDaughterDef::MatchesDaughter()
[u/mrichter/AliRoot.git] / PWG2 / RESONANCES / AliRsnCutPID.cxx
1 //
2 // Class AliRsnCutPID
3 //
4 // General implementation of a single cut strategy, which can be:
5 // - a value contained in a given interval  [--> IsBetween()   ]
6 // - a value equal to a given reference     [--> MatchesValue()]
7 //
8 // In all cases, the reference value(s) is (are) given as data members
9 // and each kind of cut requires a given value type (Int, UInt, Double),
10 // but the cut check procedure is then automatized and chosen thanks to
11 // an enumeration of the implemented cut types.
12 // At the end, the user (or any other point which uses this object) has
13 // to use the method IsSelected() to check if this cut has been passed.
14 //
15 // authors: Martin Vala (martin.vala@cern.ch)
16 //          Alberto Pulvirenti (alberto.pulvirenti@ct.infn.it)
17 //
18 #include "TMath.h"
19
20 #include "AliExternalTrackParam.h"
21
22 #include "AliRsnDaughter.h"
23 #include "AliRsnCutPID.h"
24
25 ClassImp(AliRsnCutPID)
26
27 //_________________________________________________________________________________________________
28 AliRsnCutPID::AliRsnCutPID() :
29    AliRsnCut(),
30    fPerfect(kFALSE),
31    fUseDefault(kTRUE)
32 {
33 //
34 // Default constructor.
35 // Sets the cut to realistic PID with default weights,
36 // and defines the 'fMinI' value of the base class as the PID
37 // to which we want to compare this object.
38 //
39
40    Int_t i;
41
42    for (i = 0; i < kDetectors; i++) {
43       fUseDetector[i] = kFALSE;
44       fPtThreshold[i] = 0.0;
45       fGoAboveThreshold[i] = kTRUE;
46    }
47
48    for (i = 0; i < AliPID::kSPECIES; i++) {
49       fWeight[i] = 0.0;
50       fPrior[i] = 1.0;
51    }
52
53    SetTargetType(AliRsnTarget::kDaughter);
54 }
55
56 //_________________________________________________________________________________________________
57 AliRsnCutPID::AliRsnCutPID(const char *name, AliPID::EParticleType pid, Double_t probMin, Bool_t perfectPID) :
58    AliRsnCut(name, AliRsnCut::kDaughter, (Int_t)pid),
59    fPerfect(perfectPID),
60    fUseDefault(kTRUE)
61 {
62 //
63 // Default constructor.
64 // Sets the cut to realistic PID with default weights,
65 // and defines the 'fMinI' value of the base class as the PID
66 // to which we want to compare this object.
67 //
68
69    Int_t i;
70
71    for (i = 0; i < kDetectors; i++) {
72       fUseDetector[i] = kFALSE;
73       fPtThreshold[i] = 0.0;
74       fGoAboveThreshold[i] = kTRUE;
75    }
76
77    for (i = 0; i < AliPID::kSPECIES; i++) {
78       fWeight[i] = 0.0;
79       fPrior[i] = 1.0;
80    }
81
82    fMinD = probMin;
83    fMaxD = 1.000001;
84
85    SetTargetType(AliRsnTarget::kDaughter);
86 }
87
88 //_____________________________________________________________________________
89 Bool_t AliRsnCutPID::CheckThreshold(EDetector det, Double_t value)
90 {
91 //
92 // Checks if the passed value (track pT) stays in the
93 // interval where the detector should be accepted
94 //
95
96    if (!CheckBounds(det)) return kFALSE;
97
98    if (fGoAboveThreshold[det]) return (value >= fPtThreshold[det]);
99    else return (value <= fPtThreshold[det]);
100 }
101
102 //_________________________________________________________________________________________________
103 Bool_t AliRsnCutPID::ComputeWeights(AliRsnDaughter *daughter)
104 {
105 //
106 // Compute the PID weights using the class settings.
107 // If the argument is an ESD track, customized weights can be computed
108 // It the argument is a track (ESD or AOD), at least default weights
109 // can be computed, otherwise, no weights can be defined.
110 // The return value tells if the operation was successful.
111 //
112
113    Int_t  i, j;
114    Bool_t useDefault = fUseDefault;
115    Bool_t perfectPID = fPerfect;
116    if (perfectPID && !daughter->GetRefMC()) return kFALSE;
117    if (!daughter->GetRefESDtrack()) useDefault = kTRUE;
118    if (!daughter->GetRefESDtrack() && !daughter->GetRefAODtrack()) return kFALSE;
119
120    // if perfect PID ise required,
121    // compare the PDG code of the type stored in 'fMinI' of the cut
122    // and that of the particle which is checked, looking at its MC
123    if (perfectPID) {
124       i = TMath::Abs(AliPID::ParticleCode(fMinI));
125       j = daughter->GetPDG();
126       return (i == j);
127    }
128
129    // if default weights are (or need to be) used,
130    // they are taken directly and function exits
131    if (useDefault) {
132       if (daughter->GetRefESDtrack())
133          daughter->GetRefESDtrack()->GetESDpid(fWeight);
134       else {
135          for (i = 0; i < AliPID::kSPECIES; i++)
136             fWeight[i] = daughter->GetRefAODtrack()->PID()[i];
137       }
138       return kTRUE;
139    }
140
141    // if we arrive here, this means that we have an ESD track
142    // and we want to customize the PID
143    AliESDtrack *track = daughter->GetRefESDtrack();
144    Double_t     w[kDetectors][AliPID::kSPECIES];
145    track->GetITSpid(w[kITS]);
146    track->GetTPCpid(w[kTPC]);
147    track->GetTRDpid(w[kTRD]);
148    track->GetTOFpid(w[kTOF]);
149    track->GetHMPIDpid(w[kHMPID]);
150
151    // if a detector is excluded or the track has not the right pT
152    // all related weights are set to 1 in order not to contribute
153    // to final multiplication
154    for (i = 0; i < kDetectors; i++) {
155       if (!fUseDetector[i] || !CheckThreshold((EDetector)i, track->Pt())) {
156          for (j = 0; j < AliPID::kSPECIES; j++) {
157             w[i][j] = 1.0;
158          }
159       }
160    }
161
162    // multiplicate all weights to compute final one
163    for (i = 0; i < AliPID::kSPECIES; i++) {
164       fWeight[i] = w[kITS][i] * w[kTPC][i] * w[kTRD][i] * w[kTOF][i] * w[kHMPID][i];
165    }
166
167    return kTRUE;
168 }
169
170 //_________________________________________________________________________________________________
171 Int_t AliRsnCutPID::RealisticPID(AliRsnDaughter * const daughter, Double_t &prob)
172 {
173 //
174 // Combines the weights collected from chosen detectors with the priors
175 // and gives the corresponding particle with the largest probability,
176 // in terms of the AliPID particle type enumeration.
177 // The argument, passed by reference, gives the corresponding probability,
178 // since the cut could require that it is larger than a threshold.
179 //
180
181    // try to compute the weights
182    if (!ComputeWeights(daughter)) {
183       prob = -1.0;
184       return AliPID::kUnknown;
185    }
186
187    // combine with priors and normalize
188    Int_t    i;
189    Double_t sum = 0.0, w[AliPID::kSPECIES];
190    for (i = 0; i < AliPID::kSPECIES; i++) {
191       w[i] = fWeight[i] * fPrior[i];
192       sum += w[i];
193    }
194    if (sum <= 0.0) {
195       AliError("Sum = 0");
196       prob = -1.0;
197       return AliPID::kUnknown;
198    }
199    for (i = 0; i < AliPID::kSPECIES; i++) w[i] /= sum;
200
201    // find the largest probability and related PID
202    Int_t ibest = 0;
203    prob = w[0];
204    for (i = 1; i < AliPID::kSPECIES; i++) {
205       if (w[i] > prob) {
206          prob = w[i];
207          ibest = i;
208       }
209    }
210
211    // return the value, while the probability
212    // will be consequentially set
213    return ibest;
214 }
215
216 //_________________________________________________________________________________________________
217 Int_t AliRsnCutPID::PerfectPID(AliRsnDaughter * const daughter)
218 {
219 //
220 // If MC is present, retrieve the particle corresponding to the used track
221 // (using the fRefMC data member) and then return the true particle type
222 // taken from the PDG code of the reference particle itself, converted
223 // into the enumeration from AliPID object.
224 //
225
226    // works only if the MC is present
227    if (!daughter->GetRefMC()) return AliPID::kUnknown;
228
229    // get the PDG code of the particle
230    Int_t pdg = daughter->GetPDG();
231
232    // loop over all species listed in AliPID to find the match
233    Int_t i;
234    for (i = 0; i < AliPID::kSPECIES; i++) {
235       if (AliPID::ParticleCode(i) == pdg) return i;
236    }
237
238    return AliPID::kUnknown;
239 }
240
241 //_________________________________________________________________________________________________
242 Bool_t AliRsnCutPID::IsSelected(TObject *object)
243 {
244 //
245 // Cut checker.
246 //
247
248    // convert the object into the unique correct type
249
250    if (!TargetOK(object)) {
251       AliError(Form("[%s]: this cut works only with AliRsnDaughter objects", GetName()));
252       return kTRUE;
253    }
254
255    // if target is OK, do a dynamic cast
256    AliRsnDaughter *daughter = fDaughter;
257
258    // depending on the PID type, recalls the appropriate method:
259    // in case of perfect PID, checks only if the PID type is
260    // corresponding to the request in the cut (fMinI)
261    // while in case of realistic PID checks also the probability
262    // to be within the required interval
263    if (fPerfect && daughter) {
264       fCutValueI = PerfectPID(daughter);
265       return OkValueI();
266    } else if (daughter) {
267       fCutValueI = RealisticPID(daughter, fCutValueD);
268       return OkValueI() && OkRangeD();
269    } else
270       return kFALSE;
271 }
272
273 //__________________________________________________________________________________________________
274 void AliRsnCutPID::IncludeDetector(EDetector det, Double_t threshold, Bool_t goAbove)
275 {
276 //
277 // Include a detector for a customized weight computing
278 // and specify also its eventual threshold and if the detector
279 // must be used above or below the threshold.
280 // By default the threshold is zero and detector is always used.
281 //
282
283    if (!CheckBounds(det)) return;
284
285    fUseDetector[det] = kTRUE;
286    fPtThreshold[det] = threshold;
287    fGoAboveThreshold[det] = goAbove;
288 }