KVEventSelector.cpp

#include "KVNucleusEvent.h"
#define KVEventSelector_cxx
#include "KVEventSelector.h"
#include <KVClassMonitor.h>
#include <TStyle.h>
#include "TPluginManager.h"
#include "TSystem.h"
#include "KVDataRepositoryManager.h"
#include "KVDataSetRepository.h"
#include "KVDataSetManager.h"
#ifdef WITH_PROOF
#include "TProof.h"
#endif
#include "KVDataSetAnalyser.h"
 
using namespace std;
 
ClassImp(KVEventSelector)
 
 
 
void KVEventSelector::Begin(TTree* /*tree*/)
{
   ParseOptions();
   // Need to parse options here for use in Terminate
   // Also, on PROOF, any KVDataAnalyser instance has to be passed to the workers
   // via the TSelector input list.
 
   if (IsOptGiven("CombinedOutputFile")) {
      fCombinedOutputFile = GetOpt("CombinedOutputFile");
   }
#ifdef WITH_PROOF
   else if (gProof) {
      // when running with PROOF, if the user calls SetCombinedOutputFile()
      // in InitAnalysis(), it will only be executed on the workers (in SlaveBegin()).
      // therefore we call InitAnalysis() here, but deactivate CreateTreeFile(),
      // AddTree() and AddHisto() in order to avoid interference with workers
      fDisableCreateTreeFile = kTRUE;
      if (gDataAnalyser) {
         gDataAnalyser->RegisterUserClass(this);
         gDataAnalyser->preInitAnalysis();
      }
      InitAnalysis();              //user initialisations for analysis
      if (gDataAnalyser) gDataAnalyser->postInitAnalysis();
      fDisableCreateTreeFile = kFALSE;
   }
#endif
   if (gDataAnalyser) {
      if (GetInputList()) {
         gDataAnalyser->AddJobDescriptionList(GetInputList());
         //GetInputList()->ls();
      }
   }
   if (IsOptGiven("AuxFiles")) {
      SetUpAuxEventChain();
      if (GetInputList()) GetInputList()->Add(fAuxChain);
   }
}
 
 
 
 
void KVEventSelector::SlaveBegin(TTree* /*tree*/)
{
   if (GetInputList() && GetInputList()->FindObject("JobDescriptionList")) {
      KVNameValueList* jdl = dynamic_cast<KVNameValueList*>(GetInputList()->FindObject("JobDescriptionList"));
      if (jdl) {
         KVDataAnalysisTask* task = nullptr;
         if (jdl->HasParameter("DataRepository")) {
            if (!gDataRepositoryManager) {
               new KVDataRepositoryManager;
               gDataRepositoryManager->Init();
            }
            gDataRepositoryManager->GetRepository(jdl->GetStringValue("DataRepository"))->cd();
            gDataSetManager->GetDataSet(jdl->GetStringValue("DataSet"))->cd();
            task = gDataSet->GetAnalysisTask(jdl->GetStringValue("AnalysisTask"));
         }
         else {
            if (!gDataSetManager) {
               gDataSetManager = new KVDataSetManager;
               gDataSetManager->Init();
            }
            task = gDataSetManager->GetAnalysisTaskAny(jdl->GetStringValue("AnalysisTask"));
         }
         gDataAnalyser = KVDataAnalyser::GetAnalyser(task->GetDataAnalyser());
         if (gDataSet && gDataAnalyser->InheritsFrom("KVDataSetAnalyser"))
            dynamic_cast<KVDataSetAnalyser*>(gDataAnalyser)->SetDataSet(gDataSet);
         gDataAnalyser->SetAnalysisTask(task);
         gDataAnalyser->RegisterUserClass(this);
         gDataAnalyser->SetProofMode((KVDataAnalyser::EProofMode)jdl->GetIntValue("PROOFMode"));
      }
   }
 
   ParseOptions();
 
   if (IsOptGiven("CombinedOutputFile")) {
      fCombinedOutputFile = GetOpt("CombinedOutputFile");
      Info("SlaveBegin", "Output file name = %s", fCombinedOutputFile.Data());
   }
 
   // tell the data analyser who we are
   if (gDataAnalyser) {
      gDataAnalyser->RegisterUserClass(this);
      gDataAnalyser->preInitAnalysis();
   }
   // make sure histo/tree creation are enabled
   fDisableCreateTreeFile = kFALSE;
   InitAnalysis();              //user initialisations for analysis
   if (gDataAnalyser) gDataAnalyser->postInitAnalysis();
 
   if (IsOptGiven("AuxFiles")) {
      // auxiliary files with PROOF
      if (GetInputList()) {
         fAuxChain = (TTree*)GetInputList()->FindObject(GetOpt("AuxTreeName"));
         InitFriendTree(fAuxChain, GetOpt("AuxBranchName"));
      }
   }
   Info("SlaveBegin", "fOutput->ls()");
   GetOutputList()->ls();
}
 
 
 
 
Bool_t KVEventSelector::Process(Long64_t entry)
{
   if (gDataAnalyser && gDataAnalyser->AbortProcessingLoop()) {
      // abort requested by external process
      Abort(Form("Job received KILL signal from batch system after %lld events - batch job probably needs more CPU time (see end of job statistics)", fEventsRead), kAbortFile);
      return kFALSE;
   }
 
   fTreeEntry = entry;
 
   if (gDataAnalyser && gDataAnalyser->CheckStatusUpdateInterval(fEventsRead))
      gDataAnalyser->DoStatusUpdate(fEventsRead);
 
   GetEntry(entry);
   if (gDataAnalyser) gDataAnalyser->preAnalysis();
   fEventsRead++;
   if (GetEvent()) {
      SetAnalysisFrame();//let user define any necessary reference frames for particles
      //apply particle selection criteria
      if (fPartCond.IsSet()) {
         for (auto& part : EventOKIterator(GetEvent())) {
            part.SetIsOK(fPartCond.Test(part));
         }
      }
 
      // initialise global variables at first event
      if (fFirstEvent) {
         for (auto p : fGlobalVariables) dynamic_cast<KVGVList*>(p)->Init();
         fFirstEvent = kFALSE;
      }
      RecalculateGlobalVariables();
   }
 
   Bool_t ok_anal = kTRUE;
   if (!fGlobalVariableAbortEventAnalysis) {
      // if global variables are defined, events are only analysed if no global variables
      // in the list have failed an event selection condition
      ok_anal = Analysis();     //user analysis
      if (gDataAnalyser) gDataAnalyser->postAnalysis();
   }
   CheckEndOfRun();
 
   return ok_anal;
}
 
 
 
 
void KVEventSelector::CheckEndOfRun()
{
   if (AtEndOfRun()) {
      Info("Process", "End of file reached after %lld events", fEventsRead);
      if (gDataAnalyser) gDataAnalyser->preEndRun();
      EndRun();
      if (gDataAnalyser) gDataAnalyser->postEndRun();
      fNotifyCalled = kFALSE;//Notify will be called when next file is opened (in TChain)
   }
}
 
 
 
 
void KVEventSelector::SlaveTerminate()
{
   if (writeFile) {
      auto sav_dir = gDirectory;//save current working directory
 
      writeFile->cd();
      // loop over all histos and trees, writing them in the output file
      if (GetHistoList().GetEntries()) {
         GetHistoList().R__FOR_EACH(TH1, Write)();
      }
      if (GetTreeList().GetEntries()) {
         GetTreeList().R__FOR_EACH(TTree, Write)(0, TObject::kOverwrite);
      }
      writeFile->Close();
#ifdef WITH_PROOF
      if (gDataAnalyser->GetProofMode() != KVDataAnalyser::EProofMode::None) {
         fOutput->Add(mergeFile);
      }
#endif
      sav_dir->cd();
   }
}
 
 
 
 
void KVEventSelector::Terminate()
{
   TDatime now;
   Info("Terminate", "Analysis ends at %s", now.AsString());
 
   if (gDataAnalyser) gDataAnalyser->preEndAnalysis();
   EndAnalysis();               //user end of analysis routine
   if (gDataAnalyser) gDataAnalyser->postEndAnalysis();
 
   if (GetInputList() && fAuxChain) GetInputList()->Remove(fAuxChain);
   SafeDelete(fAuxChain);
}
 
 
 
 
 
 
 
