8f7a1bf6c5b135aeca953dcba052b15e71631f83
[u/mrichter/AliRoot.git] / PWG / muon / AliAODMuonReplicator.cxx
1 /**************************************************************************
2 * Copyright(c) 1998-2013, 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 // $Id$
17
18 ///
19 /// Implementation of an AliAODBranchReplicator to produce slim muon and dimuon aods.
20 ///
21 /// This replicator is in charge of replicating the tracks,vertices,dimuons
22 /// (vzero, tzero, zdc, and, optionally, the SPD tracklets) branches
23 /// of the standard AOD into muon AODs (AliAOD.Muons.root and AliAOD.Dimuons.root)
24 ///
25 /// The tracks are filtered so that only muon tracks (and only muon tracks
26 /// that pass the trackCut if present) make it to the output aods
27 ///
28 /// The vertices are filtered so that only the primary (and pileup) vertices make it
29 /// to the output aods.
30 ///
31 /// The dimuons are recreated here, according to the set of tracks
32 /// that pass the trackCut (that set may be the same as the input set,
33 /// but to be 100% safe, we always recreate the dimuons).
34 ///
35 /// \author L. Aphecetche (Subatech)
36
37 #include "AliAODMuonReplicator.h"
38
39 #include "AliAODDimuon.h"
40 #include "AliAODEvent.h"
41 #include "AliAODMCHeader.h"
42 #include "AliAODMCParticle.h"
43 #include "AliAODTZERO.h"
44 #include "AliAODTrack.h"
45 #include "AliAODVZERO.h"
46 #include "AliAnalysisCuts.h"
47 #include <cassert>
48
49 /// \cond CLASSIMP
50 ClassImp(AliAODMuonReplicator)
51 /// \endcond
52
53 //_____________________________________________________________________________
54 AliAODMuonReplicator::AliAODMuonReplicator(const char* name, const char* title,
55                                            AliAnalysisCuts* trackCut,
56                                            AliAnalysisCuts* vertexCut,
57                                            Int_t mcMode,
58                                            Bool_t replicateHeader,
59                                            Bool_t replicateTracklets)
60 :AliAODBranchReplicator(name,title), 
61 fTrackCut(trackCut), fTracks(0x0), 
62 fVertexCut(vertexCut), fVertices(0x0), 
63 fDimuons(0x0),
64 fVZERO(0x0),
65 fTZERO(0x0),
66 fHeader(0x0),
67 fTracklets(0x0),
68 fZDC(0x0),
69 fList(0x0),
70 fMCParticles(0x0),
71 fMCHeader(0x0),
72 fMCMode(mcMode),
73 fLabelMap(),
74 fParticleSelected(),
75 fReplicateHeader(replicateHeader),
76 fReplicateTracklets(replicateTracklets)
77 {
78   /// default ctor
79   ///
80   /// \param trackCut if present will filter out tracks
81   /// \param vertexCut if present will filter out vertices
82   /// \param mcMode what to do with MC information (0: skip it, 1: copy all,
83   /// 2: copy only for events with at least one muon )
84   /// \param replicateHeader whether or not to handle the replication of the
85   /// AOD header branch
86   /// \param replicateTracklets whether or not to include the SPD tracklets branch
87   ///
88 }
89
90 //_____________________________________________________________________________
91 AliAODMuonReplicator::~AliAODMuonReplicator()
92 {
93   /// dtor
94   delete fTrackCut;
95   delete fVertexCut;
96   delete fList;
97 }
98
99 //_____________________________________________________________________________
100 void AliAODMuonReplicator::SelectParticle(Int_t i)
101 {
102   /// taking the absolute values here, need to take care
103   /// of negative daughter and mother
104   /// IDs when setting!
105   
106   if (!IsParticleSelected(TMath::Abs(i)))
107   {
108     fParticleSelected.Add(TMath::Abs(i),1);    
109   }
110 }
111
112 //_____________________________________________________________________________
113 Bool_t AliAODMuonReplicator::IsParticleSelected(Int_t i)  
114 {
115   /// taking the absolute values here, need to take
116   /// care with negative daughter and mother
117   /// IDs when setting!
118   return (fParticleSelected.GetValue(TMath::Abs(i))==1);
119 }
120
121
122 //_____________________________________________________________________________
123 void AliAODMuonReplicator::CreateLabelMap(const AliAODEvent& source)
124 {  
125   //
126   // this should be called once all selections are done 
127   //
128   
129   fLabelMap.Delete();
130   
131   TClonesArray* mcParticles = static_cast<TClonesArray*>(source.FindListObject(AliAODMCParticle::StdBranchName()));
132   
133   Int_t i(0);
134   Int_t j(0);
135   
136   TIter next(mcParticles);
137   
138   while ( next() ) 
139   {
140     if (IsParticleSelected(i))
141     {
142       fLabelMap.Add(i,j++);
143     }
144     ++i;
145   }
146 }
147
148 //_____________________________________________________________________________
149 Int_t AliAODMuonReplicator::GetNewLabel(Int_t i) 
150 {
151   /// Gets the label from the new created Map
152   /// Call CreatLabelMap before
153   /// otherwise only 0 returned
154   return fLabelMap.GetValue(TMath::Abs(i));
155 }
156
157 //_____________________________________________________________________________
158 void AliAODMuonReplicator::FilterMC(const AliAODEvent& source)
159 {
160   /// Filter MC information
161
162   fMCHeader->Reset();
163   fMCParticles->Clear("C");
164
165   AliAODMCHeader* mcHeader(0x0);
166   TClonesArray* mcParticles(0x0);
167   
168   fParticleSelected.Delete();
169   
170   if ( fMCMode>=2 && !fTracks->GetEntries() ) return;
171   // for fMCMode>=2 we only copy MC information for events where there's at least one muon track
172     
173   mcHeader = static_cast<AliAODMCHeader*>(source.FindListObject(AliAODMCHeader::StdBranchName()));
174   
175   if ( mcHeader ) 
176   {
177     *fMCHeader = *mcHeader;
178   }
179   
180   mcParticles = static_cast<TClonesArray*>(source.FindListObject(AliAODMCParticle::StdBranchName()));
181   
182   if ( mcParticles && fMCMode>=2 )
183   {
184     // loop on (kept) muon tracks to find their ancestors
185     TIter nextMT(fTracks);
186     AliAODTrack* mt;
187     
188     while ( ( mt = static_cast<AliAODTrack*>(nextMT()) ) )
189     {
190       Int_t label = mt->GetLabel();
191       
192       while ( label >= 0 ) 
193       {
194         SelectParticle(label);
195         AliAODMCParticle* mother = static_cast<AliAODMCParticle*>(mcParticles->UncheckedAt(label));
196         if (!mother)
197         {
198           AliError("Got a null mother ! Check that !");
199           label = -1;
200         }
201         else
202         {
203           label = mother->GetMother();
204         }
205       }
206     }
207     
208     CreateLabelMap(source);
209     
210     // Actual filtering and label remapping (shamelessly taken for the implementation of AliAODHandler::StoreMCParticles)
211     TIter nextMC(mcParticles);
212     AliAODMCParticle* p;
213     Int_t nmc(0);
214     Int_t nmcout(0);
215     
216     while ( ( p = static_cast<AliAODMCParticle*>(nextMC()) ) )
217     {
218       AliAODMCParticle c(*p);
219       
220       if ( IsParticleSelected(nmc) )
221       {
222         // 
223         Int_t d0 =  p->GetDaughter(0);
224         Int_t d1 =  p->GetDaughter(1);
225         Int_t m =   p->GetMother();
226         
227         // other than for the track labels, negative values mean
228         // no daughter/mother so preserve it
229         
230         if(d0<0 && d1<0)
231         {
232           // no first daughter -> no second daughter
233           // nothing to be done
234           // second condition not needed just for sanity check at the end
235           c.SetDaughter(0,d0);
236           c.SetDaughter(1,d1);
237         } else if(d1 < 0 && d0 >= 0) 
238         {
239           // Only one daughter
240           // second condition not needed just for sanity check at the end
241           if(IsParticleSelected(d0))
242           {
243             c.SetDaughter(0,GetNewLabel(d0));
244           } else 
245           {
246             c.SetDaughter(0,-1);
247           }
248           c.SetDaughter(1,d1);
249         }
250         else if (d0 > 0 && d1 > 0 )
251         {
252           // we have two or more daughters loop on the stack to see if they are
253           // selected
254           Int_t d0tmp = -1;
255           Int_t d1tmp = -1;
256           for (int id = d0; id<=d1;++id)
257           {
258             if (IsParticleSelected(id))
259             {
260               if(d0tmp==-1)
261               {
262                 // first time
263                 d0tmp = GetNewLabel(id);
264                 d1tmp = d0tmp; // this is to have the same schema as on the stack i.e. with one daugther d0 and d1 are the same 
265               }
266               else d1tmp = GetNewLabel(id);
267             }
268           }
269           c.SetDaughter(0,d0tmp);
270           c.SetDaughter(1,d1tmp);
271         } else 
272         {
273           AliError(Form("Unxpected indices %d %d",d0,d1));
274         }
275         
276         if ( m < 0 )
277         {
278           c.SetMother(m);
279         } else 
280         {
281           if (IsParticleSelected(m)) 
282           {
283             c.SetMother(GetNewLabel(m));              
284           }
285           else 
286           {
287             AliError(Form("PROBLEM Mother not selected %d", m));              
288           }
289         }
290         
291         new ((*fMCParticles)[nmcout++]) AliAODMCParticle(c);
292       }
293       
294       ++nmc;        
295     }      
296     
297     // now remap the tracks...
298     
299     TIter nextTrack(fTracks);
300     AliAODTrack* t;
301     
302     while ( ( t = static_cast<AliAODTrack*>(nextTrack()) ) )
303     {
304       t->SetLabel(GetNewLabel(t->GetLabel()));
305     }
306     
307   }
308   else if ( mcParticles ) 
309   {
310     // simple copy of input MC particles to ouput MC particles
311     TIter nextMC(mcParticles);
312     AliAODMCParticle* p;
313     Int_t nmcout(0);
314     
315     while ( ( p = static_cast<AliAODMCParticle*>(nextMC()) ) )
316     {
317       new ((*fMCParticles)[nmcout++]) AliAODMCParticle(*p);
318     }
319   }
320   
321   AliDebug(1,Form("input mc %d output mc %d",
322                   mcParticles ? mcParticles->GetEntries() : 0,
323                   fMCParticles ? fMCParticles->GetEntries() : 0));
324   
325 }
326
327 //_____________________________________________________________________________
328 TList* AliAODMuonReplicator::GetList() const
329 {
330   /// return (and build if not already done) our internal list of managed objects
331   if (!fList)
332   {
333     if ( fReplicateHeader )
334     {
335       fHeader = new AliAODHeader;
336     }
337     
338     if ( fReplicateTracklets )
339     {
340       fTracklets = new AliAODTracklets;
341       fTracklets->SetName("tracklets");
342     }
343
344     fTracks = new TClonesArray("AliAODTrack",30);
345     fTracks->SetName("tracks");    
346     
347     fVertices = new TClonesArray("AliAODVertex",2);
348     fVertices->SetName("vertices");    
349     
350     fDimuons = new TClonesArray("AliAODDimuon",20);
351     fDimuons->SetName("dimuons");
352     
353     fVZERO = new AliAODVZERO;
354     
355     fTZERO = new AliAODTZERO;
356     
357     fZDC = new AliAODZDC;
358     
359     fList = new TList;
360     fList->SetOwner(kTRUE);
361     
362     fList->Add(fHeader);
363     fList->Add(fTracks);
364     fList->Add(fVertices);
365     fList->Add(fTracklets);
366     fList->Add(fDimuons);
367     fList->Add(fVZERO);
368     fList->Add(fTZERO);
369     fList->Add(fZDC);
370     
371     if ( fMCMode > 0 )
372     {
373       fMCHeader = new AliAODMCHeader;    
374       fMCParticles = new TClonesArray("AliAODMCParticle",1000);
375       fMCParticles->SetName(AliAODMCParticle::StdBranchName());
376       fList->Add(fMCHeader);
377       fList->Add(fMCParticles);
378     }
379   }
380   return fList;
381 }
382
383 //_____________________________________________________________________________
384 void AliAODMuonReplicator::ReplicateAndFilter(const AliAODEvent& source)
385 {
386   /// Replicate (and filter if filters are there) the relevant
387   /// parts we're interested in AODEvent
388   
389   assert(fTracks!=0x0);
390   
391   if (fReplicateHeader)
392   {
393     *fHeader = *(source.GetHeader());
394   }
395
396   if (fReplicateTracklets)
397   {
398     *fTracklets = *(source.GetTracklets());
399   }
400   
401   *fVZERO = *(source.GetVZEROData());
402
403   *fTZERO = *(source.GetTZEROData());
404
405   *fZDC = *(source.GetZDCData());
406   
407   fTracks->Clear("C");
408   TIter next(source.GetTracks());
409   AliAODTrack* t;
410   Int_t nMuons=0;
411   Int_t inputMuons=0;
412   
413   while (( t = static_cast<AliAODTrack*>(next()) )) {
414
415     if (t->IsMuonTrack() || t->IsMuonGlobalTrack()) ++inputMuons;    // just a counter: MUON and MUON+MFT tracks before track cuts are applied
416      
417     // MUON and MUON+MFT tracks selected                    // AU
418     if (!fTrackCut || fTrackCut->IsSelected(t)) {
419       new ((*fTracks)[nMuons++]) AliAODTrack(*t);
420     }
421
422   }
423   
424   assert(fVertices!=0x0);
425   fVertices->Clear("C");
426   TIter nextV(source.GetVertices());
427   AliAODVertex* v;
428   Int_t nvertices(0);
429   
430   while ( ( v = static_cast<AliAODVertex*>(nextV()) ) ) {
431     if ( !fVertexCut || fVertexCut->IsSelected(v) ) {
432       AliAODVertex* tmp = v->CloneWithoutRefs();
433       AliAODVertex* copiedVertex = new((*fVertices)[nvertices++]) AliAODVertex(*tmp);
434       // to insure the main vertex retains the ncontributors information
435       // (which is otherwise computed dynamically from
436       // references to tracks, which we do not keep in muon aods...)
437       // we set it here
438       copiedVertex->SetNContributors(v->GetNContributors()); 
439       delete tmp;
440     }
441   }
442   
443   fDimuons->Clear("C");
444   
445   // as there might be a track cut (different from the one of the main filtering),
446   // we must recreate the dimuon completely from scratch to be 100% safe...
447
448   Int_t nDimuons=0;
449
450   // Dimuons built of 2 MUON tracks or 2 MUON+MFT tracks                   // AU
451   for (Int_t i=0; i<nMuons; ++i) {
452     for (Int_t j=i+1; j<nMuons; ++j) {
453       if ( (((AliAODTrack*) fTracks->At(i))->IsMuonTrack()       && ((AliAODTrack*) fTracks->At(j))->IsMuonTrack()) ||
454            (((AliAODTrack*) fTracks->At(i))->IsMuonGlobalTrack() && ((AliAODTrack*) fTracks->At(j))->IsMuonGlobalTrack()) ) {
455         new((*fDimuons)[nDimuons++]) AliAODDimuon(fTracks->At(i), fTracks->At(j));
456       }
457     }
458   }
459
460   AliDebug(1,Form("input mu tracks=%d tracks=%d vertices=%d ndimuons=%d",
461                   inputMuons,fTracks->GetEntries(),fVertices->GetEntries(),fDimuons->GetEntries()));
462
463   // Finally, deal with MC information, if needed
464
465   if (fMCMode>0) FilterMC(source);
466
467 }
468