Adding the task to add additional track or MC branches to the AOD, FastEmbedding...
[u/mrichter/AliRoot.git] / JETAN / AliAnalysisTaskFastEmbedding.cxx
1 /*************************************************************************
2  *                                                                       *
3  * Task for fast embedding                                               *
4  * read extra input from AOD                                             *
5  *                                                                       *
6  *************************************************************************/
7
8
9 /**************************************************************************
10  * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
11  *                                                                        *
12  * Author: The ALICE Off-line Project.                                    *
13  * Contributors are mentioned in the code where appropriate.              *
14  *                                                                        *
15  * Permission to use, copy, modify and distribute this software and its   *
16  * documentation strictly for non-commercial purposes is hereby granted   *
17  * without fee, provided that the above copyright notice appears in all   *
18  * copies and that both the copyright notice and this permission notice   *
19  * appear in the supporting documentation. The authors make no claims     *
20  * about the suitability of this software for any purpose. It is          *
21  * provided "as is" without express or implied warranty.                  *
22  **************************************************************************/
23
24 /* $Id: */
25
26 #include <TFile.h>
27 #include <TTree.h>
28 #include <TChain.h>
29 #include <TClonesArray.h>
30 #include <TDirectory.h>
31 #include <TSystem.h>
32 #include <TRef.h>
33 #include <TRandom3.h>
34 #include <TH1F.h>
35 #include <TH2F.h>
36
37
38 #include "AliAnalysisTaskFastEmbedding.h"
39 #include "AliAnalysisManager.h"
40 #include "AliAODEvent.h"
41 #include "AliAODTrack.h"
42 #include "AliAODJet.h"
43 #include "AliAODMCParticle.h"
44
45 #include "AliLog.h"
46
47 ClassImp(AliAnalysisTaskFastEmbedding)
48
49 //__________________________________________________________________________
50 AliAnalysisTaskFastEmbedding::AliAnalysisTaskFastEmbedding()
51     : AliAnalysisTaskSE()
52       ,fAODout(0)
53       ,fAODevent(0)
54       ,fAODtree(0)
55       ,fAODfile(0)
56       ,rndm(new TRandom3())
57       ,fAODPathArray(0)
58       ,fAODPath("AliAOD.root")
59       ,fTrackBranch("aodExtraTracks")
60       ,fMCparticlesBranch("aodExtraMCparticles")
61       ,fEntry(0)
62       ,fJobId(0)
63       ,fNEntriesPerJob(1000)
64       ,fEmbedMode(0)
65       ,fEvtSelecMode(0)
66       ,fEvtSelMinJetPt(-1)
67       ,fEvtSelMaxJetPt(-1)
68       ,fToyMinNbOfTracks(1)
69       ,fToyMaxNbOfTracks(1)
70       ,fToyMinTrackPt(50.)
71       ,fToyMaxTrackPt(50.)
72       ,fToyDistributionTrackPt(0.)
73       ,fToyMinTrackEta(-.5)
74       ,fToyMaxTrackEta(.5)
75       ,fToyMinTrackPhi(0.)
76       ,fToyMaxTrackPhi(2*TMath::Pi())
77       ,fToyFilterMap(0)
78       ,fHistList(0)
79       ,fh1TrackPt(0)
80       ,fh2TrackEtaPhi(0)
81       ,fh1TrackN(0)
82       ,fh1MCTrackPt(0)
83       ,fh2MCTrackEtaPhi(0)
84       ,fh1MCTrackN(0)
85
86
87 {
88     // default constructor
89
90 }
91
92
93 //__________________________________________________________________________
94 AliAnalysisTaskFastEmbedding::AliAnalysisTaskFastEmbedding(const char *name)
95     : AliAnalysisTaskSE(name)
96       ,fAODout(0)
97       ,fAODevent(0)
98       ,fAODtree(0)
99       ,fAODfile(0)
100       ,rndm(new TRandom3())
101       ,fAODPathArray(0)
102       ,fAODPath("AliAOD.root")
103       ,fTrackBranch("aodExtraTracks")
104       ,fMCparticlesBranch("aodExtraMCparticles")
105       ,fEntry(0)
106       ,fJobId(0)
107       ,fNEntriesPerJob(1000)
108       ,fEmbedMode(0)
109       ,fEvtSelecMode(0)
110       ,fEvtSelMinJetPt(-1)
111       ,fEvtSelMaxJetPt(-1)
112       ,fToyMinNbOfTracks(1)
113       ,fToyMaxNbOfTracks(1)
114       ,fToyMinTrackPt(50.)
115       ,fToyMaxTrackPt(50.)
116       ,fToyDistributionTrackPt(0.)
117       ,fToyMinTrackEta(-.5)
118       ,fToyMaxTrackEta(.5)
119       ,fToyMinTrackPhi(0.)
120       ,fToyMaxTrackPhi(2*TMath::Pi())
121       ,fToyFilterMap(0)
122       ,fHistList(0)
123       ,fh1TrackPt(0)
124       ,fh2TrackEtaPhi(0)
125       ,fh1TrackN(0)
126       ,fh1MCTrackPt(0)
127       ,fh2MCTrackEtaPhi(0)
128       ,fh1MCTrackN(0)
129 {
130     // constructor
131     DefineOutput(1, TList::Class());
132 }
133
134
135 //__________________________________________________________________________
136 AliAnalysisTaskFastEmbedding::AliAnalysisTaskFastEmbedding(const AliAnalysisTaskFastEmbedding &copy)
137     : AliAnalysisTaskSE()
138       ,fAODout(copy.fAODout)
139       ,fAODevent(copy.fAODevent)
140       ,fAODtree(copy.fAODtree)
141       ,fAODfile(copy.fAODfile)
142       ,rndm(copy.rndm)
143       ,fAODPathArray(copy.fAODPathArray)
144       ,fAODPath(copy.fAODPath)
145       ,fTrackBranch(copy.fTrackBranch)
146       ,fMCparticlesBranch(copy.fMCparticlesBranch)
147       ,fEntry(copy.fEntry)
148       ,fJobId(copy.fJobId)
149       ,fNEntriesPerJob(copy.fNEntriesPerJob)
150       ,fEmbedMode(copy.fEmbedMode)
151       ,fEvtSelecMode(copy.fEvtSelecMode)
152       ,fEvtSelMinJetPt(copy.fEvtSelMinJetPt)
153       ,fEvtSelMaxJetPt(copy.fEvtSelMaxJetPt)
154       ,fToyMinNbOfTracks(copy.fToyMinNbOfTracks)
155       ,fToyMaxNbOfTracks(copy.fToyMaxNbOfTracks)
156       ,fToyMinTrackPt(copy.fToyMinTrackPt)
157       ,fToyMaxTrackPt(copy.fToyMaxTrackPt)
158       ,fToyDistributionTrackPt(copy.fToyDistributionTrackPt)
159       ,fToyMinTrackEta(copy.fToyMinTrackEta)
160       ,fToyMaxTrackEta(copy.fToyMaxTrackEta)
161       ,fToyMinTrackPhi(copy.fToyMinTrackPhi)
162       ,fToyMaxTrackPhi(copy.fToyMaxTrackPhi)
163       ,fToyFilterMap(copy.fToyFilterMap)
164       ,fHistList(copy.fHistList)
165       ,fh1TrackPt(copy.fh1TrackPt)
166       ,fh2TrackEtaPhi(copy.fh2TrackEtaPhi)
167       ,fh1TrackN(copy.fh1TrackN)
168       ,fh1MCTrackPt(copy.fh1MCTrackPt)
169       ,fh2MCTrackEtaPhi(copy.fh2MCTrackEtaPhi)
170       ,fh1MCTrackN(copy.fh1MCTrackN)
171 {
172     // copy constructor
173 }
174
175
176 //__________________________________________________________________________
177 AliAnalysisTaskFastEmbedding& AliAnalysisTaskFastEmbedding::operator=(const AliAnalysisTaskFastEmbedding& o)
178 {
179     // assignment
180
181     if(this!=&o){
182         AliAnalysisTaskSE::operator=(o);
183         fAODout            = o.fAODout;
184         fAODevent          = o.fAODevent;
185         fAODtree           = o.fAODtree;
186         fAODfile           = o.fAODfile;
187         rndm               = o.rndm;
188         fAODPathArray       = o.fAODPathArray;
189         fAODPath           = o.fAODPath;
190         fTrackBranch       = o.fTrackBranch;
191         fMCparticlesBranch = o.fMCparticlesBranch;
192         fEntry             = o.fEntry;
193         fJobId             = o.fJobId;
194         fNEntriesPerJob    = o.fNEntriesPerJob;
195         fEmbedMode         = o.fEmbedMode;
196         fEvtSelecMode      = o.fEvtSelecMode;
197         fEvtSelMinJetPt    = o.fEvtSelMinJetPt;
198         fEvtSelMaxJetPt    = o.fEvtSelMaxJetPt;
199         fToyMinNbOfTracks  = o.fToyMinNbOfTracks;
200         fToyMaxNbOfTracks  = o.fToyMaxNbOfTracks;
201         fToyMinTrackPt     = o.fToyMinTrackPt;
202         fToyMaxTrackPt     = o.fToyMaxTrackPt;
203         fToyDistributionTrackPt = o.fToyDistributionTrackPt;
204         fToyMinTrackEta    = o.fToyMinTrackEta;
205         fToyMaxTrackEta    = o.fToyMaxTrackEta;
206         fToyMinTrackPhi    = o.fToyMinTrackPhi;
207         fToyMaxTrackPhi    = o.fToyMaxTrackPhi;
208         fToyFilterMap      = o.fToyFilterMap;
209         fHistList          = o.fHistList;
210         fh1TrackPt         = o.fh1TrackPt;
211         fh2TrackEtaPhi     = o.fh2TrackEtaPhi;
212         fh1TrackN          = o.fh1TrackN;
213         fh1MCTrackPt       = o.fh1MCTrackPt;
214         fh2MCTrackEtaPhi   = o.fh2MCTrackEtaPhi;
215         fh1MCTrackN        = o.fh1MCTrackN;
216     }
217
218     return *this;
219 }
220
221
222 //__________________________________________________________________________
223 AliAnalysisTaskFastEmbedding::~AliAnalysisTaskFastEmbedding()
224 {
225     // destructor
226     delete rndm;
227 }
228
229
230 //__________________________________________________________________________
231 void AliAnalysisTaskFastEmbedding::UserCreateOutputObjects()
232 {
233     // create output objects
234     if(fDebug > 1) Printf("AliAnalysisTaskFastEmbedding::UserCreateOutputObjects()");
235
236
237
238     // embed mode with AOD
239     if(fEmbedMode==kAODFull || fEmbedMode==kAODJetTracks || fEmbedMode==kAODJet4Mom){
240
241        // open input AOD
242        if(fAODPathArray){
243            Int_t nFiles = fAODPathArray->GetEntries();
244            Int_t n = rndm->Integer(nFiles);
245            AliInfo(Form("Select file %d", n));
246            TObjString *objStr = (TObjString*) fAODPathArray->At(n);
247            if(!objStr){
248               AliError("Could not get path of aod file.");
249               return;
250            }
251            fAODPath = objStr->GetString();
252        }
253
254        TDirectory *owd = gDirectory;
255        fAODfile = TFile::Open(fAODPath.Data());
256        owd->cd();
257        if(!fAODfile){
258                AliError("Could not open AOD file.");
259                return;
260        }
261
262        fAODtree = (TTree*)fAODfile->Get("aodTree");
263            
264        if(!fAODtree){
265            AliError("AOD tree not found.");
266            return;
267        }
268        fAODevent = new AliAODEvent();
269        fAODevent->ReadFromTree(fAODtree);
270        if(!fAODevent){
271            AliError("AOD event not found.");
272            return;
273        }
274     } //end: embed mode with AOD
275
276
277     // connect output aod
278     // create a new branch for extra tracks
279     fAODout = AODEvent();
280     if(!fAODout){
281         AliError("Output AOD not found.");
282         return;
283     }
284     if(!fAODout->FindListObject(fTrackBranch.Data()) && strlen(fTrackBranch.Data())){
285         AliInfo(Form("Add AOD branch %s", fTrackBranch.Data()));
286         TClonesArray *tracks = new TClonesArray("AliAODTrack",0);
287         tracks->SetName(fTrackBranch.Data());
288         AddAODBranch("TClonesArray", &tracks);
289     }
290     // create new branch for extra mcparticle if available as input
291     if(fAODevent && fAODevent->FindListObject("mcparticles") && strlen(fMCparticlesBranch.Data())){
292        AliInfo(Form("Add AOD branch %s", fMCparticlesBranch.Data()));
293        TClonesArray *mcparticles = new TClonesArray("AliAODMCParticle",0);
294        mcparticles->SetName(fMCparticlesBranch.Data());
295        AddAODBranch("TClonesArray", &mcparticles);
296     }
297
298
299
300
301     //qa histograms
302
303     OpenFile(1);
304     fHistList = new TList();
305
306     Bool_t oldStatus = TH1::AddDirectoryStatus();
307     TH1::AddDirectory(kFALSE);
308
309     fh1TrackPt      =  new TH1F("fh1TrackPt","pT of extra tracks;p_{T};entries", 300, 0., 300.);
310     fh2TrackEtaPhi  =  new TH2F("fh2TrackEtaPhi","eta-phi distribution of extra tracks;#eta;#phi", 20, -1., 1., 60, 0., 2*TMath::Pi());
311     fh1TrackN       =  new TH1F("fh1TrackN", "nb. of extra tracks per event;nb. of tracks;entries",300, 0., 3000.);
312
313     fHistList->Add(fh1TrackPt);
314     fHistList->Add(fh2TrackEtaPhi);
315     fHistList->Add(fh1TrackN);
316
317  
318     if(fAODevent && fAODevent->FindListObject("mcparticles") && strlen(fMCparticlesBranch.Data())){ 
319
320        fh1MCTrackPt      =  new TH1F("fh1MCTrackPt","pT of MC extra tracks;p_{T};entries", 300, 0., 300.);
321        fh2MCTrackEtaPhi  =  new TH2F("fh2MCTrackEtaPhi","eta-phi distribution of MC extra tracks;#eta;#phi", 20, -1., 1., 60, 0., 2*TMath::Pi());
322        fh1MCTrackN       =  new TH1F("fh1MCTrackN", "nb. of MC extra tracks per event;nb. of tracks;entries",300, 0., 3000.);
323
324        fHistList->Add(fh1MCTrackPt);
325        fHistList->Add(fh2MCTrackEtaPhi);
326        fHistList->Add(fh1MCTrackN);
327
328     }
329
330     TH1::AddDirectory(oldStatus);
331
332 }
333
334
335 //__________________________________________________________________________
336 void AliAnalysisTaskFastEmbedding::Init()
337 {
338     // Initialization
339     if(fDebug > 1) Printf("AliAnalysisTaskFastEmbedding::Init()");
340
341
342     // set seed for rndm according to sub-job id  *not implemented yet*
343 }
344
345
346 //__________________________________________________________________________
347 void AliAnalysisTaskFastEmbedding::UserExec(Option_t *)
348 {
349     if(fDebug > 1) Printf("AliAnalysisTaskFastEmbedding::UserExec()");
350
351     if(!fAODout){
352        AliError("Need output AOD, but is not connected."); 
353        return;
354     }
355
356     // connect aod out
357     TClonesArray *tracks = (TClonesArray*)(fAODout->FindListObject(fTrackBranch.Data()));
358     if(!tracks){
359         AliError("Extra track branch not found in output.");
360         return;
361     }
362     tracks->Delete();
363     Int_t nAODtracks=0;
364
365
366     TRef dummy;
367
368     // === embed mode with AOD ===
369     if(fEmbedMode==kAODFull || fEmbedMode==kAODJetTracks || fEmbedMode==kAODJet4Mom){
370        if(!fAODevent){
371           AliError("Need input AOD, but is not connected."); 
372           return;
373        }
374
375        Int_t maxEntry = fEntry+fNEntriesPerJob-1;
376        Int_t nEvents = fAODtree->GetEntries();
377        if(maxEntry>nEvents) maxEntry=nEvents;
378
379        Bool_t useEntry = kFALSE;
380        while(!useEntry){  // protection need, if no event fulfills requierment
381           if(fEntry>maxEntry) fEntry=fJobId*fNEntriesPerJob;
382           fAODtree->GetEvent(fEntry);
383
384           // jet pt selection
385           if(fEvtSelecMode==kEventsJetPt){
386              Int_t nJets = fAODevent->GetNJets();
387              for(Int_t ij=0; ij<nJets; ++ij){
388                  AliAODJet *jet = fAODevent->GetJet(ij);
389
390                  if(   (fEvtSelMinJetPt==-1. || jet->Pt()>=fEvtSelMinJetPt)
391                     && (fEvtSelMaxJetPt==-1. || jet->Pt()<=fEvtSelMaxJetPt)){
392                     useEntry = kTRUE;
393                     break;
394                  } 
395              }
396           }
397
398           // no selection
399           if(fEvtSelecMode==kEventsAll){
400              useEntry = kTRUE;
401           }
402
403           fEntry++;
404        }
405        AliInfo(Form("Use entry %d from extra AOD.", fEntry-1));
406
407
408        TClonesArray *mcpartIN  = (TClonesArray*)(fAODevent->FindListObject("mcparticles"));
409        TClonesArray *mcpartOUT = 0x0;
410        if(mcpartIN){
411            mcpartOUT = (TClonesArray*)(fAODout->FindListObject(fMCparticlesBranch.Data()));
412            mcpartOUT->Delete();
413        } else {
414            AliInfo("No extra MC particles found.");
415        }
416     
417
418        if(fEmbedMode==kAODFull || fEmbedMode==kAODJetTracks){ // take all tracks or jet tracks
419            // loop over input tracks
420            // add to output aod
421            Int_t nTracks = fAODevent->GetNumberOfTracks();
422            fh1TrackN->Fill((Float_t)nTracks);
423
424            for(Int_t it=0; it<nTracks; ++it){
425                AliAODTrack *tmpTr = fAODevent->GetTrack(it);
426                // ?? test filter bit, or something else??
427
428                if(fEmbedMode==kAODJetTracks){
429                   // jet track? else continue
430                   // not implemented yet
431                   continue;
432                } 
433
434                new ((*tracks)[nAODtracks++]) AliAODTrack(*tmpTr); 
435                dummy = (*tracks)[nAODtracks-1];
436
437                fh1TrackPt->Fill(tmpTr->Pt());
438                fh2TrackEtaPhi->Fill(tmpTr->Eta(), tmpTr->Phi());
439            }
440
441            if(mcpartIN){
442                Int_t nMCpart = mcpartIN->GetEntriesFast();
443
444                Int_t nPhysicalPrimary=0;
445                Int_t nAODmcpart=0;
446                for(Int_t ip=0; ip<nMCpart; ++ip){
447                    AliAODMCParticle *tmpPart = (AliAODMCParticle*) mcpartIN->At(ip);
448
449                    if(fEmbedMode==kAODJetTracks){
450                       // jet track? else continue
451                       // not implemented yet
452                       continue;
453                    } 
454
455                    new((*mcpartOUT)[nAODmcpart++]) AliAODMCParticle(*tmpPart);
456                    dummy = (*mcpartOUT)[nAODmcpart-1];
457
458                    if(tmpPart->IsPhysicalPrimary() && tmpPart->Charge()!=0. && tmpPart->Charge()!=-99. ){
459                       fh1MCTrackPt->Fill(tmpPart->Pt());
460                       fh2MCTrackEtaPhi->Fill(tmpPart->Eta(), tmpPart->Phi());
461                       nPhysicalPrimary++;
462                    }
463                }
464                fh1MCTrackN->Fill((Float_t)nPhysicalPrimary);
465                
466            }
467        } // end: embed all tracks or jet tracks
468
469
470        if(fEmbedMode==kAODJet4Mom){
471
472            // loop over jets
473            Int_t nJets = fAODevent->GetNJets();
474            fh1TrackN->Fill((Float_t)nJets);
475            for(Int_t ij=0; ij<nJets; ++ij){
476                AliAODJet *jet = fAODevent->GetJet(ij);
477                AliAODTrack *tmpTr = (AliAODTrack*)jet;
478
479                new ((*tracks)[nAODtracks++]) AliAODTrack(*tmpTr);
480                dummy = (*tracks)[nAODtracks-1]; 
481
482                fh1TrackPt->Fill(tmpTr->Pt());
483                fh2TrackEtaPhi->Fill(tmpTr->Eta(), tmpTr->Phi());
484            }
485
486        } // end: embed jets as 4-momenta
487
488
489     } //end: embed mode with AOD
490
491
492     // === embed mode with toy ===
493     if(fEmbedMode==kToyTracks){
494         Int_t nT = (Int_t)(rndm->Uniform(fToyMinNbOfTracks, fToyMaxNbOfTracks)+0.5);
495
496         fh1TrackN->Fill((Float_t)nT);
497
498         for(Int_t i=0; i<nT; ++i){
499
500            Double_t pt = -1.;
501            if(fToyMinTrackPt!=fToyMaxTrackPt){
502               if(fToyDistributionTrackPt==0){
503                  pt = rndm->Uniform(fToyMinTrackPt, fToyMaxTrackPt);
504               } else {
505                  while(pt<fToyMinTrackPt||pt>fToyMaxTrackPt){
506                     pt = rndm->Exp(fToyDistributionTrackPt);   // no power law yet!!
507                     pt += fToyMinTrackPt;
508                  }
509               }
510            } else {
511               pt = fToyMinTrackPt;
512            }
513            Double_t eta = rndm->Uniform(fToyMinTrackEta,fToyMaxTrackEta);
514            Double_t phi = rndm->Uniform(fToyMinTrackPhi,fToyMaxTrackPhi);
515            phi = TVector2::Phi_0_2pi(phi);
516
517            Printf("Add track #%d: pt %.2f, eta %.2f, phi %.2f", i, pt, eta, phi);
518
519            Double_t theta = 2*TMath::ATan(1/TMath::Exp(eta));
520            Float_t mom[3];
521            mom[0] = pt;
522            mom[1] = phi;
523            mom[2] = theta;
524
525            Float_t xyz[3];
526            xyz[0] = -999.;
527            xyz[1] = -999.;
528            xyz[2] = -999.;
529         
530            AliAODTrack *tmpTr = new AliAODTrack( -999,   // id
531                                               -999,   // label
532                                               mom,    // momentum[3]
533                                               kFALSE, // cartesian
534                                               xyz,    // position
535                                               kFALSE, // DCA
536                                               NULL,   // covMatrix,
537                                               -99,    // charge
538                                               0,      // itsClusMap
539                                               NULL,   // pid 
540                                               NULL,   // prodVertex
541                                               kFALSE, // used for vertex fit
542                                               kFALSE, // used for prim vtx fit
543                                               AliAODTrack::kUndef, // type
544                                               fToyFilterMap,  // select info
545                                               -999.    // chi2 per NDF
546                                             );
547            tmpTr->SetFlags(1);
548
549            new((*tracks)[nAODtracks++]) AliAODTrack(*tmpTr);
550            dummy = (*tracks)[nAODtracks-1];
551
552            fh1TrackPt->Fill(pt);
553            fh2TrackEtaPhi->Fill(eta,phi);
554
555            delete tmpTr;
556         }
557     } //end: embed mode with toy
558
559
560     PostData(1, fHistList);
561 }
562
563
564 //__________________________________________________________________________
565 void AliAnalysisTaskFastEmbedding::Terminate(Option_t *)
566 {
567     if(fDebug > 1) Printf("AliAnalysisTaskFastEmbedding::Terminate()");
568
569     if(fAODfile && fAODfile->IsOpen())
570     fAODfile->Close();  
571
572 }
573
574 //__________________________________________________________________________
575 /* NEEDS TO BE TESTED
576 Int_t AliAnalysisTaskFastEmbedding::GetJobID()
577 {
578    Int_t id=0;
579
580    const char* env = gSystem->Getenv("ALIENCOUNTER"); // GRID
581    if(!env || !strlen(env)) env = gSystem->Getenv("LSB_JOBINDEX"); // GSI
582
583    if(env && strlen(env)){
584        id= atoi(env);
585        AliInfo(Form("Job index %d", id));
586    }
587    else{
588        AliInfo("Job index not found. Okay if running locally.");
589        Int_t nEvents = fAODtree->GetEntries();
590        fNEntriesPerJob = nEvents;
591        AliInfo(Form("Asuming single job, set entries per job to maximum %d", fNEntriesPerJob));
592    }
593
594    return id;
595 }*/
596
597 //__________________________________________________________________________
598
599
600