def _rd_getitem_(rdir, name): """Get the object from ROOT-file/directory >>> t = f['A/tup'] >>> h = f['histo'] """ ## if not rdir: raise IOError("TDirectory is invalid") ## with ROOTCWD(): ## dirname, sep, fname = name.partition('/') ## while sep: rdir.cd() subdir = rdir.GetDirectory(dirname) if not subdir: raise KeyError("TDirectory[]: unknown directory name '%s%s'" % (rdir.GetPath(), dirname)) rdir = subdir dirname, sep, fname = fname.partition('/') rdir.cd() obj = rdir.Get(dirname) if not obj: raise KeyError("TDirectory[]: unknown object '%s%s'" % (rdir.GetPath(), dirname)) return obj
def _rd_keys_(rdir, recursive=True, no_dir=True): """Get all keys from ROOT file/directory >>> keys = rfile.keys() """ _res = [] if not rdir: return _res ## with ROOTCWD(): ## rdir.cd() _lst = rdir.GetListOfKeys() if not _lst: return _res for i in _lst: inam = i.GetName() idir = rdir.GetDirectory(inam) if not idir or not no_dir: _res.append(inam) if recursive and idir and not idir is rdir: ikeys = _rd_keys_(idir, recursive, no_dir) for k in ikeys: _res.append(inam + '/' + k) return _res
def __exit__(self, *_): """Exit the context: return mode to READ if needed """ if self.read: with ROOTCWD(): r = self.tfile.ReOpen('READ') if r < 0: logger.error("Can't reopen the file for READ!")
def _rd_iteritems_(rdir, fun=lambda k, t, o: True, recursive=True, no_dir=True): """Iterate over the content of ROOT directory/file: >>> for key,obj in rfile.iteritems() : print key , obj >>> for key,hist in rfile.iteritems( ROOT.TH1 ) : print key , hist >>> for key,obj in rfile.iteritems( lambda name,tkey,obj : name[0]=='M' ) : print key,obj """ ## if isinstance(fun, type) and issubclass(fun, (ROOT.TObject, cpp.TObject)): tobj = fun fun = lambda k, t, o: isinstance(o, tobj) ## with ROOTCWD(): ## rdir.cd() _lst = rdir.GetListOfKeys() if _lst: for i in _lst: inam = i.GetName() idir = rdir.GetDirectory(inam) if not idir or not no_dir: obj = rdir.Get(inam) if fun(inam, i, obj): yield inam, obj if recursive and idir and not idir is rdir: for k, o in _rd_iteritems_(idir, fun, recursive, no_dir): yield k, o
def _rd_walk_(rdir, topdown=True): """``walk'' through the directory and its content >>> rfile = ... >>> for r,d,o in rfile.walk() : ... print 'root dir: %s, %d subdirs, %d objects' % ( r , len ( d ) , len( o ) ) """ with ROOTCWD(): ## rdir.cd() _subdirs = [] _objects = [] _idirs = [] _lst = rdir.GetListOfKeys() if _lst: for i in _lst: inam = i.GetName() idir = rdir.GetDirectory(inam) if idir: _subdirs.append(inam) _idirs.append(idir) else: _objects.append(inam) _subdirs = tuple(_subdirs) _objects = tuple(_objects) if topdown: yield rdir.GetName(), _subdirs, _objects for isub in _idirs: for jj in _rd_walk_(isub, topdown): yield jj if not topdown: yield rdir.GetName(), _subdirs, _objects
def _rd_setitem_(rdir, name, tobj): """Write the object to ROOT-file/directory >>> rfile['myhisto'] = h1 >>> rfile['MyDir/MyHisto'] = h2 """ ## if not rdir: raise IOError("TDirectory is invalid") elif not rdir.IsWritable(): raise IOError("TDirectory '%s' is not writable" % rdir.GetPath()) ## with ROOTCWD(): ## dirname, sep, fname = name.partition('/') ## while sep: rdir.cd() subdir = rdir.GetDirectory(dirname) if not subdir: rdir.cd() subdir = rdir.mkdir(dirname, rdir.GetPath() + '/' + dirname) rdir.cd() subdir = rdir.GetDirectory(dirname) if not subdir: raise KeyError("TDirectory[]: can't create '%s%s'" % (rdir.GetPath(), dirname)) rdir = subdir dirname, sep, fname = fname.partition('/') rdir.cd() return rdir.WriteTObject(tobj, dirname, 'WriteDelete')
def _rd_key_object_(rdir, name): """Get the (key.,object) pair from ROOT-file/directory >>> t = f.get_key_object ('A/tup') >>> h = f.get_key_object ('histo') """ ## if not rdir: raise IOError("TDirectory is invalid") ## with ROOTCWD(): ## dirname, sep, fname = name.partition('/') ## while sep: rdir.cd() subdir = rdir.GetDirectory(dirname) if not subdir: raise KeyError("TDirectory[]: unknown directory name '%s%s'" % (rdir.GetPath(), dirname)) rdir = subdir dirname, sep, fname = fname.partition('/') rdir.cd() obj = rdir.Get(dirname) if not obj: raise KeyError("TDirectory[]: unknown object '%s%s'" % (rdir.GetPath(), dirname)) return rdir.FindKey(dirname), obj
def _rd_delitem_(rdir, name, cycle=';*'): """Delete an object from ROOT file/directory >>> del rfile['h1'] >>> del rfile['mydir/h1'] """ ## if not rdir: raise IOError("TDirectory is invalid") elif not rdir.IsWritable(): raise IOError("TDirectory '%s' is not writable" % rdir.GetPath()) ## with ROOTCWD(): ## dirname, sep, fname = name.partition('/') ## while sep: rdir.cd() subdir = rdir.GetDirectory(dirname) if not subdir: raise KeyError("TDirectory, can't delete %s" % name) rdir = subdir dirname, sep, fname = fname.partition('/') rdir.cd() if not rdir.FindKey(dirname): raise KeyError("TDirectory, can't delete %s" % dirname) return rdir.Delete(dirname + cycle)
def create_tree(fname, nentries=1000): """Create a file with a tree >>> create_tree ( 'file.root' , 1000 ) """ import ROOT, random import ostap.io.root_file from array import array var0 = array('i', [0]) var1 = array('d', [0]) var2 = array('d', [0]) var3 = array('d', [0]) NV = 200 evars = [array('d', [0]) for i in range(NV)] NM = 100 mvars = [array('d', 4 * [0]) for i in range(NM)] from ostap.core.core import ROOTCWD with ROOTCWD(), ROOT.TFile.Open(fname, 'new') as root_file: root_file.cd() tree = ROOT.TTree('S', 'tree') tree.SetDirectory(root_file) tree.Branch('evt', var0, 'evt/I') tree.Branch('mass', var1, 'mass/D') tree.Branch('pt', var2, 'pt/D') tree.Branch('eta', var3, 'eta/D') for i in range(1, NV): tree.Branch("sv%d" % i, evars[i], 'sv%d/D' % i) for i in range(1, NM): tree.Branch("vv%d" % i, mvars[i], 'vv%d[4]/D' % i) for i in range(nentries): m = random.gauss(3.1, 0.015) pt = random.uniform(0, 8) eta = random.uniform(2, 4) var0[0] = i var1[0] = m var2[0] = pt var3[0] = eta for j in range(NV): evars[j] = j for j in range(NM): for k in range(4): mvars[j][k] = 4 * j + k tree.Fill() root_file.Write()
def _rd_iter_(rdir): """Iterate (recursive) over the content in ROOT file/directory >>> for obj in rfile : print obj """ ## with ROOTCWD(): rdir.cd() for obj in _rd_itervalues_(rdir, recursive=True): yield obj
def __enter__(self): """Enter the context: make a try to reopen the file with `ReOpen('UPDATE')` """ self.read = self.tfile.GetOption() == 'READ' if self.read: with ROOTCWD(): r = self.tfile.ReOpen('UPDATE') if r < 0: logger.error("Can't reopen the file for UPDATE!") return self.tfile
def _rf_new_close_(rfile, options=''): """Close ROOT-file without making it a current working directory >>> print ROOT.gROOT.CurrentDirectory() >>> f = ROOT.TFile.Open('test_file.root','recreate') >>> print ROOT.gROOT.CurrentDirectory() >>> f.Close() >>> print ROOT.gROOT.CurrentDirectory() """ if rfile and rfile.IsOpen(): with ROOTCWD(): logger.debug("Close ROOT file %s" % rfile.GetName()) rfile._old_close_(options)
def _rd_itervalues_ ( rdir , fun = lambda k,t,o : True , recursive = True , no_dir = True ) : """Iterate over the content of ROOT directory/file: >>> for obj in rfile.itervalues () : print obj >>> for hist in rfile.itervalues ( ROOT.TH1 ) : print hist >>> for obj in rfile.itervalues ( lambda k,t,o : k[0]=='M' ) : print obj """ ## with ROOTCWD() : ## rdir.cd() for key , obj in _rd_iteritems_ ( rdir , fun , recursive , no_dir ) : yield obj
def process(self, jobid, item): """The actual processing ``params'' is assumed to be a tuple-like entity: - the file name - the tree name in the file - the variable/expression/expression list of quantities to project - the selection/weighting criteria - the first entry in tree to process - number of entries to process """ import ROOT from ostap.logger.utils import logWarning with logWarning(): import ostap.core.pyrouts import ostap.trees.trees import ostap.histos.histos import ostap.frames.frames from ostap.trees.trees import Chain, Tree input = Chain(name=item.name, files=item.files, first=item.first, nevents=item.nevents) chain = input.chain first = input.first nevents = input.nevents ## use the regular projection from ostap.trees.trees import _tt_project_ ## Create the output histogram NB! (why here???) from ostap.core.core import ROOTCWD with ROOTCWD(): ROOT.gROOT.cd() histo = self.histo.Clone() self.__output = 0, histo from ostap.trees.trees import _tt_project_ self.__output = _tt_project_(tree=chain, histo=histo, what=self.what, cuts=self.cuts, options='', nentries=nevents, firstentry=first) del item return self.__output
def _rf_new_init_(rfile, fname, mode='', *args): """Open/create ROOT-file without making it a current working directory >>> print ROOT.gROOT.CurrentDirectory() >>> f = ROOT.TFile('test_file.root','recreate') >>> print ROOT.gROOT.CurrentDirectory() Attention: IOError exception is raised for invalid file/open_mode """ with ROOTCWD(): logger.debug("Open ROOT file %s/'%s'" % (fname, mode)) result = rfile._old_init_(fname, open_mode(mode), *args) if not rfile or not rfile.IsOpen(): raise IOError("Can't open ROOT file %s/'%s'" % (fname, mode)) return result
def _rf_new_open_ ( fname , mode = '' , *args ) : """Open/create ROOT-file without making it a current working directory >>> print ROOT.gROOT.CurrentDirectory() >>> f = ROOT.TFile.Open('test_file.root','recreate') >>> print ROOT.gROOT.CurrentDirectory() ATTENTION: No exceptions are raised for invalid file/open_mode """ with ROOTCWD() : logger.debug ( "Open ROOT file %s/'%s'" % ( fname , mode ) ) fopen = ROOT.TFile._old_open_ ( fname , open_mode ( mode ) , *args ) if not fopen or not fopen.IsOpen() : logger.error ( "Can't open ROOT file %s/'%s'" % ( fname , mode ) ) return fopen
def _rd_getattr_ ( rdir , name ) : """Get the object from ROOT-file >>> t = f.A.tup >>> h = f.myhisto """ ## if not rdir : return None ## with ROOTCWD() : ## dirname , sep , fname = name.partition('/') if sep : raise AttributeError('TDirectory[]: invalid attribute %s.%s' % rdir.GetPath() , dirname ) rdir.cd() obj = rdir.Get ( dirname ) if not obj : raise AttributeError ( 'TDirectory[]: unknown attribute %s.%s' % ( rdir.GetPath(), dirname ) ) return obj
def top_dir(rdir): """```topdir'': get the top directory for the given directory""" if not rdir: return None with ROOTCWD(): top = rdir if isinstance(rdir, ROOT.TDirectory): top = rdir elif hasattr(rdir, 'GetDirectory'): top = rdir.GetDirectory() while top: moth = top.GetMotherDir() if not moth: return top top = moth else: return None
def _rd_iterkeys_ ( rdir , typ = None , recursive = True , no_dir = True ) : """Iterate over the keys in ROOT file/directory >>> for key,obj in rfile.iteritems() : print key , obj >>> for key,hist in rfile.iteritems( ROOT.TH1 ) : print key , hist """ ## with ROOTCWD() : ## rdir.cd() _lst = rdir.GetListOfKeys() for i in _lst : inam = i.GetName() idir = rdir.GetDirectory ( inam ) if not idir or not no_dir : if typ is None or isinstance ( rdir.Get ( inam ) , typ ) : yield inam if recursive and idir and not idir is rdir : for k in _rd_iterkeys_ ( idir , typ , recursive , no_dir ) : yield k
def dump_root(objects, rfile=''): if not rfile: rfile = 'ROOT_Objects_%s.root' % root_version_int if isinstance(objects, ROOT.TObject): objects = [objects] import ostap.io.root_file from ostap.core.core import ROOTCWD with ROOTCWD(), ROOT.TFile(rfile, 'RECREATE') as f: f.cd() for o in objects: o.Write() logger.info('Dump object of type %s' % type(o)) f.ls() return rfile
def _rd_contains_(rdir, name): """Check the existence of key in the root-file/directory >>> if not 'myhisto' in rfile : break >>> if 'MyDir/MyHisto' in rfile : print 'OK!' """ ## if not rdir: return False ## with ROOTCWD(): ## dirname, sep, fname = name.partition('/') ## while sep: rdir.cd() subdir = rdir.GetDirectory(dirname) if not subdir: return False ## RETURN rdir = subdir dirname, sep, fname = fname.partition('/') rdir.cd() return rdir.FindKey(dirname) ## RETURN
def _rd_ikeyskeys_(rdir, recursive=True, no_dir=True): """Iterator over keyname/key pairs from ROOT file/directory >>> for kname, key in rfile.ikeyskeys() : ... print kname , key.GetClassName() """ ## with ROOTCWD(): ## rdir.cd() klst = rdir.GetListOfKeys() for key in klst: kname = key.GetName() kdir = rdir.GetDirectory(kname) if not kdir or not no_dir: yield kname, key if recursive and kdir and not kdir is rdir: for kn, kk in _rd_ikeyskeys_(kdir, recursive, no_dir): yield kname + '/' + kn, kk
def _rd_len_(rdir): """Length of the directory : (recursive) number of keys >>> rdir = ... >>> len(rdir) """ if not rdir: return 0 with ROOTCWD(): ## rdir.cd() lst = rdir.GetListOfKeys() nkeys = lst.GetSize() if valid_pointer(lst) else 0 for item in lst: inam = item.GetName() idir = rdir.GetDirectory(inam) if idir and not idir is rdir: nkeys += _rd_len_(idir) return nkeys
def create_tree(fname, nentries=1000): """Create a file with a tree >>> create_tree ( 'file.root' , 1000 ) """ import ROOT, random import ostap.io.root_file from array import array var1 = array('d', [0]) var2 = array('d', [0]) var3 = array('d', [0]) from ostap.core.core import ROOTCWD with ROOTCWD(), ROOT.TFile.Open(fname, 'new') as root_file: root_file.cd() tree = ROOT.TTree('S', 'tree') tree.SetDirectory(root_file) tree.Branch('mass', var1, 'mass/D') tree.Branch('pt', var2, 'pt/D') tree.Branch('eta', var3, 'eta/D') for i in range(nentries): m = random.gauss(3.1, 0.015) pt = random.uniform(0, 10) eta = random.uniform(2, 5) var1[0] = m var2[0] = pt var3[0] = eta tree.Fill() root_file.Write()
def _rf_new_open_(fname, mode='', args=(), exception=False): """Open/create ROOT-file without making it a current working directory >>> print ROOT.gROOT.CurrentDirectory() >>> f = ROOT.TFile.Open('test_file.root','recreate') >>> print ROOT.gROOT.CurrentDirectory() ATTENTION: No exceptions are raised for invalid file/open_mode, unless specified """ with ROOTCWD(): logger.debug("Open ROOT file %s/'%s'" % (fname, mode)) try: fopen = ROOT.TFile._old_open_(fname, open_mode(mode), *args) except (OSError, IOError): if exception: raise else: logger.error("Can't open ROOT file %s/'%s'" % (fname, mode)) return None if not fopen or not fopen.IsOpen(): if exception: raise IOError("Can't open ROOT file %s/'%s'" % (fname, mode)) else: logger.error("Can't open ROOT file %s/'%s'" % (fname, mode)) return fopen
def __init__ ( self , tree = None , name = None , files = [] , first = 0 , nevents = -1 ) : assert ( name and files ) or valid_pointer ( tree ) , 'Invalid arguments %s/%s/%s' % ( tree , name , files ) assert isinstance ( first , int ) and 0<= first , "Invalid ``first'' %s" % first self.__first = int ( first ) self.__nevents = nevents if 0 <= nevents < ROOT.TChain.kMaxEntries else -1 if files and isinstance ( files , str ) : files = files, if name and files : self.__name = name self.__files = files chain = self.__create_chain() assert valid_pointer ( chain ), 'Invalid TChain!' assert len ( files ) == len( chain.files() ) , 'Invalid length of files' self.__chain = chain elif valid_pointer ( tree ) : self.__name = tree.GetName() if isinstance ( tree , ROOT.TChain ) : self.__files = tree.files() self.__files = tuple ( self.__files ) self.__chain = tree elif isinstance ( tree , ROOT.TTree ) : topdir = tree.topdir if isinstance ( topdir , ROOT.TFile ) : self.__files = topdir.GetName() , else : fname = CleanUp.tempfile ( suffix = '.root' , prefix = 'tree_' ) from ostap.core.core import ROOTCWD with ROOTCWD() : import ostap.io.root_file with ROOT.TFile ( fname , 'NEW') as rfile : rfile.cd() tname = tree.GetName() rfile[ tname ] = tree rfile.ls() self.__files = fname , self.tmpfiles += [ fname ] chain = ROOT.TChain( tree.GetName() ) chain.Add ( self.__files[0] ) tmp = chain assert len ( chain ) == len ( tree ) , 'Somethnig wron happens here :-( ' self.__chain = chain ## last adjustment ll = len ( self ) if ll < self.__first : logger.warning ( 'Tree/Chain will be empty %s/%s' % ( self.__first , ll ) ) self.__first = ll self.__nevents = 0 ## if number of events is specified: if 0 < self.__nevents : self.__nevents = min ( self.__nevents , ll - self.__first )
def _tt_project_ ( tree , histo , what , cuts = '' , options = '' , nentries = -1 , firstentry = 0 , silent = False ) : """Helper project method >>> tree = ... >>> h1 = ROOT.TH1D(... ) >>> tree.Project ( h1.GetName() , 'm', 'chi2<10' ) ## standart ROOT >>> h1 = ROOT.TH1D(... ) >>> tree.project ( h1.GetName() , 'm', 'chi2<10' ) ## ditto >>> h1 = ROOT.TH1D(... ) >>> tree.project ( h1 , 'm', 'chi2<10' ) ## use histo ## make invididual projections of m1 and m2 and make a sum of distributions >>> h1 = ROOT.TH1D(... ) >>> tree.project ( h1 , ('m1','m2') , 'chi2<10' ) ## two variables >>> tree.project ( h1 , 'm1,m2' , 'chi2<10' ) ## ditto >>> tree.project ( h1 , 'm1;m2' , 'chi2<10' ) ## ditto - tree : the tree - histo : the histogram (or histogram name) - what : variable/expression to project. It can be expression or list/tuple of expression or comma (or semicolumn) separated expression - cuts : selection criteria/weights """ # ## if nentries < 0 : nentries = ROOT.TTree.kMaxEntries args = options , nentries , firstentry, silent ## hname = histo if hasattr ( histo , 'GetName' ) : hname = histo.GetName() ## elif isinstance ( histo , str ) : ## h = ROOT.gROOT.FindObject ( hname ) ## if h : histo = h ## reset it! if histo and isinstance ( histo , ROOT.TH1 ) : histo.Reset() # if isinstance ( cuts , ROOT.TCut ) : cuts = str(cuts) if not what : return 0, histo # ## trivial 1-item list if hasattr ( what , '__len__' ) and 1 == len ( what ) and not isinstance ( what , (str, ROOT.TCut) ): what = what[0] ## check for comma-separated list of expressions: if isinstance ( what , str ) : what = what.split(',') if 1 == len(what) : what = what[0] ## check for semicolumn-separated list of expressions: if isinstance ( what , str ) : what = what.split(';') if 1 == len(what) : what = what[0] # if isinstance ( what , str ) : what = what elif isinstance ( what , ROOT.TCut ) : what = str(what) elif isinstance ( histo , ROOT.TH1 ) : rr = 0 hh = histo.clone() for v in what : r , h = _tt_project_ ( tree , hh , v , cuts , options , *args ) rr += r histo += h hh.Delete() del hh return rr , histo elif isinstance ( histo , str ) : ## process the head of the list: the first call creates the histo... rr, hh = _tt_project_ ( tree , histo , what[0] , cuts , *args ) histo = hh if 1 == len ( what ) : return rr , histo # normal processing of the tail of the list using created historgam hh = histo.clone() r1 , h1 = _tt_project_ ( tree , hh , what[1:] , cuts , *args ) rr += r1 histo += h1 hh.Delete() del hh, h1 return rr , histo ## the basic case from ostap.core.core import ROOTCWD with ROOTCWD() : ROOT.gROOT.cd () ## make projection result = tree.Project ( hname , what , cuts , *args[:-1] ) if isinstance ( histo , ROOT.TH1 ) : return result, histo elif isinstance ( histo , str ) : h = ROOT.gROOT.FindObject ( hname ) if h : return result, h return result, histo
def treat_file(f): """Perform treatment of the file """ fok = os.path.exists(f) and os.path.isfile(f) name, dot, ext = f.rpartition('.') if name and dot and ext in ('root', 'ROOT'): logger.debug("Try to open file '%s'" % f) try: from ostap.core.core import ROOTCWD with ROOTCWD(): _f = ROOT.TFile.Open(f, 'READ') root_files[name] = _f logger.info("Open ROOT file '%s'" % f) _f.ls() except: # - Batch mode: just re-raise exception if arguments.batch: raise # - Interactive mode: print traceback and continue logger.error('Failure to open ROOT file "%s"' % f, exc_info=True) ## execute python file elif fok and name and dot and 'py' == ext: logger.debug("Try to execute '%s'" % f) try: import runpy globs = runpy.run_path(f, globals() if arguments.WithContext else {}, run_name='__main__') globs = dict(((k, v) for k, v in globs.items() if not k.startswith('__') and not k.endswith('__'))) logger.debug('Symbols from %s: %s' % (f, globs.keys())) globals().update(globs) del globs logger.info("Executed '%s'" % f) python_scripts.append(f) except: # - Batch mode: just re-raise exception if arguments.batch: raise # - Interactive mode: print traceback and continue logger.error('Failure to execute "%s"' % f, exc_info=True) elif name and dot and ext.upper() in ('C', 'C+', 'C++', 'CPP', 'CPP+', 'CPP++', 'CXX', 'CXX+', 'CXX++'): if not load_macro(f): if arguments.batch: raise RuntimeError("No macros are loaded for '%s' pattern" % f) else: logger.error("No macros are loaded for '%s' pattern" % f) ## execute ostap-script elif fok and name and dot and ext in ('ost', 'ostp', 'ostap', 'opy'): logger.debug("Try to execute '%s'" % f) try: ## this one is always executed with the context!!! import runpy globs = runpy.run_path(f, globals(), run_name='__main__') globs = dict(((k, v) for k, v in globs.items() if not k.startswith('__') and not k.endswith('__'))) logger.debug('Symbols from %s: %s' % (f, globs.keys())) globals().update(globs) del globs python_scripts.append(f) logger.info("Executed '%s'" % f) except: # - Batch mode: just re-raise exception if arguments.batch: raise # - Interactive mode: print traceback and continue logger.error('Failure to execute "%s"' % f, exc_info=True) else: ## collect the argument as parameters PARAMETERS.append(f) logger.debug('Add argument %s to PARAMETERS' % f)