enable branches of new esd
[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     //tree->SetBranchStatus("*Calo*", kFALSE);
47
48     tree->SetBranchStatus("fTracks.*", kTRUE);
49     tree->SetBranchStatus("Tracks.*", kTRUE);
50
51     tree->SetBranchStatus("fTriggerMask", kTRUE);
52     tree->SetBranchStatus("AliESDHeader", kTRUE);
53
54     tree->SetBranchStatus("fSPDVertex*", kTRUE);
55
56     // unclear how to enable vertex branch with new ESD format
57     tree->SetBranchStatus("SPDVertex.*", kTRUE);
58     tree->SetBranchStatus("*fPosition*", kTRUE);
59     tree->SetBranchStatus("*fPosition[3]", kTRUE);
60
61     AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
62
63     if (!esdH) {
64       Printf("ERROR: Could not get ESDInputHandler");
65     } else
66       fESD = esdH->GetEvent();
67   }
68 }
69
70 //________________________________________________________________________
71 void AliCutTask::CreateOutputObjects()
72 {
73   // Create histograms
74   // Called once
75
76   fOutput = new TList;
77   fOutput->SetOwner();
78
79   fOutput->Add(fTrackCuts);
80
81   fVertex = new TH1F("fVertex", "fVertex;z vtx (cm);Count", 201, -20, 20);
82   fOutput->Add(fVertex);
83 }
84
85 //________________________________________________________________________
86 void AliCutTask::Exec(Option_t *) 
87 {
88   // Main loop
89   // Called for each event
90
91   if (!fESD) {
92     Printf("ERROR: fESD not available");
93     return;
94   }
95
96   // Post output data.
97   PostData(0, fOutput);
98
99   fESD->GetVertex()->Print();
100
101   if (!AliPWG0Helper::IsVertexReconstructed(fESD->GetVertex()))
102     return;
103
104   Printf("There are %d tracks in this event", fESD->GetNumberOfTracks());
105   fTrackCuts->CountAcceptedTracks(fESD);
106
107   // get the ESD vertex
108   fVertex->Fill(fESD->GetVertex()->GetZv());
109 }
110
111 //________________________________________________________________________
112 void AliCutTask::Terminate(Option_t *) 
113 {
114   // Draw result to the screen
115   // Called once at the end of the query
116
117   fOutput = dynamic_cast<TList*> (GetOutputData(0));
118   if (!fOutput) {
119     Printf("ERROR: fOutput not available");
120     return;
121   }
122
123   fTrackCuts = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("AliESDtrackCuts"));
124   if (!fTrackCuts) {
125     Printf("ERROR: fTrackCuts not available");
126     return;
127   }
128
129   fVertex = dynamic_cast<TH1F*> (fOutput->FindObject("fVertex"));
130   if (!fVertex) {
131     Printf("ERROR: fVertex not available");
132     return;
133   }
134
135   TFile* file = TFile::Open("trackCuts.root", "RECREATE");
136
137   fTrackCuts->SaveHistograms();
138   fVertex->Write();
139
140   file->Close();
141
142         fTrackCuts->DrawHistograms();
143
144   new TCanvas;
145   fVertex->Draw();
146 }