Fixed warning
[u/mrichter/AliRoot.git] / ANALYSIS / AliPriorsTask.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 // Example of task (running locally) to calculate 
18 // the a priori concentration within the new class
19 // AliEventPoolLoop
20 //-----------------------------------------------------------------------
21
22
23 #include <TParticle.h>
24 #include <TLine.h>
25 #include <TLegend.h>
26 #include <TH1I.h>
27 #include <TChain.h>
28 #include <TH2D.h>
29
30 #include "AliPriorsTask.h"
31 #include "AliStack.h"
32 #include "AliMCEvent.h"
33 #include "AliAnalysisManager.h"
34 #include "AliESDEvent.h"
35 #include "AliESDtrack.h"
36 #include "AliLog.h"
37
38 ClassImp(AliPriorsTask)
39
40 //__________________________________________________________________________
41 AliPriorsTask::AliPriorsTask() :
42   fHistEventsProcessed(0x0),
43   fCalcPriors(0x0),
44   fNiterations(0x0)
45 {
46   //
47   //Default ctor
48   //
49   for(Int_t i=0; i< AliPID::kSPECIES; i++){
50   fPriors[i]=0;
51   fRecId[i]=0;
52   fMCId[i]=0;
53   }
54 }
55 //___________________________________________________________________________
56 AliPriorsTask::AliPriorsTask(const Char_t* name) :
57   AliAnalysisTaskSE(name),
58   fHistEventsProcessed(0x0),
59   fCalcPriors(0x0),
60   fNiterations(0)
61 {
62   //
63   // Constructor. Initialization of Inputs and Outputs
64   //
65   
66   for(Int_t i=0; i< AliPID::kSPECIES; i++){
67   fPriors[i]=0;
68   fRecId[i]=0;
69   fMCId[i]=0;
70   }
71   
72   
73   Info("AliPriorsTask","Calling Constructor");
74
75   /*
76     DefineInput(0) and DefineOutput(0)
77     are taken care of by AliAnalysisTaskSE constructor
78   */
79   DefineOutput(1,TH1I::Class());
80   DefineOutput(2,TH2D::Class());
81 }
82
83 //___________________________________________________________________________
84 AliPriorsTask& AliPriorsTask::operator=(const AliPriorsTask& c) 
85 {
86   //
87   // Assignment operator
88   //
89   if (this!=&c) {
90     AliAnalysisTaskSE::operator=(c) ;
91     fHistEventsProcessed = c.fHistEventsProcessed;
92     fCalcPriors = c.fCalcPriors;
93     for(Int_t i=0; i< AliPID::kSPECIES; i++){
94       fPriors[i]=c.fPriors[i];
95       fRecId[i]=c.fRecId[i];
96       fMCId[i]=c.fMCId[i];
97     }
98     fNiterations=c.fNiterations;
99   }
100   return *this;
101 }
102
103 //___________________________________________________________________________
104 AliPriorsTask::AliPriorsTask(const AliPriorsTask& c) :
105   AliAnalysisTaskSE(c),
106   fHistEventsProcessed(c.fHistEventsProcessed),
107   fCalcPriors(c.fCalcPriors),
108   fNiterations(c.fNiterations)
109 {
110   //
111   // Copy Constructor
112   //
113     for(Int_t i=0; i< AliPID::kSPECIES; i++){
114       fPriors[i]=c.fPriors[i];
115       fRecId[i]=c.fRecId[i];
116       fMCId[i]=c.fMCId[i];
117     }
118 }
119
120 //___________________________________________________________________________
121 AliPriorsTask::~AliPriorsTask() {
122   //
123   // Destructor
124   //
125   if (fHistEventsProcessed) delete fHistEventsProcessed ;
126   if (fCalcPriors) delete fCalcPriors ;
127 }
128
129 //_________________________________________________
130 void AliPriorsTask::UserExec(Option_t *)
131 {
132   //
133   // Main loop function
134   //
135   AliDebug(1,"UserExec") ;
136   
137   if(!fInputEvent) {
138     Printf("No input event!");
139     return;
140   }
141
142   AliESDEvent* fESD = dynamic_cast<AliESDEvent*>(fInputEvent);
143   if (!fESD) {
144     Error("UserExec","NO ESD FOUND!");
145     return;
146   }
147   
148   if (!fMCEvent) {
149     Error("UserExec","NO MC INFO FOUND");
150     return;
151   }
152   
153   //loop on tracks
154   for (Int_t iTrack = 0; iTrack<fESD->GetNumberOfTracks(); iTrack++) {
155
156     AliESDtrack* track = fESD->GetTrack(iTrack);
157     if (!track) continue;
158
159     //some quality/pid requirements
160     if (!track->IsOn(AliESDtrack::kITSrefit)) continue;
161     if (!track->IsOn(AliESDtrack::kESDpid))   continue;
162     if (!track->IsOn(AliESDtrack::kTPCpid))   continue;
163     if (!track->IsOn(AliESDtrack::kTOFpid))   continue;
164     
165     if(track->GetLabel()<0) continue;
166
167     Double_t esdpid[AliPID::kSPECIES]; track->GetESDpid(esdpid);
168
169     AliPID pid(esdpid,kTRUE); 
170     pid.SetPriors(fPriors);
171     AliPID::EParticleType id = pid.GetMostProbable();
172     
173     if ((Int_t)id < (Int_t)AliPID::kSPECIES)   fRecId[id]++; 
174     
175     if(fNiterations == 0 ){
176       Int_t pdgcode = TMath::Abs(((fMCEvent->GetTrack(track->GetLabel()))->Particle())->GetPdgCode());
177       switch(pdgcode){
178       case 11  :         fMCId[AliPID::kElectron]++; break;
179       case 13  :         fMCId[AliPID::kMuon]++;     break;
180       case 211 :         fMCId[AliPID::kPion]++;     break;
181       case 321 :         fMCId[AliPID::kKaon]++;     break;
182       case 2212:         fMCId[AliPID::kProton]++;   break;
183       default  :                                     break;
184       }
185     }
186   } 
187   
188
189   fHistEventsProcessed->Fill(0);
190
191   PostData(1,fHistEventsProcessed) ;
192   PostData(2,fCalcPriors) ;
193 }
194
195 //___________________________________________________________________________
196 void AliPriorsTask::Terminate(Option_t*) 
197 {
198   // The Terminate() function is the last function to be called during
199   // a query. It always runs on the client, it can be used to present
200   // the results graphically or save the results to file.
201   
202   AliAnalysisTaskSE::Terminate();
203   
204   Double_t MCsum =0; for(Int_t i=0; i< AliPID::kSPECIES; i++ ) MCsum+=fMCId[i];
205   if(MCsum>0) Printf(" MC Priors :  %f  %f  %f  %f  %f",fMCId[0]/MCsum,fMCId[1]/MCsum,fMCId[2]/MCsum,fMCId[3]/MCsum,fMCId[4]/MCsum);
206   Printf("    Priors :  %f  %f  %f  %f  %f \n",fPriors[0],fPriors[1],fPriors[2],fPriors[3],fPriors[4]);
207   
208   TH2D *histo= dynamic_cast<TH2D*> (GetOutputData(2));
209   for(Int_t ip = 0; ip < AliPID::kSPECIES; ip++) histo->Fill(fNiterMax-1,fPriors[ip]);
210   histo->SetMarkerStyle(8);
211   histo->DrawCopy();
212   
213   Int_t ltype[AliPID::kSPECIES]  = {1,2,3,2,1}; 
214   Int_t lcolor[AliPID::kSPECIES] = {1,2,8,38,46}; 
215   
216   TLine *l[AliPID::kSPECIES];
217   for(Int_t isp =0; isp < AliPID::kSPECIES; isp++) {
218     if(MCsum>0)l[isp]= new TLine(0.,fMCId[isp]/MCsum,fNiterMax,fMCId[isp]/MCsum);
219     l[isp]->SetLineWidth(isp+1);
220     l[isp]->SetLineColor(lcolor[isp]);
221     l[isp]->SetLineStyle(ltype[isp]);  
222     l[isp]->Draw("same");
223   }
224   
225   const char *part[AliPID::kSPECIES]={"e","#mu","#pi","K","p"};
226   TLegend *mcpriors = new TLegend(0.4,0.6,0.6,0.8);
227   for(Int_t ip=0; ip < AliPID::kSPECIES; ip++) mcpriors->AddEntry(l[ip],part[ip],"l");
228   mcpriors->Draw("same");
229   
230 }
231 //___________________________________________________________________________
232 void AliPriorsTask::UserCreateOutputObjects() 
233 {
234   //
235   //HERE ONE CAN CREATE OUTPUT OBJECTS 
236   //BEFORE THE EXECUTION OF THE TASK
237   //
238   OpenFile(1);
239   fHistEventsProcessed = new TH1I("fHistEventsProcessed","",1,0,1) ;
240   
241   OpenFile(2);
242   fCalcPriors = new TH2D("priors",Form("priors in %i iterations",fNiterMax),fNiterMax*10,-0.5,fNiterMax+0.5,100,0,1);  
243   fCalcPriors->SetXTitle("Iteration step");
244   fCalcPriors->SetYTitle("concentrations");
245 }
246 //________________________________________________________________________
247
248 Bool_t AliPriorsTask::NotifyBinChange()
249 {
250   //
251   //METHOD CALLED AT THE END OF THE EVENT LOOP. 
252   //HERE CONCENTRATIONS ARE CALCULATED, STORED
253   //AND USED IN THE NEXT EVENT LOOP
254   //
255   
256   Double_t sum =0;   
257   for(Int_t i=0; i< AliPID::kSPECIES; i++) {
258     sum+=fRecId[i];
259   }
260   if(sum == 0) return kFALSE;
261   
262   Double_t conc[AliPID::kSPECIES];
263   for(Int_t i=0; i< AliPID::kSPECIES; i++) {
264     conc[i]=fRecId[i]/sum; 
265     fCalcPriors->Fill(fNiterations,conc[i]);
266     fRecId[i]=0;
267   }
268   
269   SetPriors(conc);
270   fNiterations++;
271   
272   return kTRUE;
273 }