core/src/SCycleBaseNTuple.cxx

Go to the documentation of this file.

// $Id: SCycleBaseNTuple.cxx 335 2012-11-21 14:11:47Z krasznaa $
/***************************************************************************
 * @Project: SFrame - ROOT-based analysis framework for ATLAS
 * @Package: Core
 *
 * @author Stefan Ask       <Stefan.Ask@cern.ch>           - Manchester
 * @author David Berge      <David.Berge@cern.ch>          - CERN
 * @author Johannes Haller  <Johannes.Haller@cern.ch>      - Hamburg
 * @author A. Krasznahorkay <Attila.Krasznahorkay@cern.ch> - CERN/Debrecen
 *
 ***************************************************************************/

// System include(s):
#include <string.h>

// STL include(s):
#include <algorithm>
#include <typeinfo>

// ROOT include(s):
#include <TFile.h>
#include <TTree.h>
#include <TChain.h>
#include <TBranch.h>
#include <TROOT.h>
#include <TList.h>
#include <TFriendElement.h>
#include <TVirtualIndex.h>
#include <TTreeFormula.h>

// Local include(s):
#include "../include/SCycleBaseNTuple.h"
#include "../include/SInputData.h"
#include "../include/SCycleConfig.h"
#include "../include/SCycleOutput.h"
#include "../include/STreeType.h"

static const Double_t EPSILON = 1e-15;

#ifndef DOXYGEN_IGNORE
ClassImp( SCycleBaseNTuple )
#endif // DOXYGEN_IGNORE

SCycleBaseNTuple::SCycleBaseNTuple()
   : SCycleBaseBase(), m_inputTrees(), m_inputBranches(),
     m_outputTrees(), m_metaInputTrees(), m_outputVarPointers(), m_output( 0 ) {

   REPORT_VERBOSE( "SCycleBaseNTuple constructed" );
}

SCycleBaseNTuple::~SCycleBaseNTuple() {

   DeleteInputVariables();
   REPORT_VERBOSE( "SCycleBaseNTuple destructed" );
}

void SCycleBaseNTuple::SetNTupleOutput( TList* output ) {

   m_output = output;
   return;
}

TList* SCycleBaseNTuple::GetNTupleOutput() const {

   return m_output;
}

TTree* SCycleBaseNTuple::GetMetadataTree( const char* name ) const throw( SError ) {

   // The result tree:
   TTree* result = 0;

   // See if this is an input tree:
   try {
      result = GetInputMetadataTree( name );
      return result;
   } catch( const SError& error ) {
      if( error.request() <= SError::SkipFile ) {
         REPORT_VERBOSE( "Input metadata tree with name \"" << name << "\" not found" );
      } else {
         REPORT_ERROR( "Exception message caught with message: " << error.what() );
         throw;
      }
   }

   // See if this is an output tree:
   try {
      result = GetOutputMetadataTree( name );
      return result;
   } catch( const SError& error ) {
      if( error.request() <= SError::SkipFile ) {
         REPORT_VERBOSE( "Output metadata tree with name \"" << name << "\" not found" );
      } else {
         REPORT_ERROR( "Exception message caught with message: " << error.what() );
         throw;
      }
   }

   //
   // Throw an exception if the tree hasn't been found:
   //
   SError error( SError::SkipFile );
   error << "Couldn't find metadata TTree with name: " << name;
   throw error;

   return result;
}

TTree* SCycleBaseNTuple::GetInputMetadataTree( const char* name ) const throw( SError ) {

   //
   // Strip off the directory name from the given tree name:
   //
   TString tname( name );
   if( tname.Contains( "/" ) ) {
      REPORT_VERBOSE( "Tokenizing the metadata tree name: " << tname );
      TObjArray* array = tname.Tokenize( "/" );
      TObjString* real_name = dynamic_cast< TObjString* >( array->Last() );
      tname = real_name->GetString();
      delete array;
   }

   REPORT_VERBOSE( "Looking for input metadata tree with name: " << tname );

   //
   // Look for such a metadata tree:
   //
   for( std::vector< TTree* >::const_iterator it = m_metaInputTrees.begin();
        it != m_metaInputTrees.end(); ++it ) {
      if( *it ) {
         if( tname == ( *it )->GetName() ) {
            REPORT_VERBOSE( "Found input metadata tree with name " << tname 
                            << " at " << ( *it ) );
            return *it;
         }
      }
   }

   //
   // Throw an exception if the tree hasn't been found:
   //
   SError error( SError::SkipFile );
   error << "Couldn't find input metadata TTree with name: " << tname;
   throw error;

   return 0;
}

TTree* SCycleBaseNTuple::GetOutputMetadataTree( const char* name ) const throw( SError ) {

   //
   // Strip off the directory name from the given tree name:
   //
   TString tname( name );
   if( tname.Contains( "/" ) ) {
      REPORT_VERBOSE( "Tokenizing the metadata tree name: " << tname );
      TObjArray* array = tname.Tokenize( "/" );
      TObjString* real_name = dynamic_cast< TObjString* >( array->Last() );
      tname = real_name->GetString();
      delete array;
   }

   REPORT_VERBOSE( "Looking for output metadata tree with name: " << tname );

   //
   // Look for such a metadata tree:
   //
   for( std::vector< TTree* >::const_iterator it = m_metaOutputTrees.begin();
        it != m_metaOutputTrees.end(); ++it ) {
      if( *it ) {
         if( tname == ( *it )->GetName() ) {
            REPORT_VERBOSE( "Found output metadata tree with name " << tname 
                            << " at " << ( *it ) );
            return *it;
         }
      }
   }

   //
   // Throw an exception if the tree hasn't been found:
   //
   SError error( SError::SkipFile );
   error << "Couldn't find input metadata TTree with name: " << tname;
   throw error;

   return 0;
}

TTree* SCycleBaseNTuple::GetInputTree( const char* treeName ) const throw( SError ) {

   //
   // Look for such input tree:
   //
   TString tname( treeName );
   for( std::vector< TTree* >::const_iterator it = m_inputTrees.begin();
        it != m_inputTrees.end(); ++it ) {
      if( *it ) {
         if( tname == ( *it )->GetName() ) {
            REPORT_VERBOSE( "Found input tree with name " << treeName 
                            << " at " << ( *it ) );
            return *it;
         }
      }
   }

   //
   // Throw an exception if the tree hasn't been found:
   //
   SError error( SError::SkipFile );
   error << "Couldn't find input TTree with name: " << treeName;
   throw error;

   return 0;
}

TTree* SCycleBaseNTuple::GetOutputTree( const char* treeName ) const throw( SError ) {

   //
   // Look for such output tree:
   //
   TString tname( treeName );
   for( std::vector< TTree* >::const_iterator it = m_outputTrees.begin();
        it != m_outputTrees.end(); ++it ) {
      if( *it ) {
         if( tname == ( *it )->GetName() ) {
            REPORT_VERBOSE( "Found output tree with name " << treeName 
                            << " at " << ( *it ) );
            return *it;
         }
      }
   }

   //
   // Throw an exception if the tree hasn't been found:
   //
   SError error( SError::SkipFile );
   error << "Couldn't find output TTree with name: " << treeName;
   throw error;

   return 0;
}

