Commits

Anonymous committed 42ffd56

support for seeking and athena-mp + add a unit test for n-tuple files with different layout

Comments (0)

Files changed (11)

+2012-03-07  Sebastien Binet  <binet@voatlas51.cern.ch>
+
+	* tagging AthenaRootComps-00-10-01
+	* support for seeking and athena-mp
+	* add a unit test for n-tuple files with different layout
+	* M python/ReadAthenaRoot.py
+	* M python/OutputStreamAthenaRoot.py
+	* M share/tests/test_athena_ntuple_dumper.py
+	* M src/RootNtupleEventSelector.cxx
+	* M src/RootNtupleEventSelector.h
+	* M src/RootConnection.cxx
+	* A share/tests/test_athena_variable_shape_ntuple.py
+	* A test/run_test_athena_variable_shape_ntuple.sh
+
 2012-03-02  Peter van Gemmeren  <gemmeren@anl.gov>
 
 	* tagging AthenaRootComps-00-10-00

python/OutputStreamAthenaRoot.py

         topSequence = AlgSequence()
         pass
 
-    # configure the AthenaRootCnvSvc to match ROOT n-tuple expectations
-    #FIXME: this customization should be less global than that
-    #       ex: somewhere in the WritingTool handle perhaps.
-    svcMgr.AthenaRootCnvSvc.SubLevelBranchName = "<key>"
-    
     # Now do standard output stream
     writingTool1 = AthenaRootOutputStreamTool( streamName + "Tool" )
     writingTool1.TupleName = tupleName
     outputStream = AthenaRootOutputStreamProtect(
         streamName,
         WritingTool = writingTool1,
-        #ItemList    = [ "EventInfo#*" ]
         ItemList = ["unsigned int#RunNumber",
-                    "unsigned int#EventNumber",]
+                    "unsigned int#EventNumber",],
+        #StrictClidMatching = False,
         )
     if asAlg:
         from AthenaCommon.AlgSequence import AlgSequence

python/ReadAthenaRoot.py

         msg.error( err )
         raise RuntimeError( err )
 
+    # disable the event loop heartbeat as it is somewhat I/O hungry for
+    # no real gain in n-tuple reading/writing scenarii
+    if not hasattr(svcMgr, theApp.EventLoop):
+        svcMgr += getattr(CfgMgr, theApp.EventLoop)()
+        evtloop = getattr(svcMgr, theApp.EventLoop)
+    try:
+        msg.info('disabling event loop heartbeat...')
+        evtloop.EventPrintoutInterval = 0
+        msg.info('disabling event loop heartbeat... [done]')
+    except Exception, err:
+        msg.info('disabling event loop heartbeat... [failed]')
+        msg.info('performances might be sub-par... sorry.')
+        pass
+    
     # Load ProxyProviderSvc
     if not hasattr (svcMgr, 'ProxyProviderSvc'):
         svcMgr += CfgMgr.ProxyProviderSvc()

share/tests/test_athena_ntuple_dumper.py

         '/tmp/binet/data/small.ntuple.0.root',
         '/tmp/binet/data/small.ntuple.1.root',
         ]
+if not 'FNAMES' in dir():
+    FNAMES = _cbntfiles[:]
+
+if not isinstance(FNAMES, (list,tuple)):
+    FNAMES = [FNAMES]
+    pass
 
 import AthenaRootComps.ReadAthenaRoot
-svcMgr.EventSelector.InputCollections = _cbntfiles
+svcMgr.EventSelector.InputCollections = FNAMES
 svcMgr.EventSelector.TupleName = "egamma"
 
 from AthenaCommon.AlgSequence import AlgSequence
         "std::vector<std::vector<float> >#el_jetcone_dr",
         ]
     out.ForceRead = True
-
-evtloop = getattr(CfgMgr, theApp.EventLoop)()
-evtloop.EventPrintoutInterval = 100000
-svcMgr += evtloop

share/tests/test_athena_variable_shape_ntuple.py

