fixes from Laurent for the MC branch in the AOD filters
[u/mrichter/AliRoot.git] / PWG3 / muon / AliAnalysisTaskESDMuonFilter.cxx
CommitLineData
1d1ea3ce 1/**************************************************************************
2 * Copyright(c) 1998-1999, 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
26ba01d4 16//
1d1ea3ce 17// Add the muon tracks to the generic AOD track branch during the
26ba01d4 18// filtering of the ESD.
19//
20// Authors: R. Arnaldi 5/5/08 and L. Aphecetche January 2011
21//
22// Note that we :
23// - completely disable all the branches that are not required by (most) the muon analyses,
24// e.g. cascades, v0s, kinks, jets, etc...
25// - filter severely the tracks (keep only muon tracks) and vertices (keep only primary -including
26// pile-up - vertices) branches
27//
28// (see AddFilteredAOD method)
29//
aba11173 30
1d1ea3ce 31#include "AliAnalysisTaskESDMuonFilter.h"
dcc1f876 32
33#include "AliAODDimuon.h"
1d1ea3ce 34#include "AliAODEvent.h"
1d1ea3ce 35#include "AliAODHandler.h"
14d6fad5 36#include "AliAODExtension.h"
dcc1f876 37#include "AliAODMCParticle.h"
38#include "AliAODMuonReplicator.h"
39#include "AliAODVertex.h"
1d1ea3ce 40#include "AliAnalysisFilter.h"
dcc1f876 41#include "AliAnalysisManager.h"
42#include "AliCodeTimer.h"
43#include "AliESDEvent.h"
44#include "AliESDInputHandler.h"
1d1ea3ce 45#include "AliESDMuonTrack.h"
46#include "AliESDVertex.h"
dcc1f876 47#include "AliESDtrack.h"
1d1ea3ce 48#include "AliLog.h"
fc3a4c45 49#include "AliMCEvent.h"
50#include "AliMCEventHandler.h"
dcc1f876 51#include "AliMultiplicity.h"
52#include "AliStack.h"
53#include <TChain.h>
54#include <TFile.h>
55#include <TParticle.h>
1d1ea3ce 56
57ClassImp(AliAnalysisTaskESDMuonFilter)
26ba01d4 58ClassImp(AliAnalysisNonMuonTrackCuts)
1d1ea3ce 59
60////////////////////////////////////////////////////////////////////////
61
26ba01d4 62AliAnalysisNonMuonTrackCuts::AliAnalysisNonMuonTrackCuts()
63{
64 // default ctor
65}
66
67Bool_t AliAnalysisNonMuonTrackCuts::IsSelected(TObject* obj)
68{
69 // Returns true if the object is a muon track
70 AliAODTrack* track = dynamic_cast<AliAODTrack*>(obj);
71 if (track && track->IsMuonTrack()) return kTRUE;
72 return kFALSE;
73}
74
75AliAnalysisNonPrimaryVertices::AliAnalysisNonPrimaryVertices()
76{
77 // default ctor
78}
79
80Bool_t AliAnalysisNonPrimaryVertices::IsSelected(TObject* obj)
81{
82 // Returns true if the object is a primary vertex
83
84 AliAODVertex* vertex = dynamic_cast<AliAODVertex*>(obj);
85 if (vertex)
86 {
87 if ( vertex->GetType() == AliAODVertex::kPrimary ||
88 vertex->GetType() == AliAODVertex::kMainSPD ||
89 vertex->GetType() == AliAODVertex::kPileupSPD ||
90 vertex->GetType() == AliAODVertex::kPileupTracks ||
91 vertex->GetType() == AliAODVertex::kMainTPC )
92 {
93 return kTRUE;
94 }
95 }
96
97// enum AODVtx_t {kUndef=-1, kPrimary, kKink, kV0, kCascade, kMulti, kMainSPD, kPileupSPD, kPileupTracks,kMainTPC};
98
99 return kFALSE;
100
101}
102
14d6fad5 103AliAnalysisTaskESDMuonFilter::AliAnalysisTaskESDMuonFilter(Bool_t onlyMuon, Bool_t keepAllEvents, Int_t mcMode):
aba11173 104 AliAnalysisTaskSE(),
f785fc59 105 fTrackFilter(0x0),
4a497116 106 fEnableMuonAOD(kFALSE),
26ba01d4 107 fEnableDimuonAOD(kFALSE),
108 fOnlyMuon(onlyMuon),
14d6fad5 109 fKeepAllEvents(keepAllEvents),
110 fMCMode(mcMode)
1d1ea3ce 111{
112 // Default constructor
113}
114
14d6fad5 115AliAnalysisTaskESDMuonFilter::AliAnalysisTaskESDMuonFilter(const char* name, Bool_t onlyMuon, Bool_t keepAllEvents, Int_t mcMode):
aba11173 116 AliAnalysisTaskSE(name),
f785fc59 117 fTrackFilter(0x0),
4a497116 118 fEnableMuonAOD(kFALSE),
26ba01d4 119 fEnableDimuonAOD(kFALSE),
120 fOnlyMuon(onlyMuon),
14d6fad5 121 fKeepAllEvents(keepAllEvents),
122 fMCMode(mcMode)
1d1ea3ce 123{
124 // Constructor
125}
126
3d270be0 127//______________________________________________________________________________
1d1ea3ce 128void AliAnalysisTaskESDMuonFilter::UserCreateOutputObjects()
129{
130 // Create the output container
aba11173 131 if (fTrackFilter) OutputTree()->GetUserInfo()->Add(fTrackFilter);
1d1ea3ce 132}
133
3d270be0 134//______________________________________________________________________________
135void AliAnalysisTaskESDMuonFilter::PrintTask(Option_t *option, Int_t indent) const
136{
137 // Specify how we are configured
138
139 AliAnalysisTaskSE::PrintTask(option,indent);
140
141 TString spaces(' ',indent+3);
142
143 if ( fOnlyMuon )
144 {
145 cout << spaces.Data() << "Keep only muon information " << endl;
146 }
147 else
148 {
149 cout << spaces.Data() << "Keep all information from standard AOD" << endl;
150 }
151
152 if ( fKeepAllEvents )
153 {
154 cout << spaces.Data() << "Keep all events, regardless of number of muons" << endl;
155 }
156 else
157 {
158 cout << spaces.Data() << "Keep only events with at least one muon" << endl;
159 }
14d6fad5 160
161 if ( fMCMode > 0 )
162 {
163 cout << spaces.Data() << "Assuming work on MC data (i.e. will transmit MC branches)" << endl;
164 }
3d270be0 165}
166
167//______________________________________________________________________________
26ba01d4 168void AliAnalysisTaskESDMuonFilter::AddFilteredAOD(const char* aodfilename, const char* title)
169{
170 // Add an output filtered and replicated aod
171
172 AliAODHandler *aodH = (AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
173 if (!aodH) Fatal("UserCreateOutputObjects", "No AOD handler");
174
175 AliAODExtension* ext = aodH->AddFilteredAOD(aodfilename,title);
176
3d270be0 177 if (!ext) return;
178
179 if ( fOnlyMuon )
180 {
14d6fad5 181
182 AliAODMuonReplicator* murep = new AliAODMuonReplicator("MuonReplicator",
183 "remove non muon tracks and non primary or pileup vertices",
184 new AliAnalysisNonMuonTrackCuts,
185 new AliAnalysisNonPrimaryVertices,
186 fMCMode);
dcc1f876 187
188 ext->DropUnspecifiedBranches(); // all branches not part of a FilterBranch call (below) will be dropped
189
26ba01d4 190 ext->FilterBranch("tracks",murep);
3d270be0 191 ext->FilterBranch("vertices",murep);
192 ext->FilterBranch("dimuons",murep);
14d6fad5 193
194 if ( fMCMode > 0 )
195 {
196 // MC branches will be copied (if present), as they are, but only
197 // for events with at least one muon.
198 // For events w/o muon, mcparticles array will be empty and mcheader will be dummy
199 // (e.g. strlen(GetGeneratorName())==0)
200
201 ext->FilterBranch("mcparticles",murep);
202 ext->FilterBranch("mcHeader",murep);
203 }
26ba01d4 204 }
205}
206
3d270be0 207//______________________________________________________________________________
1d1ea3ce 208void AliAnalysisTaskESDMuonFilter::Init()
209{
210 // Initialization
26ba01d4 211 if(fEnableMuonAOD) AddFilteredAOD("AliAOD.Muons.root", "MuonEvents");
212 if(fEnableDimuonAOD) AddFilteredAOD("AliAOD.Dimuons.root", "DimuonEvents");
1d1ea3ce 213}
214
215
3d270be0 216//______________________________________________________________________________
1d1ea3ce 217void AliAnalysisTaskESDMuonFilter::UserExec(Option_t */*option*/)
218{
219 // Execute analysis for current event
26ba01d4 220
1d1ea3ce 221 Long64_t ientry = Entry();
edee4062 222 if(fDebug)printf("Muon Filter: Analysing event # %5d\n", (Int_t) ientry);
aba11173 223
1d1ea3ce 224 ConvertESDtoAOD();
225}
226
3d270be0 227//______________________________________________________________________________
1d1ea3ce 228void AliAnalysisTaskESDMuonFilter::ConvertESDtoAOD()
229{
230 // ESD Muon Filter analysis task executed for each event
dcc1f876 231
232 AliCodeTimerAuto("",0);
233
1d1ea3ce 234 AliESDEvent* esd = dynamic_cast<AliESDEvent*>(InputEvent());
aba11173 235
fc3a4c45 236 // Fetch Stack for debuggging if available
237 AliStack *pStack = 0;
238 AliMCEventHandler *mcH = 0;
239 if(MCEvent()){
240 pStack = MCEvent()->Stack();
241 mcH = (AliMCEventHandler*) ((AliAnalysisManager::GetAnalysisManager())->GetMCtruthEventHandler());
242 }
243
1d1ea3ce 244 // Define arrays for muons
245 Double_t pos[3];
246 Double_t p[3];
247 Double_t pid[10];
aba11173 248
249 // has to be changed once the muon pid is provided by the ESD
a6e0ebfe 250 for (Int_t i = 0; i < 10; pid[i++] = 0.) {}
aba11173 251 pid[AliAODTrack::kMuon]=1.;
252
1d1ea3ce 253 AliAODHeader* header = AODEvent()->GetHeader();
254 AliAODTrack *aodTrack = 0x0;
aba11173 255 AliESDMuonTrack *esdMuTrack = 0x0;
256
1d1ea3ce 257 // Access to the AOD container of tracks
258 TClonesArray &tracks = *(AODEvent()->GetTracks());
259 Int_t jTracks = header->GetRefMultiplicity();
260
261 // Read primary vertex from AOD event
aba11173 262 AliAODVertex *primary = AODEvent()->GetPrimaryVertex();
edee4062 263 if(fDebug)primary->Print();
aba11173 264
1d1ea3ce 265 // Loop on muon tracks to fill the AOD track branch
266 Int_t nMuTracks = esd->GetNumberOfMuonTracks();
866d8d78 267
daf100e1 268 for (Int_t iTrack=0; iTrack<nMuTracks; ++iTrack) esd->GetMuonTrack(iTrack)->SetESDEvent(esd);
aba11173 269
270 // Update number of positive and negative tracks from AOD event (M.G.)
1d1ea3ce 271 Int_t nPosTracks = header->GetRefMultiplicityPos();
272 Int_t nNegTracks = header->GetRefMultiplicityNeg();
aba11173 273
866d8d78 274 // Access to the AOD container of dimuons
275 TClonesArray &dimuons = *(AODEvent()->GetDimuons());
276 AliAODDimuon *aodDimuon = 0x0;
277
866d8d78 278 Int_t nMuons=0;
29e7e51b 279 Int_t nDimuons=0;
866d8d78 280 Int_t jDimuons=0;
397081d5 281 Int_t nMuonTrack[100];
87ac315c 282
397081d5 283 for(int imuon=0;imuon<100;imuon++) nMuonTrack[imuon]=0;
26ba01d4 284
285 for (Int_t nMuTrack = 0; nMuTrack < nMuTracks; ++nMuTrack)
286 {
aba11173 287 esdMuTrack = esd->GetMuonTrack(nMuTrack);
288
289 if (!esdMuTrack->ContainTrackerData()) continue;
daf100e1 290
aba11173 291 UInt_t selectInfo = 0;
292 // Track selection
293 if (fTrackFilter) {
294 selectInfo = fTrackFilter->IsSelected(esdMuTrack);
295 if (!selectInfo) {
296 continue;
297 }
87ac315c 298 }
26ba01d4 299
aba11173 300 p[0] = esdMuTrack->Px();
301 p[1] = esdMuTrack->Py();
302 p[2] = esdMuTrack->Pz();
303
304 pos[0] = esdMuTrack->GetNonBendingCoor();
305 pos[1] = esdMuTrack->GetBendingCoor();
306 pos[2] = esdMuTrack->GetZ();
307
fc3a4c45 308 if(mcH)mcH->SelectParticle(esdMuTrack->GetLabel());
309
aba11173 310 aodTrack = new(tracks[jTracks++]) AliAODTrack(esdMuTrack->GetUniqueID(), // ID
26ba01d4 311 esdMuTrack->GetLabel(), // label
312 p, // momentum
313 kTRUE, // cartesian coordinate system
314 pos, // position
315 kFALSE, // isDCA
316 0x0, // covariance matrix
317 esdMuTrack->Charge(), // charge
318 0, // ITSClusterMap
319 pid, // pid
320 primary, // primary vertex
321 kFALSE, // used for vertex fit?
322 kFALSE, // used for primary vertex fit?
323 AliAODTrack::kPrimary,// track type
324 selectInfo);
aba11173 325
326 aodTrack->SetXYAtDCA(esdMuTrack->GetNonBendingCoorAtDCA(), esdMuTrack->GetBendingCoorAtDCA());
327 aodTrack->SetPxPyPzAtDCA(esdMuTrack->PxAtDCA(), esdMuTrack->PyAtDCA(), esdMuTrack->PzAtDCA());
f43586f0 328 aodTrack->SetRAtAbsorberEnd(esdMuTrack->GetRAtAbsorberEnd());
aba11173 329 aodTrack->ConvertAliPIDtoAODPID();
330 aodTrack->SetChi2perNDF(esdMuTrack->GetChi2() / (2.*esdMuTrack->GetNHit() - 5.));
331 aodTrack->SetChi2MatchTrigger(esdMuTrack->GetChi2MatchTrigger());
332 aodTrack->SetHitsPatternInTrigCh(esdMuTrack->GetHitsPatternInTrigCh());
333 aodTrack->SetMuonClusterMap(esdMuTrack->GetMuonClusterMap());
334 aodTrack->SetMatchTrigger(esdMuTrack->GetMatchTrigger());
5c15a68b 335 aodTrack->Connected(esdMuTrack->IsConnected());
aba11173 336 primary->AddDaughter(aodTrack);
337
338 if (esdMuTrack->Charge() > 0) nPosTracks++;
339 else nNegTracks++;
866d8d78 340
87ac315c 341 nMuonTrack[nMuons]= jTracks-1.;
26ba01d4 342 ++nMuons;
87ac315c 343 }
26ba01d4 344
345 if(nMuons>=2)
346 {
347 for(int i=0;i<nMuons;i++){
348 Int_t index0 = nMuonTrack[i];
349 for(int j=i+1;j<nMuons;j++){
350 Int_t index1 = nMuonTrack[j];
351 aodDimuon = new(dimuons[jDimuons++]) AliAODDimuon(tracks.At(index0),tracks.At(index1));
352 ++nDimuons;
353 if (fDebug > 1){
354 AliAODDimuon *dimuon0 = (AliAODDimuon*)dimuons.At(jDimuons-1);
355 printf("Dimuon: mass = %f, px=%f, py=%f, pz=%f\n",dimuon0->M(),dimuon0->Px(),dimuon0->Py(),dimuon0->Pz());
356 AliAODTrack *mu0 = (AliAODTrack*) dimuon0->GetMu(0);
357 AliAODTrack *mu1 = (AliAODTrack*) dimuon0->GetMu(1);
358 printf("Muon0 px=%f py=%f pz=%f\n",mu0->Px(),mu0->Py(),mu0->Pz());
359 printf("Muon1 px=%f py=%f pz=%f\n",mu1->Px(),mu1->Py(),mu1->Pz());
360 }
87ac315c 361 }
866d8d78 362 }
26ba01d4 363 }
364
aba11173 365
1d1ea3ce 366 header->SetRefMultiplicity(jTracks);
367 header->SetRefMultiplicityPos(nPosTracks);
368 header->SetRefMultiplicityNeg(nNegTracks);
29e7e51b 369 header->SetNumberOfMuons(nMuons);
370 header->SetNumberOfDimuons(nDimuons);
866d8d78 371
26ba01d4 372 if ( fEnableMuonAOD && ( (nMuons>0) || fKeepAllEvents ) )
373 {
f785fc59 374 AliAODExtension *extMuons = dynamic_cast<AliAODHandler*>
375 ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler())->GetFilteredAOD("AliAOD.Muons.root");
376 extMuons->SelectEvent();
866d8d78 377 }
26ba01d4 378
379 if ( fEnableDimuonAOD && ( (nMuons>1) || fKeepAllEvents ) )
380 {
866d8d78 381 AliAODExtension *extDimuons = dynamic_cast<AliAODHandler*>
26ba01d4 382 ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler())->GetFilteredAOD("AliAOD.Dimuons.root");
866d8d78 383 extDimuons->SelectEvent();
f785fc59 384 }
26ba01d4 385
1d1ea3ce 386}
387
388void AliAnalysisTaskESDMuonFilter::Terminate(Option_t */*option*/)
389{
aba11173 390 // Terminate analysis
391 //
392 if (fDebug > 1) printf("AnalysisESDfilter: Terminate() \n");
1d1ea3ce 393}
aba11173 394
26ba01d4 395void AliAnalysisTaskESDMuonFilter::PrintMCInfo(AliStack *pStack,Int_t label)
396{
397 // print mc info
fc3a4c45 398 if(!pStack)return;
399 label = TMath::Abs(label);
400 TParticle *part = pStack->Particle(label);
401 Printf("########################");
402 Printf("%s:%d %d UniqueID %d PDG %d P %3.3f",(char*)__FILE__,__LINE__,label,part->GetUniqueID(),part->GetPdgCode(),part->P());
403 part->Print();
404 TParticle* mother = part;
405 Int_t imo = part->GetFirstMother();
406 Int_t nprim = pStack->GetNprimary();
407 // while((imo >= nprim) && (mother->GetUniqueID() == 4)) {
408 while((imo >= nprim)) {
409 mother = pStack->Particle(imo);
410 Printf("Mother %s:%d Label %d UniqueID %d PDG %d P %3.3f",(char*)__FILE__,__LINE__,imo,mother->GetUniqueID(),mother->GetPdgCode(),mother->P());
411 mother->Print();
412 imo = mother->GetFirstMother();
413 }
414 Printf("########################");
415}