AliAODEvent::GetHeader() returns AliVHeader
[u/mrichter/AliRoot.git] / PWG / muon / AliAnalysisTaskCreateMixedDimuons.cxx
1 #include "TChain.h"
2 #include "TTree.h"
3
4 #include "AliAnalysisTaskME.h"
5 #include "AliAnalysisManager.h"
6
7 #include "AliAODEvent.h"
8 #include "AliAODTrack.h"
9 #include "AliAODVertex.h"
10 #include "AliMultiEventInputHandler.h"
11 #include "AliAODHandler.h"
12
13 #include "AliAnalysisTaskCreateMixedDimuons.h"
14 #include "AliEventPoolMuon.h"
15 #include "TDatabasePDG.h"
16 #include "TRandom.h"
17
18 // Example of an analysis task creating aod events filled with mixed muon pairs
19 //
20 // Authors Alessandro De Falco and Antonio Uras, INFN Cagliari
21 // alessandro.de.falco@ca.infn.it  antonio.uras@ca.infn.it
22
23 #define AliAnalysisTaskCreateMixedDimuons_CXX
24
25 ClassImp(AliAnalysisTaskCreateMixedDimuons)
26
27 //=================================================================================
28
29 AliAnalysisTaskCreateMixedDimuons::AliAnalysisTaskCreateMixedDimuons(const char *name) 
30 : AliAnalysisTaskME(name),
31   fBufferSize(0),
32   fOutputUserHandler(0x0),
33   fOutputUserAOD(0X0),
34   fOutputUserAODTree(0X0),
35   fPoolMuon(0X0),
36   fDebug(0X0)
37 {
38
39   // Constructor
40
41   // Default input and output containers
42   DefineInput (0, TChain::Class());
43
44   // User-defined input and output containers
45   DefineOutput(1, TTree::Class());
46
47   // ---------------------------------------------------------------
48
49   fDebug = kFALSE;
50   fBufferSize = 0;
51
52   RequireFreshBuffer();
53
54   for (Int_t i=0; i<100; i++) fInputAOD[i] = 0;
55
56   fOutputUserHandler = 0;
57   fOutputUserAOD     = 0;
58   fOutputUserAODTree = 0;  
59
60 }
61
62 //=================================================================================
63
64 void AliAnalysisTaskCreateMixedDimuons::ConnectInputData(Option_t *) {
65
66   // Connect ESD or AOD here
67   // Called once
68
69   printf("-> AliAnalysisTaskCreateMixedDimuons::ConnectInputData\n");
70
71   TTree* tree = (TTree*) GetInputData(0);
72
73   if (!tree) {
74     Printf("ERROR: Could not read chain from input slot 0");
75   } 
76   else {
77
78     fInputHandler = (AliMultiEventInputHandler*) AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler();
79     fPoolMuon     = (AliEventPoolMuon*)          AliAnalysisManager::GetAnalysisManager()->GetEventPool();
80
81     if (!fInputHandler) {
82       Printf("ERROR: Could not get AliMultiAODInputHandler");
83     }
84     else {  
85       fBufferSize = fInputHandler->GetBufferSize();
86       if (fBufferSize>100) {
87         printf("\n*** WARNING AliAnalysisTaskCreateMixedDimuons::ConnectInputData -> Trying to set fBufferSize>100, forcing fBufferSize=100 ***\n\n");
88         fBufferSize = 100;
89       }
90       for (Int_t i=0; i<fBufferSize; i++) fInputAOD[i] = (AliAODEvent*) fInputHandler->GetEvent(i);
91     }
92   }
93
94   printf("<- AliAnalysisTaskCreateMixedDimuons::ConnectInputData\n");
95
96 }
97
98 //=================================================================================
99
100 void AliAnalysisTaskCreateMixedDimuons::UserCreateOutputObjects() {
101
102   // Here the user-defined output containers should be created!!!
103   // Called once
104
105   fOutputUserHandler = new AliAODHandler();
106
107   fOutputUserHandler -> Init("");
108
109   fOutputUserAOD = fOutputUserHandler -> GetAOD();
110
111   fOutputUserAODTree = fOutputUserHandler -> GetTree();
112
113 }
114
115 //=================================================================================
116
117 void AliAnalysisTaskCreateMixedDimuons::UserExec(Option_t *) {
118
119   if (!fOutputUserAOD) {
120     Printf("ERROR: fOutputUserAOD not available\n");
121     return;
122   }
123
124   printf("Calling USER EXEC\n\n");
125
126   for (Int_t iEv=0; iEv<fBufferSize; iEv++) {
127     for (Int_t jEv=0; jEv<iEv; jEv++) {
128
129       Int_t nTracksEv[2]  = {0};
130       Int_t nFWMuonsEv[2] = {0};
131       
132       nTracksEv[0] = fInputAOD[iEv]->GetNumberOfTracks();
133       nTracksEv[1] = fInputAOD[jEv]->GetNumberOfTracks();
134       
135       for (Int_t i=0; i<nTracksEv[0]; i++) if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[0]++;
136       for (Int_t i=0; i<nTracksEv[1]; i++) if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[1]++;
137       
138       // Muon track mixing to fill a mass spectrum
139       
140       if (nFWMuonsEv[0] && nFWMuonsEv[1]) {
141         
142         Int_t rndMuonTrack[2] = {0};
143         rndMuonTrack[0] = gRandom->Integer(nFWMuonsEv[0]);
144         rndMuonTrack[1] = gRandom->Integer(nFWMuonsEv[1]);
145         
146         Int_t nFWMUonsAdded = 0;
147         Int_t nPosTracksAdded = 0;
148         Int_t nNegTracksAdded = 0;
149         
150         AliAODVertex *vertex = new AliAODVertex();
151         vertex -> SetX(0.0);
152         vertex -> SetY(0.0);
153         vertex -> SetZ(fPoolMuon->GetMeanPrimaryVertexZ());
154         
155         Int_t muonCounter[2] = {0};
156         
157         // adding tracks and vertex to the output event...
158         
159         for (Int_t i=0; i<nTracksEv[0]; i++) {
160           if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) {
161             if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p    pt = %f     uniqueID = %d\n",
162                                iEv,i,fInputAOD[iEv]->GetTrack(i),fInputAOD[iEv]->GetTrack(i)->Pt(),
163                                fInputAOD[iEv]->GetTrack(i)->GetUniqueID());
164             if (muonCounter[0]==rndMuonTrack[0]) {
165               fOutputUserAOD->AddTrack(fInputAOD[iEv]->GetTrack(i));
166               nFWMUonsAdded++;
167               if (fInputAOD[iEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
168               else nNegTracksAdded++;
169             }
170             muonCounter[0]++;
171           }
172         }
173         
174         for (Int_t i=0; i<nTracksEv[1]; i++) {
175           if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) {
176             if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p    pt = %f     uniqueID = %d\n",
177                                jEv,i,fInputAOD[jEv]->GetTrack(i),fInputAOD[jEv]->GetTrack(i)->Pt(),
178                                fInputAOD[jEv]->GetTrack(i)->GetUniqueID());
179             if (muonCounter[1]==rndMuonTrack[1]) {
180               fOutputUserAOD->AddTrack(fInputAOD[jEv]->GetTrack(i));
181               nFWMUonsAdded++;
182               if (fInputAOD[jEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
183               else nNegTracksAdded++;
184             }
185             muonCounter[1]++;   
186           }
187         }
188         
189         fOutputUserAOD->AddVertex(vertex);
190         
191         // ... done!
192         
193         if (fDebug) {
194           for (Int_t i=0; i<nFWMUonsAdded; i++) {
195             AliAODTrack *tr = (AliAODTrack*) fOutputUserAOD->GetTrack(i);
196             printf("fOutputUserAOD->GetTrack(%d) = %p    pt = %f\n",i,tr,tr->Pt());
197           }
198         }
199         
200
201         AliAODHeader * header = dynamic_cast<AliAODHeader*>(fOutputUserAOD->GetHeader());
202         if(!header) AliFatal("Not a standard AOD");
203         header->SetRefMultiplicity(nFWMUonsAdded); 
204         header->SetRefMultiplicityPos(nPosTracksAdded);
205         header->SetRefMultiplicityNeg(nNegTracksAdded);
206         
207         fOutputUserHandler -> FinishEvent();
208         
209       }
210
211       PostData(1, fOutputUserAODTree);
212       
213     }
214     
215   }
216   
217 }      
218
219 //===================================================================================
220
221 void AliAnalysisTaskCreateMixedDimuons::Terminate(Option_t *) {
222
223   // Called once at the end of the query
224
225   printf("\n\nCalling TERMINATE \n\n\n");
226
227   fOutputUserHandler -> Terminate();
228
229 }
230
231 //===================================================================================