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