Update of the Xiaoming code for pp900 first muon analysis. Fixing wanirngs and violti...
[u/mrichter/AliRoot.git] / PWG3 / muon / AliAnalysisTaskSEMuonsHF.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-2008, 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 //
18 // AliAnalysisTaskSE for the single muon and dimuon from HF analysis,
19 // using the classes AliMuonsHFHeader,
20 //                   AliMuonInfoStoreRD,
21 //                   AliDimuInfoStoreRD,
22 //                   AliMuonInfoStoreMC,
23 //                   AliDimuInfoStoreMC.
24 //
25 // Author: X-M. Zhang, zhang@clermont.in2p3.fr
26 //                     zhangxm@iopp.ccnu.edu.cn
27 /////////////////////////////////////////////////////////////
28
29 #include <TList.h>
30 #include <TClonesArray.h>
31 #include <TFile.h>
32
33 #include "AliAnalysisManager.h"
34 #include "AliMCEventHandler.h"
35 #include "AliAODMCParticle.h"
36 #include "AliVEvent.h"
37 #include "AliESDEvent.h"
38 #include "AliAODEvent.h"
39 #include "AliESDMuonTrack.h"
40 #include "AliAODTrack.h"
41 #include "AliMuonsHFHeader.h"
42 #include "AliMuonInfoStoreRD.h"
43 #include "AliMuonInfoStoreMC.h"
44 #include "AliDimuInfoStoreRD.h"
45 #include "AliDimuInfoStoreMC.h"
46 #include "AliAnalysisTaskSEMuonsHF.h"
47
48 class AliAnalysisTaskSE;
49
50 ClassImp(AliAnalysisTaskSEMuonsHF)
51
52 //_____________________________________________________________________________
53 AliAnalysisTaskSEMuonsHF::AliAnalysisTaskSEMuonsHF() :
54 AliAnalysisTaskSE(),
55 fAnaMode(0),
56 fIsOutputTree(kFALSE),
57 fIsUseMC(kFALSE),
58 fHeader(0),
59 fMuonClArr(0),
60 fDimuClArr(0),
61 fListHisHeader(0),
62 fListHisMuon(),
63 fListHisDimu(0)
64 {
65   //
66   // Default constructor
67   //
68 }
69
70 //_____________________________________________________________________________
71 AliAnalysisTaskSEMuonsHF::AliAnalysisTaskSEMuonsHF(const char *name) :
72 AliAnalysisTaskSE(name),
73 fAnaMode(0),
74 fIsOutputTree(kFALSE),
75 fIsUseMC(kFALSE),
76 fHeader(0),
77 fMuonClArr(0),
78 fDimuClArr(0),
79 fListHisHeader(0),
80 fListHisMuon(0),
81 fListHisDimu(0)
82 {
83   //
84   // Constructor
85   //
86 }
87
88 //_____________________________________________________________________________
89 AliAnalysisTaskSEMuonsHF::~AliAnalysisTaskSEMuonsHF()
90 {
91   //
92   // Default destructor
93   //
94   if (fHeader)    { delete fHeader;     fHeader    =NULL; }
95   if (fMuonClArr) { delete fMuonClArr;  fMuonClArr =NULL; }
96   if (fDimuClArr) { delete fDimuClArr;  fDimuClArr =NULL; }
97
98   if (fListHisHeader) { delete    fListHisHeader; fListHisHeader=NULL; }
99   if (fListHisMuon)   { delete [] fListHisMuon;   fListHisMuon  =NULL; }
100   if (fListHisDimu)   { delete [] fListHisDimu;   fListHisDimu  =NULL; }
101 }
102
103 //_____________________________________________________________________________
104 void AliAnalysisTaskSEMuonsHF::Init()
105 {
106   // Initialization
107   // Setting and initializing the running mode and status
108
109   AliMuonsHFHeader::SetAnaMode(fAnaMode);
110   AliMuonsHFHeader::SetIsMC(fIsUseMC);
111   if (!fHeader) {
112     fHeader = new AliMuonsHFHeader();
113     fHeader->SetName(AliMuonsHFHeader::StdBranchName());
114   }
115
116   if (!fMuonClArr) {
117     if (fIsUseMC) { 
118       fMuonClArr = new TClonesArray("AliMuonInfoStoreMC", 0);
119       fMuonClArr->SetName(AliMuonInfoStoreMC::StdBranchName());
120     } else {
121       fMuonClArr = new TClonesArray("AliMuonInfoStoreRD", 0);
122       fMuonClArr->SetName(AliMuonInfoStoreRD::StdBranchName());
123     }
124   }
125
126   if (fAnaMode!=1 && !fDimuClArr) {
127     if (fIsUseMC) {
128       fDimuClArr = new TClonesArray("AliDimuInfoStoreMC", 0);
129       fDimuClArr->SetName(AliDimuInfoStoreMC::StdBranchName());
130     } else {
131       fDimuClArr = new TClonesArray("AliDimuInfoStoreRD", 0);
132       fDimuClArr->SetName(AliDimuInfoStoreRD::StdBranchName());
133     }
134   }
135
136   return;
137 }
138
139 //_____________________________________________________________________________
140 void AliAnalysisTaskSEMuonsHF::UserCreateOutputObjects()
141 {
142   // Create the output container
143
144   if (!fListHisHeader) fListHisHeader = new TList();
145   if (fAnaMode!=2 && !fListHisMuon) {
146     if (fIsUseMC) fListHisMuon = new TList[AliMuonInfoStoreMC::NSources()];
147     else fListHisMuon = new TList();
148   }
149   if (fAnaMode!=1 && !fListHisDimu) {
150     if (fIsUseMC) fListHisDimu = new TList[AliDimuInfoStoreMC::NSources()];
151     else fListHisDimu = new TList();
152   }
153   fHeader->CreateHistograms(fListHisHeader, fListHisMuon, fListHisDimu);
154
155   if (fIsOutputTree) {
156     AddAODBranch("AliMuonsHFHeader", &fHeader);
157     AddAODBranch("TClonesArray", &fMuonClArr);
158     if (fAnaMode!=1) AddAODBranch("TClonesArray", &fDimuClArr);
159   }
160
161   return;
162 }
163
164 //_____________________________________________________________________________
165 void AliAnalysisTaskSEMuonsHF::UserExec(Option_t *)
166 {
167   // Execute analysis for current event:
168   // muon event header & (di)muon info store
169
170   AliVEvent *event = dynamic_cast<AliVEvent*>(InputEvent());
171   //if (AODEvent() && IsStandardAOD()) event = dynamic_cast<AliVEvent*>(AODEvent());
172   TString evName = event->IsA()->GetName();
173   Bool_t isAOD = ((evName=="AliAODEvent") ? kTRUE : kFALSE);
174   event = 0x0;
175
176   Int_t ntrks = 0;
177   AliAODEvent *aod = 0;
178   AliESDEvent *esd = 0;
179   TClonesArray *mcClArr = 0;
180   AliMCEventHandler *mcH = 0;
181   if (isAOD) {
182     aod = dynamic_cast<AliAODEvent*>(InputEvent());
183     //if (AODEvent() && IsStandardAOD()) aod = dynamic_cast<AliAODEvent*>(AODEvent());
184     if (!aod) { AliError("AOD event not found. Nothing done!"); return; }
185     if (fIsUseMC) {
186       mcClArr = (TClonesArray*)aod->GetList()->FindObject(AliAODMCParticle::StdBranchName());
187       if (!mcClArr) { AliError("MC Array not found. Nothing done!"); return; }
188     }
189     fHeader->SetEvent(aod);
190     ntrks = aod->GetNTracks();
191   } else {
192     esd = dynamic_cast<AliESDEvent*>(InputEvent());
193     if (!esd) { AliError("ESD event not found. Nothing done!"); return; }
194     if (fIsUseMC) {
195       mcH = (AliMCEventHandler*)((AliAnalysisManager::GetAnalysisManager())->GetMCtruthEventHandler());
196       if (!mcH) { AliError("MC Handler not found. Nothing done!"); return; }
197     }
198     fHeader->SetEvent(esd);
199     ntrks = esd->GetNumberOfMuonTracks();
200   }
201
202   fHeader->FillHistosEventH(fListHisHeader);
203
204   fMuonClArr->Delete();
205   TClonesArray &muonRef = *fMuonClArr;
206   Int_t countN = fMuonClArr->GetEntriesFast();
207
208   AliAODTrack        *trkAOD = 0;
209   AliESDMuonTrack    *trkESD = 0;
210   AliMuonInfoStoreRD *trkRD  = 0;
211   AliMuonInfoStoreMC *trkMC  = 0;
212   for (Int_t itrk=0; itrk<ntrks; itrk++) {  // loop over all tracks
213     if (isAOD) {
214       trkAOD = (AliAODTrack*)aod->GetTrack(itrk);
215       if (!trkAOD->IsMuonTrack()) { trkAOD=0; trkRD=0; trkMC=0; continue; }
216       if (fIsUseMC) trkMC = new AliMuonInfoStoreMC(trkAOD, mcClArr);
217       else trkRD = new AliMuonInfoStoreRD(trkAOD);
218       trkAOD = 0;
219     } else {
220       trkESD = (AliESDMuonTrack*)esd->GetMuonTrack(itrk);
221       if (!trkESD->ContainTrackerData()) { trkESD=0; trkRD=0; trkMC=0; continue; }
222       if (fIsUseMC) trkMC = new AliMuonInfoStoreMC(trkESD, esd, mcH);
223       else trkRD = new AliMuonInfoStoreRD(trkESD);
224       trkESD = 0;
225     }
226
227     if (trkRD) {
228       if (fAnaMode!=2) fHeader->FillHistosMuonRD(fListHisMuon, trkRD);
229       new(muonRef[countN++]) AliMuonInfoStoreRD(*trkRD);
230     }
231     if (trkMC) {
232       if (fAnaMode!=2) fHeader->FillHistosMuonMC(fListHisMuon, trkMC);
233       new(muonRef[countN++]) AliMuonInfoStoreMC(*trkMC);
234     }
235
236     if (trkRD) { delete trkRD; trkRD=0; }
237     if (trkMC) { delete trkMC; trkMC=0; }
238   }  // end loop of all tracks
239
240   if (fAnaMode==1) return;
241   fDimuClArr->Delete();
242   countN = fDimuClArr->GetEntriesFast();
243   TClonesArray &dimuRef = *fDimuClArr;
244
245   AliDimuInfoStoreRD *dimuRD = 0;
246   AliDimuInfoStoreMC *dimuMC = 0;
247   ntrks = fMuonClArr->GetEntriesFast();
248   for (Int_t itrk=0; itrk<ntrks-1; itrk++) {  // 1st loop over muon tracks
249     for (Int_t jtrk=itrk+1; jtrk<ntrks; jtrk++) {  // 2nd loop ofver muon tracks
250       if (fIsUseMC)
251         dimuMC = new AliDimuInfoStoreMC((AliMuonInfoStoreMC*)fMuonClArr->At(itrk), (AliMuonInfoStoreMC*)fMuonClArr->At(jtrk));
252       else
253         dimuRD = new AliDimuInfoStoreRD((AliMuonInfoStoreRD*)fMuonClArr->At(itrk), (AliMuonInfoStoreRD*)fMuonClArr->At(jtrk));
254
255       if (dimuRD) {
256         fHeader->FillHistosDimuRD(fListHisDimu, dimuRD);
257         if (fIsOutputTree) new(dimuRef[countN++]) AliDimuInfoStoreRD(*dimuRD);
258       }
259       if (dimuMC) {
260         fHeader->FillHistosDimuMC(fListHisDimu, dimuMC);
261         if (fIsOutputTree) new(dimuRef[countN++]) AliDimuInfoStoreMC(*dimuMC);
262       }
263
264       if (dimuRD) { delete dimuRD; dimuRD=0; }
265       if (dimuMC) { delete dimuMC; dimuMC=0; }
266     }  // end 2nd loop of muon tracks
267   }  // end 1st loop of muon tracks
268
269   return;
270 }
271
272 //_____________________________________________________________________________
273 void AliAnalysisTaskSEMuonsHF::Terminate(Option_t *)
274 {
275   // Terminate analysis
276
277   TFile *fout = new TFile("muonsHF.root", "RECREATE");
278   gDirectory->mkdir("header");
279   gDirectory->cd("header");
280   fListHisHeader->Write();
281   gDirectory->cd("..");
282
283   if (fAnaMode!=2) {
284     gDirectory->mkdir("muon");
285     gDirectory->cd("muon");
286     if (fIsUseMC) {
287       char *muonName[] = {"BottomMu", "CharmMu", "PrimaryMu", "SecondaryMu", "NotMu", "Undentified", "All"};
288       for (Int_t i=AliMuonInfoStoreMC::NSources(); i--;) {
289         gDirectory->mkdir(muonName[i]);
290         gDirectory->cd(muonName[i]);
291         fListHisMuon[i].Write();
292         gDirectory->cd("..");
293       }
294     } else fListHisMuon->Write();
295     gDirectory->cd("..");
296   }
297
298   if (fAnaMode!=1) {
299     gDirectory->mkdir("dimu");
300     gDirectory->cd("dimu");
301     if (fIsUseMC) {
302       char *dimuName[] = {"BBdiff", "Bchain", "DDdiff", "Dchain", "Resonance", "Background", "All"};
303       for (Int_t i=AliDimuInfoStoreMC::NSources(); i--;) {
304         gDirectory->mkdir(dimuName[i]);
305         gDirectory->cd(dimuName[i]);
306         fListHisDimu[i].Write();
307         gDirectory->cd("..");
308       }
309     } else fListHisDimu->Write();
310     gDirectory->cd("..");
311   }
312
313   fout->Write();
314   fout->Close();
315   return;
316 }