void SCycleBaseNTuple::CreateOutputTrees( const SInputData& iD,
                                          std::vector< TTree* >& outTrees,
                                          TFile* outputFile ) throw( SError ) {

   // sanity checks
   if( outTrees.size() ) {
      m_logger << WARNING << "Vector of output trees is not empty in "
               << "\"CreateOutputTrees\"!"  << SLogger::endmsg;
   }

   // Clear the vector of output trees:
   m_outputTrees.clear();
   m_metaOutputTrees.clear();

   // Clear the vector of output variable pointers:
   m_outputVarPointers.clear();

   const std::vector< STree >* sOutTree = iD.GetTrees( STreeType::OutputSimpleTree );

   // Open output file / create output trees
   gROOT->cd();

   //
   // Create all the regular output trees, but don't create any branches in them
   // just yet.
   //
   const Int_t branchStyle = 1;
   const Int_t autoSave = 10000000;
   if( sOutTree ) {
      for( std::vector< STree >::const_iterator st = sOutTree->begin();
           st != sOutTree->end(); ++st ) {

         m_logger << DEBUG << "Creating output event tree with name: "
                  << st->treeName << SLogger::endmsg;

         TTree* tree = new TTree( st->treeName.Data(), TString( "Format: User" ) +
                                  ", data type: " + iD.GetType() );

         tree->SetAutoSave( autoSave );
         TTree::SetBranchStyle( branchStyle );

         outTrees.push_back( tree );
         m_outputTrees.push_back( tree );

         if( outputFile ) {
            tree->SetDirectory( outputFile );
            REPORT_VERBOSE( "Attached TTree \"" << st->treeName.Data()
                            << "\" to file: " << outputFile->GetName() );
         } else {
            SCycleOutput* out = new SCycleOutput( tree, st->treeName );
            m_output->Add( out );
            REPORT_VERBOSE( "Keeping TTree \"" << st->treeName.Data()
                            << "\" in memory" );
         }
      }
   }

   //
   // Create the metadata output trees. These don't have to be reported back to the
   // caller, since the user will be responsible for filling them.
   //
   const std::vector< STree >* sMetaTrees = iD.GetTrees( STreeType::OutputMetaTree );
   if( sMetaTrees ) {
      for( std::vector< STree >::const_iterator mt = sMetaTrees->begin();
           mt != sMetaTrees->end(); ++mt ) {

         //
         // Split the name into the name of the tree and the name of the
         // directory:
         //
         TString tname( mt->treeName ), dirname( "" );
         if( tname.Contains( "/" ) ) {
            REPORT_VERBOSE( "Tokenizing the metadata tree name: " << tname );
            TObjArray* array = tname.Tokenize( "/" );
            TObjString* real_name = dynamic_cast< TObjString* >( array->Last() );
            tname = real_name->GetString();
            if( array->GetSize() > 1 ) {
               for( Int_t i = 0; i < array->GetSize() - 1; ++i ) {
                  TObjString* dirletname = dynamic_cast< TObjString* >( array->At( i ) );
                  if( ! dirletname ) continue;
                  if( dirletname->GetString() == "" ) continue;
                  if( dirletname->GetString() == tname ) break;
                  dirname += ( dirname == "" ? dirletname->GetString() : "/" + dirletname->GetString() );
               }
            }
            delete array;
         }

         m_logger << DEBUG << "Creating output metadata tree with name: "
                  << tname  << " in directory: " << dirname << SLogger::endmsg;

         TTree* tree = new TTree( tname, TString( "Format: User" ) +
                                  ", data type: " + iD.GetType() );

         tree->SetAutoSave( autoSave );
         TTree::SetBranchStyle( branchStyle );

         m_metaOutputTrees.push_back( tree );

         if( outputFile ) {
            tree->SetDirectory( MakeSubDirectory( dirname, outputFile ) );
            REPORT_VERBOSE( "Attached TTree \"" << mt->treeName
                            << "\" to file: " << outputFile->GetName() );
         } else {
            SCycleOutput* out = new SCycleOutput( tree, tname, dirname );
            m_output->Add( out );
            REPORT_VERBOSE( "Keeping TTree \"" << mt->treeName
                            << "\" in memory" );
         }
      }
   }

   return;
}

