ca358bdcef914673fb1e5e052e99c25197bafd93
[u/mrichter/AliRoot.git] / PWG3 / muon / AliAnalysisTaskSingleMu.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-2007, 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 /// \class AliAnalysisTaskSingleMu
18 /// Analysis task for single muons in the spectrometer.
19 /// The output is a tree with:
20 ///  - pt, y and phi of the muon
21 ///  - z position of primary vertex
22 ///  - transverse distance at vertex (DCA)
23 ///  - matched trigger
24 ///
25 /// \author Diego Stocco
26 //-----------------------------------------------------------------------------
27
28 //----------------------------------------------------------------------------
29 //    Implementation of the single muon analysis class
30 // An example of usage can be found in the macro RunSingleMuAnalysisFromAOD.C.
31 //----------------------------------------------------------------------------
32
33 #define AliAnalysisTaskSingleMu_cxx
34
35 // ROOT includes
36 #include "TChain.h"
37 #include "TROOT.h"
38 #include "TLorentzVector.h"
39 #include "TCanvas.h"
40
41 // STEER includes
42 #include "AliLog.h"
43
44 #include "AliAODEvent.h"
45 #include "AliAODTrack.h"
46 #include "AliAODVertex.h"
47 #include "AliAODInputHandler.h"
48
49 #include "AliAnalysisTask.h"
50 #include "AliAnalysisDataSlot.h"
51 #include "AliAnalysisManager.h"
52 #include "AliAnalysisTaskSingleMu.h"
53
54 /// \cond CLASSIMP
55 ClassImp(AliAnalysisTaskSingleMu) // Class implementation in ROOT context
56 /// \endcond
57
58 //________________________________________________________________________
59 AliAnalysisTaskSingleMu::AliAnalysisTaskSingleMu(const char *name) :
60   AliAnalysisTask(name,""), 
61   fAOD(0),
62   fResults(0),
63   fVarFloat(0),
64   fVarInt(0),
65   fFloatVarName(0),
66   fIntVarName(0)
67 {
68   //
69   /// Constructor.
70   //
71   InitVariables();
72   // Input slot #0 works with an Ntuple
73   DefineInput(0, TChain::Class());
74   // Output slot #0 writes into a TTree container
75   DefineOutput(0,  TTree::Class());
76 }
77
78 //___________________________________________________________________________
79 void AliAnalysisTaskSingleMu::ConnectInputData(Option_t *) 
80 {
81   //
82   /// Connect AOD here
83   /// Called once
84   //
85
86   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
87   if (!tree) {
88     Printf("ERROR: Could not read chain from input slot 0");
89   } else {
90     /*
91     // Disable all branches and enable only the needed ones
92     // The next two lines are different when data produced as AliAODEvent is read
93     tree->SetBranchStatus("*", kFALSE);
94     tree->SetBranchStatus("fTracks.*", kTRUE);
95     */
96
97     AliAODInputHandler *aodH = dynamic_cast<AliAODInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
98
99     if (!aodH) {
100       Printf("ERROR: Could not get AODInputHandler");
101     } else
102       printf("   ConnectInputData of task %s\n", GetName());
103     fAOD = aodH->GetEvent();
104   }
105 }
106
107 //___________________________________________________________________________
108 void AliAnalysisTaskSingleMu::CreateOutputObjects() 
109 {
110   //
111   /// Create output histograms
112   //
113   printf("   CreateOutputObjects of task %s\n", GetName());
114
115   // initialize tree
116   if(!fResults) fResults = new TTree("Results", "Single mu selection results");
117
118   TString baseName, suffixName;
119
120   suffixName="/F";
121   for(Int_t iVar=0; iVar<kNfloatVars; iVar++){
122     baseName = fFloatVarName[iVar];
123     if(iVar==0) baseName += suffixName;
124     fResults->Branch(fFloatVarName[iVar].Data(), &fVarFloat[iVar], baseName.Data());
125   }
126
127   suffixName="/I";
128   for(Int_t iVar=0; iVar<kNintVars; iVar++){
129     baseName = fIntVarName[iVar];
130     if(iVar==0) baseName += suffixName;
131     fResults->Branch(fIntVarName[iVar].Data(), &fVarInt[iVar], baseName.Data());
132   }
133 }
134
135 //________________________________________________________________________
136 void AliAnalysisTaskSingleMu::Exec(Option_t *) 
137 {
138   //
139   /// Main loop
140   /// Called for each event
141   //
142
143   TTree *tinput = (TTree*)GetInputData(0);
144   tinput->GetReadEntry();
145
146   if (!fAOD) {
147     Printf("ERROR: fAOD not available");
148     return;
149   }
150
151
152   // Object declaration
153   AliAODTrack *muonTrack = 0x0;
154
155   Int_t nTracks = fAOD->GetNumberOfTracks();
156   for (Int_t itrack = 0; itrack < nTracks; itrack++) {
157     muonTrack = fAOD->GetTrack(itrack);
158
159     // Apply cuts
160     if(!FillTrackVariables(*muonTrack)) continue;
161     fResults->Fill();
162   }
163
164   // Post final data. It will be written to a file with option "RECREATE"
165   PostData(0, fResults);
166 }
167
168 //________________________________________________________________________
169 void AliAnalysisTaskSingleMu::Terminate(Option_t *) {
170   //
171   /// Draw some histogram at the end.
172   //
173   if (!gROOT->IsBatch()) {
174     TCanvas *c1 = new TCanvas("c1","Vz vs Pt",10,10,310,310);
175     c1->SetFillColor(10); c1->SetHighLightColor(10);
176     c1->SetLeftMargin(0.15); c1->SetBottomMargin(0.15);  
177     fResults->Draw("pt:vz","","COLZ");
178   }
179 }
180
181 //________________________________________________________________________
182 void AliAnalysisTaskSingleMu::InitVariables() 
183 {
184   //
185   /// Reset histograms
186   //
187
188   fVarFloat = new Float_t[kNfloatVars];
189   fVarInt = new Int_t[kNintVars];
190
191   fFloatVarName = new TString[kNfloatVars];
192   fFloatVarName[kVarPt]     = "pt";
193   fFloatVarName[kVarY]      = "y";
194   fFloatVarName[kVarPhi]    = "phi";
195   fFloatVarName[kVarVz]     = "vz";
196   fFloatVarName[kVarDCA]    = "dca";
197
198   fIntVarName = new TString[kNintVars];
199   fIntVarName[kVarTrig]     = "matchTrig";
200 }
201
202
203 //________________________________________________________________________
204 Bool_t AliAnalysisTaskSingleMu::FillTrackVariables(AliAODTrack &muonTrack)
205 {
206   //
207   /// Fill lorentz vector and check cuts
208   //
209
210   TLorentzVector lorVec;
211
212   // Check if track is a muon
213   if(muonTrack.GetMostProbablePID()!=AliAODTrack::kMuon) return kFALSE;
214
215   // Check if track is triggered
216   fVarInt[kVarTrig] = (muonTrack.GetMatchTrigger() && 0x3);
217   
218   // Fill track parameters
219   Double_t px = muonTrack.Px();
220   Double_t py = muonTrack.Py();
221   Double_t pz = muonTrack.Pz();
222   Double_t p  = muonTrack.P();
223
224   const Double_t kMuonMass = 0.105658369;
225
226   Double_t energy = TMath::Sqrt(p*p + kMuonMass*kMuonMass);
227   lorVec.SetPxPyPzE(px,py,pz,energy);
228
229   fVarFloat[kVarPt]  = lorVec.Pt();
230   fVarFloat[kVarY]   = lorVec.Rapidity();
231   fVarFloat[kVarPhi] = lorVec.Phi();
232
233   fVarFloat[kVarVz] = fAOD->GetPrimaryVertex()->GetZ();
234
235   Double_t xDca = muonTrack.XAtDCA();
236   Double_t yDca = muonTrack.YAtDCA();
237
238   fVarFloat[kVarDCA] = TMath::Sqrt(xDca*xDca + yDca*yDca);
239
240   return kTRUE;
241 }