preparation for recalculation of rho
[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   return kTRUE;
95 }
96
97 void AliAnalysisTaskJetBackgroundSubtract::UserCreateOutputObjects()
98 {
99
100   //
101   // Create the output container
102   //
103   // Connect the AOD
104
105
106   if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::UserCreateOutputObjects() \n");
107   if(fNonStdFile.Length()!=0){
108     // 
109     // case that we have an AOD extension we need to fetch the jets from the extended output
110     // we identifay the extension aod event by looking for the branchname
111     AliAODHandler *aodH = dynamic_cast<AliAODHandler*>(AliAnalysisManager::GetAnalysisManager()->GetOutputEventHandler());
112     fAODExtension = aodH->GetExtension(fNonStdFile.Data());
113     
114     if(!fAODExtension){
115       if(fDebug>1)Printf("AODExtension found for %s",fNonStdFile.Data());
116     }
117   }
118   fAODOut = AODEvent();
119   fAODIn = dynamic_cast<AliAODEvent*>(InputEvent());
120
121
122   // collect the jet branches 
123
124   if(!fInJetArrayList)fInJetArrayList =new TList();
125   if(!fOutJetArrayList)fOutJetArrayList =new TList();
126
127   for(int iJB = 0;iJB<fJBArray->GetEntries();iJB++){
128     TClonesArray* jarray = 0;      
129     TObjString *ostr = (TObjString*)fJBArray->At(iJB);
130     if(!jarray&&fAODOut){
131       jarray = (TClonesArray*)(fAODOut->FindListObject(ostr->GetString().Data()));
132     }
133     if(!jarray&&fAODExtension){
134       jarray = (TClonesArray*)(fAODExtension->GetAOD()->FindListObject(ostr->GetString().Data()));
135     }
136     if(!jarray){
137       if(fDebug)Printf("%s:%d jet branch %s not found",(char*)__FILE__,__LINE__,ostr->GetString().Data());
138       continue;
139     }
140
141     TString newName(jarray->GetName());
142     if(!newName.Contains(fReplaceString1.Data())){
143       Printf("%s:%d cannot replace string %s in %s",(char*)__FILE__,__LINE__,fReplaceString1.Data(),
144              jarray->GetName());
145       continue;
146     }
147
148     fInJetArrayList->Add(jarray);
149
150     // add a new branch to the output for the background subtracted jets take the names from
151     // the input jets and replace the background flag names
152     TClonesArray *tca = new TClonesArray("AliAODJet", 0);
153
154     newName.ReplaceAll(fReplaceString1.Data(),Form(fReplaceString2.Data(),fSubtraction));
155     if(fDebug){
156       Printf("%s:%d created branch \n %s from \n %s",(char*)__FILE__,__LINE__,newName.Data(),
157              jarray->GetName());
158     }
159     tca->SetName(newName.Data());
160     AddAODBranch("TClonesArray",&tca,fNonStdFile.Data());
161     fOutJetArrayList->Add(tca);
162   }
163   
164
165
166
167   if(!fHistList)fHistList = new TList();
168   fHistList->SetOwner();
169
170   for(int ij = 0;ij <  fInJetArrayList->GetEntries();ij++){
171     TH2F *hTmp = new TH2F(Form("h2PtInPtOut_%d",ij),Form(";%s p_{T};%s p_{T}",fInJetArrayList->At(ij)->GetName(),fOutJetArrayList->At(ij)->GetName()),200,0,200.,200,0.,200.);
172     fHistList->Add(hTmp);
173   }
174
175   Bool_t oldStatus = TH1::AddDirectoryStatus();
176   TH1::AddDirectory(kFALSE);
177
178   //
179   //  Histogram booking, add som control histograms here
180   //    
181
182   // =========== Switch on Sumw2 for all histos ===========
183   for (Int_t i=0; i<fHistList->GetEntries(); ++i) {
184     TH1 *h1 = dynamic_cast<TH1*>(fHistList->At(i));
185     if (h1){
186       h1->Sumw2();
187       continue;
188     }
189     THnSparse *hn = dynamic_cast<THnSparse*>(fHistList->At(i));
190     if(hn)hn->Sumw2();
191   }
192   TH1::AddDirectory(oldStatus);
193
194   if(fBackgroundBranch.Length()==0)
195     AliError(Form("%s:%d No BackgroundBranch defined",(char*)__FILE__,__LINE__));
196   if(fJBArray->GetEntries()==0)
197     AliError(Form("%s:%d No Jet Branches defined defined",(char*)__FILE__,__LINE__));
198 }
199
200 void AliAnalysisTaskJetBackgroundSubtract::Init()
201 {
202   //
203   // Initialization
204   //
205   if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::Init() \n");
206 }
207
208 void AliAnalysisTaskJetBackgroundSubtract::UserExec(Option_t */*option*/)
209 {
210
211   ResetOutJets();
212   if(fBackgroundBranch.Length()==0||fJBArray->GetEntries()==0){
213     PostData(1,fHistList);
214   }
215     if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::UserExec() \n");
216
217
218   static AliAODJetEventBackground*  evBkg = 0;
219   static TClonesArray*              bkgClusters = 0;
220   static TString bkgClusterName(fBackgroundBranch.Data());
221   if(!evBkg&&!bkgClusters&&fAODOut){
222     bkgClusterName.ReplaceAll("jeteventbackground_","");
223     evBkg = (AliAODJetEventBackground*)(fAODOut->FindListObject(fBackgroundBranch.Data()));
224     bkgClusters = (TClonesArray*)(fAODOut->FindListObject(bkgClusterName.Data()));
225   }
226   if(!evBkg&&!bkgClusters&&fAODExtension){
227     evBkg = (AliAODJetEventBackground*)(fAODExtension->GetAOD()->FindListObject(fBackgroundBranch.Data()));
228     bkgClusters = (TClonesArray*)(fAODExtension->GetAOD()->FindListObject(bkgClusterName.Data()));
229   }
230   if(!evBkg&&!bkgClusters&&fAODIn){
231     evBkg = (AliAODJetEventBackground*)(fAODIn->FindListObject(fBackgroundBranch.Data()));
232     bkgClusters = (TClonesArray*)(fAODOut->FindListObject(bkgClusterName.Data()));
233   }
234
235   if(!evBkg){
236     if(fDebug)Printf("%s:%d Backroundbranch %s not found",(char*)__FILE__,__LINE__,fBackgroundBranch.Data());
237     PostData(1,fHistList);
238     return;
239   }
240
241   if(!bkgClusters){
242     if(fDebug)Printf("%s:%d Backround cluster branch %s not found",(char*)__FILE__,__LINE__,bkgClusterName.Data());
243     PostData(1,fHistList);
244     return;
245   }
246
247
248   // LOOP over all jet branches and subtract the background
249
250   Float_t rho = 0;
251   if(fSubtraction==kArea)rho =  evBkg->GetBackground(2);
252   
253   for(int iJB = 0;iJB<fInJetArrayList->GetEntries();iJB++){
254     TClonesArray* jarray = (TClonesArray*)fInJetArrayList->At(iJB);
255     TClonesArray* jarrayOut = (TClonesArray*)fOutJetArrayList->At(iJB);
256     if(!jarray||!jarrayOut)continue;
257     TH2F* h2PtInOut = (TH2F*)fHistList->FindObject(Form("h2PtInPtOut_%d",iJB));
258     // loop over all jets
259     Int_t nOut = 0;
260     for(int i = 0;i < jarray->GetEntriesFast();i++){
261       AliAODJet *jet = (AliAODJet*)jarray->At(i);
262       AliAODJet tmpNewJet(*jet);
263       Bool_t bAdd = false;
264       if(fSubtraction==kArea){  
265         Double_t background = rho * jet->EffectiveAreaCharged();
266         Float_t ptSub = jet->Pt() - background; 
267         if(fDebug>2){
268           Printf("%s:%d Jet %d %3.3f %3.3f",(char*)__FILE__,__LINE__,i,jet->Pt(),ptSub);
269         }
270         if(ptSub<0){
271           // optionally rescale it and keep??
272           bAdd = RescaleJetMomentum(&tmpNewJet,0.1);
273           if(h2PtInOut)h2PtInOut->Fill(jet->Pt(),0.1);
274         }
275         else{
276           bAdd = RescaleJetMomentum(&tmpNewJet,ptSub);
277           if(h2PtInOut)h2PtInOut->Fill(jet->Pt(),ptSub);
278         }
279         // add background estimates to the new jet object
280         // allows to recover old p_T and rho...
281         tmpNewJet.SetBgEnergy(background,0);
282       }// kAREA
283       else if(fSubtraction==kRhoRecalc1){
284         // Put a function call to calculate rho here
285         // * exclude edges
286         // * exclude clusters with small area
287         // * exclude areas around leading jets
288       }
289
290       if(bAdd){
291         AliAODJet *newJet = new ((*jarrayOut)[nOut++]) AliAODJet(tmpNewJet);
292         // what about track references, clear for now...
293         newJet->GetRefTracks()->Clear();
294       }
295
296    }
297
298
299
300     // subtract the background
301     
302
303     // remove jets??
304
305     // sort jets...
306
307   }
308   PostData(1, fHistList);
309 }
310
311 void AliAnalysisTaskJetBackgroundSubtract::Terminate(Option_t */*option*/)
312 {
313   // Terminate analysis
314   //
315   if (fDebug > 1) printf("AnalysisJetBackgroundSubtract: Terminate() \n");
316 }
317
318 Bool_t AliAnalysisTaskJetBackgroundSubtract::RescaleJetMomentum(AliAODJet *jet,Float_t pT){
319   // keep the direction and the jet mass
320   if(pT<=0)return kFALSE;
321   Double_t pTold = jet->Pt();
322   Double_t scale  = pT/pTold;
323   Double_t mass  = jet->M();
324   Double_t pNew = jet->P() * scale;
325   jet->SetPxPyPzE(scale*jet->Px(),scale*jet->Py(),scale*jet->Pz(),TMath::Sqrt(mass*mass+pNew*pNew));
326   return kTRUE;
327 }
328
329 void AliAnalysisTaskJetBackgroundSubtract::ResetOutJets(){
330   if(!fOutJetArrayList)return;
331   for(int iJB = 0;iJB<fOutJetArrayList->GetEntries();iJB++){
332     TClonesArray* jarray = (TClonesArray*)fOutJetArrayList->At(iJB);
333     if(jarray)jarray->Delete();
334   }
335 }