fixes from Laurent for the MC branch in the AOD filters
[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 "AliAnalysisTaskESDMuonFilter.h"
32
33 #include "AliAODDimuon.h"
34 #include "AliAODEvent.h"
35 #include "AliAODHandler.h"
36 #include "AliAODExtension.h"
37 #include "AliAODMCParticle.h"
38 #include "AliAODMuonReplicator.h"
39 #include "AliAODVertex.h"
40 #include "AliAnalysisFilter.h"
41 #include "AliAnalysisManager.h"
42 #include "AliCodeTimer.h"
43 #include "AliESDEvent.h"
44 #include "AliESDInputHandler.h"
45 #include "AliESDMuonTrack.h"
46 #include "AliESDVertex.h"
47 #include "AliESDtrack.h"
48 #include "AliLog.h"
49 #include "AliMCEvent.h"
50 #include "AliMCEventHandler.h"
51 #include "AliMultiplicity.h"
52 #include "AliStack.h"
53 #include <TChain.h>
54 #include <TFile.h>
55 #include <TParticle.h>
56
57 ClassImp(AliAnalysisTaskESDMuonFilter)
58 ClassImp(AliAnalysisNonMuonTrackCuts)
59
60 ////////////////////////////////////////////////////////////////////////
61
62 AliAnalysisNonMuonTrackCuts::AliAnalysisNonMuonTrackCuts()
63 {
64   // default ctor 
65 }
66
67 Bool_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
75 AliAnalysisNonPrimaryVertices::AliAnalysisNonPrimaryVertices()
76 {
77   // default ctor   
78 }
79
80 Bool_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
103 AliAnalysisTaskESDMuonFilter::AliAnalysisTaskESDMuonFilter(Bool_t onlyMuon, Bool_t keepAllEvents, Int_t mcMode):
104   AliAnalysisTaskSE(),
105   fTrackFilter(0x0),
106   fEnableMuonAOD(kFALSE),
107   fEnableDimuonAOD(kFALSE),
108   fOnlyMuon(onlyMuon),
109   fKeepAllEvents(keepAllEvents),
110   fMCMode(mcMode)
111 {
112   // Default constructor
113 }
114
115 AliAnalysisTaskESDMuonFilter::AliAnalysisTaskESDMuonFilter(const char* name, Bool_t onlyMuon, Bool_t keepAllEvents, Int_t mcMode):
116   AliAnalysisTaskSE(name),
117   fTrackFilter(0x0),
118   fEnableMuonAOD(kFALSE),
119   fEnableDimuonAOD(kFALSE),
120   fOnlyMuon(onlyMuon),
121   fKeepAllEvents(keepAllEvents),
122   fMCMode(mcMode)
123 {
124   // Constructor
125 }
126
127 //______________________________________________________________________________
128 void AliAnalysisTaskESDMuonFilter::UserCreateOutputObjects()
129 {
130   // Create the output container
131   if (fTrackFilter) OutputTree()->GetUserInfo()->Add(fTrackFilter);
132 }
133
134 //______________________________________________________________________________
135 void 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   }
160   
161   if ( fMCMode > 0 ) 
162   {
163     cout << spaces.Data() << "Assuming work on MC data (i.e. will transmit MC branches)" << endl;
164   }
165 }
166
167 //______________________________________________________________________________
168 void 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
177   if (!ext) return;
178   
179   if ( fOnlyMuon ) 
180   {    
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);
187     
188     ext->DropUnspecifiedBranches(); // all branches not part of a FilterBranch call (below) will be dropped
189     
190     ext->FilterBranch("tracks",murep);    
191     ext->FilterBranch("vertices",murep);  
192     ext->FilterBranch("dimuons",murep);
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     }
204   }  
205 }
206
207 //______________________________________________________________________________
208 void AliAnalysisTaskESDMuonFilter::Init()
209 {
210   // Initialization
211   if(fEnableMuonAOD) AddFilteredAOD("AliAOD.Muons.root", "MuonEvents");
212   if(fEnableDimuonAOD) AddFilteredAOD("AliAOD.Dimuons.root", "DimuonEvents");    
213 }
214
215
216 //______________________________________________________________________________
217 void AliAnalysisTaskESDMuonFilter::UserExec(Option_t */*option*/)
218 {
219   // Execute analysis for current event                                     
220   
221   Long64_t ientry = Entry();
222   if(fDebug)printf("Muon Filter: Analysing event # %5d\n", (Int_t) ientry);
223   
224   ConvertESDtoAOD();
225 }
226
227 //______________________________________________________________________________
228 void AliAnalysisTaskESDMuonFilter::ConvertESDtoAOD() 
229 {
230   // ESD Muon Filter analysis task executed for each event
231   
232   AliCodeTimerAuto("",0);
233   
234   AliESDEvent* esd = dynamic_cast<AliESDEvent*>(InputEvent());
235   
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     
244   // Define arrays for muons
245   Double_t pos[3];
246   Double_t p[3];
247   Double_t pid[10];
248   
249   // has to be changed once the muon pid is provided by the ESD
250   for (Int_t i = 0; i < 10; pid[i++] = 0.) {}
251   pid[AliAODTrack::kMuon]=1.;
252   
253   AliAODHeader* header = AODEvent()->GetHeader();
254   AliAODTrack *aodTrack = 0x0;
255   AliESDMuonTrack *esdMuTrack = 0x0;
256   
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 
262   AliAODVertex *primary = AODEvent()->GetPrimaryVertex();
263   if(fDebug)primary->Print();
264   
265   // Loop on muon tracks to fill the AOD track branch
266   Int_t nMuTracks = esd->GetNumberOfMuonTracks();
267
268   for (Int_t iTrack=0; iTrack<nMuTracks; ++iTrack) esd->GetMuonTrack(iTrack)->SetESDEvent(esd);
269   
270   // Update number of positive and negative tracks from AOD event (M.G.)
271   Int_t nPosTracks = header->GetRefMultiplicityPos();
272   Int_t nNegTracks = header->GetRefMultiplicityNeg();
273   
274   // Access to the AOD container of dimuons
275   TClonesArray &dimuons = *(AODEvent()->GetDimuons());
276   AliAODDimuon *aodDimuon = 0x0;
277   
278   Int_t nMuons=0;
279   Int_t nDimuons=0;
280   Int_t jDimuons=0;
281   Int_t nMuonTrack[100];
282   
283   for(int imuon=0;imuon<100;imuon++) nMuonTrack[imuon]=0;
284   
285   for (Int_t nMuTrack = 0; nMuTrack < nMuTracks; ++nMuTrack)
286   {
287     esdMuTrack = esd->GetMuonTrack(nMuTrack);
288     
289     if (!esdMuTrack->ContainTrackerData()) continue;
290     
291     UInt_t selectInfo = 0;
292     // Track selection
293     if (fTrackFilter) {
294         selectInfo = fTrackFilter->IsSelected(esdMuTrack);
295         if (!selectInfo) {
296           continue;
297         }  
298     }
299     
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     
308     if(mcH)mcH->SelectParticle(esdMuTrack->GetLabel());
309     
310     aodTrack = new(tracks[jTracks++]) AliAODTrack(esdMuTrack->GetUniqueID(), // ID
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); 
325     
326     aodTrack->SetXYAtDCA(esdMuTrack->GetNonBendingCoorAtDCA(), esdMuTrack->GetBendingCoorAtDCA());
327     aodTrack->SetPxPyPzAtDCA(esdMuTrack->PxAtDCA(), esdMuTrack->PyAtDCA(), esdMuTrack->PzAtDCA());
328     aodTrack->SetRAtAbsorberEnd(esdMuTrack->GetRAtAbsorberEnd());
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());
335     aodTrack->Connected(esdMuTrack->IsConnected());
336     primary->AddDaughter(aodTrack);
337     
338     if (esdMuTrack->Charge() > 0) nPosTracks++;
339     else nNegTracks++;
340     
341     nMuonTrack[nMuons]= jTracks-1.;
342     ++nMuons;
343   }
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         }  
361       }
362     }
363   }
364   
365   
366   header->SetRefMultiplicity(jTracks); 
367   header->SetRefMultiplicityPos(nPosTracks);
368   header->SetRefMultiplicityNeg(nNegTracks);
369   header->SetNumberOfMuons(nMuons);
370   header->SetNumberOfDimuons(nDimuons);
371   
372   if ( fEnableMuonAOD && ( (nMuons>0) || fKeepAllEvents ) )
373   {
374     AliAODExtension *extMuons = dynamic_cast<AliAODHandler*>
375     ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler())->GetFilteredAOD("AliAOD.Muons.root");
376     extMuons->SelectEvent();
377   }
378   
379   if ( fEnableDimuonAOD && ( (nMuons>1) || fKeepAllEvents )  )
380   {
381     AliAODExtension *extDimuons = dynamic_cast<AliAODHandler*>
382     ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler())->GetFilteredAOD("AliAOD.Dimuons.root");
383     extDimuons->SelectEvent();
384   }
385   
386 }
387
388 void AliAnalysisTaskESDMuonFilter::Terminate(Option_t */*option*/)
389 {
390   // Terminate analysis
391   //
392   if (fDebug > 1) printf("AnalysisESDfilter: Terminate() \n");
393 }
394
395 void  AliAnalysisTaskESDMuonFilter::PrintMCInfo(AliStack *pStack,Int_t label)
396 {
397   // print mc info
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 }