+from glob import glob
+if 1:
+    _cbntfiles = [
+        'root://eosatlas//eos/atlas/user/b/binet/utests/utests/filter-d3pd/ntuple.0.root',
+        'root://eosatlas//eos/atlas/user/b/binet/utests/utests/filter-d3pd/ntuple.1.root',
+        ]
+    # _cbntfiles = [
+    #     '/tmp/binet/data/ntuple.0.root',
+    #     '/tmp/binet/data/ntuple.1.root',
+    #     ]
+    
+if 0:
+    _cbntfiles = [
+        'root://eosatlas//eos/atlas/user/b/binet/utests/utests/filter-d3pd/small.ntuple.0.root',
+        'root://eosatlas//eos/atlas/user/b/binet/utests/utests/filter-d3pd/small.ntuple.1.root',
+        ]
+    _cbntfiles = [
+        '/tmp/binet/data/small.ntuple.0.root',
+        '/tmp/binet/data/small.ntuple.1.root',
+        ]
+if not 'FNAMES' in dir():
+    FNAMES = _cbntfiles[:]
+
+if not isinstance(FNAMES, (list,tuple)):
+    FNAMES = [FNAMES]
+    pass
+
+import AthenaRootComps.ReadAthenaRoot
+svcMgr.EventSelector.InputCollections = FNAMES
+svcMgr.EventSelector.TupleName = "egamma"
+
+from AthenaCommon.AlgSequence import AlgSequence
+job = AlgSequence()
+
+from AthenaPython import PyAthena
+StatusCode = PyAthena.StatusCode
+
+class MyAlg( PyAthena.Alg ):
+
+    def __init__(self, name='MyAlg', **kw):
+        kw['name'] = name
+        self.activeBranches = kw.get(
+            'activeBranches',
+            ['RunNumber',
+             'EventNumber',
+             'el_n',
+             'el_eta',
+             'el_jetcone_dr',
+             ])
+        super(MyAlg, self).__init__(**kw)
+        return
+
+    def initialize(self):
+        self.evtstore = PyAthena.py_svc('StoreGateSvc')
+        
+        return StatusCode.Success
+
+    def execute(self):
+        self.msg.info('running execute...')
+        #keys = self.evtstore.keys()
+        for br in self.activeBranches:
+            try:
+                o = self.evtstore[br]
+                if hasattr(o, 'at'):
+                    o = list(o)
+                    for i,v in enumerate(o):
+                        if hasattr(v, 'at') and not isinstance(v, (basestring,)):
+                            o[i] = list(v)
+                self.msg.info('%s: %r', br, o)
+            except Exception, err:
+                self.msg.info(' --> err for [%s]: %s' % (br, err))
+                pass
+        return StatusCode.Success
+
+    def finalize(self):
+        return StatusCode.Success
+
+    pass # MyAlg
+
+from AthenaCommon.AlgSequence import AlgSequence
+job = AlgSequence()
+job += MyAlg('py_alg')
+
+if 'BRANCHES' in dir():
+    job.py_alg.activeBranches = BRANCHES
+    pass
+
+if not 'EVTMAX' in dir():
+    EVTMAX=-1
+theApp.EvtMax = EVTMAX
+
+if not 'DOWRITE' in dir():
+    DOWRITE=1
+if DOWRITE:
+    svcMgr += CfgMgr.DecisionSvc()
+    import AthenaRootComps.WriteAthenaRoot as arcw
+    out = arcw.createNtupleOutputStream("StreamD3PD", "d3pd.root", "egamma")
+    if not 'OUTBRANCHES' in dir():
+        OUTBRANCHES=[
+            "int#el_n",
+            "std::vector<float>#el_eta",
+            "std::vector<std::vector<float> >#el_jetcone_dr",
+            ]
+    out.ItemList += OUTBRANCHES
+    out.ForceRead = True
+

