adding support for AliESDEvent in AliESDtrackCuts
[u/mrichter/AliRoot.git] / PWG0 / esdTrackCuts / AliCutTask.cxx
1 #include "TChain.h"
2 #include "TTree.h"
3 #include "TH1F.h"
4 #include "TCanvas.h"
5 #include "TList.h"
6 #include "TFile.h"
7
8 #include "esdTrackCuts/AliESDtrackCuts.h"
9 #include "AliPWG0Helper.h"
10
11 #include "AliAnalysisTask.h"
12 #include "AliAnalysisManager.h"
13 #include "AliESDEvent.h"
14 #include "AliESDInputHandler.h"
15 #include "AliESDVertex.h"
16
17 #include "AliCutTask.h"
18
19 // simple task that runs the esd track cuts to evaluate the basic plots created during the cuts
20
21 ClassImp(AliCutTask)
22
23 //________________________________________________________________________
24 AliCutTask::AliCutTask(const char *name) 
25   : AliAnalysisTask(name, ""), fESD(0), fTrackCuts(0), fVertex(0), fOutput(0)
26 {
27   // Constructor
28
29   // Define input and output slots here
30   DefineInput(0, TChain::Class());
31   DefineOutput(0, TList::Class());
32 }
33
34 //________________________________________________________________________
35 void AliCutTask::ConnectInputData(Option_t *) 
36 {
37   // Connect ESD or AOD here
38   // Called once
39
40   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
41   if (!tree) {
42     Printf("ERROR: Could not read chain from input slot 0");
43   } else {
44     // Disable all branches and enable only the needed ones
45     //tree->SetBranchStatus("*", kFALSE);
46
47     tree->SetBranchStatus("fTracks.*", kTRUE);
48     tree->SetBranchStatus("Tracks.*", kTRUE);
49
50     tree->SetBranchStatus("fTriggerMask", kTRUE);
51     tree->SetBranchStatus("AliESDHeader", kTRUE);
52
53     tree->SetBranchStatus("fSPDVertex*", kTRUE);
54     tree->SetBranchStatus("SPDVertex", kTRUE);
55     //tree->SetBranchStatus("fPosition[3]", kTRUE);
56
57     AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
58
59     if (!esdH) {
60       Printf("ERROR: Could not get ESDInputHandler");
61     } else
62       fESD = esdH->GetEvent();
63   }
64 }
65
66 //________________________________________________________________________
67 void AliCutTask::CreateOutputObjects()
68 {
69   // Create histograms
70   // Called once
71
72   fOutput = new TList;
73   fOutput->SetOwner();
74
75   fOutput->Add(fTrackCuts);
76
77   fVertex = new TH1F("fVertex", "fVertex;z vtx (cm);Count", 201, -20, 20);
78   fOutput->Add(fVertex);
79 }
80
81 //________________________________________________________________________
82 void AliCutTask::Exec(Option_t *) 
83 {
84   // Main loop
85   // Called for each event
86
87   if (!fESD) {
88     Printf("ERROR: fESD not available");
89     return;
90   }
91
92   // Post output data.
93   PostData(0, fOutput);
94
95   if (!AliPWG0Helper::IsVertexReconstructed(fESD->GetVertex()))
96     return;
97
98   Printf("There are %d tracks in this event", fESD->GetNumberOfTracks());
99   fTrackCuts->CountAcceptedTracks(fESD);
100
101   // get the ESD vertex
102   fVertex->Fill(fESD->GetVertex()->GetZv());
103 }
104
105 //________________________________________________________________________
106 void AliCutTask::Terminate(Option_t *) 
107 {
108   // Draw result to the screen
109   // Called once at the end of the query
110
111   fOutput = dynamic_cast<TList*> (GetOutputData(0));
112   if (!fOutput) {
113     Printf("ERROR: fOutput not available");
114     return;
115   }
116
117   fTrackCuts = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("AliESDtrackCuts"));
118   if (!fTrackCuts) {
119     Printf("ERROR: fTrackCuts not available");
120     return;
121   }
122
123   fVertex = dynamic_cast<TH1F*> (fOutput->FindObject("fVertex"));
124   if (!fVertex) {
125     Printf("ERROR: fVertex not available");
126     return;
127   }
128
129   TFile* file = TFile::Open("trackCuts.root", "RECREATE");
130
131   fTrackCuts->SaveHistograms();
132   fVertex->Write();
133
134   file->Close();
135
136         fTrackCuts->DrawHistograms();
137
138   new TCanvas;
139   fVertex->Draw();
140 }