removed eff-c++ warnings
[u/mrichter/AliRoot.git] / PWG2 / RESONANCES / AliRsnReader.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 AliRsnReader
18 //
19 //   Reader for conversion of ESD or Kinematics output into AliRsnEvent
20 //   .....
21 //   .....
22 //   .....
23 //   .....
24 //   .....
25 //   .....
26 //   .....
27 // 
28 // author: A. Pulvirenti             (email: alberto.pulvirenti@ct.infn.it)
29 //-------------------------------------------------------------------------
30
31 #include <Riostream.h>
32
33 #include <TH1.h>
34 #include <TH3.h>
35 #include <TFile.h>
36 #include <TTree.h>
37 #include <TArrayF.h>
38 #include <TParticle.h>
39 #include <TRandom.h>
40
41 #include "AliRun.h"
42 #include "AliESD.h"
43 #include "AliStack.h"
44 #include "AliHeader.h"
45 #include "AliESDtrack.h"
46 #include "AliRunLoader.h"
47
48 #include "AliRsnDaughter.h"
49 #include "AliRsnEvent.h"
50 #include "AliRsnReader.h"
51
52 ClassImp(AliRsnReader)
53
54 //--------------------------------------------------------------------------------------------------------
55 AliRsnReader::AliRsnReader() :
56   TObject(),
57   fPIDMethod(kESDPID),
58   fPtLimit4PID(4.0),
59   fProbThreshold(0.0),
60   fMaxRadius(3.0),
61   fUseKineInfo(kFALSE),
62   fEvents(0x0)
63 {
64 //
65 // Constructor.
66 // Initializes all working parameters to default values:
67 // - ESD particle identification
68 // - rejection of non-ITS-refitted tracks
69 // - maximum distance allowed from primary vertex = 3.0 cm (beam pipe)
70 //
71         Int_t i;
72         for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = 1.0;
73 }
74 //--------------------------------------------------------------------------------------------------------
75 AliRsnReader::AliRsnReader(const AliRsnReader &copy) : 
76   TObject(copy),
77   fPIDMethod(copy.fPIDMethod),
78   fPtLimit4PID(copy.fPtLimit4PID),
79   fProbThreshold(copy.fProbThreshold),
80   fMaxRadius(copy.fMaxRadius),
81   fUseKineInfo(copy.fUseKineInfo),
82   fEvents(0x0)
83 {
84 //
85 // Copy constructor.
86 // Initializes all working parameters to same values of another simlar object.
87 // TTree data member is not created.
88 //
89         Int_t i;
90         for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = copy.fPrior[i];
91 }
92 //--------------------------------------------------------------------------------------------------------
93 AliRsnReader& AliRsnReader::operator=(const AliRsnReader &copy)
94 {
95 //
96 // Assignment operator.
97 // Initializes all working parameters to same values of another simlar object.
98 // TTree data member is not created.
99 //
100         fPIDMethod = copy.fPIDMethod;
101         Int_t i;
102         for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = copy.fPrior[i];
103         fPtLimit4PID = copy.fPtLimit4PID;
104         fProbThreshold = copy.fProbThreshold;
105         fMaxRadius = copy.fMaxRadius;
106         
107         fUseKineInfo = copy.fUseKineInfo;
108
109         fEvents = 0;
110         
111         return (*this);
112 }
113 //--------------------------------------------------------------------------------------------------------
114 void AliRsnReader::Clear(Option_t *option)
115 {
116 //
117 // Clear collection of filenames.
118 // If requested with the option "DELETELIST", 
119 // the collection object is also deleted.
120 //
121         TString opt(option);
122         
123         if (!opt.CompareTo("TREE", TString::kIgnoreCase)) {
124                 fEvents->Reset();
125                 if (!opt.CompareTo("DELTREE", TString::kIgnoreCase)) {
126                         delete fEvents;
127                         fEvents = 0;
128                 }
129         }
130 }
131 //--------------------------------------------------------------------------------------------------------
132 Double_t* AliRsnReader::GetPIDprobabilities(AliRsnDaughter track) const
133 {
134 //
135 // Computes PID probabilites starting from priors and weights
136 //
137         Double_t *prob = new Double_t[AliPID::kSPECIES];
138         
139         Int_t i;
140         
141         // step 1 - compute the normalization factor
142         Double_t sum = 0.0;
143         for (i = 0; i < AliPID::kSPECIES; i++) {
144                 prob[i] = fPrior[i] * track.GetPIDweight(i);
145                 sum += prob[i];
146         }
147         if (sum <= 0.0) return 0;
148         
149         // step 2 - normalize PID weights by the relative prior probability
150         for (Int_t i = 0; i < AliPID::kSPECIES; i++) {
151                 prob[i] /= sum;
152         }
153         
154         return prob;
155 }
156 //--------------------------------------------------------------------------------------------------------
157 void AliRsnReader::Identify(AliRsnDaughter &track)
158 {
159 //
160 // Identifies a track according to method selected
161 //
162         Bool_t doESDPID = (fPIDMethod == kESDPID);
163         Bool_t keepRecSign = (fPIDMethod == kPerfectPIDwithRecSign);
164         
165         if (doESDPID) {
166                 // when doing ESD PID it is supposed that charge sign
167                 // comes from ESD track and it is not modified
168                 Double_t pt = track.GetPt();
169                 if (pt <= fPtLimit4PID) {
170                         Double_t *prob = GetPIDprobabilities(track);
171                         if (!prob) track.SetPDG(0);
172                         Int_t idx[AliPID::kSPECIES];
173                         TMath::Sort(AliPID::kSPECIES, prob, idx);
174                         Int_t imax = idx[0];
175                         Double_t maxprob = prob[imax];
176                         if (maxprob >= fProbThreshold) {
177                                 track.SetPDG((UShort_t)AliPID::ParticleCode(imax));
178                         }
179                         delete [] prob;
180                 }
181                 else {
182                         track.SetPDG(0);
183                 }
184         }
185         else {
186                 Short_t truePDG = track.GetTruePDG();
187                 track.SetPDG((UShort_t)TMath::Abs(truePDG));
188                 if (!keepRecSign) {
189                         if (TMath::Abs(truePDG) <= 13) {
190                                 if (truePDG > 0) track.SetSign(-1); else track.SetSign(1);
191                         }
192                         else {
193                                 if (truePDG > 0) track.SetSign(1); else track.SetSign(-1);
194                         }
195                 }
196         }
197 }
198 //--------------------------------------------------------------------------------------------------------
199 TTree* AliRsnReader::ReadTracks(const char *path, Option_t *option)
200 {
201 //
202 // Reads AliESD in a given path, with and "experimental" method.
203 // When using this method, no Kinematics information is assumed
204 // and particle identification is always done with the bayesian method.
205 // No Kinematics informations are stored.
206 // Allowed options (actually):
207 // - "R" : reject tracks which do not have the flag "kITSRefit" set to TRUE
208 // - "F" : reject 'fake' tracks (negative label)
209 //
210         fPIDMethod = kESDPID;
211
212         TTree *events = new TTree("selection", "AliRsnEvents");
213         TTree::SetBranchStyle(1);
214         AliRsnEvent *event = new AliRsnEvent;
215         TBranch *branch = events->Branch("events", "AliRsnEvent", &event, 32000, 1);
216         branch->SetAutoDelete(kFALSE);
217         
218         // get path
219         TString strPath(path);
220         if (strPath.Length() < 1) return 0;
221         strPath += '/';
222         
223         // evaluate options
224         TString opt(option);
225         opt.ToUpper();
226         Bool_t checkITSrefit = (opt.Contains("R"));
227         Bool_t rejectFakes = (opt.Contains("F"));
228         
229         // opens ESD file
230         TFile *fileESD = TFile::Open(Form("%s/AliESDs.root", strPath.Data()));
231         if (!fileESD) return 0;
232         if (fileESD->IsOpen() == kFALSE) return 0;
233         TTree* treeESD = (TTree*)fileESD->Get("esdTree");
234         AliESD *esd = 0;
235         treeESD->SetBranchAddress("ESD", &esd);
236         Int_t nev = (Int_t)treeESD->GetEntries();
237         
238         // loop on events
239         Int_t i, nSelTracks = 0; 
240         Double_t vertex[3];
241         for (i = 0; i < nev; i++) {
242         
243                 // message
244 //              cout << "\rEvent " << i << flush;
245                 treeESD->GetEntry(i);
246                 
247                 // primary vertex
248                 vertex[0] = esd->GetVertex()->GetXv();
249                 vertex[1] = esd->GetVertex()->GetYv();
250                 vertex[2] = esd->GetVertex()->GetZv();
251                 
252                 // create new AliRsnEvent object
253                 event->Clear("DELETE");
254                 event->Init();
255                                 
256                 // get number of tracks
257                 Int_t ntracks = esd->GetNumberOfTracks();
258                 if (!ntracks) continue;
259                 
260                 // store tracks from ESD
261                 Int_t index, label;
262                 Double_t vtot, v[3];
263                 for (index = 0; index < ntracks; index++) {
264                         
265                         // get track
266                         AliESDtrack *esdTrack = esd->GetTrack(index);
267                         
268                         // check against vertex constraint
269                         esdTrack->GetXYZ(v);
270                         vtot  = (v[0] - vertex[0])*(v[0] - vertex[0]);
271                         vtot += (v[1] - vertex[1])*(v[1] - vertex[1]);
272                         vtot += (v[2] - vertex[2])*(v[2] - vertex[2]);
273                         vtot  = TMath::Sqrt(vtot);
274                         if (vtot > fMaxRadius) continue;
275                         
276                         // check for ITS refit
277                         if (checkITSrefit) {
278                                 if (!(esdTrack->GetStatus() & AliESDtrack::kITSrefit)) continue;
279                         }
280                         
281                         // check for fakes
282                         label = esdTrack->GetLabel();
283                         if (rejectFakes && (label < 0)) continue;
284                         
285                         // create AliRsnDaughter (and make Bayesian PID)
286                         AliRsnDaughter track;
287                         if (!track.Adopt(esdTrack, checkITSrefit)) continue;
288                         track.SetIndex(index);
289                         track.SetLabel(label);
290                         Identify(track);
291                         
292                         // store in TClonesArray
293                         event->AddTrack(track);
294                         nSelTracks++;
295                 }
296                 
297                 // compute total multiplicity
298                 event->ComputeMultiplicity();
299         
300                 // link to events tree and fill
301                 events->Fill();
302         }
303         
304         fileESD->Close();
305         return events;
306 }
307 //--------------------------------------------------------------------------------------------------------
308 TTree* AliRsnReader::ReadTracksAndParticles(const char *path, Option_t *option)
309 {
310 //
311 // Reads AliESD in a given path, getting also informations from Kinematics.
312 // In this case, the PID method used is the one selected with apposite setter.
313 // Allowed options (actually):
314 // - "R" : reject tracks which do not have the flag "kITSRefit" set to TRUE
315 // - "F" : reject 'fake' tracks (negative label)
316 // - "M" : use 'true' momentum instead of reconstructed one
317 //
318         TTree *events = new TTree("selection", "AliRsnEvents");
319         TTree::SetBranchStyle(1);
320         AliRsnEvent *event = new AliRsnEvent;
321         TBranch *branch = events->Branch("events", "AliRsnEvent", &event, 32000, 1);
322         branch->SetAutoDelete(kFALSE);
323         
324         // get path
325         TString strPath(path);
326         if (strPath.Length() < 1) return 0;
327         
328         // evaluate options
329         TString opt(option);
330         opt.ToUpper();
331         Bool_t checkITSrefit = (opt.Contains("R"));
332         Bool_t rejectFakes = (opt.Contains("F"));
333         Bool_t copyMomentum = (opt.Contains("M"));
334         
335         // opens ESD file
336         TFile *fileESD = TFile::Open(Form("%s/AliESDs.root", strPath.Data()));
337         if (!fileESD) return 0;
338         if (fileESD->IsOpen() == kFALSE) return 0;
339         TTree* treeESD = (TTree*)fileESD->Get("esdTree");
340         AliESD *esd = 0;
341         treeESD->SetBranchAddress("ESD", &esd);
342         Int_t nevRec = (Int_t)treeESD->GetEntries();
343         
344         // initialize run loader
345         AliRunLoader *runLoader = OpenRunLoader(path);
346         if (!runLoader) return 0;
347         Int_t nevSim = runLoader->GetNumberOfEvents();
348         
349         // check number of reconstructed and simulated events
350         if ( (nevSim != 0 && nevRec != 0) && (nevSim != nevRec) ) {
351                 cerr << "Count mismatch: sim = " << nevSim << " -- rec = " << nevRec << endl;
352                 return 0;
353         }
354         else if (nevSim == 0 && nevRec == 0) {
355                 cerr << "Count error: sim = rec = 0" << endl;
356                 return 0;
357         }
358         
359         // loop on events
360         Int_t i /*, procType*/, ntracks, nSelTracks = 0; 
361         Double_t vertex[3];
362         for (i = 0; i < nevRec; i++) {
363         
364                 // get event
365 //              cout << "\rEvent " << i << " " << flush;
366                 treeESD->GetEntry(i);
367                 runLoader->GetEvent(i);
368                                 
369                 // reject event if it is diffractive
370                 //procType = ((AliGenPythiaEventHeader*)runLoader->GetHeader()->GenEventHeader())->ProcessType();
371                 //if (procType == 92 || procType == 93 || procType == 94) {
372                 //      cout << "Skipping diffractive event" << endl;
373                 //      continue;
374                 //}
375                 
376                 // get particle stack
377                 AliStack *stack = runLoader->Stack();
378                                 
379                 // primary vertex
380                 vertex[0] = esd->GetVertex()->GetXv();
381                 vertex[1] = esd->GetVertex()->GetYv();
382                 vertex[2] = esd->GetVertex()->GetZv();
383                 
384                 // multiplicity
385                 ntracks = esd->GetNumberOfTracks();
386                 if (!ntracks) {
387                         Warning("ReadTracksAndParticles", "Event %d has no tracks!", i);
388                         continue;
389                 }
390                 
391                 // create new AliRsnEvent object
392                 event->Clear("DELETE");
393                 event->Init();
394                                 
395                 // store tracks from ESD
396                 Int_t index, label;
397                 Double_t vtot, v[3];
398                 for (index = 0; index < ntracks; index++) {
399                         
400                         // get track
401                         AliESDtrack *esdTrack = esd->GetTrack(index);
402                         
403                         // check against vertex constraint
404                         esdTrack->GetXYZ(v);
405                         vtot  = (v[0] - vertex[0])*(v[0] - vertex[0]);
406                         vtot += (v[1] - vertex[1])*(v[1] - vertex[1]);
407                         vtot += (v[2] - vertex[2])*(v[2] - vertex[2]);
408                         vtot  = TMath::Sqrt(vtot);
409                         if (vtot > fMaxRadius) continue;
410                         
411                         // check for fakes
412                         label = esdTrack->GetLabel();
413                         if (rejectFakes) {
414                                 if (label < 0) continue;
415                         }
416                         
417                         // create AliRsnDaughter (and make Bayesian PID)
418                         AliRsnDaughter track;
419                         if (!track.Adopt(esdTrack, checkITSrefit)) continue;
420                         track.SetIndex(index);
421                         
422                         // retrieve particle and get Kine info
423                         TParticle *part = stack->Particle(TMath::Abs(label));
424                         track.SetTruePDG(part->GetPdgCode());
425                         Int_t mother = part->GetFirstMother();
426                         track.SetMother(mother);
427                         if (mother >= 0) {
428                                 TParticle *mum = stack->Particle(mother);
429                                 track.SetMotherPDG(mum->GetPdgCode());
430                         }
431                         if (copyMomentum) track.SetPxPyPz(part->Px(), part->Py(), part->Pz());
432                         
433                         // identification
434                         Identify(track);
435                         
436                         // store in TClonesArray
437 //                      track.Print();
438                         event->AddTrack(track);
439                         nSelTracks++;
440                 }
441                 
442                 // compute total multiplicity
443                 event->ComputeMultiplicity();
444         
445                 // link to events tree and fill
446                 events->Fill();
447         }
448         
449         runLoader->UnloadKinematics();
450         delete runLoader;
451         fileESD->Close();
452         
453         return events;
454 }
455 //--------------------------------------------------------------------------------------------------------
456 void AliRsnReader::SetPriorProbabilities(Double_t *prior)
457 {
458 //
459 // Set prior probabilities to be used in case of ESD PID.
460 //
461         if (!prior) return;
462         
463         Int_t i = 0;
464         for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = prior[i];
465 }
466 //--------------------------------------------------------------------------------------------------------
467 void AliRsnReader::SetPriorProbability(AliPID::EParticleType type, Double_t value)
468 {
469 //
470 // Sets prior probability referred to a single particle species.
471 //
472         if (type >= AliPID::kElectron && type < AliPID::kPhoton) fPrior[type] = value;
473 }
474 //--------------------------------------------------------------------------------------------------------
475 AliPID::EParticleType AliRsnReader::FindType(Int_t pdg)
476 {
477 //
478 // Finds the enum value corresponding to a PDG code
479 //
480         pdg = TMath::Abs(pdg);
481         switch (pdg) {
482                 case   11: return AliPID::kElectron; break;
483                 case   13: return AliPID::kMuon; break;
484                 case  211: return AliPID::kPion; break;
485                 case  321: return AliPID::kKaon; break;
486                 case 2212: return AliPID::kProton; break;
487                 default  : return AliPID::kPhoton;
488         }
489 }
490 //--------------------------------------------------------------------------------------------------------
491 AliRunLoader* AliRsnReader::OpenRunLoader(const char *path)
492 {
493 //
494 // Open the Run loader with events in a given path
495 //
496         // clear gALICE
497         if (gAlice) {
498                 delete gAlice;
499                 gAlice = 0;
500         }
501         
502         // initialize run loader
503         TString name(path);
504         name += "/galice.root";
505         AliRunLoader *runLoader = AliRunLoader::Open(name.Data());
506         if (runLoader) {
507                 runLoader->LoadgAlice();
508                 gAlice = runLoader->GetAliRun();
509                 runLoader->LoadKinematics();
510                 runLoader->LoadHeader();
511         }
512         
513         return runLoader;
514 }