void KVEventSelector::RecalculateGlobalVariables()
{
   // Calculates all global variables in all global variable lists
 
   fGlobalVariableAbortEventAnalysis = false;
   for (auto p : fGlobalVariables) {
      auto gvl = dynamic_cast<KVGVList*>(p);
      gvl->CalculateGlobalVariables(GetEvent());
      // check if a variable event selection failed
      if (gvl->AbortEventAnalysis()) {
         fGlobalVariableAbortEventAnalysis = true;
         return;
      }
   }
}
 
 
 
 
 
void KVEventSelector::add_histo(TH1* histo)
{
   // Declare a histogram to be used in analysis.
   // This method must be called when using PROOF.
 
   if (fDisableCreateTreeFile) return;
 
   GetHistoList().Add(histo);
 
   if (writeFile) histo->SetDirectory(writeFile); // if output file is initialized, associate histo with it
   else create_output_file();// try to initialize output file
}
 
 
 
 
void KVEventSelector::add_tree(TTree* tree)
{
   // Declare a TTree to be used in analysis.
   // This method must be called when using PROOF.
 
   if (fDisableCreateTreeFile) return;
 
   tree->SetAutoSave();
   GetTreeList().Add(tree);
 
   if (writeFile) tree->SetDirectory(writeFile); // if output file is initialized, associate tree with it
   else create_output_file();// try to initialize output file
}
 
 
 
 
void KVEventSelector::create_output_file()
{
   // Create the file for saving histos and/or trees created during analysis.
   //
   // The name of the file must first be set using SetJobOutputFileName()
 
   if (fCombinedOutputFile == "") return;
 
   auto sav_dir = gDirectory;//save current working directory
   if (gDataAnalyser->GetProofMode() == KVDataAnalyser::EProofMode::None) {
      // Analysis running in 'normal' (non-PROOF) mode
      writeFile = TFile::Open(fCombinedOutputFile, "RECREATE");
   }
#ifdef WITH_PROOF
   else {
      // Analysis running with PROOF(Lite)
      mergeFile = new TProofOutputFile(fCombinedOutputFile, "M");
      mergeFile->SetOutputFileName(fCombinedOutputFile);
      writeFile = mergeFile->OpenFile("RECREATE");
      if (writeFile && writeFile->IsZombie()) SafeDelete(writeFile);
   }
#endif
   sav_dir->cd();//return to previous working directory
 
   if (!writeFile) {
      TString amsg = TString::Format("%s::create_output_file: could not create output ROOT file '%s'!",
                                     ClassName(), fCombinedOutputFile.Data());
      Abort(amsg, kAbortProcess);
   }
 
   // any previously declared histograms or trees must now be associated with the output file
   if (GetHistoList().GetEntries()) {
      GetHistoList().R__FOR_EACH(TH1, SetDirectory)(writeFile);
   }
   if (GetTreeList().GetEntries()) {
      GetTreeList().R__FOR_EACH(TTree, SetDirectory)(writeFile);
   }
}
 
 
 
 
void KVEventSelector::SetUpAuxEventChain()
{
   // Called by SlaveBegin() when user gives the following options:
   //
   //~~~~~~~~~~~~~~~~
   //    AuxFiles:               list of files containing "friend" TTrees to be made available during analysis
   //    AuxDir:                 directory in which to find AuxFiles
   //    AuxTreeName:            name of tree in AuxFiles containing KVEvent objects
   //    AuxBranchName:          name of branch in AuxFiles containing KVEvent objects
   //~~~~~~~~~~~~~~~~
 
   if (!IsOptGiven("AuxDir") || !IsOptGiven("AuxTreeName") || !IsOptGiven("AuxBranchName")) {
      KVError::Error(this, "SetUpAuxEventChain", "if AuxFiles option given, you must define AuxDir, AuxTreeName and AuxBranchName");
      return;
   }
   KVString filelist = GetOpt("AuxFiles");
   KVString filedir = GetOpt("AuxDir");
   if (!filedir.EndsWith("/")) filedir += "/";
   TChain* auxchain = new TChain(GetOpt("AuxTreeName"));
   filelist.Begin("|");
   while (!filelist.End()) {
      KVString path = filedir + filelist.Next();
      auxchain->Add(path);
   }
   InitFriendTree(auxchain, GetOpt("AuxBranchName"));
}
 
 
 
 
void KVEventSelector::ParseOptions()
{
   // Analyse comma-separated list of options given to TTree::Process
   // and store all `"option=value"` pairs in fOptionList.
   // Options can then be accessed using IsOptGiven(), GetOptString(), etc.
   //
   //~~~~~~~~~~~~~~~
   //     BranchName=xxxx  :  change name of branch in TTree containing data
   //     EventsReadInterval=N: print "+++ 12345 events processed +++" every N events
   //~~~~~~~~~~~~~~~
   //
 
   fOptionList.ParseOptions(GetOption());
 
   // check for branch name
   if (IsOptGiven("BranchName")) SetBranchName(GetOpt("BranchName"));
   // check for events read interval
   if (IsOptGiven("EventsReadInterval")) SetEventsReadInterval(GetOpt("EventsReadInterval").Atoi());
}
 
 
 
 
void KVEventSelector::Init(TTree* tree)
{
   Event = nullptr;
   // Set branch addresses and branch pointers
   if (!tree) return;
   fChain = tree;
   fChain->SetMakeClass(1);
   // When using PROOF, need to set tree pointer in KVDataAnalyser
   if (gDataAnalyser->GetProofMode() != KVDataAnalyser::EProofMode::None) {
      gDataAnalyser->SetAnalysedTree(tree);
   }
 
   if (strcmp(GetBranchName(), "")  && fChain->GetBranch(GetBranchName())) {
      Info("Init", "Analysing data in branch : %s", GetBranchName());
      fChain->SetBranchAddress(GetBranchName(), &Event, &b_Event);
   }
   else {
      KVError::Error(this, "Init", "Failed to link KVEvent object with a branch. Expected branch name=%s",
            GetBranchName());
   }
   //user additional branches addressing
   SetAdditionalBranchAddress();
   fEventsRead = 0;
 
}
 
 
 
 
void KVEventSelector::InitFriendTree(TTree* tree, const TString& branchname)
{
   //      @param tree
   //      @param branchname
   //
   //    Set up a "friend" TTree/TChain containing KVEvent-derived objects in branch 'branchname'
   //    N.B. this is not a "friend" in the sense of TTree::AddFriend, the main TTree and the
   //    "friend" TTree can have different numbers of entries
   //
   //    After calling this method at the beginning of the analysis, you can
   //    access any of the events stored in the "friend" by doing:
   //
   //~~~{.cpp}
   //   GetFriendTreeEntry(entry_number);
   //   KVEvent* friend_event = GetFriendEvent();
   //~~~
 
   AuxEvent = 0;
   fAuxChain = tree;
   fAuxChain->SetBranchAddress(branchname, &AuxEvent);
   fAuxChain->Print();
   fAuxChain->GetEntry(0);
   fAuxChain->GetTree()->GetEntry(0);
}
 
 
 
 
Bool_t KVEventSelector::Notify()
{
   if (fNotifyCalled) return kTRUE; // avoid multiple calls at beginning of analysis
   fNotifyCalled = kTRUE;
 
   Info("Notify", "Beginning analysis of file %s (%lld events)", fChain->GetCurrentFile()->GetName(), fChain->GetTree()->GetEntries());
 
   if (gDataAnalyser) gDataAnalyser->preInitRun();
   InitRun();                   //user initialisations for run
   if (gDataAnalyser) gDataAnalyser->postInitRun();
 
   KVClassMonitor::GetInstance()->SetInitStatistics();
 
   return kTRUE;
}