Derive from AliAnalysisTaskSE.
[u/mrichter/AliRoot.git] / JETAN / AliJetESDReader.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 // Jet ESD Reader 
18 // ESD reader for jet analysis
19 // Authors: Mercedes Lopez Noriega (mercedes.lopez.noriega@cern.ch)
20 //          Magali Estienne <magali.estienne@IReS.in2p3.fr>
21 //------------------------------------------------------------------------- 
22
23 // --- Standard library ---
24 #include <Riostream.h>
25
26 // --- ROOT system ---
27 #include <TSystem.h>
28 #include <TStopwatch.h>
29 #include <TLorentzVector.h>
30 #include <TVector3.h>
31 #include <TTask.h>
32 #include <TGeoManager.h>
33 #include <assert.h>
34 #include <TRefArray.h>
35 #include <TMath.h>
36 #include <TChain.h>
37
38
39 // --- AliRoot header files ---
40 #include "AliJetESDReader.h"
41 #include "AliJetESDReaderHeader.h"
42 #include "AliESDEvent.h"
43 #include "AliESD.h"
44 #include "AliESDtrack.h"
45 #include "AliJetDummyGeo.h"
46 #include "AliJetFillUnitArrayTracks.h"
47 #include "AliJetFillUnitArrayEMCalDigits.h"
48 #include "AliJetUnitArray.h"
49
50 ClassImp(AliJetESDReader)
51
52 AliJetESDReader::AliJetESDReader():
53   AliJetReader(),  
54   fGeom(0),
55   fChain(0x0),
56   fESD(0x0),
57   fHadCorr(0x0),
58   fTpcGrid(0x0),
59   fEmcalGrid(0x0),
60   fGrid0(0),
61   fGrid1(0),
62   fGrid2(0),
63   fGrid3(0),
64   fGrid4(0),
65   fPtCut(0),
66   fHCorrection(0),
67   fNumUnits(0),
68   fDebug(0),
69   fMass(0),
70   fSign(0),
71   fNIn(0),
72   fOpt(0),
73   fDZ(0),
74   fNeta(0),
75   fNphi(0),
76   fArrayInitialised(0)
77 {
78   // Constructor 
79 }
80
81 //____________________________________________________________________________
82 AliJetESDReader::~AliJetESDReader()
83 {
84   // Destructor
85     delete fChain;
86     delete fESD;
87     delete fTpcGrid;
88     delete fEmcalGrid;
89     if(fDZ)
90       {
91         delete fGrid0;
92         delete fGrid1;
93         delete fGrid2;
94         delete fGrid3;
95         delete fGrid4;
96       }
97 }
98
99 //____________________________________________________________________________
100 void AliJetESDReader::OpenInputFiles()
101 {
102   // chain for the ESDs
103   fChain   = new TChain("esdTree");
104
105   // get directory and pattern name from the header
106    const char* dirName=fReaderHeader->GetDirectory();
107    const char* pattern=fReaderHeader->GetPattern();
108     
109    // Add files matching patters to the chain
110    void *dir  = gSystem->OpenDirectory(dirName);
111    const char *name = 0x0;
112    int nesd = ((AliJetESDReaderHeader*) fReaderHeader)->GetNesd();
113    int a = 0;
114    while ((name = gSystem->GetDirEntry(dir))){
115        if (a>=nesd) continue;
116        
117        if (strstr(name,pattern)){
118          printf("Adding %s\n",name);
119          char path[256];
120          //        sprintf(path,"%s/%s/AliESDs.root",dirName,name);
121          sprintf(path,"%s/%s/AliESDs.root",dirName,name);
122          fChain->AddFile(path);
123          a++;
124        }
125    }
126   
127   gSystem->FreeDirectory(dir);
128   
129   fESD = new AliESDEvent();
130   fESD->ReadFromTree(fChain);
131
132   int nMax = fChain->GetEntries(); 
133
134   printf("\n AliJetESDReader: Total number of events in chain= %d \n",nMax);
135   
136   // set number of events in header
137   if (fReaderHeader->GetLastEvent() == -1)
138     fReaderHeader->SetLastEvent(nMax);
139   else {
140     Int_t nUsr = fReaderHeader->GetLastEvent();
141     fReaderHeader->SetLastEvent(TMath::Min(nMax,nUsr));
142   }
143 }
144
145 //____________________________________________________________________________
146 void AliJetESDReader::SetInputEvent(TObject* esd, TObject* /*aod*/, TObject* /*mc*/) {
147     // Connect the tree
148      fESD   = (AliESDEvent*) esd;
149 }
150
151 //____________________________________________________________________________
152 Bool_t AliJetESDReader::FillMomentumArray()
153 {
154   // Fill momentum array
155
156   Int_t goodTrack = 0;
157   Int_t nt = 0;
158   Float_t pt, eta;
159   TVector3 p3;
160   
161   // clear momentum array
162   ClearArray();
163   fDebug = fReaderHeader->GetDebug();
164   fOpt = fReaderHeader->GetDetector();
165
166   if (!fESD) {
167       return kFALSE;
168   }
169
170   // get number of tracks in event (for the loop)
171   nt = fESD->GetNumberOfTracks();
172   printf("Fill Momentum Array %5d  \n", nt);
173   
174   // temporary storage of signal and pt cut flag
175   Int_t* sflag  = new Int_t[nt];
176   Int_t* cflag  = new Int_t[nt];
177   
178   // get cuts set by user
179   Float_t ptMin =  fReaderHeader->GetPtCut();
180   Float_t etaMin = fReaderHeader->GetFiducialEtaMin();
181   Float_t etaMax = fReaderHeader->GetFiducialEtaMax();  
182   
183   //loop over tracks
184   for (Int_t it = 0; it < nt; it++) {
185       AliESDtrack *track = fESD->GetTrack(it);
186       UInt_t status = track->GetStatus();
187       
188       Double_t mom[3];
189       track->GetPxPyPz(mom);
190       p3.SetXYZ(mom[0],mom[1],mom[2]);
191       pt = p3.Pt();
192       if ((status & AliESDtrack::kTPCrefit) == 0)    continue;      // quality check
193       if ((status & AliESDtrack::kITSrefit) == 0)    continue;      // quality check
194       if (((AliJetESDReaderHeader*) fReaderHeader)->ReadSignalOnly() 
195           && TMath::Abs(track->GetLabel()) > 10000)  continue;      // quality check
196       if (((AliJetESDReaderHeader*) fReaderHeader)->ReadBkgdOnly() 
197           && TMath::Abs(track->GetLabel()) < 10000)  continue;      // quality check
198       eta = p3.Eta();
199       if ( (eta > etaMax) || (eta < etaMin))         continue;      // checking eta cut
200       
201       new ((*fMomentumArray)[goodTrack]) TLorentzVector(p3,p3.Mag());
202       sflag[goodTrack]=0;
203       if (TMath::Abs(track->GetLabel()) < 10000) sflag[goodTrack]=1;
204       cflag[goodTrack]=0;
205       if (pt > ptMin) cflag[goodTrack]=1;                           // pt cut
206       goodTrack++;
207   }
208   // set the signal flags
209   fSignalFlag.Set(goodTrack,sflag);
210   fCutFlag.Set(goodTrack,cflag);
211
212   delete[] sflag;
213   delete[] cflag;
214
215   return kTRUE;
216 }
217
218 //____________________________________________________________________________
219 void AliJetESDReader::CreateTasks()
220 {
221   fDebug = fReaderHeader->GetDebug();
222   fDZ = fReaderHeader->GetDZ();
223
224   // Init EMCAL geometry and create UnitArray object
225   SetEMCALGeometry();
226   InitParameters();
227   InitUnitArray();
228
229   fFillUnitArray = new TTask("fFillUnitArray","Fill unit array jet finder");
230   fFillUAFromTracks = new AliJetFillUnitArrayTracks(); 
231   fFillUAFromTracks->SetReaderHeader(fReaderHeader);
232   fFillUAFromTracks->SetGeom(fGeom);
233   fFillUAFromTracks->SetTPCGrid(fTpcGrid);
234   fFillUAFromTracks->SetEMCalGrid(fEmcalGrid);
235
236   if(fDZ)
237     {
238       fFillUAFromTracks->SetGrid0(fGrid0);
239       fFillUAFromTracks->SetGrid1(fGrid1);
240       fFillUAFromTracks->SetGrid2(fGrid2);
241       fFillUAFromTracks->SetGrid3(fGrid3);
242       fFillUAFromTracks->SetGrid4(fGrid4);
243     }
244   fFillUAFromTracks->SetHadCorrection(fHCorrection);
245   fFillUAFromTracks->SetHadCorrector(fHadCorr);
246   fFillUAFromEMCalDigits = new AliJetFillUnitArrayEMCalDigits();
247   fFillUAFromEMCalDigits->SetReaderHeader(fReaderHeader);
248   fFillUAFromEMCalDigits->SetGeom(fGeom);
249   fFillUAFromEMCalDigits->SetTPCGrid(fTpcGrid);
250   fFillUAFromEMCalDigits->SetEMCalGrid(fEmcalGrid);
251   //      fFillUnitArray->Add(fFillUAFromTracks);
252   fFillUnitArray->Add(fFillUAFromEMCalDigits);
253   fFillUAFromTracks->SetActive(kFALSE);
254   fFillUAFromEMCalDigits->SetActive(kFALSE);
255
256   cout << "Tasks instantiated at that stage ! " << endl;
257   cout << "You can loop over events now ! " << endl;
258    
259 }
260
261 //____________________________________________________________________________
262 //void AliJetESDReader::ExecTasks(Int_t event)
263 Bool_t AliJetESDReader::ExecTasks(Int_t /*event*/)
264 {
265   // clear momentum array
266   Int_t nEntRef = fRefArray->GetEntries();
267
268   for(Int_t i=0; i<nEntRef; i++)
269     {      
270       ((AliJetUnitArray*)fRefArray->At(i))->SetUnitTrackID(0);
271       ((AliJetUnitArray*)fRefArray->At(i))->SetUnitEnergy(0.);
272       ((AliJetUnitArray*)fRefArray->At(i))->SetUnitCutFlag(kPtSmaller);
273       ((AliJetUnitArray*)fRefArray->At(i))->SetUnitDetectorFlag(kTpc);
274       ((AliJetUnitArray*)fRefArray->At(i))->SetUnitFlag(kOutJet);
275     }
276
277   ClearArray();
278
279   fDebug = fReaderHeader->GetDebug();
280   fOpt = fReaderHeader->GetDetector();
281   //  InitParameters();
282
283   if(!fESD) {
284     return kFALSE;
285   }
286   
287   /*
288   // get event from chain
289   // For TSelectors
290   //  fChain->GetTree()->GetEntry(event);
291   // For interactive process
292   //  fChain->GetEntry(event);
293   fChain->GetEvent(event);
294   */
295
296   // TPC only or Digits+TPC or Clusters+TPC
297   if(fOpt%2==!0 && fOpt!=0){ 
298     fFillUAFromTracks->SetESD(fESD);
299     fFillUAFromTracks->SetActive(kTRUE);
300     fFillUAFromTracks->SetUnitArray(fUnitArray);
301     fFillUAFromTracks->SetRefArray(fRefArray);
302     //    fFillUAFromTracks->ExecuteTask("tpc"); // => Temporarily changed  !!!
303                                                  // Incompatibility with Andrei's analysis framework
304     fFillUAFromTracks->Exec("tpc");
305     if(fOpt==1){
306       fNumCandidate = fFillUAFromTracks->GetMult();
307       fNumCandidateCut = fFillUAFromTracks->GetMultCut();
308     }
309   }
310
311   // Digits only or Digits+TPC  
312   if(fOpt>=2 && fOpt<=3){
313     fFillUAFromEMCalDigits->SetESD(fESD);
314     fFillUAFromEMCalDigits->SetActive(kTRUE);
315     fFillUAFromEMCalDigits->SetUnitArray(fUnitArray);
316     fFillUAFromEMCalDigits->SetRefArray(fRefArray);
317     fFillUAFromEMCalDigits->SetInitMult(fFillUAFromTracks->GetMult());
318     fFillUAFromEMCalDigits->SetInitMultCut(fFillUAFromTracks->GetMultCut());
319     fFillUAFromEMCalDigits->Exec("digits"); // => Temporarily changed !!!
320     fNumCandidate = fFillUAFromEMCalDigits->GetMult();
321     fNumCandidateCut = fFillUAFromEMCalDigits->GetMultCut();
322   }
323
324   //  fFillUnitArray->ExecuteTask(); // => Temporarily commented
325
326   return kTRUE;
327 }
328
329 //____________________________________________________________________________
330 void AliJetESDReader::SetEMCALGeometry()
331 {
332   // Define EMCAL geometry to be able to read ESDs
333     fGeom = AliJetDummyGeo::GetInstance();
334     if (fGeom == 0)
335         fGeom = AliJetDummyGeo::GetInstance("SHISH_77_TRD1_2X2_FINAL_110DEG","EMCAL");
336
337     // To be setted to run some AliEMCALGeometry functions
338     TGeoManager::Import("geometry.root");
339     //    fGeom->GetTransformationForSM();  
340     printf("\n EMCal Geometry set ! \n");
341
342 }
343
344
345 //____________________________________________________________________________  
346 void AliJetESDReader::InitParameters()
347 {
348     // Initialise parameters
349     fHCorrection    = 0;                 // For hadron correction
350     fHadCorr        = 0;                 // For hadron correction
351     fNumUnits       = fGeom->GetNCells();      // Number of cells in EMCAL
352     if(fDebug>1) printf("\n EMCal parameters initiated ! \n");
353 }
354
355 //____________________________________________________________________________
356 void AliJetESDReader::InitUnitArray()
357 {
358   //Initialises unit arrays
359   Int_t nElements = fTpcGrid->GetNEntries();
360   Float_t eta = 0., phi = 0., Deta = 0., Dphi = 0.;
361   if(fArrayInitialised) fUnitArray->Delete();
362
363   if(fTpcGrid->GetGridType()==0)
364     { // Fill the following quantities :
365       // Good track ID, (Eta,Phi) position ID, eta, phi, energy, px, py, pz, Deta, Dphi, 
366       // detector flag, in/out jet, pt cut, mass, cluster ID)
367       for(Int_t nBin = 1; nBin < nElements+1; nBin++)
368         {
369           //      fTpcGrid->GetEtaPhiFromIndex2(nBin,eta,phi);
370           fTpcGrid->GetEtaPhiFromIndex2(nBin,phi,eta);
371           phi = ((phi < 0) ? phi + 2. * TMath::Pi() : phi);
372           Deta = fTpcGrid->GetDeta();
373           Dphi = fTpcGrid->GetDphi();
374           new ((*fUnitArray)[nBin-1]) AliJetUnitArray(nBin-1,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
375         }
376     }
377
378   if(fTpcGrid->GetGridType()==1)
379     {
380       Int_t nGaps = 0;
381       Int_t n0 = 0, n1 = 0, n2 = 0, n3 = 0, n4 = 0;
382
383       if(fDZ)
384         {
385           // Define a grid of cell for the gaps between SM
386           Double_t phimin0 = 0., phimin1 = 0., phimin2 = 0., phimin3 = 0., phimin4 = 0.;
387           Double_t phimax0 = 0., phimax1 = 0., phimax2 = 0., phimax3 = 0., phimax4 = 0.;
388           fGeom->GetPhiBoundariesOfSMGap(0,phimin0,phimax0);
389           fGrid0 = new AliJetGrid(0,95,phimin0,phimax0,-0.7,0.7); // 0.015 x 0.015
390           fGrid0->SetGridType(0);
391           fGrid0->SetMatrixIndexes();
392           fGrid0->SetIndexIJ();
393           n0 = fGrid0->GetNEntries();
394           fGeom->GetPhiBoundariesOfSMGap(1,phimin1,phimax1);
395           fGrid1 = new AliJetGrid(0,95,phimin1,phimax1,-0.7,0.7); // 0.015 x 0.015
396           fGrid1->SetGridType(0);
397           fGrid1->SetMatrixIndexes();
398           fGrid1->SetIndexIJ();
399           n1 = fGrid1->GetNEntries();
400           fGeom->GetPhiBoundariesOfSMGap(2,phimin2,phimax2);
401           fGrid2 = new AliJetGrid(0,95,phimin2,phimax2,-0.7,0.7); // 0.015 x 0.015
402           fGrid2->SetGridType(0);
403           fGrid2->SetMatrixIndexes();
404           fGrid2->SetIndexIJ();
405           n2 = fGrid2->GetNEntries();
406           fGeom->GetPhiBoundariesOfSMGap(3,phimin3,phimax3);
407           fGrid3 = new AliJetGrid(0,95,phimin3,phimax3,-0.7,0.7); // 0.015 x 0.015
408           fGrid3->SetGridType(0);  
409           fGrid3->SetMatrixIndexes();
410           fGrid3->SetIndexIJ();
411           n3 = fGrid3->GetNEntries();
412           fGeom->GetPhiBoundariesOfSMGap(4,phimin4,phimax4);
413           fGrid4 = new AliJetGrid(0,95,phimin4,phimax4,-0.7,0.7); // 0.015 x 0.015
414           fGrid4->SetGridType(0);
415           fGrid4->SetMatrixIndexes();
416           fGrid4->SetIndexIJ();
417           n4 = fGrid4->GetNEntries();
418
419           if(fDebug>1) 
420             {
421               cout << "n0 cells: " << n0 << "phimin0: " << phimin0 << ", phimax0: " << phimax0 << endl;
422               cout << "n1 cells: " << n1 << "phimin1: " << phimin1 << ", phimax1: " << phimax1 << endl;
423               cout << "n2 cells: " << n2 << "phimin2: " << phimin2 << ", phimax2: " << phimax2 << endl;
424               cout << "n3 cells: " << n3 << "phimin3: " << phimin3 << ", phimax3: " << phimax3 << endl;
425               cout << "n4 cells: " << n4 << "phimin4: " << phimin4 << ", phimax4: " << phimax4 << endl;
426             }
427           
428           nGaps = n0+n1+n2+n3+n4;
429
430         }
431
432       for(Int_t nBin = 0; nBin < fNumUnits+nElements+nGaps; nBin++) 
433         {
434           if(nBin<fNumUnits)
435             {
436               fGeom->EtaPhiFromIndex(nBin, eta, phi); // From EMCal geometry 
437               // fEmcalGrid->GetEtaPhiFromIndex2(nBin,phi,eta); // My function from Grid
438               phi = ((phi < 0) ? phi + 2. * TMath::Pi() : phi);
439               Deta = fEmcalGrid->GetDeta(); // Modify with the exact detector values
440               Dphi = fEmcalGrid->GetDphi(); // Modify with the exact detector values
441               new ((*fUnitArray)[nBin]) AliJetUnitArray(nBin,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
442             } 
443           else {
444             if(nBin>=fNumUnits && nBin<fNumUnits+nElements){
445               fTpcGrid->GetEtaPhiFromIndex2(nBin+1-fNumUnits,phi,eta);
446               phi = ((phi < 0) ? phi + 2. * TMath::Pi() : phi);
447               Deta = fTpcGrid->GetDeta();
448               Dphi = fTpcGrid->GetDphi();
449               new ((*fUnitArray)[nBin]) AliJetUnitArray(nBin,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
450             }
451             else {
452               if(fDZ) {
453                 if(nBin>=fNumUnits+nElements && nBin<fNumUnits+nElements+nGaps){
454                   if(nBin<fNumUnits+nElements+n0)
455                     {
456                       Float_t phi = eta = 0.;
457                       fGrid0->GetEtaPhiFromIndex2(nBin+1-(fNumUnits+nElements),phi,eta);
458                       Deta = fGrid0->GetDeta(); 
459                       Dphi = fGrid0->GetDphi(); 
460                       new ((*fUnitArray)[nBin]) AliJetUnitArray(nBin,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
461                     }
462                   else if(nBin>=fNumUnits+nElements+n0 && nBin<fNumUnits+nElements+n0+n1)
463                     {
464                       Float_t phi = eta = 0.;
465                       fGrid1->GetEtaPhiFromIndex2(nBin+1-(fNumUnits+nElements+n0),phi,eta);
466                       Deta = fGrid1->GetDeta(); 
467                       Dphi = fGrid1->GetDphi(); 
468                       new ((*fUnitArray)[nBin]) AliJetUnitArray(nBin,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
469                     }
470                   else if(nBin>=fNumUnits+nElements+n0+n1 && nBin<fNumUnits+nElements+n0+n1+n2)
471                     {
472                       Float_t phi = eta = 0.;
473                       fGrid2->GetEtaPhiFromIndex2(nBin+1-(fNumUnits+nElements+n0+n1),phi,eta);
474                       Deta = fGrid2->GetDeta(); 
475                       Dphi = fGrid2->GetDphi(); 
476                       new ((*fUnitArray)[nBin]) AliJetUnitArray(nBin,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
477                     }
478                   else if(nBin>=fNumUnits+nElements+n0+n1+n2 && nBin<fNumUnits+nElements+n0+n1+n2+n3)
479                     {
480                       Float_t phi = eta = 0.;
481                       fGrid3->GetEtaPhiFromIndex2(nBin+1-(fNumUnits+nElements+n0+n1+n2),phi,eta);
482                       Deta = fGrid3->GetDeta(); 
483                       Dphi = fGrid3->GetDphi(); 
484                       new ((*fUnitArray)[nBin]) AliJetUnitArray(nBin,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
485                     }
486                   else if(nBin>=fNumUnits+nElements+n0+n1+n2+n3 && nBin<fNumUnits+nElements+nGaps)
487                     {
488                       Float_t phi = eta = 0.;
489                       fGrid4->GetEtaPhiFromIndex2(nBin+1-(fNumUnits+nElements+n0+n1+n2+n3),phi,eta);
490                       Deta = fGrid4->GetDeta(); 
491                       Dphi = fGrid4->GetDphi(); 
492                       new ((*fUnitArray)[nBin]) AliJetUnitArray(nBin,0,eta,phi,0.,0.,0.,0.,Deta,Dphi,kTpc,kOutJet,kPtSmaller,0.,-1);
493                     }
494                 }
495               } // end if(fDZ)
496             } // end else 2
497           } // end else 1
498         } // end loop on nBin
499     } // end grid type == 1
500   fArrayInitialised = 1;
501 }
502
503
504