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