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