Adding static_cast to keep the compiler happy
[u/mrichter/AliRoot.git] / PWG2 / RESONANCES / AliRsnSelectorRL.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 AliRsnSelectorRL
18 // ------------------------
19 // Reader for conversion of ESD output into the internal format
20 // used for resonance study.
21 // ---
22 // author: A. Pulvirenti             (email: alberto.pulvirenti@ct.infn.it)
23 // ---
24 // adapted for TSelector compliance
25 // by    : R. Vernet                 (email: renaud.vernet@cern.ch)
26 //-------------------------------------------------------------------------
27
28 #include <Riostream.h>
29
30 #include <TFile.h>
31 #include <TChain.h>
32 #include <TParticle.h>
33 #include <TRandom.h>
34 #include <TObjString.h>
35 #include <TObjectTable.h>
36 #include <TOrdCollection.h>
37 #include "AliRun.h"
38 #include "AliESD.h"
39 #include "AliStack.h"
40 #include "AliHeader.h"
41 #include "AliESDtrack.h"
42 #include "AliRunLoader.h"
43 #include "AliRsnDaughter.h"
44 #include "AliRsnEvent.h"
45 #include "AliRsnSelectorRL.h"
46
47 ClassImp(AliRsnSelectorRL)
48
49 //--------------------------------------------------------------------------------------------------------
50 AliRsnSelectorRL::AliRsnSelectorRL(TTree*) :
51   AliSelectorRL(),
52   fOutputPath(0),
53   fDebugFlag(0),
54   fStoreKineInfo(0),
55   fCheckITSRefit(0),
56   fRejectFakes(0),
57   fCopyMomentum(0),
58   fIsRunLoaderOpen(0),
59   fRsnEventTree(0),
60   fRsnEvent(0),
61   fRsnEventBranch(0),
62   fPIDMethod(kESDPID),
63   fPtLimit4PID(4.0),
64   fProbThreshold(0.0),
65   fMaxRadius(3.0)
66 {
67 //
68 // Constructor.
69 // Initializes all working parameters to default values:
70 // - ESD particle identification
71 // - rejection of non-ITS-refitted tracks
72 // - maximum distance allowed from primary vertex = 3.0 cm (beam pipe)
73 //
74         Int_t i;
75         for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = 1.0;
76 }
77
78 AliRsnSelectorRL::~AliRsnSelectorRL() {
79   Clear();
80 }
81 //--------------------------------------------------------------------------------------------------------
82 AliRsnSelectorRL::AliRsnSelectorRL(const AliRsnSelectorRL& obj) :
83   AliSelectorRL(), // not implemented a copy constructor for AliRsnSelectorRL
84   fOutputPath(obj.fOutputPath),
85   fDebugFlag(obj.fDebugFlag),
86   fStoreKineInfo(obj.fStoreKineInfo),
87   fCheckITSRefit(obj.fCheckITSRefit),
88   fRejectFakes(obj.fRejectFakes),
89   fCopyMomentum(obj.fCopyMomentum),
90   fIsRunLoaderOpen(obj.fIsRunLoaderOpen),
91   fRsnEventTree(obj.fRsnEventTree),
92   fRsnEvent(obj.fRsnEvent),
93   fRsnEventBranch(obj.fRsnEventBranch),
94   fPIDMethod(obj.fPIDMethod),
95   fPtLimit4PID(obj.fPtLimit4PID),
96   fProbThreshold(obj.fProbThreshold),
97   fMaxRadius(obj.fMaxRadius)
98 {
99 //
100 // Copy constructor
101 //
102         Int_t i;
103         for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = obj.fPrior[i];
104 }
105 //--------------------------------------------------------------------------------------------------------
106 AliRsnSelectorRL& AliRsnSelectorRL::operator=(const AliRsnSelectorRL& obj) 
107 {
108 //
109 // Assignment operator
110 // works in the same way as the copy constructor
111 //
112         if (this!=&obj) {
113                 fDebugFlag = obj.fDebugFlag;
114                 fOutputPath = obj.fOutputPath;
115                 fStoreKineInfo = obj.fStoreKineInfo;
116                 fCheckITSRefit = obj.fCheckITSRefit;
117                 fRejectFakes = obj.fRejectFakes;
118                 fCopyMomentum = obj.fCopyMomentum;
119                 fIsRunLoaderOpen = obj.fIsRunLoaderOpen;
120                 fRsnEventTree = obj.fRsnEventTree;
121                 fRsnEvent = obj.fRsnEvent;
122                 fRsnEventBranch = obj.fRsnEventBranch;
123                 fPIDMethod=obj.fPIDMethod;
124                 for (Int_t i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = obj.fPrior[i];
125                 fPtLimit4PID = obj.fPtLimit4PID;
126                 fProbThreshold = obj.fProbThreshold;
127                 fMaxRadius = obj.fMaxRadius;
128         }
129         
130         return *this;
131 }
132 //--------------------------------------------------------------------------------------------------------
133 void AliRsnSelectorRL::Clear(Option_t * /*option*/)
134 {
135   //
136   // Does nothing.
137   //
138 }
139 //--------------------------------------------------------------------------------------------------------
140 Double_t* AliRsnSelectorRL::GetPIDprobabilities(AliRsnDaughter track) const
141 {
142 //
143 // Computes PID probabilites starting from priors and weights
144 //
145         Int_t i;
146         Double_t *prob = new Double_t[AliPID::kSPECIES];
147         
148         // step 1 - compute the normalization factor
149         Double_t sum = 0.0;
150         for (i = 0; i < AliPID::kSPECIES; i++) {
151                 prob[i] = fPrior[i] * track.GetPIDweight(i);
152                 sum += prob[i];
153         }
154         if (sum <= 0.0) return 0;
155         
156         // step 2 - normalize PID weights by the relative prior probability
157         for (Int_t i = 0; i < AliPID::kSPECIES; i++) {
158                 prob[i] /= sum;
159         }
160         
161         return prob;
162 }
163 //--------------------------------------------------------------------------------------------------------
164 void AliRsnSelectorRL::Identify(AliRsnDaughter &track)
165 {
166 //
167 // Identifies a track according to method selected
168 //
169         Bool_t doESDPID = (fPIDMethod == kESDPID);
170         Bool_t keepRecSign = (fPIDMethod == kPerfectPIDwithRecSign);
171         
172         if (doESDPID) {
173                 // when doing ESD PID it is supposed that charge sign
174                 // comes from ESD track and it is not modified
175                 Double_t pt = track.GetPt();
176                 if (pt <= fPtLimit4PID) {
177               Double_t *prob = GetPIDprobabilities(track);
178           if (!prob) track.SetPDG(0);
179               Int_t idx[AliPID::kSPECIES];
180               TMath::Sort(static_cast<Int_t>(AliPID::kSPECIES), prob, idx);
181               Int_t imax = idx[0];
182           Double_t maxprob = prob[imax];
183               if (maxprob >= fProbThreshold) {
184                         track.SetPDG((UShort_t)AliPID::ParticleCode(imax));
185               }
186          delete [] prob;
187         }
188             else {
189           track.SetPDG(0);
190         }
191         }
192         else {
193             Short_t truePDG = track.GetTruePDG();
194         track.SetPDG((UShort_t)TMath::Abs(truePDG));
195                 if (!keepRecSign) {
196                         if (TMath::Abs(truePDG) <= 13) {
197                                 if (truePDG > 0) track.SetSign(-1); else track.SetSign(1);
198                         }
199                         else {
200                                 if (truePDG > 0) track.SetSign(1); else track.SetSign(-1);
201                         }
202                 }
203         }
204 }
205 //--------------------------------------------------------------------------------------------------------
206 void AliRsnSelectorRL::SetPriorProbabilities(Double_t *prior)
207 {
208 //
209 // Set prior probabilities to be used in case of ESD PID.
210 //
211         if (!prior) return;
212         Int_t i = 0;
213         for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = prior[i];
214 }
215 //--------------------------------------------------------------------------------------------------------
216 void AliRsnSelectorRL::SetPriorProbability(AliPID::EParticleType type, Double_t value)
217 {
218 //
219 // Sets prior probability referred to a single particle species.
220 //
221         if (type >= AliPID::kElectron && type < AliPID::kPhoton) fPrior[type] = value;
222 }
223 //--------------------------------------------------------------------------------------------------------
224 AliPID::EParticleType AliRsnSelectorRL::FindType(Int_t pdg)
225 {
226 //
227 // Finds the enum value corresponding to a PDG code
228 //
229         pdg = TMath::Abs(pdg);
230         switch (pdg) {
231                 case   11: return AliPID::kElectron; break;
232                 case   13: return AliPID::kMuon; break;
233                 case  211: return AliPID::kPion; break;
234                 case  321: return AliPID::kKaon; break;
235                 case 2212: return AliPID::kProton; break;
236                 default  : return AliPID::kPhoton;
237         }
238 }
239 //--------------------------------------------------------------------------------------------------------
240
241 //--------------------------------------------------------
242 //             The following is Selector stuff
243 //--------------------------------------------------------
244
245 //--------------------------------------------------------------------------------------------------------
246 void AliRsnSelectorRL::Begin(TTree *) const
247 {
248 //
249 // Implementation of BEGIN method
250 //
251         Info("Begin", "");
252         TString option = GetOption();
253 }
254 //--------------------------------------------------------------------------------------------------------
255 void AliRsnSelectorRL::SlaveBegin(TTree * tree)
256 {
257 //
258 // Implementation of secondary BEGIN which 
259 // is called by separate process managers
260 //
261         Info("SlaveBegin", "");
262         Init(tree);
263         TString option = GetOption();
264 }
265 //--------------------------------------------------------------------------------------------------------
266 void AliRsnSelectorRL::Init(TTree *tree)
267 {
268 //
269 // Initializer
270 // Connects the selector to a TTree and links the branch
271 // which is used to make translation ESD --> Rsn
272 //
273         Info("Init","");
274         
275         if (!tree) return;
276         fTree = tree;
277         if (fDebugFlag) {
278                 Info("Init", "fTree=%p   fTree->GetCurrentFile()=%p", fTree, fTree->GetCurrentFile());
279         }
280         fTree->SetBranchAddress("ESD", &fESD);
281         fRsnEventTree = new TTree("selection", "AliRsnEvents");
282         TTree::SetBranchStyle(1);
283         fRsnEvent = new AliRsnEvent;
284         fRsnEventBranch = fRsnEventTree->Branch("events", "AliRsnEvent", &fRsnEvent, 32000, 1);
285         fRsnEventBranch->SetAutoDelete(kFALSE);
286 }
287 //--------------------------------------------------------------------------------------------------------
288 Bool_t AliRsnSelectorRL::Process(Long64_t entry)
289 {
290 //
291 // Main core of the Selector processing
292 // Reads the ESD input and creates a TTree or AliRsnEvent objects
293 //
294         if (fDebugFlag) Info("Process", "Processing event %d", entry);
295         if (!AliSelectorRL::Process(entry)) return kFALSE;
296         
297         AliStack* stack = GetStack();
298         if (!stack) {
299                 Warning("Process", "NULL stack: cannot get kinematics info");
300                 fStoreKineInfo = kFALSE;
301         }
302                 
303         Int_t ntracks, nSelTracks = 0;
304                 
305         // primary vertex
306         Double_t vertex[3];
307         vertex[0] = fESD->GetVertex()->GetXv();
308         vertex[1] = fESD->GetVertex()->GetYv();
309         vertex[2] = fESD->GetVertex()->GetZv();
310         
311         // multiplicity
312         ntracks = fESD->GetNumberOfTracks();
313         if (!ntracks) {
314                 Warning("Process", "Event %d has no tracks!", entry);
315                 return kFALSE;
316         }
317         if (fDebugFlag) Info("Process", "Number of ESD tracks : %d", ntracks);
318         
319         // create new AliRsnEvent object
320         fRsnEvent->Clear("DELETE");
321         fRsnEvent->Init();
322         
323         // store tracks from ESD
324         Int_t index, label;
325         Double_t vtot, v[3];
326         AliRsnDaughter track;
327         for (index = 0; index < ntracks; index++) {
328                 if (fDebugFlag) Info("Process","Track # %d", index);
329                 AliESDtrack *esdTrack = fESD->GetTrack(index);
330                 
331                 // check against vertex constraint
332                 esdTrack->GetXYZ(v);
333                 vtot  = (v[0] - vertex[0])*(v[0] - vertex[0]);
334                 vtot += (v[1] - vertex[1])*(v[1] - vertex[1]);
335                 vtot += (v[2] - vertex[2])*(v[2] - vertex[2]);
336                 vtot  = TMath::Sqrt(vtot);
337                 if (vtot > fMaxRadius) continue;
338                 
339                 // check for fakes
340                 label = esdTrack->GetLabel();
341                 if (fRejectFakes && (label < 0)) continue;
342                 
343                 // copy ESDtrack data into RsnDaughter (and make Bayesian PID)
344                 if (!track.Adopt(esdTrack, fCheckITSRefit)) continue;
345                 track.SetIndex(index);
346                 
347                 // if requested, get Kine info
348                 if (fStoreKineInfo) {
349                         if (fDebugFlag) Info("Process", "Getting part label=%d stack=%p", label, stack);
350                         TParticle *part = stack->Particle(TMath::Abs(label));
351                         if (fDebugFlag) Info("Process", "part=%p", part);
352                         track.SetTruePDG(part->GetPdgCode());
353                         Int_t mother = part->GetFirstMother();
354                         track.SetMother(mother);
355                         if (mother >= 0) {
356                                 TParticle *mum = stack->Particle(mother);
357                                 track.SetMotherPDG(mum->GetPdgCode());
358                         }
359                         // if requested, reconstructed momentum is replaced with true momentum
360                         if (fCopyMomentum) track.SetPxPyPz(part->Px(), part->Py(), part->Pz());
361                 }
362                 
363                 // identification
364                 Identify(track);
365                 
366                 // store in TClonesArray
367                 if (fDebugFlag) track.Print("SLITVPMNW");
368                 fRsnEvent->AddTrack(track);
369                 nSelTracks++;
370         }
371         
372         // compute total multiplicity
373         fRsnEvent->ComputeMultiplicity();
374         
375         // link to events tree and fill
376         fRsnEventTree->Fill();
377         
378         return kTRUE;
379 }
380 //--------------------------------------------------------------------------------------------------------
381 void AliRsnSelectorRL::SlaveTerminate()
382 {
383 //
384 // SlaveTerminate
385 // Partial termination method
386 //
387 // test have been performed only on AliEn, please contact the author in case of merging problem under CAF/PROOF.
388 //
389 //
390         Info("SlaveTerminate", "");
391         
392         // Add the histograms to the output on each slave server
393         fOutput->Add(fRsnEventTree);
394 }
395 //--------------------------------------------------------------------------------------------------------
396 void AliRsnSelectorRL::Terminate()
397 {
398 //
399 // Global termination method
400 //
401         Info("Terminate","");
402         // fRsnEventTree = dynamic_cast<TTree*>(fOutput->FindObject("fRsnEventTree"));
403         
404         //AliSelector::Terminate();
405         cout << fOutputPath << endl;
406         Info("Terminate", Form("Saving in: %s", fOutputPath->Data()));
407         TFile* file = TFile::Open(fOutputPath->Data(), "RECREATE");
408         fRsnEventTree->Write();
409         file->Close();
410         
411         delete fRsnEventTree;
412         delete file;
413         delete fOutputPath;
414 }