src/RootConnection.cxx

 StatusCode
 RootConnection::connectWrite(IIoSvc::IoType mode)
 {
+  // std::cerr << "::RootConnection::connectWrite(" << m_pfn
+  //           << "|" << m_treeName
+  //           << ")...\n";
   if (m_file == 0) {
     delete m_tree; m_tree = 0;
     m_file = TFile::Open(m_pfn.c_str(),

src/RootNtupleEventSelector.cxx

 #include <stdint.h>
 
 // ROOT includes
-#include "TTree.h"
+#include "TBranchElement.h"
 #include "TClass.h"
 #include "TClassEdit.h"
-#include "TBranchElement.h"
+#include "TFile.h"
 #include "TLeaf.h"
-#include "TFile.h"
+#include "TROOT.h"
+#include "TTree.h"
 
 // Framework includes
+//#include "GaudiKernel/GenericAddress.h"
+#include "GaudiKernel/IIoComponentMgr.h"
 #include "GaudiKernel/ISvcLocator.h"
-//#include "GaudiKernel/GenericAddress.h"
+#include "GaudiKernel/ITHistSvc.h"
+#include "GaudiKernel/MsgStream.h"
 #include "GaudiKernel/StatusCode.h"
-#include "GaudiKernel/MsgStream.h"
-#include "GaudiKernel/ITHistSvc.h"
 #include "GaudiKernel/System.h"
 #include "AthenaKernel/IClassIDSvc.h"
 #include "AthenaKernel/IDictLoaderSvc.h"
 
 // StoreGate includes
+
+#include "SGTools/BuiltinsClids.h"   // to make sure we have their clids
+#include "SGTools/CLASS_DEF.h"
+#include "SGTools/StlMapClids.h"     // to make sure we have their clids
+#include "SGTools/StlVectorClids.h"  // to make sure we have their clids
+#include "SGTools/TransientAddress.h"
 #include "StoreGate/StoreGate.h" 
-#include "SGTools/TransientAddress.h"
-#include "SGTools/BuiltinsClids.h"   // to make sure we have their clids
-#include "SGTools/StlVectorClids.h"  // to make sure we have their clids
-#include "SGTools/StlMapClids.h"     // to make sure we have their clids
-
-#include "SGTools/CLASS_DEF.h"
 
 #include "TObject.h"
 CLASS_DEF( TObject,    74939790 , 1 )
   const RootNtupleEventSelector* m_evtsel;
 
   /// the file container managed by this context
-  FileNames_t m_files;
+  //FileNames_t m_files;
 
   /// the index into the container of file names
-  std::size_t m_fidx;
+  //std::size_t m_fidx;
 
   /// current entry of current file
-  int64_t m_entry;
+  //int64_t m_entry;
 
   /// reference to top-level tree
-  TTree* m_tree;
+  //TTree* m_tree;
 
   /// connection FID
   std::string m_fid;
   /// standard c-tor with initialization
   RootNtupleEventContext(const RootNtupleEventSelector* sel) :
     m_evtsel(sel),
-    m_files(),
-    m_fidx(0),
-    m_entry(-1),
-    m_tree(NULL),
+    //m_files(),
+    //m_fidx(0),
+    //m_entry(-1),
+    //m_tree(NULL),
     m_fid("")
   {}
 
   virtual ~RootNtupleEventContext() {}
 
   // access to the container of files
-  const FileNames_t& files() const { return m_files; }
+  const FileNames_t& files() const 
+  { return m_evtsel->m_inputCollectionsName.value(); }
 
-  /// set the container of files
-  void setFiles(const FileNames_t& fnames) 
-  {
-    m_files = fnames;
-    m_fidx = 0;
-  }
+  // /// set the container of files
+  // void setFiles(const FileNames_t& fnames) 
+  // {
+  //   //m_files = fnames;
+  //   //m_fidx = 0;
+  // }
 
   /// context identifier
   virtual void* identifier() const 
 
   /// access to the file iterator
   std::size_t fileIndex() const
-  { return m_fidx; }
+  { return m_evtsel->m_collIdx; }
 
   /// set file iterator
-  void setFileIndex(std::size_t idx) 
-  { m_fidx = idx; }
+  // void setFileIndex(std::size_t idx) 
+  // { m_fidx = idx; }
 
   /// access to the current event entry number
-  int64_t entry() const { return m_entry; }
+  int64_t entry() const { return m_evtsel->m_curEvt; }
 
-  /// set the current event entry number
-  void setEntry(int64_t entry) { m_entry = entry; }
+  // /// set the current event entry number
+  // void setEntry(int64_t entry) { m_entry = entry; }
 
   /// set connection FID
   void setFID(const std::string& fid) { m_fid = fid; }
   const std::string& fid() const { return m_fid; }
 
   /// access to the tree used to iterate
-  TTree* tree() const { return m_tree; }
+  TTree* tree() const { return m_evtsel->m_tuple; }
 
   /// set the tree used to iterate
-  void setTree(TTree* tree) { m_tree = tree; }
+  void setTree(TTree* tree) { m_evtsel->m_tuple = tree; }
 };
 
   typedef RootNtupleEventContext::FileNames_t FileNames_t;
   m_clidsvc  ( "ClassIDSvc",   name ),
   m_dictsvc  ( "AthDictLoaderSvc", name ),
   m_nbrEvts  ( 0 ),
+  m_curEvt   ( 0 ),
+  m_collIdx  ( 0 ),
+  m_collEvts (   ),
   m_tuple    (NULL),
   m_rootAddresses ()
 {
        << "                      tuple [" << m_tupleName.value() << "]");
   }
 
-  m_tuple = fetchNtuple(m_inputCollectionsName.value()[0]);
-  if (!m_tuple) {
+  {
+    // register this service for 'I/O' events
+    ServiceHandle<IIoComponentMgr> iomgr("IoComponentMgr", name());
+    if (!iomgr.retrieve().isSuccess()) {
+      ATH_MSG_FATAL("Could not retrieve IoComponentMgr !");
+      return StatusCode::FAILURE;
+    }
+    if (!iomgr->io_register(this).isSuccess()) {
+      ATH_MSG_FATAL("Could not register myself with the IoComponentMgr !");
+      return StatusCode::FAILURE;
+    }
+    // register input file's names with the I/O manager
+    const std::vector<std::string>& incol = m_inputCollectionsName.value();
+    bool allGood = true;
+    for (std::size_t icol = 0, imax = incol.size(); icol < imax; icol++) {
+      if (!iomgr->io_register(this, 
+                              IIoComponentMgr::IoMode::Input, 
+                              incol[icol]).isSuccess()) {
+        ATH_MSG_FATAL("could not register [" << incol[icol] << "] for output !");
+        allGood = false;
+      } else {
+        ATH_MSG_VERBOSE("io_register[" << this->name() << "](" << incol[icol] << ") [ok]");
+      }
+    }
+    if (!allGood) {
+      return StatusCode::FAILURE;
+    }
+  }
+
+  if (!do_init_io().isSuccess()) {
     return StatusCode::FAILURE;
   }
 
   // retrieve event store
-  // this needs to happen *after* having retrieved the n-tuple
+  // this needs to happen *after* having initialized the i/o
   // as our branches (which need a valid m_ntuple pointer)
   // may be asked to be registered as we are a ProxyProvider.
   // retrieving the event store will poke the ProxyProviderSvc...
     return StatusCode::FAILURE;
   }
 
-  // reset the list of branches
-  m_rootAddresses.clear();
-
-  // skip events we are asked to skip
-  m_nbrEvts += m_skipEvts;
-
   return StatusCode::SUCCESS;
 }
 
     *ppvInterface = dynamic_cast<IEvtSelector*>(this);
   } else if ( IEventSeek::interfaceID().versionMatch(riid) ) {
     *ppvInterface = dynamic_cast<IEventSeek*>(this);
+  } else if ( IIoComponent::interfaceID().versionMatch(riid) ) {
+    *ppvInterface = dynamic_cast<IIoComponent*>(this);
   } else {
     // Interface is not directly available : try out a base class
     return AthService::queryInterface(riid, ppvInterface);
 StatusCode
 RootNtupleEventSelector::next( IEvtSelector::Context& ctx ) const
 {
-  ATH_MSG_DEBUG ("next() : iEvt " << m_nbrEvts);
+  // std::cout << "::next(fidx=" << m_collIdx << ", eidx=" << m_curEvt << ")"
+  //           << std::endl;
+  ATH_MSG_DEBUG ("next() : iEvt " << m_curEvt);
 
   // get evt context
   RootNtupleEventContext* rctx = dynamic_cast<RootNtupleEventContext*>(&ctx);
     const FileNames_t& fnames = rctx->files();
     std::size_t fidx = rctx->fileIndex();
     rctx->setTree(NULL);
-    rctx->setEntry(-1);
+    //rctx->setEntry(-1);
 
     if (fidx < rctx->files().size()) {
       const std::string& fname = fnames[fidx];
         return StatusCode::FAILURE;
       }
       rctx->setTree(tree);
-      rctx->setFileIndex(++fidx);
+      //rctx->setFileIndex(++fidx);
       //rctx->setFileIterator(++fitr);
       // FIXME: close file holding this m_tuple ?
       // what if that file holds other still-used objects ?
       // if (m_tuple) {
       //   m_tuple->GetCurrentFile()->Close();
       // }
-      m_tuple = tree;
+      //m_tuple = tree;
 
       // notify everybody a new file has been connected.
       for (Addrs_t::const_iterator 
         }
       }
 
+      const bool forceRemove = false;
+      m_storeGate->clearStore(forceRemove).ignore();
+
     } else {
       // end of collections
       return StatusCode::FAILURE;
     }
   }