void SCycleBaseNTuple::SaveOutputTrees( TDirectory* /*output*/ ) throw( SError ) {

   // Remember which directory we were in:
   TDirectory* savedir = gDirectory;

   // Save each regular output tree:
   for( std::vector< TTree* >::iterator tree = m_outputTrees.begin();
        tree != m_outputTrees.end(); ++tree ) {
      TDirectory* dir = ( *tree )->GetDirectory();
      if( dir ) dir->cd();
      ( *tree )->Write();
      ( *tree )->AutoSave();
      ( *tree )->SetDirectory( 0 );
      delete ( *tree );
   }

   // Save each metadata output tree:
   for( std::vector< TTree* >::iterator tree = m_metaOutputTrees.begin();
        tree != m_metaOutputTrees.end(); ++tree ) {
      TDirectory* dir = ( *tree )->GetDirectory();
      if( dir ) dir->cd();
      ( *tree )->Write();
      ( *tree )->AutoSave();
      ( *tree )->SetDirectory( 0 );
      delete ( *tree );
   }

   // Go back to the original directory:
   gDirectory = savedir;

   return;
}

void SCycleBaseNTuple::LoadInputTrees( const SInputData& iD,
                                       TTree* main_tree,
                                       TFile*& inputFile ) throw( SError ) {

   REPORT_VERBOSE( "Loading/accessing the event-level input trees" );

   //
   // Initialize some variables:
   //
   const std::vector< STree >* sInTree = iD.GetTrees( STreeType::InputSimpleTree );
   const std::vector< STree >* sMetaTree = iD.GetTrees( STreeType::InputMetaTree );
   Bool_t firstPassed = kFALSE;
   Int_t nEvents = 0;
   m_inputTrees.clear();
   m_inputBranches.clear();
   DeleteInputVariables();
   m_metaInputTrees.clear();

   //
   // Access the physical file that is currently being opened:
   //
   inputFile = 0;
   if( GetConfig().GetRunMode() == SCycleConfig::LOCAL ) {
      TChain* chain = dynamic_cast< TChain* >( main_tree );
      if( ! chain ) {
         throw SError( "In LOCAL running the input TTree is not a TChain!",
                       SError::StopExecution );
      }
      inputFile = chain->GetFile();
   } else if( GetConfig().GetRunMode() == SCycleConfig::PROOF ) {
      inputFile = main_tree->GetCurrentFile();
   } else {
      throw SError( "Running mode not recongnised", SError::SkipCycle );
      return;
   }
   if( ! inputFile ) {
      throw SError( "Couldn't get the input file pointer!", SError::SkipFile );
      return;
   }
   REPORT_VERBOSE( "Accessed the pointer to the input file: " << inputFile );

   //
   // Handle the regular input trees:
   //
   if( sInTree ) {
      for( std::vector< STree >::const_iterator st = sInTree->begin(); st != sInTree->end();
           ++st ) {

         REPORT_VERBOSE( "Now trying to access TTree: " << st->treeName );

         TTree* tree = dynamic_cast< TTree* >( inputFile->Get( st->treeName ) );
         if( ! tree ) {
            SError error( SError::SkipFile );
            error << "Tree " << st->treeName << " doesn't exist in File "
                  << inputFile->GetName();
            throw error;
         }

         // do we need this at all now that we loop over the branches
         // one-by-one?
         
         // Remove friends if any, for better performance
         bool skipFriends = true; // can be made configurable
         if( skipFriends ) {
            TList* flist = tree->GetListOfFriends();
            TIter nextf( flist );
            TFriendElement* fe = 0;
            while( ( fe = ( TFriendElement* ) nextf() ) ) {
               m_logger << DEBUG << "Remove friend " << fe->GetName() << " from tree " 
                        << tree->GetName() << SLogger::endmsg;
               flist->Remove( fe );
               delete fe;
               fe = 0;
            }
         }
         // Delete index if any, for better performance
         bool deleteIndex = true; // can be made configurable
         if( deleteIndex ) {
            if( tree->GetTreeIndex() ) {
               m_logger << DEBUG << "Delete index from tree " 
                        << tree->GetName() << SLogger::endmsg;
               tree->SetTreeIndex( 0 );
               delete tree->GetTreeIndex();
            }
         }

         m_inputTrees.push_back( tree );
         if( firstPassed && tree->GetEntries() != nEvents ) {
            SError error( SError::SkipFile );
            error << "Conflict in number of entries - Tree " << tree->GetName()
                  << " has " << tree->GetEntries() << ", NOT "
                  << nEvents;
            throw error;
         } else if( ! firstPassed ) {
            firstPassed = kTRUE;
            nEvents = tree->GetEntries();
         }
      }
   }

   //
   // Handle the metadata trees:
   //
   if( sMetaTree ) {
      for( std::vector< STree >::const_iterator mt = sMetaTree->begin();
           mt != sMetaTree->end(); ++mt ) {

         TTree* tree = dynamic_cast< TTree* >( inputFile->Get( mt->treeName ) );
         if( ! tree ) {
            SError error( SError::SkipFile );
            error << "Tree " << mt->treeName << " doesn't exist in File "
                  << inputFile->GetName();
            throw error;
         } else {
            m_metaInputTrees.push_back( tree );
         }
      }
   }

   return;
}

