Skip to content
Snippets Groups Projects
Commit 7d88b65b authored by Nils Erik Krumnack's avatar Nils Erik Krumnack
Browse files

Merge branch 'RootNtupleEventSelector_multiTree' into '21.2'

Support for using wildcards in tree name for multi-tree processing in AthAnalysis

See merge request atlas/athena!12809

Former-commit-id: 7215e63abb67747132e443a2718dbf5e194484ae
parents fed6cf35 be006a65
No related branches found
No related tags found
No related merge requests found
......@@ -49,6 +49,7 @@
#include "TObject.h"
#include "TTree.h"
#include "TRegexp.h"
CLASS_DEF( TObject, 74939790 , 1 )
#include "AthenaRootComps/TransferTree.h"
......@@ -289,6 +290,8 @@ RootNtupleEventSelector::RootNtupleEventSelector( const std::string& name,
"List of branch names to activate" );
declareProperty( "CreateEventInfo", m_createEventInfo=false, "Are EventInfo objects created per event or not? .. there is overhead to their creation" );
declareProperty( "IgnoreMissingTrees", m_ignoreMissingTrees=false, "If True, missing trees will not terminate the event loop" );
}
// Destructor
......@@ -296,6 +299,21 @@ RootNtupleEventSelector::RootNtupleEventSelector( const std::string& name,
RootNtupleEventSelector::~RootNtupleEventSelector()
{}
void getTreeNames( TList& keys , std::set<std::string>& treeNames, const std::string& path ) {
TIter keyIter(&keys);
TKey* key;
while( (key = (TKey*)keyIter()) ) {
if(strcmp(key->GetClassName(),"TTree")==0) {
treeNames.insert(path+key->GetName());
} else if(strcmp(key->GetClassName(),"TDirectoryFile")==0) {
TDirectory* dir = static_cast<TDirectory*>(key->ReadObj());
TList* dirKeys = dir->GetListOfKeys();
if(dirKeys) getTreeNames( *dirKeys , treeNames, path + key->GetName() + "/" );
}
}
}
StatusCode RootNtupleEventSelector::initialize()
{
ATH_MSG_INFO ("Enter RootNtupleEventSelector initialization...");
......@@ -342,6 +360,57 @@ StatusCode RootNtupleEventSelector::initialize()
<< endmsg
<< " TTree [" << m_tupleName.value() << "]");
}
bool hasWildcards(false);
for(auto& tupleName : m_tupleNames) {
if(tupleName.find("*")!=std::string::npos) {
hasWildcards=true; break;
}
}
if(hasWildcards) {
//use the all files to resolve the tree names ...
std::set<std::string> treeNames;
for(unsigned int i=0;i<m_inputCollectionsName.value().size();i++) {
std::unique_ptr<TFile> theFile( TFile::Open( m_inputCollectionsName.value()[i].c_str() ) );
if(!theFile) {
ATH_MSG_FATAL("Could not open the first input file: " << m_inputCollectionsName.value()[i]);
return StatusCode::FAILURE;
}
TList* keys = theFile->GetListOfKeys();
if(keys) getTreeNames( *keys, treeNames, "" );
theFile->Close();
}
std::vector<std::string> foundTupleNames;
for(auto& tupleName : m_tupleNames) {
TRegexp re(tupleName.c_str(),true);
for(auto& treeName : treeNames) {
Ssiz_t tmp;
if(re.Index(treeName,&tmp)==-1) continue;
foundTupleNames.push_back(treeName);
ATH_MSG_INFO("Will read tree: " << treeName);
}
}
m_tupleNames = foundTupleNames;
if ( m_tupleNames.empty() ) {
ATH_MSG_ERROR
("No matching trees found");
return StatusCode::FAILURE;
}
}
{
// register this service for 'I/O' events
......@@ -533,6 +602,10 @@ RootNtupleEventSelector::next( IEvtSelector::Context& ctx ) const
const std::string& fname = fnames[fidx];
tree = fetchNtuple(fname, m_tupleNames[rctx->tupleIndex()]);
if (!tree) {
if(m_ignoreMissingTrees) {
ATH_MSG_WARNING("Unable to get tree: " << m_tupleNames[rctx->tupleIndex()] << " from file: " << fname);
fidx++;continue;
}
throw "RootNtupleEventSelector: Unable to get tree";
}
rctx->setTree(tree);
......@@ -1124,7 +1197,7 @@ RootNtupleEventSelector::fetchNtuple(const std::string& fname,
}
}
if (!f || f->IsZombie()) {
ATH_MSG_ERROR("could not open next file in input collection ["
if(!m_ignoreMissingTrees) ATH_MSG_ERROR("could not open next file in input collection ["
<< fname << "]");
if (f) {
f->Close();
......@@ -1132,9 +1205,9 @@ RootNtupleEventSelector::fetchNtuple(const std::string& fname,
return tree;
}
// std::cout << "::TFile::GetTree(" << m_tupleName << ")..." << std::endl;
tree = (TTree*)f->Get(tupleName.c_str());
tree = dynamic_cast<TTree*>(f->Get(tupleName.c_str()));
if (!tree) {
ATH_MSG_ERROR("could not retrieve tree [" << tupleName << "]"
if(!m_ignoreMissingTrees) ATH_MSG_ERROR("could not retrieve tree [" << tupleName << "]"
<< " from file [" << fname << "]");
f->Close();
return tree;
......@@ -1159,7 +1232,7 @@ RootNtupleEventSelector::fetchNtuple(const std::string& fname,
void RootNtupleEventSelector::addMetadataFromDirectoryName(const std::string &metadirname, TFile *fileObj, const std::string &prefix) const
{
TDirectoryFile *metadir = (TDirectoryFile*)fileObj->Get(metadirname.c_str());
TDirectoryFile *metadir = dynamic_cast<TDirectoryFile*>(fileObj->Get(metadirname.c_str()));
if (!metadir) return;
addMetadataFromDirectory(metadir, prefix);
}
......@@ -1247,7 +1320,7 @@ RootNtupleEventSelector::do_init_io()
m_tuple = fetchNtuple(m_inputCollectionsName.value()[m_collIdx],
m_tupleNames[0]);
if (!m_tuple) {
if (!m_tuple && !m_ignoreMissingTrees) {
throw "RootNtupleEventSelector: Unable to fetch Ntuple";
}
......@@ -1282,7 +1355,7 @@ RootNtupleEventSelector::find_coll_idx (int evtidx,
if (itr.min_entries == -1) {
TTree *tree = fetchNtuple(m_inputCollectionsName.value()[icoll],
m_tupleNames[ituple]);
if (tree) {
if (tree || m_ignoreMissingTrees) {
long offset = 0;
if (icoll > 0) {
offset = m_collEvts[ituple][icoll-1].max_entries;
......@@ -1290,7 +1363,7 @@ RootNtupleEventSelector::find_coll_idx (int evtidx,
else if (ituple > 0) {
offset = m_collEvts[ituple-1].back().max_entries;
}
itr.entries = tree->GetEntriesFast();
itr.entries = (tree) ? tree->GetEntriesFast() : 0;
itr.min_entries = offset;
itr.max_entries = offset + itr.entries;
} else {
......
......@@ -216,6 +216,9 @@ class RootNtupleEventSelector :
/// List of branches to activate in the @c TTree
StringArrayProperty m_activeBranchNames;
/// Should missing trees be ignored (only prints a warning)
bool m_ignoreMissingTrees;
/// Number of events to skip at the beginning
long m_skipEvts;
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment