1b136a273c620896578f82e93c35c8824b596bba
[u/mrichter/AliRoot.git] / JETAN / AliAnalysisTaskJetBackgroundSubtract.cxx
1 // **************************************
2 // Task used for the correction of determiantion of reconstructed jet spectra
3 // Compares input (gen) and output (rec) jets   
4 // *******************************************
5
6
7 /**************************************************************************
8  * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
9  *                                                                        *
10  * Author: The ALICE Off-line Project.                                    *
11  * Contributors are mentioned in the code where appropriate.              *
12  *                                                                        *
13  * Permission to use, copy, modify and distribute this software and its   *
14  * documentation strictly for non-commercial purposes is hereby granted   *
15  * without fee, provided that the above copyright notice appears in all   *
16  * copies and that both the copyright notice and this permission notice   *
17  * appear in the supporting documentation. The authors make no claims     *
18  * about the suitability of this software for any purpose. It is          *
19  * provided "as is" without express or implied warranty.                  *
20  **************************************************************************/
21
22  
23 #include <TROOT.h>
24 #include <TH1F.h>
25 #include <TH2F.h>
26 #include <THnSparse.h>
27 #include <TSystem.h>
28 #include <TInterpreter.h>
29 #include <TList.h>
30 #include <TLorentzVector.h>
31 #include <TRefArray.h>
32 #include  "TDatabasePDG.h"
33
34 #include "AliAnalysisTaskJetBackgroundSubtract.h"
35 #include "AliAnalysisManager.h"
36 #include "AliAODHandler.h"
37 #include "AliAODTrack.h"
38 #include "AliAODJet.h"
39 #include "AliAODEvent.h"
40 #include "AliInputEventHandler.h"
41 #include "AliAODJetEventBackground.h"
42
43
44 ClassImp(AliAnalysisTaskJetBackgroundSubtract)
45
46 AliAnalysisTaskJetBackgroundSubtract::~AliAnalysisTaskJetBackgroundSubtract(){
47   delete fJBArray;
48   delete fOutJetArrayList;
49   delete fInJetArrayList;
50 }
51
52 AliAnalysisTaskJetBackgroundSubtract::AliAnalysisTaskJetBackgroundSubtract(): 
53   AliAnalysisTaskSE(),
54   fAODOut(0x0),
55   fAODIn(0x0),  
56   fAODExtension(0x0),
57   fJBArray(new TObjArray()),
58   fBackgroundBranch(""),
59   fNonStdFile(""),
60   fReplaceString1("B0"),
61   fReplaceString2("B%d"),
62   fSubtraction(kArea),
63   fInJetArrayList(0x0),
64   fOutJetArrayList(0x0),
65   fHistList(0x0)  
66 {
67
68 }
69
70 AliAnalysisTaskJetBackgroundSubtract::AliAnalysisTaskJetBackgroundSubtract(const char* name):
71
72   AliAnalysisTaskSE(name),
73   fAODOut(0x0),
74   fAODIn(0x0),  
75   fAODExtension(0x0),
76   fJBArray(new TObjArray()),
77   fBackgroundBranch(""),
78   fNonStdFile(""),
79   fReplaceString1("B0"),
80   fReplaceString2("B%d"),
81   fSubtraction(kArea),
82   fInJetArrayList(0x0),
83   fOutJetArrayList(0x0),
84   fHistList(0x0)  
85 {
86  DefineOutput(1, TList::Class());  
87 }
88
89
90
91 Bool_t AliAnalysisTaskJetBackgroundSubtract::Notify()
92 {
93   //
94   fAODIn = dynamic_cast<AliAODEvent*>(InputEvent());
95
96   ResetOutJets();
97
98   // Now we also have the Input Event Available! Fillvthe pointers in the list
99
100   fInJetArrayList->Clear();
101   fOutJetArrayList->Clear();
102
103   for(int iJB = 0;iJB<fJBArray->GetEntries();iJB++){
104     TObjString *ostr = (TObjString*)fJBArray->At(iJB);
105
106     TClonesArray* jarray = 0;      
107     if(!jarray&&fAODOut){
108       jarray = (TClonesArray*)(fAODOut->FindListObject(ostr->GetString().Data()));
109     }
110     if(!jarray&&fAODExtension){
111       jarray = (TClonesArray*)(fAODExtension->GetAOD()->FindListObject(ostr->GetString().Data()));
112     }
113     if(!jarray&&fAODIn){
114       jarray = (TClonesArray*)(fAODIn->FindListObject(ostr->GetString().Data()));
115     }
116
117     if(!jarray){
118       if(fDebug){
119         Printf("%s:%d Input jet branch %s not found",(char*)__FILE__,__LINE__,ostr->GetString().Data());
120       }
121       continue;
122     }
123     
124     TString newName(ostr->GetString().Data());
125     newName.ReplaceAll(fReplaceString1.Data(),Form(fReplaceString2.Data(),fSubtraction));
126     TClonesArray* jarrayOut = 0;      
127     if(!jarrayOut&&fAODOut){
128       jarrayOut = (TClonesArray*)(fAODOut->FindListObject(newName.Data()));
129     }
130     if(!jarrayOut&&fAODExtension){
131       jarrayOut = (TClonesArray*)(fAODExtension->GetAOD()->FindListObject(newName.Data()));
132     }
133
134     if(!jarrayOut){
135       if(fDebug){
136         Printf("%s:%d Output jet branch %s not found",(char*)__FILE__,__LINE__,newName.Data());
137         PrintAODContents();
138       }
139       continue;
140     }
141     if(jarrayOut&&jarray){
142       fOutJetArrayList->Add(jarrayOut);
143       fInJetArrayList->Add(jarray);
144     }
145   }
146   return kTRUE;
147 }
148
149 void AliAnalysisTaskJetBackgroundSubtract::UserCreateOutputObjects()
150 {
151
152   //
153   // Create the output container
154   //
155   // Connect the AOD
156
157   if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::UserCreateOutputObjects() \n");
158   if(fNonStdFile.Length()!=0){
159     // 
160     // case that we have an AOD extension we need to fetch the jets from the extended output
161     // we identifay the extension aod event by looking for the branchname
162     AliAODHandler *aodH = dynamic_cast<AliAODHandler*>(AliAnalysisManager::GetAnalysisManager()->GetOutputEventHandler());
163     fAODExtension = aodH->GetExtension(fNonStdFile.Data());
164     
165     if(!fAODExtension){
166       if(fDebug>1)Printf("AODExtension found for %s",fNonStdFile.Data());
167     }
168   }
169   fAODOut = AODEvent();
170
171   // usually we do not have the input already here
172
173   if(!fInJetArrayList)fInJetArrayList =new TList();
174   if(!fOutJetArrayList)fOutJetArrayList =new TList();
175
176   for(int iJB = 0;iJB<fJBArray->GetEntries();iJB++){
177     TObjString *ostr = (TObjString*)fJBArray->At(iJB);
178     TString newName(ostr->GetString().Data());
179     if(!newName.Contains(fReplaceString1.Data())){
180       Printf("%s:%d cannot replace string %s in %s",(char*)__FILE__,__LINE__,fReplaceString1.Data(),
181              newName.Data());
182       continue;
183     }
184
185     // add a new branch to the output for the background subtracted jets take the names from
186     // the input jets and replace the background flag names
187     TClonesArray *tca = new TClonesArray("AliAODJet", 0);
188     newName.ReplaceAll(fReplaceString1.Data(),Form(fReplaceString2.Data(),fSubtraction));
189     if(fDebug){
190       Printf("%s:%d created branch \n %s from \n %s",(char*)__FILE__,__LINE__,newName.Data(),
191              ostr->GetString().Data());
192     }
193     tca->SetName(newName.Data());
194     AddAODBranch("TClonesArray",&tca,fNonStdFile.Data());
195   }
196   
197
198   if(!fHistList)fHistList = new TList();
199   fHistList->SetOwner();
200
201   for(int iJB = 0;iJB<fJBArray->GetEntries();iJB++){
202     TObjString *ostr = (TObjString*)fJBArray->At(iJB);
203     TString oldName(ostr->GetString().Data()); 
204     TString newName(ostr->GetString().Data()); 
205     if(!newName.Contains(fReplaceString1.Data())){
206       Printf("%s:%d cannot replace string %s in %s",(char*)__FILE__,__LINE__,fReplaceString1.Data(),
207              newName.Data());
208       continue;
209     }
210     newName.ReplaceAll(fReplaceString1.Data(),Form(fReplaceString2.Data(),fSubtraction));
211     TH2F *hTmp = new TH2F(Form("h2PtInPtOut_%d",iJB),Form(";%s p_{T}; %s p_{T}",oldName.Data(),newName.Data()),200,0,200.,200,0.,200.);
212     fHistList->Add(hTmp);
213   }
214
215   Bool_t oldStatus = TH1::AddDirectoryStatus();
216   TH1::AddDirectory(kFALSE);
217
218   //
219   //  Histogram booking, add som control histograms here
220   //    
221
222   // =========== Switch on Sumw2 for all histos ===========
223   for (Int_t i=0; i<fHistList->GetEntries(); ++i) {
224     TH1 *h1 = dynamic_cast<TH1*>(fHistList->At(i));
225     if (h1){
226       h1->Sumw2();
227       continue;
228     }
229     THnSparse *hn = dynamic_cast<THnSparse*>(fHistList->At(i));
230     if(hn)hn->Sumw2();
231   }
232   TH1::AddDirectory(oldStatus);
233
234   if(fBackgroundBranch.Length()==0)
235     AliError(Form("%s:%d No BackgroundBranch defined",(char*)__FILE__,__LINE__));
236   if(fJBArray->GetEntries()==0)
237     AliError(Form("%s:%d No Jet Branches defined defined",(char*)__FILE__,__LINE__));
238 }
239
240 void AliAnalysisTaskJetBackgroundSubtract::Init()
241 {
242   //
243   // Initialization
244   //
245   if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::Init() \n");
246 }
247
248 void AliAnalysisTaskJetBackgroundSubtract::UserExec(Option_t */*option*/)
249 {
250
251   if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::UserExec() \n");
252   ResetOutJets();
253   if(fBackgroundBranch.Length()==0||fJBArray->GetEntries()==0){
254     if(fDebug)Printf("%s:%d No background subtraction done",(char*)__FILE__,__LINE__);
255     PostData(1,fHistList);
256   }
257   if(fJBArray->GetEntries()!=fInJetArrayList->GetEntries()){
258     if(fDebug)Printf("%s:%d different Array  sizes %d %d %d",(char*)__FILE__,__LINE__,fJBArray->GetEntries(),fInJetArrayList->GetEntries(),fOutJetArrayList->GetEntries());
259     PostData(1,fHistList);
260   }
261
262
263
264   static AliAODJetEventBackground*  evBkg = 0;
265   static TClonesArray*              bkgClusters = 0;
266   static TString bkgClusterName(fBackgroundBranch.Data());
267   bkgClusterName.ReplaceAll(Form("%s_",AliAODJetEventBackground::StdBranchName()),"");
268   if(!evBkg&&!bkgClusters&&fAODOut){
269     evBkg = (AliAODJetEventBackground*)(fAODOut->FindListObject(fBackgroundBranch.Data()));
270     bkgClusters = (TClonesArray*)(fAODOut->FindListObject(bkgClusterName.Data()));
271     if(fDebug&&bkgClusters)Printf("%s:%d Background cluster branch %s found",(char*)__FILE__,__LINE__,bkgClusterName.Data());
272     if(fDebug&&evBkg)Printf("%s:%d Backgroundbranch %s found",(char*)__FILE__,__LINE__,fBackgroundBranch.Data());
273   }
274   if(!evBkg&&!bkgClusters&&fAODExtension){
275     evBkg = (AliAODJetEventBackground*)(fAODExtension->GetAOD()->FindListObject(fBackgroundBranch.Data()));
276     bkgClusters = (TClonesArray*)(fAODExtension->GetAOD()->FindListObject(bkgClusterName.Data()));
277     if(fDebug&&bkgClusters)Printf("%s:%d Background cluster branch %s found",(char*)__FILE__,__LINE__,bkgClusterName.Data());
278     if(fDebug&&evBkg)Printf("%s:%d Backgroundbranch %s found",(char*)__FILE__,__LINE__,fBackgroundBranch.Data());
279   }
280
281   if(!evBkg&&!bkgClusters&&fAODIn){
282     evBkg = (AliAODJetEventBackground*)(fAODIn->FindListObject(fBackgroundBranch.Data()));
283     bkgClusters = (TClonesArray*)(fAODIn->FindListObject(bkgClusterName.Data()));
284     if(fDebug&&bkgClusters)Printf("%s:%d Background cluster branch %s found",(char*)__FILE__,__LINE__,bkgClusterName.Data());
285     if(fDebug&&evBkg)Printf("%s:%d Backgroundbranch %s found",(char*)__FILE__,__LINE__,fBackgroundBranch.Data());
286   }
287
288   if(!evBkg){
289     if(fDebug){
290       Printf("%s:%d Backgroundbranch %s not found",(char*)__FILE__,__LINE__,fBackgroundBranch.Data());
291       PrintAODContents();
292     }
293     PostData(1,fHistList);
294     return;
295   }
296
297   if(!bkgClusters){
298     if(fDebug){
299       Printf("%s:%d Background cluster branch %s not found",(char*)__FILE__,__LINE__,bkgClusterName.Data());
300       PrintAODContents();
301     }
302     PostData(1,fHistList);
303     return;
304   }
305
306
307   // LOOP over all jet branches and subtract the background
308
309   Float_t rho = 0;
310   if(fSubtraction==kArea)rho =  evBkg->GetBackground(2);
311   
312
313   for(int iJB = 0;iJB<fInJetArrayList->GetEntries();iJB++){
314     TClonesArray* jarray = (TClonesArray*)fInJetArrayList->At(iJB);
315     TClonesArray* jarrayOut = (TClonesArray*)fOutJetArrayList->At(iJB);
316     
317     if(!jarray||!jarrayOut){
318       Printf("%s:%d Array not found %d: %p %p",(char*)__FILE__,__LINE__,iJB,jarray,jarrayOut);
319       continue;
320     }
321     TH2F* h2PtInOut = (TH2F*)fHistList->FindObject(Form("h2PtInPtOut_%d",iJB));
322     // loop over all jets
323     Int_t nOut = 0;
324     for(int i = 0;i < jarray->GetEntriesFast();i++){
325       AliAODJet *jet = (AliAODJet*)jarray->At(i);
326       AliAODJet tmpNewJet(*jet);
327       Bool_t bAdd = false;
328       if(fSubtraction==kArea){  
329         Double_t background = rho * jet->EffectiveAreaCharged();
330         Float_t ptSub = jet->Pt() - background; 
331         if(fDebug>2){
332           Printf("%s:%d Jet %d %3.3f %3.3f",(char*)__FILE__,__LINE__,i,jet->Pt(),ptSub);
333         }
334         if(ptSub<0){
335           // optionally rescale it and keep??
336           bAdd = RescaleJetMomentum(&tmpNewJet,0.1);
337           if(h2PtInOut)h2PtInOut->Fill(jet->Pt(),0.1);
338         }
339         else{
340           bAdd = RescaleJetMomentum(&tmpNewJet,ptSub);
341           if(h2PtInOut)h2PtInOut->Fill(jet->Pt(),ptSub);
342         }
343         // add background estimates to the new jet object
344         // allows to recover old p_T and rho...
345         tmpNewJet.SetBgEnergy(background,0);
346       }// kAREA
347       else if(fSubtraction==kRhoRecalc1){
348         // Put a function call to calculate rho here
349         // * exclude edges
350         // * exclude clusters with small area
351         // * exclude areas around leading jets
352       }
353
354       if(bAdd){
355         AliAODJet *newJet = new ((*jarrayOut)[nOut++]) AliAODJet(tmpNewJet);
356         // what about track references, clear for now...
357         newJet->GetRefTracks()->Clear();
358       }
359
360    }
361
362
363
364     // subtract the background
365     
366
367     // remove jets??
368
369     // sort jets...
370
371   }
372   PostData(1, fHistList);
373 }
374
375 void AliAnalysisTaskJetBackgroundSubtract::Terminate(Option_t */*option*/)
376 {
377   // Terminate analysis
378   //
379   if (fDebug > 1) printf("AnalysisJetBackgroundSubtract: Terminate() \n");
380 }
381
382 Bool_t AliAnalysisTaskJetBackgroundSubtract::RescaleJetMomentum(AliAODJet *jet,Float_t pT){
383   // keep the direction and the jet mass
384   if(pT<=0)return kFALSE;
385   Double_t pTold = jet->Pt();
386   Double_t scale  = pT/pTold;
387   Double_t mass  = jet->M();
388   Double_t pNew = jet->P() * scale;
389   jet->SetPxPyPzE(scale*jet->Px(),scale*jet->Py(),scale*jet->Pz(),TMath::Sqrt(mass*mass+pNew*pNew));
390   return kTRUE;
391 }
392
393 void AliAnalysisTaskJetBackgroundSubtract::ResetOutJets(){
394   if(!fOutJetArrayList)return;
395   for(int iJB = 0;iJB<fOutJetArrayList->GetEntries();iJB++){
396     TClonesArray* jarray = (TClonesArray*)fOutJetArrayList->At(iJB);
397     if(jarray)jarray->Delete();
398   }
399 }
400
401
402 void AliAnalysisTaskJetBackgroundSubtract::PrintAODContents(){
403   if(fAODIn){
404     Printf("%s:%d >>>>>> Input",(char*)__FILE__,__LINE__);
405     fAODIn->Print();
406   }
407   if(fAODExtension){
408     Printf("%s:%d >>>>>> Extenstion",(char*)__FILE__,__LINE__);
409     fAODExtension->GetAOD()->Print();
410   }
411   if(fAODOut){
412     Printf("%s:%d >>>>>> Output",(char*)__FILE__,__LINE__);
413     fAODOut->Print();
414   }
415 }