unigen/UManager.cxx

#include <iostream>
using namespace std;

#include "TFile.h"
#include "TTree.h"
#include "TList.h"

#include "UTask.h"
#include "URun.h"
#include "UEvent.h"
#include "UManager.h"

ClassImp(UManager)

//____________________________________________________________________
//
// UManager
//
// This class is the singleton manager for organising I/O and task
// execution.
// SetInputFile and SetOutputFile are used to set the input and
// output files.
// User tasks, which are derived from the UTask class, with
// implemented Init and Exec functions, are to be added using
// AddTask function.
// Init and Run functions make initialisation and execution
//


UManager *UManager::fgInstance = NULL;


//--------------------------------------------------------------------
UManager::UManager()
{
  // Default constructor
  if(NULL != fgInstance) {
    Fatal("UManager", "Singleton class instantiated twice!");
  }
  fgInstance = this;
  fInFile = nullptr;
  fOutFile = nullptr;
  fInTree = nullptr;
  fOutTree = nullptr;
  fEvent = nullptr;
  fRun = nullptr;
  fTaskList = new TList();
  fWrite = kTRUE;
}
//--------------------------------------------------------------------


//--------------------------------------------------------------------
UManager::~UManager()
{
  // Destructor
  fgInstance = nullptr;
  CloseInputFile();
  CloseOutputFile();
  fInTree = nullptr;
  fOutTree = nullptr;
  fEvent = nullptr;
  fRun = nullptr;
  delete fTaskList;
}
//--------------------------------------------------------------------


//--------------------------------------------------------------------
UManager* UManager::Instance()
{
  // Get instance of singleton manager
  return fgInstance;
}
//--------------------------------------------------------------------


//--------------------------------------------------------------------
void UManager::SetInputFile(const char *name)
{
  // Open the input file and get the run description and tree of events,
  // which later can be accessed by GetRun and GetEvent respectively.
  // name - name of the input file
  fInFile = TFile::Open(name);
  if(nullptr == fInFile) {
    Fatal("SetInputFile", "Input file does not exist!");
  }

  fRun = (URun*) fInFile->Get("run");

  fInTree = (TTree*) fInFile->Get("events");
  if(nullptr == fInTree) {
    Fatal("SetInputFile", "Input file has no events tree!");
  }
  cout << "-I- UManager : Input file has " << fInTree->GetEntries() << " events"
       << endl;

  if(nullptr == fEvent) {
      fEvent = new UEvent();
  }
  fInTree->SetBranchAddress("event", &fEvent);
}
//--------------------------------------------------------------------


//--------------------------------------------------------------------
void UManager::CloseInputFile()
{
  // Close the input file
  if(nullptr != fInFile) {
    fInFile->Close();
    fInFile = nullptr;
  }
}
//--------------------------------------------------------------------


//--------------------------------------------------------------------
void UManager::SetOutputFile(const char *name, Bool_t writeTree)
{
  // Open the output file. Create the output tree with event branch.
  // name - name of the output file
  fOutFile = TFile::Open(name, "RECREATE");
  if(fOutFile == nullptr) {
    Fatal("SetOutputFile", "Output file can not be created!");
  }
  if(writeTree) {
    // Added by Grigory Nigmatkulov
    fOutFile->SetCompressionLevel(9);
    int bufsize =  65536;
    int split = 99;

    //fOutTree = new TTree("events", "Event data");
    fOutTree = new TTree("events", "Event data", split);
    fOutTree->SetAutoSave(1000000);
    if(fEvent == nullptr) {
      fEvent = new UEvent();
    }

    //fOutTree->Branch("event", &fEvent);
    fOutTree->Branch("event", &fEvent, bufsize, split);
  }
}

//--------------------------------------------------------------------
void UManager::CloseOutputFile()
{
  // Close the output file
  if(nullptr != fOutFile) {
    fOutFile->Close();
    fOutFile = nullptr;
  }
}

//--------------------------------------------------------------------
Int_t UManager::GetEntries()
{
  // Get the number of entries in input tree
  if(nullptr == fInTree) return 0;
  return fInTree->GetEntries();
}

//--------------------------------------------------------------------
void UManager::GetEntry(Int_t i)
{
  // Get entry from the input tree.
  // i - index of entry
  if(nullptr == fInTree) return;
  fInTree->GetEntry(i);
}

//--------------------------------------------------------------------
void UManager::Fill()
{
  // Fill the output tree
  if(nullptr != fOutTree) {
    if(fWrite) {
      fOutTree->Fill();
    } else {
      fWrite = kTRUE;
    }
  }
}

//--------------------------------------------------------------------
void UManager::AddTask(UTask *task) {
  // Add task to the task list.
  // task - pointer to the task
  fTaskList->Add(task);
}

//--------------------------------------------------------------------
void UManager::createRun(const char* generator, const char* comment,
			 const Int_t& aProj, const Int_t& zProj, const Double_t& pProj,
			 const Int_t& aTarg, const Int_t& zTarg, const Double_t& pTarg,
			 const Double_t& bMin, const Double_t& bMax, const Int_t& bWeight,
			 const Double_t& phiMin, const Double_t& phiMax,
			 const Double_t& sigma, const Int_t& nEvents) {
  // Creates the run description, which later can be accessed by GetRun method
  fRun = new URun(generator, comment,
		  aProj, zProj, pProj, aTarg, zTarg, pTarg,
		  bMin, bMax, bWeight, phiMin, phiMax, sigma, nEvents);
}

//--------------------------------------------------------------------
void UManager::Init()
{
  // Run initialisation of the tasks
  UTask *task;
  for(Int_t i = 0; i < fTaskList->GetSize(); i++) {
    task = (UTask*) fTaskList->At(i);
    task->Init();
  }
}

//--------------------------------------------------------------------
void UManager::Run(Int_t a, Int_t b)
{
  // Execute the tasks.
  // Run()    - process all events from input tree.
  // Run(i)   - process only i-th event
  // Run(i,j) - process events from i-th to j-th
  if(nullptr==fInTree && nullptr==fOutTree) return;

  UTask *task;
  if(nullptr==fInTree) {

    for(Int_t i = 0; i < a; i++) {
      for(Int_t i = 0; i < fTaskList->GetSize(); i++) {
	task = (UTask*) fTaskList->At(i);
	task->Exec();
      }
      Fill();
    }

  } else {

    if(-1==a && -1==b) {
      for(Int_t iEv = 0; iEv < GetEntries(); iEv++) {
	GetEntry(iEv);
	for(Int_t i = 0; i < fTaskList->GetSize(); i++) {
	  task = (UTask*) fTaskList->At(i);
	  task->Exec();
	}
	Fill();
      }
    } else if(a>-1 && -1==b) {
      GetEntry(a);
      for(Int_t i = 0; i < fTaskList->GetSize(); i++) {
	task = (UTask*) fTaskList->At(i);
	task->Exec();
      }
      Fill();
    } else {
      for(Int_t iEv = a; iEv < b; iEv++) {
	GetEntry(iEv);
	for(Int_t i = 0; i < fTaskList->GetSize(); i++) {
	  task = (UTask*) fTaskList->At(i);
	  task->Exec();
	}
	Fill();
      }
    }

  }

  for(Int_t i = 0; i < fTaskList->GetSize(); i++) {
    task = (UTask*) fTaskList->At(i);
    task->Finish();
  }

  if(nullptr != fOutFile) {
    if(nullptr != fOutTree) {
      fOutFile->cd();
      fOutTree->Write();
    }
    if(nullptr != fRun) {
      fOutFile->cd();
      fRun->Write();
    }
  }
}