// RooTreeDataStore is the abstract base class for data collection that
// use a TTree as internal storage mechanism
// END_HTML
#include "RooFit.h"
#include "RooMsgService.h"
#include "RooTreeDataStore.h"
#include "Riostream.h"
#include "TTree.h"
#include "TChain.h"
#include "TDirectory.h"
#include "TROOT.h"
#include "RooFormulaVar.h"
#include "RooRealVar.h"
#include <iomanip>
using namespace std ;
ClassImp(RooTreeDataStore)
;
Int_t RooTreeDataStore::_defTreeBufSize = 4096 ;
RooTreeDataStore::RooTreeDataStore() :
_tree(0),
_cacheTree(0),
_defCtor(kTRUE)
{
}
RooTreeDataStore::RooTreeDataStore(TTree* t, const RooArgSet& vars) :
RooAbsDataStore("blah","blah",vars),
_tree(t),
_cacheTree(0),
_defCtor(kTRUE)
{
}
RooTreeDataStore::RooTreeDataStore(const char* name, const char* title, const RooArgSet& vars) :
RooAbsDataStore(name,title,vars),
_tree(0),
_cacheTree(0),
_defCtor(kFALSE)
{
initialize() ;
}
RooTreeDataStore::RooTreeDataStore(const char* name, const char* title, const RooArgSet& vars, TTree& t, const RooFormulaVar& select) :
RooAbsDataStore(name,title,vars),
_tree(0),
_cacheTree(0),
_defCtor(kFALSE)
{
initialize() ;
loadValues(&t,&select) ;
}
RooTreeDataStore::RooTreeDataStore(const char* name, const char* title, const RooArgSet& vars, TTree& t, const char* selExpr) :
RooAbsDataStore(name,title,vars),
_tree(0),
_cacheTree(0),
_defCtor(kFALSE)
{
initialize() ;
if (selExpr && *selExpr) {
RooFormulaVar select(selExpr,selExpr,_vars) ;
loadValues(&t,&select);
} else {
loadValues(&t);
}
}
RooTreeDataStore::RooTreeDataStore(const char* name, const char* title, const RooArgSet& vars, const RooAbsDataStore& tds, const RooFormulaVar& select) :
RooAbsDataStore(name,title,vars),
_tree(0),
_cacheTree(0),
_defCtor(kFALSE)
{
initialize() ;
loadValues(&tds,&select) ;
}
RooTreeDataStore::RooTreeDataStore(const char* name, const char* title, const RooArgSet& vars, const RooAbsDataStore& ads, const char* selExpr) :
RooAbsDataStore(name,title,vars),
_tree(0),
_cacheTree(0),
_defCtor(kFALSE)
{
initialize() ;
if (selExpr && *selExpr) {
RooFormulaVar select(selExpr,selExpr,_vars) ;
loadValues(&ads,&select);
} else {
loadValues(&ads);
}
}
RooTreeDataStore::RooTreeDataStore(const char *name, const char *title, RooAbsDataStore& tds,
const RooArgSet& vars, const RooFormulaVar* cutVar, const char* cutRange,
Int_t nStart, Int_t nStop, Bool_t ) :
RooAbsDataStore(name,title,vars), _defCtor(kFALSE)
{
_tree = 0 ;
_cacheTree = 0 ;
createTree(name,title) ;
RooFormulaVar* cloneVar = 0;
if (cutVar) {
cloneVar = (RooFormulaVar*) cutVar->cloneTree() ;
cloneVar->attachDataStore(tds) ;
}
initialize();
initCache(((RooTreeDataStore&)tds)._cachedVars) ;
loadValues(&tds,cloneVar,cutRange,nStart,nStop);
if (cloneVar) delete cloneVar ;
}
void RooTreeDataStore::initCache(const RooArgSet& cachedVars)
{
_cachedVars.removeAll() ;
TIterator* iter = cachedVars.createIterator() ;
RooAbsArg *var;
while((0 != (var= (RooAbsArg*)iter->Next()))) {
var->attachToTree(*_cacheTree,_defTreeBufSize) ;
_cachedVars.add(*var) ;
}
delete iter ;
}
RooTreeDataStore::RooTreeDataStore(const RooTreeDataStore& other, const char* newname) :
RooAbsDataStore(other,newname),
_tree(0),
_cacheTree(0),
_defCtor(kFALSE)
{
initialize() ;
loadValues(&other) ;
}
RooTreeDataStore::RooTreeDataStore(const RooTreeDataStore& other, const RooArgSet& vars, const char* newname) :
RooAbsDataStore(other,vars,newname),
_tree(0),
_cacheTree(0),
_defCtor(kFALSE)
{
initialize() ;
loadValues(&other) ;
}
RooTreeDataStore::~RooTreeDataStore()
{
if (_tree) {
delete _tree ;
}
if (_cacheTree) {
delete _cacheTree ;
}
}
void RooTreeDataStore::initialize()
{
createTree(GetName(),GetTitle()) ;
_iterator->Reset() ;
RooAbsArg *var;
while((0 != (var= (RooAbsArg*)_iterator->Next()))) {
var->attachToTree(*_tree,_defTreeBufSize) ;
}
}
void RooTreeDataStore::createTree(const char* name, const char* title)
{
TString pwd(gDirectory->GetPath()) ;
TString memDir(gROOT->GetName()) ;
memDir.Append(":/") ;
Bool_t notInMemNow= (pwd!=memDir) ;
if (notInMemNow) {
gDirectory->cd(memDir) ;
}
if (!_tree) {
_tree = new TTree(name,title) ;
_tree->SetDirectory(0) ;
}
if (!_cacheTree) {
_cacheTree = new TTree(name,title) ;
_cacheTree->SetDirectory(0) ;
}
gDirectory->RecursiveRemove(_tree) ;
gDirectory->RecursiveRemove(_cacheTree) ;
if (notInMemNow) {
gDirectory->cd(pwd) ;
}
}
void RooTreeDataStore::loadValues(const TTree *t, const RooFormulaVar* select, const char* , Int_t , Int_t )
{
TString pwd(gDirectory->GetPath()) ;
TString memDir(gROOT->GetName()) ;
memDir.Append(":/") ;
Bool_t notInMemNow= (pwd!=memDir) ;
if (notInMemNow) {
gDirectory->cd(memDir) ;
}
TTree* tClone ;
if (dynamic_cast<const TChain*>(t)) {
tClone = (TTree*) t->Clone() ;
} else {
tClone = ((TTree*)t)->CloneTree() ;
}
tClone->SetDirectory(0) ;
if (notInMemNow) {
gDirectory->cd(pwd) ;
}
RooArgSet *sourceArgSet = (RooArgSet*) _vars.snapshot(kFALSE) ;
TIterator* sourceIter = sourceArgSet->createIterator() ;
RooAbsArg* sourceArg = 0;
while ((sourceArg=(RooAbsArg*)sourceIter->Next())) {
sourceArg->attachToTree(*tClone,_defTreeBufSize) ;
}
RooFormulaVar* selectClone(0) ;
if (select) {
selectClone = (RooFormulaVar*) select->cloneTree() ;
selectClone->recursiveRedirectServers(*sourceArgSet) ;
RooArgSet branchList ;
selectClone->branchNodeServerList(&branchList) ;
TIterator* iter = branchList.createIterator() ;
RooAbsArg* arg ;
while((arg=(RooAbsArg*)iter->Next())) {
arg->setOperMode(RooAbsArg::ADirty) ;
}
delete iter ;
}
RooAbsArg* destArg = 0;
Int_t numInvalid(0) ;
Int_t nevent= (Int_t)tClone->GetEntries();
for(Int_t i=0; i < nevent; ++i) {
Int_t entryNumber=tClone->GetEntryNumber(i);
if (entryNumber<0) break;
tClone->GetEntry(entryNumber,1);
_iterator->Reset() ;
sourceIter->Reset() ;
Bool_t allOK(kTRUE) ;
while ((destArg = (RooAbsArg*)_iterator->Next())) {
sourceArg = (RooAbsArg*) sourceIter->Next() ;
destArg->copyCache(sourceArg) ;
if (!destArg->isValid()) {
numInvalid++ ;
allOK=kFALSE ;
break ;
}
}
if (!allOK || (selectClone && selectClone->getVal()==0)) {
continue ;
}
fill() ;
}
if (numInvalid>0) {
coutI(Eval) << "RooTreeData::loadValues(" << GetName() << ") Ignored " << numInvalid << " out of range events" << endl ;
}
SetTitle(t->GetTitle());
delete sourceIter ;
delete sourceArgSet ;
delete selectClone ;
delete tClone ;
}
void RooTreeDataStore::loadValues(const RooAbsDataStore *ads, const RooFormulaVar* select,
const char* rangeName, Int_t nStart, Int_t nStop)
{
RooFormulaVar* selectClone(0) ;
if (select) {
selectClone = (RooFormulaVar*) select->cloneTree() ;
selectClone->recursiveRedirectServers(*ads->get()) ;
RooArgSet branchList ;
selectClone->branchNodeServerList(&branchList) ;
TIterator* iter = branchList.createIterator() ;
RooAbsArg* arg ;
while((arg=(RooAbsArg*)iter->Next())) {
arg->setOperMode(RooAbsArg::ADirty) ;
}
delete iter ;
}
ads->get(0) ;
RooAbsArg* arg = 0;
Int_t nevent = nStop < ads->numEntries() ? nStop : ads->numEntries() ;
Bool_t allValid ;
for(Int_t i=nStart; i < nevent ; ++i) {
ads->get(i) ;
if (selectClone && selectClone->getVal()==0) {
continue ;
}
_vars = ((RooTreeDataStore*)ads)->_vars ;
_iterator->Reset() ;
allValid=kTRUE ;
while((arg=(RooAbsArg*)_iterator->Next())) {
if (!arg->isValid() || (rangeName && !arg->inRange(rangeName))) {
allValid=kFALSE ;
break ;
}
}
if (!allValid) continue ;
_cachedVars = ((RooTreeDataStore*)ads)->_cachedVars ;
fill() ;
}
delete selectClone ;
SetTitle(ads->GetTitle());
}
Bool_t RooTreeDataStore::valid() const
{
return kTRUE ;
}
Int_t RooTreeDataStore::fill()
{
return _tree->Fill() ;
}
const RooArgSet* RooTreeDataStore::get(Int_t index) const
{
checkInit() ;
Int_t ret = ((RooTreeDataStore*)this)->GetEntry(index, 1) ;
if(!ret) return 0;
if (_doDirtyProp) {
_iterator->Reset() ;
RooAbsArg* var = 0;
while ((var=(RooAbsArg*)_iterator->Next())) {
var->setValueDirty() ;
}
_cacheIter->Reset() ;
while ((var=(RooAbsArg*)_cacheIter->Next())) {
var->setValueDirty() ;
var->clearValueDirty() ;
}
}
return &_vars;
}
Bool_t RooTreeDataStore::changeObservableName(const char* from, const char* to)
{
RooAbsArg* var = _vars.find(from) ;
if (!var) {
coutE(InputArguments) << "RooTreeDataStore::changeObservableName(" << GetName() << " no observable " << from << " in this dataset" << endl ;
return kTRUE ;
}
TString oldBranchName = var->cleanBranchName() ;
var->SetName(to) ;
if (_tree->GetBranch(oldBranchName.Data())) {
_tree->GetBranch(oldBranchName.Data())->SetName(var->cleanBranchName().Data()) ;
if (_tree->GetBranch(Form("%s_err",oldBranchName.Data()))) {
_tree->GetBranch(Form("%s_err",oldBranchName.Data()))->SetName(Form("%s_err",var->cleanBranchName().Data())) ;
}
if (_tree->GetBranch(Form("%s_aerr_lo",oldBranchName.Data()))) {
_tree->GetBranch(Form("%s_aerr_lo",oldBranchName.Data()))->SetName(Form("%s_aerr_lo",var->cleanBranchName().Data())) ;
}
if (_tree->GetBranch(Form("%s_aerr_hi",oldBranchName.Data()))) {
_tree->GetBranch(Form("%s_aerr_hi",oldBranchName.Data()))->SetName(Form("%s_aerr_hi",var->cleanBranchName().Data())) ;
}
} else {
if (_tree->GetBranch(Form("%s_idx",oldBranchName.Data()))) {
_tree->GetBranch(Form("%s_idx",oldBranchName.Data()))->SetName(Form("%s_idx",var->cleanBranchName().Data())) ;
}
if (_tree->GetBranch(Form("%s_lbl",oldBranchName.Data()))) {
_tree->GetBranch(Form("%s_lbl",oldBranchName.Data()))->SetName(Form("%s_lb",var->cleanBranchName().Data())) ;
}
}
return kFALSE ;
}
RooAbsArg* RooTreeDataStore::addColumn(RooAbsArg& newVar, Bool_t adjustRange)
{
checkInit() ;
RooAbsArg* valHolder= newVar.createFundamental();
if(!valHolder->isFundamental()) {
coutE(InputArguments) << GetName() << "::addColumn: holder argument is not fundamental: \""
<< valHolder->GetName() << "\"" << endl;
return 0;
}
RooAbsArg* newVarClone = newVar.cloneTree() ;
newVarClone->recursiveRedirectServers(_vars,kFALSE) ;
((RooAbsArg*)valHolder)->attachToTree(*_tree,_defTreeBufSize) ;
_vars.add(*valHolder) ;
for (int i=0 ; i<GetEntries() ; i++) {
get(i) ;
newVarClone->syncCache(&_vars) ;
valHolder->copyCache(newVarClone) ;
valHolder->fillTreeBranch(*_tree) ;
}
if (adjustRange) {
}
delete newVarClone ;
return valHolder ;
}
RooArgSet* RooTreeDataStore::addColumns(const RooArgList& varList)
{
TIterator* vIter = varList.createIterator() ;
RooAbsArg* var ;
checkInit() ;
TList cloneSetList ;
RooArgSet cloneSet ;
RooArgSet* holderSet = new RooArgSet ;
while((var=(RooAbsArg*)vIter->Next())) {
RooAbsArg* valHolder= var->createFundamental();
holderSet->add(*valHolder) ;
if(!valHolder->isFundamental()) {
coutE(InputArguments) << GetName() << "::addColumn: holder argument is not fundamental: \""
<< valHolder->GetName() << "\"" << endl;
return 0;
}
RooArgSet* newVarCloneList = (RooArgSet*) RooArgSet(*var).snapshot() ;
if (!newVarCloneList) {
coutE(InputArguments) << "RooTreeDataStore::RooTreeData(" << GetName()
<< ") Couldn't deep-clone variable " << var->GetName() << ", abort." << endl ;
return 0 ;
}
RooAbsArg* newVarClone = newVarCloneList->find(var->GetName()) ;
newVarClone->recursiveRedirectServers(_vars,kFALSE) ;
newVarClone->recursiveRedirectServers(*holderSet,kFALSE) ;
cloneSetList.Add(newVarCloneList) ;
cloneSet.add(*newVarClone) ;
((RooAbsArg*)valHolder)->attachToTree(*_tree,_defTreeBufSize) ;
_vars.addOwned(*valHolder) ;
}
delete vIter ;
TIterator* cIter = cloneSet.createIterator() ;
TIterator* hIter = holderSet->createIterator() ;
RooAbsArg *cloneArg, *holder ;
for (int i=0 ; i<GetEntries() ; i++) {
get(i) ;
cIter->Reset() ;
hIter->Reset() ;
while((cloneArg=(RooAbsArg*)cIter->Next())) {
holder = (RooAbsArg*)hIter->Next() ;
cloneArg->syncCache(&_vars) ;
holder->copyCache(cloneArg) ;
holder->fillTreeBranch(*_tree) ;
}
}
delete cIter ;
delete hIter ;
cloneSetList.Delete() ;
return holderSet ;
}
RooAbsDataStore* RooTreeDataStore::merge(const RooArgSet& allVars, list<RooAbsDataStore*> dstoreList)
{
RooTreeDataStore* mergedStore = new RooTreeDataStore("merged","merged",allVars) ;
Int_t nevt = dstoreList.front()->numEntries() ;
for (int i=0 ; i<nevt ; i++) {
mergedStore->_vars = *get(i) ;
for (list<RooAbsDataStore*>::iterator iter = dstoreList.begin() ; iter!=dstoreList.end() ; iter++) {
const RooArgSet* partSet = (*iter)->get(i) ;
mergedStore->_vars = *partSet ;
}
mergedStore->fill() ;
}
return mergedStore ;
}
void RooTreeDataStore::append(RooAbsDataStore& other)
{
Int_t nevt = other.numEntries() ;
for (int i=0 ; i<nevt ; i++) {
_vars = *other.get(i) ;
fill() ;
}
}
Int_t RooTreeDataStore::numEntries() const
{
return _tree->GetEntries() ;
}
void RooTreeDataStore::reset()
{
Reset() ;
}
void RooTreeDataStore::cacheArgs(RooArgSet& newVarSet, const RooArgSet* nset)
{
TIterator *iter = newVarSet.createIterator() ;
RooAbsArg *arg ;
Bool_t doTreeFill = (_cachedVars.getSize()==0) ;
while ((arg=(RooAbsArg*)iter->Next())) {
arg->attachToTree(*_cacheTree,_defTreeBufSize) ;
arg->redirectServers(_vars) ;
_cachedVars.add(*arg) ;
}
for (int i=0 ; i<GetEntries() ; i++) {
get(i) ;
iter->Reset() ;
while ((arg=(RooAbsArg*)iter->Next())) {
arg->setValueDirty() ;
arg->syncCache(nset) ;
if (!doTreeFill) {
arg->fillTreeBranch(*_cacheTree) ;
}
}
if (doTreeFill) {
_cacheTree->Fill() ;
}
}
delete iter ;
}
void RooTreeDataStore::setArgStatus(const RooArgSet& set, Bool_t active)
{
TIterator* iter = set.createIterator() ;
RooAbsArg* arg ;
while ((arg=(RooAbsArg*)iter->Next())) {
RooAbsArg* depArg = _vars.find(arg->GetName()) ;
if (!depArg) {
coutE(InputArguments) << "RooTreeDataStore::setArgStatus(" << GetName()
<< ") dataset doesn't contain variable " << arg->GetName() << endl ;
continue ;
}
depArg->setTreeBranchStatus(*_tree,active) ;
}
delete iter ;
}
void RooTreeDataStore::resetCache()
{
_cachedVars.removeAll() ;
delete _cacheTree ;
_cacheTree = 0 ;
createTree(GetName(),GetTitle()) ;
return ;
}
void RooTreeDataStore::checkInit() const
{
if (_defCtor) {
const_cast<RooTreeDataStore*>(this)->initialize() ;
_defCtor = kFALSE ;
}
}
Stat_t RooTreeDataStore::GetEntries() const
{
return _tree->GetEntries() ;
}
void RooTreeDataStore::Reset(Option_t* option)
{
_tree->Reset(option) ;
}
Int_t RooTreeDataStore::Fill()
{
return _tree->Fill() ;
}
Int_t RooTreeDataStore::GetEntry(Int_t entry, Int_t getall)
{
Int_t ret1 = _tree->GetEntry(entry,getall) ;
if (!ret1) return 0 ;
_cacheTree->GetEntry(entry,getall) ;
return ret1 ;
}