fixes from Laurent for the MC branch in the AOD filters
[u/mrichter/AliRoot.git] / STEER / AliAODExtension.cxx
1 #include "AliAODExtension.h"
2
3 //-------------------------------------------------------------------------
4 //     Support class for AOD extensions. This is created by the user analysis
5 //     that requires a separate file for some AOD branches. The name of the 
6 //     AliAODExtension object is the file name where the AOD branches will be
7 //     stored.
8 //-------------------------------------------------------------------------
9
10 #include "AliAODBranchReplicator.h"
11 #include "AliAODEvent.h"
12 #include "AliCodeTimer.h"
13 #include "AliLog.h"
14 #include "Riostream.h"
15 #include "TDirectory.h"
16 #include "TFile.h"
17 #include "TList.h"
18 #include "TMap.h"
19 #include "TMap.h"
20 #include "TObjString.h"
21 #include "TROOT.h"
22 #include "TString.h"
23 #include "TTree.h"
24
25 ClassImp(AliAODExtension)
26
27 //______________________________________________________________________________
28 AliAODExtension::AliAODExtension() : TNamed(), 
29 fAODEvent(0), fTreeE(0), fFileE(0), fNtotal(0), fNpassed(0), 
30 fSelected(kFALSE), fRepFiMap(0x0), fRepFiList(0x0), fEnableReferences(kTRUE), fObjectList(0x0)
31 {
32   // default ctor
33 }
34
35 //______________________________________________________________________________
36 AliAODExtension::AliAODExtension(const char* name, const char* title, Bool_t isfilter)
37 :TNamed(name,title), 
38 fAODEvent(0), 
39 fTreeE(0), 
40 fFileE(0), 
41 fNtotal(0), 
42 fNpassed(0),
43 fSelected(kFALSE),
44 fRepFiMap(0x0),
45 fRepFiList(0x0),
46 fEnableReferences(kTRUE),
47 fObjectList(0x0)
48 {
49   // Constructor.
50   if (isfilter) {
51     TObject::SetBit(kFilteredAOD);
52     printf("####### Added AOD filter %s\n", name);
53   } else printf("####### Added AOD extension %s\n", name);
54   KeepUnspecifiedBranches();
55 }   
56
57 //______________________________________________________________________________
58 AliAODExtension::~AliAODExtension()
59 {
60   // Destructor.
61   if(fFileE){
62     // is already handled in TerminateIO
63     fFileE->Close();
64     delete fFileE;
65     fTreeE = 0;
66     fAODEvent = 0;
67   }
68   if (fTreeE) delete fTreeE;
69   if (fRepFiMap) fRepFiMap->DeleteAll();
70   delete fRepFiMap; // the map is owner
71   delete fRepFiList; // the list is not
72   delete fObjectList; // not owner
73 }
74
75 //______________________________________________________________________________
76 void AliAODExtension::AddBranch(const char* cname, void* addobj)
77 {
78   // Add a new branch to the aod 
79   
80   if (!fAODEvent) {
81     char type[20];
82     gROOT->ProcessLine(Form("TString s_tmp; AliAnalysisManager::GetAnalysisManager()->GetAnalysisTypeString(s_tmp); sprintf((char*)%p, \"%%s\", s_tmp.Data());", type));
83     Init(type);
84   }
85   TDirectory *owd = gDirectory;
86   if (fFileE) {
87     fFileE->cd();
88   }
89   char** apointer = (char**) addobj;
90   TObject* obj = (TObject*) *apointer;
91   
92   fAODEvent->AddObject(obj);
93   
94   TString bname(obj->GetName());
95   
96   if (!fTreeE->FindBranch(bname.Data())) 
97   {
98     Bool_t acceptAdd(kTRUE);
99     
100     if ( TestBit(kDropUnspecifiedBranches) )
101     {
102       // check that this branch is in our list of specified ones...
103       // otherwise do not add it !
104       TIter next(fRepFiMap);
105       TObjString* p;
106       
107       acceptAdd=kFALSE;
108       
109       while ( ( p = static_cast<TObjString*>(next()) ) && !acceptAdd )
110       {
111         if ( p->String() == bname ) acceptAdd=kTRUE;
112       }
113     }
114     
115     if ( acceptAdd ) 
116     {
117       // Do the same as if we book via 
118       // TTree::Branch(TCollection*)
119       
120       fObjectList->Add(obj);
121
122       const Int_t kSplitlevel = 99; // default value in TTree::Branch()
123       const Int_t kBufsize = 32000; // default value in TTree::Branch()
124       
125       fTreeE->Bronch(bname.Data(), cname, 
126                      fAODEvent->GetList()->GetObjectRef(obj),
127                      kBufsize, kSplitlevel - 1);
128     }
129   }
130   owd->cd();
131 }
132
133 //______________________________________________________________________________
134 Bool_t AliAODExtension::FinishEvent()
135 {
136   // Fill current event.
137   fNtotal++;
138   if (!IsFilteredAOD()) {
139     fAODEvent->MakeEntriesReferencable();
140     fTreeE->Fill();
141     return kTRUE;
142   }  
143   // Filtered AOD. Fill only if event is selected.
144   if (!fSelected) return kTRUE;
145   
146   TIter next(fRepFiList);
147   
148   AliAODBranchReplicator* repfi;
149   
150   while ( ( repfi = static_cast<AliAODBranchReplicator*>(next()) ) )
151   {
152     repfi->ReplicateAndFilter(*fAODEvent);
153   }
154   fNpassed++;
155   fTreeE->Fill();
156   fSelected = kFALSE; // so that next event will not be selected unless demanded
157   return kTRUE;
158 }  
159
160 //______________________________________________________________________________
161 Bool_t AliAODExtension::Init(Option_t *option)
162 {
163   // Initialize IO.
164   
165   AliCodeTimerAuto(GetName(),0);
166   
167   if(!fAODEvent) 
168   {
169     fAODEvent = new AliAODEvent();    
170   }
171   
172   TDirectory *owd = gDirectory;
173   TString opt(option);
174   opt.ToLower();
175   
176   if (opt.Contains("proof")) 
177   {
178     // proof
179     // Merging via files. Need to access analysis manager via interpreter.
180     gROOT->ProcessLine(Form("AliAnalysisDataContainer *c_common_out = AliAnalysisManager::GetAnalysisManager()->GetCommonOutputContainer();"));
181     gROOT->ProcessLine(Form("AliAnalysisManager::GetAnalysisManager()->OpenProofFile(c_common_out, \"RECREATE\", \"%s\");", fName.Data()));
182     fFileE = gFile;
183   } 
184   else 
185   {
186     fFileE = new TFile(GetName(), "RECREATE");
187   }  
188   fTreeE = new TTree("aodTree", "AliAOD tree");
189   
190   delete fObjectList;
191   fObjectList = new TList;
192   fObjectList->SetOwner(kFALSE); // be explicit we're not the owner...
193   TList* inputList = fAODEvent->GetList();
194   TIter next(inputList);
195   TObject* o;
196   
197   while ( ( o = next() ) )
198   {
199     // Loop on the objects that are within the main AOD, and see what to do with them :
200     // - transmit them to our AOD as they are
201     // - filter them (by means of an AliAODBranchReplicator)
202     // - drop them completely
203     
204     Bool_t mustKeep(kFALSE);
205     
206     TString test(o->ClassName());
207     test.ToUpper();
208     if (test.Contains("HEADER"))
209     {
210       // do not allow to drop header branch
211       mustKeep=kTRUE;
212     }
213     
214     if ( fRepFiMap && !mustKeep )
215     {
216       // we have some replicators, so let's see what the relevant one decides about this object
217       TObject* specified = fRepFiMap->FindObject(o->GetName()); // FindObject finds key=o->GetName() in the map
218       if (specified)
219       {
220         AliAODBranchReplicator* repfi = dynamic_cast<AliAODBranchReplicator*>(fRepFiMap->GetValue(o->GetName())); // GetValue gets the replicator corresponding to key=o->GetName()
221         if ( repfi ) 
222         {        
223           TList* replicatedList = repfi->GetList();
224           if (replicatedList)
225           {
226             AliAODEvent::AssignIDtoCollection(replicatedList);
227             TIter nextRep(replicatedList);
228             TObject* objRep;
229             while ( ( objRep = nextRep() ) )
230             {
231               if ( !fObjectList->FindObject(objRep) ) // insure we're not adding several times the same object
232               {                
233                 fObjectList->Add(objRep);                  
234               }
235             }
236           }
237           else
238           {
239             AliError(Form("replicatedList from %s is null !",repfi->GetName()));
240           }
241         }
242       }
243       else
244       {
245         if ( !TestBit(kDropUnspecifiedBranches) ) 
246         {
247           // object o will be transmitted to the output AOD, unchanged
248           fObjectList->Add(o);
249         }
250       }
251     } 
252     else
253     {
254       // no replicator, so decide based on the policy about dropping unspecified branches
255       if ( mustKeep || !TestBit(kDropUnspecifiedBranches) )
256       {
257         // object o will be transmitted to the output AOD, unchanged
258         fObjectList->Add(o);
259       }
260     }
261   }
262     
263   if (fEnableReferences) 
264   {
265     fTreeE->BranchRef();    
266   }
267     
268   fTreeE->Branch(fObjectList);
269   
270   owd->cd();
271   
272   return kTRUE;
273 }
274
275 //______________________________________________________________________________
276 void AliAODExtension::Print(Option_t* opt) const
277 {
278   // Print info about this extension
279   
280   cout << opt << Form("%s - %s - %s - aod %p",IsFilteredAOD() ? "FilteredAOD" : "Extension",
281                       GetName(),GetTitle(),GetAOD()) << endl;
282   if ( !fEnableReferences ) 
283   {
284     cout << opt << opt << "References are disabled ! Hope you know what you are doing !" << endl;
285   }
286   if ( TestBit(kDropUnspecifiedBranches) )
287   {
288     cout << opt << opt << "All branches not explicitely specified will be dropped" << endl;
289   }
290   
291   TIter next(fRepFiMap);
292   TObjString* s;
293   
294   while ( ( s = static_cast<TObjString*>(next()) ) )
295   {
296     AliAODBranchReplicator* br = static_cast<AliAODBranchReplicator*>(fRepFiMap->GetValue(s->String().Data()));
297     
298     cout << opt << opt << "Branch " << s->String();
299     if (br)
300     {
301       cout << " will be filtered by class " << br->ClassName();
302     }
303     else
304     {
305       cout << " will be transmitted as is";
306     }
307     cout << endl;
308   }
309 }
310
311 //______________________________________________________________________________
312 void AliAODExtension::SetEvent(AliAODEvent* event)
313 {
314   // Connects to an external event
315   if (!IsFilteredAOD()) {
316     Error("SetEvent", "Not allowed to set external event for non filtered AOD's");   
317     return;
318   }
319   fAODEvent = event;
320 }
321
322 //______________________________________________________________________________
323 void AliAODExtension::AddAODtoTreeUserInfo()
324 {
325   // Add aod event to tree user info
326   
327   if (!fTreeE) return;
328   
329   AliAODEvent* aodEvent(fAODEvent);
330   
331   if ( IsFilteredAOD() )
332   {
333     // cannot attach fAODEvent (which is shared with our AliAODHandler mother)
334     // so we create a custom (empty) AliAODEvent 
335     // Has also the advantage we can specify only the list of objects
336     // that are actually there in this filtered aod
337     //
338     aodEvent = new AliAODEvent;
339     TIter nextObj(fObjectList);
340     TObject* o;
341     while ( ( o = nextObj() ) ) 
342     {
343       aodEvent->AddObject(o);
344     }    
345   }
346   
347   fTreeE->GetUserInfo()->Add(aodEvent);
348 }
349
350 //______________________________________________________________________________
351 Bool_t AliAODExtension::TerminateIO()
352 {
353   // Terminate IO
354   if (TObject::TestBit(kFilteredAOD))
355     printf("AOD Filter %s: events processed: %d   passed: %d\n", GetName(), fNtotal, fNpassed);
356   else
357     printf("AOD extension %s: events processed: %d\n", GetName(), fNtotal);
358   if (fFileE) 
359   {
360     fFileE->Write();
361     fFileE->Close();
362     delete fFileE;
363     fFileE = 0;
364     fTreeE = 0;
365     fAODEvent = 0;
366   }
367   return kTRUE;
368 }
369
370 //______________________________________________________________________________
371 void AliAODExtension::FilterBranch(const char* branchName, AliAODBranchReplicator* repfi)
372 {
373   // Specify a filter/replicator for a given branch
374   //
375   // If repfi=0x0, this will disable the branch (in the output) completely.
376   //
377   // repfi is adopted by this class, i.e. user should not delete it.
378   //
379   // WARNING : branch name must be exact.
380   //
381   // See also the documentation for AliAODBranchReplicator class.
382   //
383   
384   if (!fRepFiMap)
385   {
386     fRepFiMap = new TMap;
387     fRepFiMap->SetOwnerKeyValue(kTRUE,kTRUE);
388     fRepFiList = new TList;
389     fRepFiList->SetOwner(kFALSE);
390   }
391   
392   fRepFiMap->Add(new TObjString(branchName),repfi);
393   
394   if (repfi && !fRepFiList->FindObject(repfi))
395   {
396     // insure we get unique and non-null replicators in this list
397     fRepFiList->Add(repfi);
398   }
399 }
400