void SCycleBaseNTuple::GetEvent( Long64_t entry ) throw( SError ) {

   // Tell all trees to update their cache:
   for( std::vector< TTree* >::const_iterator it = m_inputTrees.begin();
        it != m_inputTrees.end(); ++it ) {
      ( *it )->LoadTree( entry );
   }

   // Load the current entry for all the regular input variables:
   for( std::vector< TBranch* >::const_iterator it = m_inputBranches.begin();
        it != m_inputBranches.end(); ++it ) {
      ( *it )->GetEntry( entry );
   }

   return;
}

Double_t SCycleBaseNTuple::CalculateWeight( const SInputData& inputData,
                                            Long64_t entry ) {

   // the type of this input data
   const TString& type    = inputData.GetType();
   const TString& version = inputData.GetVersion();

   Double_t weight = 0.;
   Double_t totlum = 0.;

   if( inputData.GetType() == "data" ) {
      REPORT_VERBOSE( "Data" );
      weight = 1.;
      return weight;
   }

   //iterate over vector of input data and addup the weight for the type of this input data
   for( std::vector< SInputData >::const_iterator iD = GetConfig().GetInputData().begin();
        iD != GetConfig().GetInputData().end(); ++iD ) {

      if( ( iD->GetType() == type ) && ( iD->GetVersion() == version ) ) {

         const std::vector< SGeneratorCut >& sgencuts = iD->GetSGeneratorCuts();
         Bool_t inside = kTRUE;

         for( std::vector< SGeneratorCut >::const_iterator sgc = sgencuts.begin();
              sgc != sgencuts.end(); ++sgc ) {
          
            // loop over the trees 
            for( std::vector< TTree* >::const_iterator it = m_inputTrees.begin();
                 it != m_inputTrees.end(); ++it ) {

               // consider the one with the correct name
               if( ( *it )->GetName() == sgc->GetTreeName() ) {
                  // check for this entry, if Formula is true  
                  TString teststring = sgc->GetFormula();
                  TTreeFormula f( "testFormula", teststring.Data(), *it );
                  // if true for all cuts, then add the Lumi of this input data to totlum
                  if( ! f.EvalInstance( entry ) ) inside = kFALSE;
                  break;
               }

            }
         }
         if( inside ) totlum += iD->GetScaledLumi();
      }
   }

   if( totlum > EPSILON ) 
      weight = ( GetConfig().GetTargetLumi() / totlum );
  
   return weight;
}

void SCycleBaseNTuple::ClearCachedTrees() {

   m_inputTrees.clear();
   m_inputBranches.clear();
   m_outputTrees.clear();
   m_metaInputTrees.clear();
   m_metaOutputTrees.clear();

   DeleteInputVariables();

   return;
}

