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