-  int64_t entry = rctx->entry();
+  int64_t global_entry = rctx->entry();
+  int64_t entry = global_entry;
+  if (m_collEvts[m_collIdx].min_entries < 0) {
+    // need to trigger collmetadata...
+    const_cast<RootNtupleEventSelector*>(this)->find_coll_idx(entry);
+  }
+  // rctx::entry is the *global* entry number.
+  // we need the local one...
+  entry = global_entry - m_collEvts[m_collIdx].min_entries;
+
   Long64_t nentries = tree->GetEntriesFast();
-  if ( nentries > (entry+1) ) {
-    rctx->setEntry(++entry);
+  // std::cout << "::entry=" << global_entry 
+  //           << ", nentries=" << nentries
+  //           << ", local=" << entry
+  //           << " (min=" << m_collEvts[m_collIdx].min_entries
+  //           << ", max=" << m_collEvts[m_collIdx].max_entries << ")"
+  //           << " (tree=" << tree << ")"
+  //           << std::endl;
+  if ( nentries > entry ) {
 
     // load data from tuple
+    // std::cout << "--load-data--" << std::endl;
     if (tree->LoadTree(entry) < 0) {
       ATH_MSG_ERROR
-        ("Problem loading tree for event [" << m_nbrEvts << "] !!");
+        ("Problem loading tree for event [" << m_curEvt << "] !!");
       return StatusCode::FAILURE;
     } else {
-      ATH_MSG_DEBUG("==> loaded-tree(" << m_nbrEvts << ")");
+      ATH_MSG_DEBUG("==> loaded-tree(" << m_curEvt << ")");
     }
 
     // load entry
+    // std::cout << "--load-entry--" << std::endl;
     if ( tree->GetEntry(entry) < 0 ) {
       ATH_MSG_ERROR
-        ("Problem retrieving data from tree for entry [" << m_nbrEvts 
+        ("Problem retrieving data from tree for entry [" << m_curEvt 
          << "] !!");
       return StatusCode::FAILURE;
     } else {
-      ATH_MSG_DEBUG("==> got-tree(" << m_nbrEvts << ")");
+      ATH_MSG_DEBUG("==> got-tree(" << m_curEvt << ")");
     }
     ++m_nbrEvts;
+    m_curEvt = global_entry + 1;
 
-    // FIXME: we should try to get the RunNumber/EventNumber from
-    //        the ntuple if these branches exist...
-
+    // std::cout << "--event-info--" << std::endl;
+    // try to get the RunNumber/EventNumber from
+    // the ntuple if these branches exist...
+    const uint32_t* runnbr = NULL;
+    const uint32_t* evtnbr = NULL;
+    if (m_storeGate->contains<uint32_t>("RunNumber")) {
+      m_storeGate->retrieve(runnbr, "RunNumber").ignore();
+    }
+    if (m_storeGate->contains<uint32_t>("EventNumber")) {
+      m_storeGate->retrieve(evtnbr, "EventNumber").ignore();
+    }
     // event info
     EventType* evtType = new EventType;
-    const std::size_t runNbr = 0;
-    EventInfo* evtInfo = new EventInfo( new EventID( runNbr, m_nbrEvts-1, 0 ),
+    const std::size_t runNbr = runnbr ? *runnbr : 0;
+    EventInfo* evtInfo = new EventInfo( new EventID( runNbr, 
+                                                     evtnbr ? 
+                                                     *evtnbr : m_curEvt-1, 0 ),
                                         evtType );
     if ( !m_storeGate->record( evtInfo, "TTreeEventInfo" ).isSuccess() ) {
       ATH_MSG_ERROR ("Could not record TTreeEventInfo !");
       delete evtInfo; evtInfo = 0;
       return StatusCode::FAILURE;
     }
+    // std::cout << "--done--" << std::endl;
     return StatusCode::SUCCESS;
   } else {
     // prepare for next file, if any...
-    std::size_t fidx = rctx->fileIndex();
-    rctx->setFileIndex(++fidx);
+    // std::cout << "=========================================================="
+    //           << std::endl;
+    // std::cout << "::switch to next file..." << std::endl;
+    //std::size_t fidx = rctx->fileIndex();
+    //rctx->setFileIndex(++fidx);
+    m_collIdx += 1;
     rctx->setTree(NULL);
-    rctx->setEntry(-1);
+    //rctx->setEntry(-1);
     return next(*rctx);
   }
   
+  // std::cout << "***end of collections***" << std::endl;
   // end of collections
   return StatusCode::FAILURE;
 }
 
 StatusCode RootNtupleEventSelector::next( Context& ctx, int jump ) const
 {
-  ATH_MSG_DEBUG ("next(" << jump << ") : iEvt " << m_nbrEvts);
+  ATH_MSG_DEBUG ("next(" << jump << ") : iEvt " << m_curEvt);
 
-  if (jump>0) {
-    for (int i=0; i<jump; ++i) {
-      if (!next(ctx).isSuccess()) {
-        return StatusCode::FAILURE;
-      }
-    }
-    return StatusCode::SUCCESS;
+  RootNtupleEventSelector *self=const_cast<RootNtupleEventSelector*>(this);
+  if (self->seek(m_curEvt + jump).isSuccess()) {
+    return StatusCode::FAILURE;
   }
-  // still here ? So this is it. EOF.
-  return StatusCode::FAILURE;
+  return next(ctx);
 }
 
 StatusCode
 StatusCode 
 RootNtupleEventSelector::rewind( Context& /*ctxt*/ ) const 
 {
-  ATH_MSG_ERROR 
-    ("............. ::rewind(Context) Not Implemented .............");
-  return StatusCode::FAILURE;
+  RootNtupleEventSelector *self=const_cast<RootNtupleEventSelector*>(this);
+  return self->seek(0);
 }
 
 StatusCode
 RootNtupleEventSelector::createContext( Context*& refCtx ) const
 {
   RootNtupleEventContext *ctx = new RootNtupleEventContext(this);
-  ctx->setFiles(m_inputCollectionsName.value());
-  ctx->setTree(m_tuple);
   refCtx = ctx;
   return StatusCode::SUCCESS;
 }
  * @param evtnum  The event number to which to seek.
  */
 StatusCode
-RootNtupleEventSelector::seek (int /*evtnum*/)
+RootNtupleEventSelector::seek (int evtnum)
 {
-#if 0
-  // jump to right position
-  if ( ( evtnum >= 0 ) &&
-       ( evtnum <  m_totalNbrEvts ) ) {
-    // adjust pointer
-    m_nbrEvts = evtnum;
-    
-    // load data from tuple
-    if ( m_tuple->LoadTree (m_nbrEvts) < 0 ) {
-      ATH_MSG_ERROR
-        ("Problem loading tree for event [" << m_nbrEvts << "] !!");
-      return StatusCode::FAILURE;
-    }
+  // std::cout << "::seek - evtnum=" << evtnum 
+  //           << " curevt=" << m_curEvt 
+  //           << " curcol=" << m_collIdx
+  //           << std::endl;
+  long coll_idx = find_coll_idx(evtnum);
+  // std::cout << "::seek - evtnum=" << evtnum 
+  //           << " curevt=" << m_curEvt 
+  //           << " curcol=" << m_collIdx
+  //           << " colidx=" << coll_idx
+  //           << std::endl;
+  if (coll_idx == -1 && evtnum < m_curEvt) {
+    coll_idx = m_collIdx;
+  }
 
-    if ( m_tuple->GetEntry( m_nbrEvts ) < 0 ) {
-      ATH_MSG_ERROR
-        ("Problem retrieving data from tree for entry [" << m_nbrEvts 
-         << "] !!");
-      return StatusCode::FAILURE;
-    }
-  } else {
-    return StatusCode::FAILURE;
+  if (coll_idx == -1) {
+    ATH_MSG_INFO("seek: reached end of input.");
+    return StatusCode::RECOVERABLE;
   }
-#endif
+
+  if (coll_idx != m_collIdx) {
+    // tell everyone we switched files...
+    m_tuple = NULL;
+
+  }
+
+  m_collIdx = coll_idx;
+  m_curEvt = evtnum;
+
   return StatusCode::SUCCESS;
 }
 
 int 
 RootNtupleEventSelector::curEvent() const
 {
-  return m_nbrEvts;
+  return m_curEvt;
 }
 
+/// Callback method to reinitialize the internal state of the component 
+/// for I/O purposes (e.g. upon @c fork(2))
+StatusCode 
+RootNtupleEventSelector::io_reinit()
+{
+  ATH_MSG_INFO("I/O reinitialization...");
+
+  ServiceHandle<IIoComponentMgr> iomgr("IoComponentMgr", name());
+  if (!iomgr.retrieve().isSuccess()) {
+    ATH_MSG_FATAL("Could not retrieve IoComponentMgr !");
+    return StatusCode::FAILURE;
+  }
+  if (!iomgr->io_hasitem(this)) {
+    ATH_MSG_FATAL("IoComponentMgr does not know about myself !");
+    return StatusCode::FAILURE;
+  }
+  std::vector<std::string> inputCollections = m_inputCollectionsName.value();
+
+  for (std::size_t 
+         i = 0, 
+         imax = m_inputCollectionsName.value().size(); 
+       i < imax; 
+       ++i) {
+    std::string &fname = inputCollections[i];
+    // std::cout << "--retrieve new name for [" << fname << "]...\n";
+    if (!iomgr->io_contains(this, fname)) {
+      ATH_MSG_ERROR("IoComponentMgr does not know about [" << fname << "] !");
+      return StatusCode::FAILURE;
+    }
+    if (!iomgr->io_retrieve(this, fname).isSuccess()) {
+      ATH_MSG_FATAL("Could not retrieve new value for [" << fname << "] !");
+      return StatusCode::FAILURE;
+    }
+    // std::cout << "--> [" << fname << "]\n" << std::flush;
+  }
+  // all good... copy over.
+  m_inputCollectionsName = inputCollections;
+
+  // remove our EventInfo if any...
+  // {
+  //   const bool force_remove = true;
+  //   if (!m_storeGate->clearStore(force_remove).isSuccess()) {
+  //     ATH_MSG_ERROR("could not clear event store!");
+  //     return StatusCode::FAILURE;
+  //   } else {
+  //     ATH_MSG_INFO("sgdump: \n" << m_storeGate->dump()); 
+  //   }
+  // }
+
+  // std::cout << "--> do_init_io...\n" << std::flush;
+  if (!do_init_io().isSuccess()) {
+    return StatusCode::FAILURE;
+  }
+  
+  ATH_MSG_INFO("I/O reinitialization... [done]");
+  return StatusCode::SUCCESS;
+}
+
+
 ///@{
 /// @c IAddressProvider interface
 ///get all addresses from Provider : Called before Begin Event
          m_tupleName.value(), 
          br_name, 
          (unsigned long)(value_ptr),
-         (unsigned long)(m_nbrEvts-1));
+         (unsigned long)(m_curEvt-1));
       SG::TransientAddress* taddr = new SG::TransientAddress
         (id, sg_key, addr);
       taddr->setProvider(this);
 TTree*
 RootNtupleEventSelector::fetchNtuple(const std::string& fname) const
 {
+  // std::cout << "::fetchNtuple(" << fname << ")..." << std::endl;
   TTree* tree = NULL;
   RootGlobalsRestore rgr;
-  TFile *f = TFile::Open(fname.c_str(), "READ");
+  // std::cout << "::TFile::Open()..." << std::endl;
+  TFile *f = (TFile*)gROOT->GetListOfFiles()->FindObject(fname.c_str());
+  if (!f) {
+    f = TFile::Open(fname.c_str(), "READ");
+  }
   if (!f || f->IsZombie()) {
     ATH_MSG_ERROR("could not open next file in input collection ["
                   << fname << "]");
     }
     return tree;
   }
+  // std::cout << "::TFile::GetTree(" << m_tupleName << ")..." << std::endl;
   tree = (TTree*)f->Get(m_tupleName.value().c_str());
   if (!tree) {
     ATH_MSG_ERROR("could not retrieve tree [" << m_tupleName << "]"
     f->Close();
     return tree;
   }
+  // std::cout << "::TTree::SetBranchStatus()..." << std::endl;
   // disable all branches
   tree->SetBranchStatus("*", 0);
 
   return tree;
 }
+
+StatusCode
+RootNtupleEventSelector::do_init_io()
+{
+  // std::cout << "::fetchNtuple..." << std::endl;
+
+  // initialize some helper structures and data
+  {
+    CollMetaData zero; 
+    zero.min_entries = -1;
+    zero.max_entries = -1;
+    m_collEvts.resize(m_inputCollectionsName.value().size(), zero);
+    m_collIdx = 0;
+  }
+
+  m_tuple = fetchNtuple(m_inputCollectionsName.value()[m_collIdx]);
+  if (!m_tuple) {
+    return StatusCode::FAILURE;
+  }
+
+  // std::cout << "::clear-root-addresses..." << std::endl;
+  // reset the list of branches
+  m_rootAddresses.clear();
+
+  // skip events we are asked to skip
+  m_curEvt  = m_skipEvts;
+  m_nbrEvts = 0;
+
+  // std::cout << "::fetchNtuple...[done]" << std::endl;
+  return StatusCode::SUCCESS;
+}
+
+/// helper method to get the collection index (into `m_inputCollectionsName`)
+/// for a given event index `evtidx`.
+/// returns -1 if not found.
+int 
+RootNtupleEventSelector::find_coll_idx(int evtidx)
+{
+  // std::cout << "--find_coll_idx(" << evtidx << ")..." << std::endl
+  //           << "--collsize: " << m_collEvts.size() << std::endl;
+  for (std::size_t i = 0, imax = m_collEvts.size();
+       i < imax;
+       ++i) {
+    // std::cout << "--[" << i << "]...\n";
+    CollMetaData &itr = m_collEvts[i];
+    if (itr.min_entries == -1) {
+      TTree *tree = fetchNtuple(m_inputCollectionsName.value()[i]);
+      if (tree) {
+        long offset = 0;
+        if (i > 0) {
+          CollMetaData &jtr = m_collEvts[i-1];
+          offset += jtr.max_entries;
+        }
+        itr.min_entries = offset;
+        itr.max_entries = offset + tree->GetEntriesFast();
+      }
+    }
+    // std::cout << "--[" << i << "] => [" << itr.min_entries << ", "
+    //           << itr.max_entries << ") evtidx=[" << evtidx << "]"
+    //           << std::endl;
+    if (itr.min_entries <= evtidx && evtidx < itr.max_entries) {
+      return i;
+    }
+  }
+  return -1;
+}
+
+
 } //> namespace Athena

src/RootNtupleEventSelector.h

 // framework includes
 #include "AthenaBaseComps/AthService.h"
 #include "GaudiKernel/IEvtSelector.h"
+#include "GaudiKernel/IIoComponent.h"
+#include "GaudiKernel/MsgStream.h"
+#include "GaudiKernel/Property.h"
 #include "GaudiKernel/ServiceHandle.h"
-#include "GaudiKernel/Property.h"
-#include "GaudiKernel/MsgStream.h"
 #include "AthenaKernel/IEventSeek.h"
 #include "AthenaKernel/IAddressProvider.h"
 
 class IDictLoaderSvc;
 //class IAthenaRootCnvSvc;
 //namespace Athena { class RootCnvSvc; }
-//namespace Athena { class RootNtupleEventContext; }
+namespace Athena { class RootNtupleEventContext; }
 
 namespace Athena {
 
   virtual public IEvtSelector,
   virtual public IEventSeek,
   virtual public IAddressProvider,
-  virtual public AthService
+  virtual public IIoComponent,
+          public ::AthService
 { 
+  friend class Athena::RootNtupleEventContext;
 
   /////////////////////////////////////////////////////////////////// 
   // Public methods: 
   virtual int curEvent () const;
   ///@}
   
+  /// Callback method to reinitialize the internal state of the component 
+  /// for I/O purposes (e.g. upon @c fork(2))
+  virtual StatusCode io_reinit();
+
   ///@{
   /// @c IAddressProvider interface
   ///get all addresses from Provider : Called before Begin Event
   /// helper method to retrieve the correct tuple
   TTree* fetchNtuple(const std::string& fname) const;
 
+  /// helper method to init the i/o components
+  StatusCode do_init_io();
+
+  /// helper method to get the collection index (into `m_inputCollectionsName`)
+  /// for a given event index `evtidx`.
+  /// returns -1 if not found.
+  int find_coll_idx(int evtidx);
+
   /////////////////////////////////////////////////////////////////// 
   // Private data: 
   /////////////////////////////////////////////////////////////////// 
 
   /// Number of Events read so far.
   mutable long m_nbrEvts;
+  
+  /// current event index
+  mutable long m_curEvt;
+
+  /// current collection index (into `m_inputCollectionsName`)
+  mutable long m_collIdx;
+
+  struct CollMetaData {
+    /// number of entries up to this collection
+    long min_entries;
+    /// number of entries after this collection
+    long max_entries;
+  };
+  /// cache of the number of entries for each collection
+  mutable std::vector<CollMetaData> m_collEvts;
 
   /// current tree being read
   mutable TTree *m_tuple;

test/AthenaRootComps.xml

       </expectations>
    </TEST>
 
+   <TEST name="athenarootcomps.variable_shape" type="script" suite="athenarootcomps">
+      <package_atn>Control/AthenaRootComps</package_atn>
+      <options_atn>run_test_athena_ntuple_dumper.sh</options_atn>
+      <timelimit>30</timelimit>
+      <author> Sebastien Binet </author>
+      <mailto> binet@cern.ch </mailto>
+      <expectations>
+         <returnValue>0</returnValue>
+      </expectations>
+   </TEST>
+
 </atn>

test/run_test_athena_ntuple_dumper.sh

     diff -urN $reffile $chkfile
     sc=$?
     echo "::: comparing ascii dumps...[$sc]"
+
+    if [ $sc -ne 0 ]; then
+	return $sc
+    fi
+
+    echo "::: reading back n-tuple file..."
+    /bin/mv d3pd.root rb.d3pd.root || return 1
+    /bin/rm -f $chkfile > /dev/null
+    ($time_cmd athena.py -c'FNAMES=["rb.d3pd.root"]' AthenaRootComps/test_athena_ntuple_dumper.py >& d3pd.log.txt) || return 1
+    /bin/cat d3pd.timing.log
+    echo "::: comparing ascii dumps..."
+    diff -urN $reffile $chkfile
+    sc=$?
+    echo "::: comparing ascii dumps...[$sc]"
+
     return $sc
 }
 

test/run_test_athena_variable_shape_ntuple.sh

+#!/bin/sh
+
+function check_disappearing_branches()
+{
+echo "::: generate f1.root..."
+athena.py \
+    -c 'EVTMAX=10; BRANCHES=["RunNumber", "EventNumber", "el_n", "el_eta"]; OUTBRANCHES=["int#el_n","std::vector<float>#el_eta"]' \
+    AthenaRootComps/test_athena_variable_shape_ntuple.py \
+    >| log.001.txt \
+    || return 1
+/bin/cp d3pd.root f1.root || return 1
+acmd.py dump-root f1.root | grep "^egamma" | egrep "RunNumber|EventNumber|el_n" | tee f1.ascii || return 1
+cat f1.ascii| cut -d. -f3 >| f1.ascii.todiff || return 1
+
+echo "::: generate f2.root..."
+athena.py \
+    -c 'EVTMAX=10; BRANCHES=["RunNumber", "EventNumber", "el_n", "el_eta"]; OUTBRANCHES=["int#el_n"]' \
+    AthenaRootComps/test_athena_variable_shape_ntuple.py \
+    >| log.002.txt \
+    || return 1
+/bin/cp d3pd.root f2.root || return 1
+acmd.py dump-root f2.root | grep "^egamma" | egrep "RunNumber|EventNumber|el_n" | tee f2.ascii || return 1
+cat f2.ascii| cut -d. -f3 >| f2.ascii.todiff || return 1
+
+echo "::: generate f3.root..."
+athena.py \
+    -c 'EVTMAX=10; BRANCHES=["RunNumber", "EventNumber", "el_n", "el_eta"]; OUTBRANCHES=["int#el_n","std::vector<float>#el_eta"]' \
+    AthenaRootComps/test_athena_variable_shape_ntuple.py \
+    >| log.003.txt \
+    || return 1
+/bin/cp d3pd.root f3.root || return 1
+acmd.py dump-root f3.root | grep "^egamma" | egrep "RunNumber|EventNumber|el_n" | tee f3.ascii || return 1
+cat f3.ascii| cut -d. -f3 >| f3.ascii.todiff || return 1
+
+echo "::: generate merged.root..."
+athena.py \
+    -c 'EVTMAX=30; BRANCHES=["RunNumber", "EventNumber", "el_n", "el_eta"]; OUTBRANCHES=["int#el_n",]; FNAMES=["f1.root","f2.root","f3.root"]' \
+    AthenaRootComps/test_athena_variable_shape_ntuple.py \
+    >| log.004.txt \
+    || return 1
+/bin/cp d3pd.root f4.root || return 1
+acmd.py dump-root f4.root | grep "^egamma" | egrep "RunNumber|EventNumber|el_n" | tee f4.ascii || return 1
+cat f4.ascii| cut -d. -f3 >| f4.ascii.todiff || return 1
+
+echo "::: compare dump-root outputs..."
+cat f[123].ascii.todiff > merged.ascii.todiff || return 1
+diff -urN merged.ascii.todiff f4.ascii.todiff || return 1
+echo "::: compare dump-root outputs... [ok]"
+
+# echo "::: compare logfiles..."
+# cat log.00[123].txt | grep "py_alg" >| log.merged.txt.todiff || return 1
+# grep "py_alg" log.004.txt >| log.004.txt.todiff || return 1
+# diff -urN log.merged.txt.todiff log.004.txt.todiff || return 1
+# echo "::: compare logfiles... [ok]"
+
+return 0
+}
+
+check_disappearing_branches
+
+