const char* SCycleBaseNTuple::RootType( const char* typeid_type ) throw( SError ) {

   if( strlen( typeid_type ) != 1 ) {
      throw SError( "SCycleBaseNTuple::RootType received complex object description",
                    SError::StopExecution );
   }

   switch( typeid_type[ 0 ] ) {

   case 'c':
      return "B";
      break;
   case 'h':
      return "b";
      break;
   case 's':
      return "S";
      break;
   case 't':
      return "s";
      break;
   case 'i':
      return "I";
      break;
   case 'j':
      return "i";
      break;
   case 'f':
      return "F";
      break;
   case 'd':
      return "D";
      break;
   case 'x':
      return "L";
      break;
   case 'y':
      return "l";
      break;
   case 'b':
      return "O";
      break;

   }

   throw SError( "Unknown primitive type encountered: " + TString( typeid_type ),
                 SError::StopExecution );
   return "";
}

const char* SCycleBaseNTuple::TypeidType( const char* root_type ) throw( SError ) {

   if( ! strcmp( root_type, "Char_t" ) ) {
      return "c";
   } else if( ! strcmp( root_type, "UChar_t" ) ) {
      return "h";
   } else if( ! strcmp( root_type, "Short_t" ) ) {
      return "s";
   } else if( ! strcmp( root_type, "UShort_t" ) ) {
      return "t";
   } else if( ! strcmp( root_type, "Int_t" ) ) {
      return "i";
   } else if( ! strcmp( root_type, "UInt_t" ) ) {
      return "j";
   } else if( ! strcmp( root_type, "Float_t" ) ) {
      return "f";
   } else if( ! strcmp( root_type, "Double_t" ) ) {
      return "d";
   } else if( ! strcmp( root_type, "Long64_t" ) ) {
      return "x";
   } else if( ! strcmp( root_type, "ULong64_t" ) ) {
      return "y";
   } else if( ! strcmp( root_type, "Bool_t" ) ) {
      return "b";
   }

   throw SError( "Unknown ROOT primitive type encountered: " + TString( root_type ),
                 SError::StopExecution );
   return "";
}

void SCycleBaseNTuple::RegisterInputBranch( TBranch* br ) throw( SError ) {

   if( std::find( m_inputBranches.begin(), m_inputBranches.end(), br ) !=
       m_inputBranches.end() ) {
      m_logger << DEBUG << "Branch '" << br->GetName() << "' already registered!"
               << SLogger::endmsg;
   } else {
      m_inputBranches.push_back( br );
   }

   return;
}

void SCycleBaseNTuple::DeleteInputVariables() {

   for( std::list< TObject* >::iterator it = m_inputVarPointers.begin();
        it != m_inputVarPointers.end(); ++it ) {
      delete ( *it );
   }
   m_inputVarPointers.clear();

   return;
}

TDirectory*
SCycleBaseNTuple::MakeSubDirectory( const TString& path,
                                    TDirectory* dir ) const throw( SError ) {

   if( ! path.Length() ) return dir;

   TDirectory* result = 0;
   if( ! ( result = dir->GetDirectory( path ) ) ) {

      REPORT_VERBOSE( "Creating directory: "
                      << dir->GetPath() << "/" << path );

      //
      // Break up the path name at the slashes:
      //
      TObjArray* directories = path.Tokenize( "/" );

      //
      // Create each necessary directory:
      //
      result = dir;
      TDirectory* tempDir = 0;
      for( Int_t i = 0; i < directories->GetSize(); ++i ) {

         TObjString* path_element = dynamic_cast< TObjString* >( directories->At( i ) );
         if( ! path_element ) continue;
         if( path_element->GetString() == "" ) continue;

         REPORT_VERBOSE( "Accessing directory: " << path_element->GetString() );
         if( ! ( tempDir = result->GetDirectory( path_element->GetString() ) ) ) {
            REPORT_VERBOSE( "Directory doesn't exist, creating it..." );
            if( ! ( tempDir = result->mkdir( path_element->GetString(), "dummy title" ) ) ) {
               SError error( SError::SkipInputData );
               error << "Couldn't create directory: " << path
                     << " in the output file!";
               throw error;
            }
         }
         result = tempDir;

      }

      // Delete the object created by TString::Tokenize(...):
      delete directories;

   }

   return result;
}