"""
The :py:mod:`bamboo.root` module collects a set of thin wrappers around ROOT
methods, and centralizes the import of the Cling interpreter global namespace
in PyROOT. For compatibility, it is recommended that user code uses
``from bamboo.root import gbl`` rather than ``import ROOT as gbl`` or
``from cppyy import gbl``.
"""
import functools
from contextlib import contextmanager
from importlib import resources
import ROOT; gbl = ROOT # noqa: E702
_rootVersion = gbl.gROOT.GetVersion()
_rootVersion_split = tuple(int(iv) for tk in _rootVersion.split("/") for iv in tk.split("."))
if _rootVersion_split[0] == 6 and _rootVersion_split[1] < 22:
gbl.PyConfig.IgnoreCommandLineOptions = True
_includePaths = set()
_headers = set()
_dynamicPaths = set()
_libraries = set()
[docs]
def addIncludePath(incPath):
""" Add an include path to the ROOT interpreter """
global _includePaths
_includePaths.add(incPath)
gbl.gInterpreter.AddIncludePath(incPath)
[docs]
def addDynamicPath(libPath):
""" Add a dynamic library path to the ROOT interpreter"""
global _dynamicPaths
_dynamicPaths.add(libPath)
gbl.gSystem.AddDynamicPath(libPath)
[docs]
def loadLibrary(libName):
""" Load a shared library in the ROOT interpreter """
global _libraries
_libraries.add(libName)
st = gbl.gSystem.Load(libName)
if st == -1:
raise RuntimeError(f"Library {libName} could not be found")
elif st == -2:
raise RuntimeError(f"Version match for library {libName}")
[docs]
def findLibrary(libName):
""" Check if a library can be found, and returns the path in that case """
ret = gbl.gSystem.FindDynamicLibrary(gbl.TString(libName), True)
if ret != "":
return str(ret)
[docs]
def loadDependency(bambooLib=None, includePath=None, headers=None, dynamicPath=None, libraries=None):
"""
Load a C++ extension
:param bambooLib: name(s) of the bamboo extension libraries, if any
:param includePath: include directory for headers
:param headers: headers to load explicitly (which can depend on other headers in the inclue path)
:param dynamicPath: dynamic library path to add
:param libraries: additional shared libraries to load
"""
def asList(arg):
if isinstance(arg, str):
return [arg]
else:
return arg
if dynamicPath:
for pth in asList(dynamicPath):
addDynamicPath(pth)
if libraries:
for lib in asList(libraries):
loadLibrary(f"lib{lib}")
if includePath:
for pth in asList(includePath):
addIncludePath(pth)
if headers:
for hdr in asList(headers):
loadHeader(hdr)
if bambooLib:
for lib in asList(bambooLib):
loadLibrary(f"lib{lib}")
def getIMTPoolSize():
return ROOT.GetThreadPoolSize() if ROOT.IsImplicitMTEnabled() else 0
@contextmanager
def returnToDirectory():
gpwd = gbl.gDirectory.GetPath()
yield
gbl.gDirectory.cd(gpwd)
[docs]
class once:
""" Function decorator to make sure things are not loaded more than once """
def __init__(self, fun):
self.fun = fun
self.has_run = False
functools.update_wrapper(self, fun)
def __call__(self, *args, **kwargs):
if not self.has_run:
self.has_run = True
return self.fun(*args, **kwargs)
_defaultHeaders = ("Math/VectorUtil.h", "bamboohelpers.h", "range.h",
"LumiMask.h", "bamboorandom.h",
"scalefactors.h", "BTagCalibrationStandalone.h")
_defaultBambooLibs = ("BambooLumiMask", "BambooRandom", "BinnedValues")
_defaultDependencies = ("ROOT::ROOTDataFrame", "ROOT::GenVector")
@once
def loadBambooExtensions():
# Add extension libraries and necessary header files to the ROOT interpreter
import os.path
with (
resources.as_file(resources.files("bamboo") / "include") as incDir,
resources.as_file(resources.files("bamboo") / "lib") as libDir
):
incDir = str(incDir)
libDir = str(libDir)
if not (os.path.exists(incDir) and os.path.exists(libDir)):
raise RuntimeError(
f"Could not find {incDir} and/or {libDir}, this may happen when running "
"from the source repository and using a virtualenv with non-editable install")
addIncludePath(incDir)
addDynamicPath(libDir)
# now load default headers and libraries
for lib in _defaultBambooLibs:
loadLibrary(f"lib{lib}")
for fname in _defaultHeaders:
loadHeader(fname)
@once
def loadJMESystematicsCalculators(backend=None):
try:
import CMSJMECalculators # noqa: F401
except ImportError:
raise RuntimeError(
"Could not load jet&MET correction and systematic variations calculators. "
"Please install them (see https://gitlab.cern.ch/cp3-cms/CMSJMECalculators "
"for more information)"
) from None
loadBambooExtensions()
loadcorrectionlib(backend)
loader = backend.addDependency if backend else loadDependency
with (
resources.as_file(resources.files("CMSJMECalculators") / "include") as inc_dir,
resources.as_file(resources.files("CMSJMECalculators") / "lib") as libDir
):
inc_dir = str(inc_dir)
libDir = str(libDir)
headers = [
"CMSCalculatorsUtilities.h",
"JetMETVariationsCalculatorBase.h",
"JetVariationsCalculator.h",
"FatJetVariationsCalculator.h",
"Type1METVariationsCalculator.h",
"EGammaVariationsCalculator.h",
"TauVariationsCalculator.h",
"MuonVariationsCalculator.h"
]
for header in headers:
loader(includePath=inc_dir, headers=header)
loader(dynamicPath=libDir, libraries="CMSJMECalculators")
if backend:
with resources.as_file(resources.files("CMSJMECalculators") / "cmake") as calc_dir:
calc_dir = str(calc_dir)
backend.addDependency(
package="CMSJMECalculators", cmakeArgs=f"-DCMSJMECalculators_DIR={calc_dir}",
)
@once
def loadRochesterCorrectionCalculator(backend=None):
loadBambooExtensions()
loader = backend.addDependency if backend else loadDependency
loader(bambooLib="RoccoR", headers="RochesterCorrectionCalculator.h")
getattr(gbl, "RochesterCorrectionCalculator::result_t") # trigger dictionary generation
@once
def loadSVfitCalculator(backend=None, pathToSVfit=""):
loadBambooExtensions()
loader = backend.addDependency if backend else loadDependency
import os
incPath, dynPath = None, None
if "VIRTUAL_ENV" in os.environ and pathToSVfit == "":
incPath = os.path.join(os.environ["VIRTUAL_ENV"], "../ClassicSVfit/interface")
dynPath = os.path.join(os.environ["VIRTUAL_ENV"], "../build-ClassicSVfit-FastMTT/src")
elif pathToSVfit != "":
incPath = os.path.join(pathToSVfit, "../ClassicSVfit/interface")
dynPath = os.path.join(pathToSVfit, "../build-ClassicSVfit-FastMTT/src")
else:
raise RuntimeError(
"Could not load SVfit."
)
loader(headers="SVfitBambooProducer.h",
includePath=incPath, dynamicPath=dynPath,
libraries="ClassicSVfit")
@once
def loadLibTorch(backend=None):
if loadTensorflowC.has_run:
raise RuntimeError("If loading both Tensorflow-C and libtorch, the latter should be loaded first")
loadBambooExtensions()
torchLib = findLibrary("libtorch")
dynPath = None
if not torchLib:
with resources.as_file(resources.files("torch") / "lib") as lib_dir:
lib_dir = str(lib_dir)
import os.path
if os.path.isdir(lib_dir):
dynPath = lib_dir
loader = backend.addDependency if backend else loadDependency
loader(bambooLib="BambooTorch", headers="bambootorch.h", dynamicPath=dynPath)
def _addVirtualEnvPaths():
import os
incPath, dynPath = None, None
if "VIRTUAL_ENV" in os.environ:
incPath = addIncludePath(os.path.join(os.environ["VIRTUAL_ENV"], "include"))
dynPath = addDynamicPath(os.path.join(os.environ["VIRTUAL_ENV"], "lib"))
return incPath, dynPath
@once
def loadTensorflowC(backend=None):
loadBambooExtensions()
tfLib = findLibrary("libtensorflow")
incPath, dynPath = None, None
if not tfLib:
incPath, dynPath = _addVirtualEnvPaths()
loader = backend.addDependency if backend else loadDependency
loader(bambooLib="BambooTensorflowC", headers="bambootensorflowc.h",
includePath=incPath, dynamicPath=dynPath, libraries="tensorflow")
@once
def loadlwtnn(backend=None):
loadBambooExtensions()
lwtnnLib = findLibrary("liblwtnn")
incPath, dynPath = None, None
if not lwtnnLib:
incPath, dynPath = _addVirtualEnvPaths()
loader = backend.addDependency if backend else loadDependency
loader(bambooLib="BambooLwtnn", headers="bamboolwtnn.h",
includePath=incPath, dynamicPath=dynPath, libraries="lwtnn")
@once
def loadONNXRuntime(backend=None):
loadBambooExtensions()
onnxRuntimeLib = findLibrary("libonnxruntime")
incPath, dynPath = None, None
if not onnxRuntimeLib:
incPath, dynPath = _addVirtualEnvPaths()
onnxRuntimeLib = findLibrary("libonnxruntime")
import os.path
incDir = os.path.join(os.path.dirname(os.path.dirname(onnxRuntimeLib)), "include")
if os.path.isdir(os.path.join(incDir, "onnxruntime")):
incDir = os.path.join(incDir, "onnxruntime")
onnx_cand_paths = [
os.path.join(incDir, "core", "session", "onnxruntime_cxx_api.h"),
os.path.join(incDir, "onnxruntime_cxx_api.h")
]
for path in onnx_cand_paths:
if os.path.isfile(path):
onnx_cxx_hdr = path
break
else:
msg = "\n".join(onnx_cand_paths)
raise RuntimeError(
f"Could not find onnxruntime header. Tried:\n{msg}"
)
loader = backend.addDependency if backend else loadDependency
loader(bambooLib="BambooONNXRuntime", headers="bambooonnxruntime.h",
includePath=os.path.dirname(onnx_cxx_hdr),
dynamicPath=dynPath, libraries="onnxruntime")
@once
def loadBambooSOFIE(backend=None):
loadBambooExtensions()
loader = backend.addDependency if backend else loadDependency
loader(headers="bamboosofie.h")
@once
def loadBambooRBDT(backend=None):
loadBambooExtensions()
loader = backend.addDependency if backend else loadDependency
loader(headers="bamboorbdt.h")
@once
def loadcorrectionlib(backend=None):
loader = backend.addDependency if backend else loadDependency
with (
resources.as_file(resources.files("correctionlib") / "include") as incDir,
resources.as_file(resources.files("correctionlib") / "lib") as libDir
):
incDir = str(incDir)
libDir = str(libDir)
loader(headers="correction.h",
includePath=incDir, dynamicPath=libDir,
libraries="correctionlib"
)