1775f4ddd800fa604b05ea76a2e8647250de2c53
[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 <THnSparse.h>
26 #include <TSystem.h>
27 #include <TInterpreter.h>
28 #include <TList.h>
29 #include <TLorentzVector.h>
30 #include  "TDatabasePDG.h"
31
32 #include "AliAnalysisTaskJetBackgroundSubtract.h"
33 #include "AliAnalysisManager.h"
34 #include "AliAODHandler.h"
35 #include "AliAODTrack.h"
36 #include "AliAODJet.h"
37 #include "AliAODEvent.h"
38 #include "AliInputEventHandler.h"
39 #include "AliAODJetEventBackground.h"
40
41
42 ClassImp(AliAnalysisTaskJetBackgroundSubtract)
43
44 AliAnalysisTaskJetBackgroundSubtract::~AliAnalysisTaskJetBackgroundSubtract(){
45   delete fJBArray;
46   delete fOutJetArrayList;
47   delete fInJetArrayList;
48 }
49
50 AliAnalysisTaskJetBackgroundSubtract::AliAnalysisTaskJetBackgroundSubtract(): 
51   AliAnalysisTaskSE(),
52   fAODOut(0x0),
53   fAODIn(0x0),  
54   fAODExtension(0x0),
55   fJBArray(new TObjArray()),
56   fBackgroundBranch(""),
57   fNonStdFile(""),
58   fReplaceString1("B0"),
59   fReplaceString2("B%d"),
60   fSubtraction(kArea),
61   fInJetArrayList(0x0),
62   fOutJetArrayList(0x0),
63   fHistList(0x0)  
64 {
65
66 }
67
68 AliAnalysisTaskJetBackgroundSubtract::AliAnalysisTaskJetBackgroundSubtract(const char* name):
69
70   AliAnalysisTaskSE(name),
71   fAODOut(0x0),
72   fAODIn(0x0),  
73   fAODExtension(0x0),
74   fJBArray(new TObjArray()),
75   fBackgroundBranch(""),
76   fNonStdFile(""),
77   fReplaceString1("B0"),
78   fReplaceString2("B%d"),
79   fSubtraction(kArea),
80   fInJetArrayList(0x0),
81   fOutJetArrayList(0x0),
82   fHistList(0x0)  
83 {
84  DefineOutput(1, TList::Class());  
85 }
86
87
88
89 Bool_t AliAnalysisTaskJetBackgroundSubtract::Notify()
90 {
91   //
92   return kTRUE;
93 }
94
95 void AliAnalysisTaskJetBackgroundSubtract::UserCreateOutputObjects()
96 {
97
98   //
99   // Create the output container
100   //
101   // Connect the AOD
102
103
104   if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::UserCreateOutputObjects() \n");
105   if(fNonStdFile.Length()!=0){
106     // 
107     // case that we have an AOD extension we need to fetch the jets from the extended output
108     // we identifay the extension aod event by looking for the branchname
109     AliAODHandler *aodH = dynamic_cast<AliAODHandler*>(AliAnalysisManager::GetAnalysisManager()->GetOutputEventHandler());
110     fAODExtension = aodH->GetExtension(fNonStdFile.Data());
111     
112     if(!fAODExtension){
113       if(fDebug>1)Printf("AODExtension found for %s",fNonStdFile.Data());
114     }
115   }
116   fAODOut = AODEvent();
117   fAODIn = dynamic_cast<AliAODEvent*>(InputEvent());
118
119
120   // collect the jet branches 
121
122   if(!fInJetArrayList)fInJetArrayList =new TList();
123   if(!fOutJetArrayList)fOutJetArrayList =new TList();
124
125   for(int iJB = 0;iJB<fJBArray->GetEntries();iJB++){
126     TClonesArray* jarray = 0;      
127     TObjString *ostr = (TObjString*)fJBArray->At(iJB);
128     if(!jarray&&fAODOut){
129       jarray = (TClonesArray*)(fAODOut->FindListObject(ostr->GetString().Data()));
130     }
131     if(!jarray&&fAODExtension){
132       jarray = (TClonesArray*)(fAODExtension->GetAOD()->FindListObject(ostr->GetString().Data()));
133     }
134     if(!jarray){
135       if(fDebug)Printf("%s:%d jet branch %s not found",(char*)__FILE__,__LINE__,ostr->GetString().Data());
136       continue;
137     }
138
139     TString newName(jarray->GetName());
140     if(!newName.Contains(fReplaceString1.Data())){
141       Printf("%s:%d cannot replace string %s in %s",(char*)__FILE__,__LINE__,fReplaceString1.Data(),
142              jarray->GetName());
143       continue;
144     }
145
146     fInJetArrayList->Add(jarray);
147
148     // add a new branch to the output for the background subtracted jets take the names from
149     // the input jets and replace the background flag names
150     TClonesArray *tca = new TClonesArray("AliAODJet", 0);
151
152     newName.ReplaceAll(fReplaceString1.Data(),Form(fReplaceString2.Data(),fSubtraction));
153     if(fDebug){
154       Printf("%s:%d created branch \n %s from \n %s",(char*)__FILE__,__LINE__,newName.Data(),
155              jarray->GetName());
156     }
157     tca->SetName(newName.Data());
158     AddAODBranch("TClonesArray",&tca,fNonStdFile.Data());
159     fOutJetArrayList->Add(tca);
160   }
161   
162
163
164
165   if(!fHistList)fHistList = new TList();
166   fHistList->SetOwner();
167
168   Bool_t oldStatus = TH1::AddDirectoryStatus();
169   TH1::AddDirectory(kFALSE);
170
171   //
172   //  Histogram booking, add som control histograms here
173   //    
174
175   // =========== Switch on Sumw2 for all histos ===========
176   for (Int_t i=0; i<fHistList->GetEntries(); ++i) {
177     TH1 *h1 = dynamic_cast<TH1*>(fHistList->At(i));
178     if (h1){
179       h1->Sumw2();
180       continue;
181     }
182     THnSparse *hn = dynamic_cast<THnSparse*>(fHistList->At(i));
183     if(hn)hn->Sumw2();
184   }
185   TH1::AddDirectory(oldStatus);
186
187   if(fBackgroundBranch.Length()==0)
188     AliError(Form("%s:%d No BackgroundBranch defined",(char*)__FILE__,__LINE__));
189   if(fJBArray->GetEntries()==0)
190     AliError(Form("%s:%d No Jet Branches defined defined",(char*)__FILE__,__LINE__));
191 }
192
193 void AliAnalysisTaskJetBackgroundSubtract::Init()
194 {
195   //
196   // Initialization
197   //
198   if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::Init() \n");
199 }
200
201 void AliAnalysisTaskJetBackgroundSubtract::UserExec(Option_t */*option*/)
202 {
203
204   ResetOutJets();
205   if(fBackgroundBranch.Length()==0||fJBArray->GetEntries()==0){
206     PostData(1,fHistList);
207   }
208     if (fDebug > 1) printf("AnalysisTaskJetBackgroundSubtract::UserExec() \n");
209
210
211   static AliAODJetEventBackground*  evBkg = 0;
212   if(!evBkg&&fAODOut){
213     evBkg = (AliAODJetEventBackground*)(fAODOut->FindListObject(fBackgroundBranch.Data()));
214   }
215   if(!evBkg&&fAODExtension){
216     evBkg = (AliAODJetEventBackground*)(fAODExtension->GetAOD()->FindListObject(fBackgroundBranch.Data()));
217   }
218   if(!evBkg&&fAODIn){
219     evBkg = (AliAODJetEventBackground*)(fAODIn->FindListObject(fBackgroundBranch.Data()));
220   }
221
222   if(!evBkg){
223     if(fDebug)Printf("%s:%d Backroundbranch %s not found",(char*)__FILE__,__LINE__,fBackgroundBranch.Data());
224     PostData(1,fHistList);
225   }
226
227
228   // LOOP over all jet branches and subtract the background
229
230   Float_t rho = 0;
231   if(fSubtraction==kArea)rho =  evBkg->GetBackground(2);
232   
233   for(int iJB = 0;iJB<fInJetArrayList->GetEntries();iJB++){
234     TClonesArray* jarray = (TClonesArray*)fInJetArrayList->At(iJB);
235     TClonesArray* jarrayOut = (TClonesArray*)fOutJetArrayList->At(iJB);
236     if(!jarray||!jarrayOut)continue;
237     
238     // loop over all jets
239     Int_t nOut = 0;
240     for(int i = 0;i < jarray->GetEntriesFast();i++){
241       AliAODJet *jet = (AliAODJet*)jarray->At(i);
242       AliAODJet tmpNewJet(*jet);
243       Bool_t bAdd = false;
244       if(fSubtraction==kArea){  
245         Float_t ptSub = jet->Pt() -  rho * jet->EffectiveAreaCharged();
246         if(fDebug>2){
247           Printf("%s:%d Jet %d %3.3f %3.3f",(char*)__FILE__,__LINE__,i,jet->Pt(),ptSub);
248         }
249         if(ptSub<0){
250           // optionally rescale it and keep??
251           bAdd = RescaleJetMomentum(&tmpNewJet,0.1);
252         }
253         else{
254           bAdd = RescaleJetMomentum(&tmpNewJet,ptSub);
255         }
256       }// kAREA
257
258       if(bAdd){
259         new ((*jarrayOut)[nOut++]) AliAODJet(tmpNewJet);
260         // optionally add background estimates to the jet object? CKB
261       }
262
263    }
264
265
266
267     // subtract the background
268     
269
270     // remove jets??
271
272     // sort jets...
273
274   }
275   PostData(1, fHistList);
276 }
277
278 void AliAnalysisTaskJetBackgroundSubtract::Terminate(Option_t */*option*/)
279 {
280   // Terminate analysis
281   //
282   if (fDebug > 1) printf("AnalysisJetBackgroundSubtract: Terminate() \n");
283 }
284
285 Bool_t AliAnalysisTaskJetBackgroundSubtract::RescaleJetMomentum(AliAODJet *jet,Float_t pT){
286   // keep the direction and the jet mass
287   if(pT<=0)return kFALSE;
288   Double_t pTold = jet->Pt();
289   Double_t scale  = pT/pTold;
290   Double_t mass  = jet->M();
291   Double_t pNew = jet->P() * scale;
292   jet->SetPxPyPzE(scale*jet->Px(),scale*jet->Py(),scale*jet->Pz(),TMath::Sqrt(mass*mass+pNew*pNew));
293   return kTRUE;
294 }
295
296 void AliAnalysisTaskJetBackgroundSubtract::ResetOutJets(){
297   if(!fOutJetArrayList)return;
298   for(int iJB = 0;iJB<fOutJetArrayList->GetEntries();iJB++){
299     TClonesArray* jarray = (TClonesArray*)fOutJetArrayList->At(iJB);
300     if(jarray)jarray->Delete();
301   }
302 }