def prepare_data(tmpdir, nfiles=100, nentries=100, ppservers=(), silent=True):
## Use generic Task from Kisa
from ostap.parallel.parallel import GenericTask as Task
task = Task(processor=create_tree)
## task = PrepareTask ()
wmgr = Parallel.WorkManager(ppservers=ppservers, silent=silent)
from ostap.utils.cleanup import CleanUp
tmpfile = CleanUp.tempfile(prefix='test_kisa_', suffix='.root', dir=tmpdir)
fname = '%s/test_kisa_%d.root'
files = [
CleanUp.tempfile(prefix='test_kisa_', suffix='.root', dir=tmpdir)
for i in range(nfiles)
]
wmgr.process(task, [(f, nentries) for f in files])
the_files = set()
for f in task.results():
if os.path.exists(f):
the_files.add(f)
from ostap.trees.data import Data
the_files = list(the_files)
the_files.sort()
return Data('S', list(the_files))
def tempdir(cls,
suffix='=db-dir',
prefix='compress-shelve-dir-',
date=True):
"""Create the temporary directory
The directory will be cleaned-up and deleted at-exit.
"""
from ostap.utils.cleanup import CleanUp as CU
return CU.tempdir(suffix=suffix, prefix=prefix, date=date)
def tempfile(cls,
suffix='-db',
prefix='compress-shelve-',
dir=None,
date=True):
"""Ccreate the name for the temproary file
The file will be deleted at-axit
"""
from ostap.utils.cleanup import CleanUp as CU
return CU.tempfile(suffix=suffix, prefix=prefix, dir=dir, date=date)
def prepare_data(nfiles=50, nentries=500):
from ostap.utils.cleanup import CleanUp
files = [
CleanUp.tempfile(prefix='ostap-test-trees-addbranch-%d-' % i,
suffix='.root') for i in range(nfiles)
]
for f in progress_bar(files):
create_tree(f, nentries)
return files
def make_build_dir ( build = None ) :
"""Create proper temporary directory for ROOT builds
"""
if not build or not writeable ( build ) :
from ostap.utils.cleanup import CleanUp
build = CleanUp.tempdir ( prefix = 'ostap-build-' )
if not os.path.exists ( build ) :
make_dir ( build )
return build
def prepare_data(nfiles=50, nentries=500):
""" prepare data for tests
"""
from ostap.utils.cleanup import CleanUp
files = [
CleanUp.tempfile(prefix='ostap-test-fitting-fill-%d-' % i,
suffix='.root') for i in range(nfiles)
]
for f in progress_bar(files):
create_tree(f, nentries)
return files
def prepare_data(nfiles=10, nentries=100):
"""Prepare data for the test
"""
files = []
for i in progress_bar(range(nfiles)):
from ostap.utils.cleanup import CleanUp
tmpfile = CleanUp.tempfile(prefix='ostap-test-selectors-',
suffix='.root')
files.append(create_tree(tmpfile, nentries))
files.sort()
return Data('S', files)
## 2) check the configuration file
import ostap.core.config as OCC
build_dir = OCC.general.get ( 'BUILD_DIR' , None )
del OCC
if build_dir and not os.path.exists ( build_dir ) :
build_dir = make_dir ( build_dir )
if build_dir and not writeable ( build_dir ) :
logger.warning ('Directory %s is not writeable!' % tmp_dir )
build_dir = None
## 3) Construct something temporary and unique
if not build_dir :
from ostap.utils.cleanup import CleanUp
build_dir = CleanUp.tempdir ( prefix = 'ostap-build-' )
# ==============================================================================
## Context manager to use certain build build directory
# (useful for multiprocessing environment)
# @code
# with UseBuildDir() :
# ...
# @endcode
class UseBuildDir ( object ) :
"""Context manager to use certain build build directory
(useful for multiprocessing environment)
>>> with UseBuildDir() :
>>> ...
"""
import ostap.parallel.parallel_fill ## ATTENTION!
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('ostap.test_parallel_kisa')
else:
logger = getLogger(__name__)
# =============================================================================
with timing('Prepare data'):
logger.info('Prepare data, it could take some time')
from ostap.utils.cleanup import CleanUp
tmpdir = CleanUp().tmpdir
import prepare_test_kisa as PTK
## data = PTK.prepare_data ( tmpdir , nfiles = 250 , nentries = 200000 , silent = False )
data = PTK.prepare_data(tmpdir, nfiles=100, nentries=50000, silent=False)
## data = PTK.prepare_data ( tmpdir , nfiles = 10 , nentries = 100 , silent = False )
logger.info('DATA %s' % data)
# =============================================================================
def test_kisa():
h1 = ROOT.TH1D('h1', '', 200, 3, 3.2)
h2 = h1.clone()
chain = data.chain
import ROOT, os, random
import ostap.core.pyrouts
import ostap.trees.param
import ostap.math.models
from ostap.core.core import hID, SE, Ostap
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('test_trees_param')
else:
logger = getLogger(__name__)
# =============================================================================
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile(suffix='.root', prefix='ostap-test-trees-param-')
if not os.path.exists(data_file):
import random
N = 200000
logger.info('Prepare input ROOT file with data %s' % data_file)
with ROOT.TFile.Open(data_file, 'recreate') as test_file:
## test_file.cd()
tree = ROOT.TTree('S', 'signal tree')
tree.SetDirectory(test_file)
from array import array
var1 = array('d', [0])
import ostap.io.root_file
from builtins import range
from ostap.core.core import ROOTCWD
from ostap.utils.progress_bar import progress_bar
from array import array
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__ :
logger = getLogger ( 'ostap.test_tools_tmva2' )
else :
logger = getLogger ( __name__ )
# ==============================================================================
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile ( suffix = '.root' , prefix = 'ostap-test-tools-tmva2-' )
if not os.path.exists( data_file ) :
import random
nB = 10000
nS = 10000
logger.info('Prepare input ROOT file with data %s' % data_file )
with ROOT.TFile.Open( data_file ,'recreate') as test_file:
## test_file.cd()
treeSignal = ROOT.TTree('S','signal tree')
treeBkg = ROOT.TTree('B','background tree')
treeSignal.SetDirectory ( test_file )
treeBkg .SetDirectory ( test_file )
from array import array
var1 = array ( 'd', [0] )
var2 = array ( 'd', [0] )
import ostap.io.root_file
from builtins import range
from ostap.core.core import ROOTCWD
from ostap.utils.progress_bar import progress_bar
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__: logger = getLogger('ostap.test_tools_tmva')
else: logger = getLogger(__name__)
# ==============================================================================
# def test_tmva () :
if 1 < 2:
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile(suffix='.root',
prefix='ostap-test-tools-tmva-')
if not os.path.exists(data_file):
import random
nB = 10000
nS = 10000
## nB = 100000
## nS = 20000
logger.info('Prepare input ROOT file with data %s' % data_file)
with ROOT.TFile.Open(data_file, 'recreate') as test_file:
test_file.cd()
treeSignal = ROOT.TTree('S', 'signal tree')
treeBkg = ROOT.TTree('B', 'background tree')
treeSignal.SetDirectory(test_file)
treeBkg.SetDirectory(test_file)
from array import array
import ostap.io.zipshelve as DBASE
from ostap.utils.timing import timing
from ostap.logger.colorized import attention, allright
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('ostap.test_tools_reweight2')
else:
logger = getLogger(__name__)
# =============================================================================
logger.info('Test for 2D-Reweighting machinery')
# ============================================================================
from ostap.utils.cleanup import CleanUp
testdata = CleanUp.tempfile(suffix='.root',
prefix='ostap-test-tools-reweight2-')
tag_data = 'DATA2_histogram'
tag_datax = 'DATAX_histogram'
tag_datay = 'DATAY_histogram'
tag_mc = 'MC2_tree'
dbname = CleanUp.tempfile(suffix='.db', prefix='ostap-test-tools-reweight2-')
if os.path.exists(testdata): os.remove(testdata)
if os.path.exists(dbname): os.remove(dbname)
import ostap.parallel.kisa
N1 = 1000000
N2 = 100000
N1 = 100000
from ostap.core.core import ROOTCWD
from ostap.utils.progress_bar import progress_bar
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('ostap.test_tools_tmva')
else:
logger = getLogger(__name__)
# ==============================================================================
# def test_tmva () :
if 1 < 2:
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile(suffix='.root', prefix='test_tools_tmva_')
if not os.path.exists(data_file):
import random
nB = 10000
nS = 10000
## nB = 100000
## nS = 20000
logger.info('Prepare input ROOT file with data %s' % data_file)
with ROOT.TFile.Open(data_file, 'recreate') as test_file:
logger.info('HERE-0!')
test_file.cd()
logger.info('HERE-1!')
treeSignal = ROOT.TTree('S', 'signal tree')
treeBkg = ROOT.TTree('B', 'background tree')
treeSignal.SetDirectory(test_file)
treeBkg.SetDirectory(test_file)
import ostap.io.root_file
from builtins import range
from ostap.core.core import ROOTCWD
from ostap.utils.progress_bar import progress_bar
from array import array
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('ostap.test_tools_chopping')
else:
logger = getLogger(__name__)
# =============================================================================
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile(suffix='.root', prefix='test_tools_chopping_')
if not os.path.exists(data_file):
import random
nB = 20000
nS = 10000
s_evt_per_run = 927
b_evt_per_run = 511
logger.info('Prepare input ROOT file with data %s' % data_file)
with ROOT.TFile.Open(data_file, 'recreate') as test_file:
## test_file.cd()
treeSignal = ROOT.TTree('S', 'signal tree')
treeBkg = ROOT.TTree('B', 'background tree')
from ostap.core.meta_info import root_info
from builtins import range
from ostap.core.core import ROOTCWD
from ostap.utils.progress_bar import progress_bar
from array import array
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('ostap.test_tools_chopping')
else:
logger = getLogger(__name__)
# =============================================================================
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile(suffix='.root',
prefix='ostap-test-tools-chopping-')
if not os.path.exists(data_file):
import random
nB = 20000
nS = 10000
if root_info < (6, 15):
nB = 2000
nS = 1000
s_evt_per_run = 927
b_evt_per_run = 511
logger.info('Prepare input ROOT file with data %s' % data_file)
import ostap.io.root_file
from builtins import range
from ostap.core.core import ROOTCWD
from ostap.utils.progress_bar import progress_bar
from array import array
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__ :
logger = getLogger ( 'ostap.test_tools_chopping' )
else :
logger = getLogger ( __name__ )
# =============================================================================
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile ( suffix = '.root' , prefix = 'ostap-test-tools-chopping-' )
if not os.path.exists( data_file ) :
import random
nB = 20000
nS = 10000
s_evt_per_run = 927
b_evt_per_run = 511
logger.info('Prepare input ROOT file with data %s' % data_file )
with ROOT.TFile.Open( data_file ,'recreate') as test_file:
## test_file.cd()
treeSignal = ROOT.TTree('S','signal tree')
treeBkg = ROOT.TTree('B','background tree')
import ostap.core.config as _CONFIG
if 'BUILD_DIR' in _CONFIG.general:
bdir = _CONFIG.general.get('BUILD_DIR', fallback='')
if good_dir(bdir): build_dir = bdir
elif bdir and not os.path.exists(bdir): make_dir(bdir)
if good_dir(bdir): build_dir = bdir
# ==============================================================================
# 3) use the temporary directory
if not build_dir:
from ostap.utils.cleanup import CleanUp as _CU
bdir = _CU.tempdir(prefix='build_')
if good_dir(bdir): build_dir = bdir
del _CU
# =============================================================================
if build_dir and good_dir(build_dir):
ROOT.gSystem.SetBuildDir(build_dir)
build_dir = ROOT.gSystem.GetBuildDir()
# =============================================================================
if '__main__' == __name__:
# =========================================================================
# logging
# =========================================================================
from ostap.trees.data_utils import Data
from ostap.logger.colorized import attention, allright
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('ostap.test_tools_gbreweight')
else:
logger = getLogger(__name__)
# =============================================================================
logger.info('Test for 2D-Reweighting machinery using GBReweighter')
# ============================================================================
from ostap.utils.cleanup import CleanUp
testdata = CleanUp.tempfile(suffix='.root',
prefix='ostap-test-tools-gbreweight-')
## dbname = CleanUp.tempfile ( suffix = '.db' , prefix ='ostap-test-tools-gbreweight-' )
## testdata = 'testdata.root'
## dbname = 'testdb.db'
tag_data = 'DATA2_histogram'
tag_datax = 'DATAX_histogram'
tag_datay = 'DATAY_histogram'
tag_mc = 'MC2_tree'
## if os.path.exists ( testdata ) : os.remove ( testdata )
## if os.path.exists ( dbname ) : os.remove ( dbname )
import ostap.parallel.kisa
N1 = 1000000
def __exit__ ( self , *_ ) :
ROOT.gSystem.SetBuildDir ( self.__prev )
if self.__created :
from ostap.utils.cleanup import CleanUp
CleanUp.remove_dir ( self.__build )
import ostap.io.zipshelve as DBASE
from ostap.utils.timing import timing
from ostap.logger.colorized import attention, allright
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('ostap.test_tools_reweight2')
else:
logger = getLogger(__name__)
# =============================================================================
logger.info('Test for 2D-Reweighting machinery')
# ============================================================================
from ostap.utils.cleanup import CleanUp
testdata = CleanUp.tempfile(suffix='.root', prefix='test_tools_reweight2_')
tag_data = 'DATA2_histogram'
tag_datax = 'DATAX_histogram'
tag_datay = 'DATAY_histogram'
tag_mc = 'MC2_tree'
dbname = CleanUp.tempfile(suffix='.db', prefix='test_tools_reweight2_')
if os.path.exists(testdata): os.remove(testdata)
if os.path.exists(dbname): os.remove(dbname)
import ostap.parallel.kisa
if not os.path.exists(testdata):
#
seed = 1234567890
random.seed(seed)
import ROOT, os, random
import ostap.core.pyrouts
import ostap.trees.param
import ostap.math.models
from ostap.core.core import hID, SE, Ostap
# =============================================================================
# logging
# =============================================================================
from ostap.logger.logger import getLogger
if '__main__' == __name__ or '__builtin__' == __name__:
logger = getLogger('test_trees_param')
else:
logger = getLogger(__name__)
# =============================================================================
from ostap.utils.cleanup import CleanUp
data_file = CleanUp.tempfile(suffix='.root', prefix='test_trees_param_')
if not os.path.exists(data_file):
import random
N = 200000
logger.info('Prepare input ROOT file with data %s' % data_file)
with ROOT.TFile.Open(data_file, 'recreate') as test_file:
## test_file.cd()
tree = ROOT.TTree('S', 'signal tree')
tree.SetDirectory(test_file)
from array import array
var1 = array('d', [0])