Source code for spey_pyhf.manager
"""Manager for pyhf integration"""
import importlib
import logging
from typing import Callable, List, Text
from spey.system.exceptions import MethodNotAvailable
[docs]class PyhfManager:
"""Manager to unify the usage of pyhf through out the package"""
pyhf = importlib.import_module("pyhf")
[docs] def __init__(self):
PyhfManager.pyhf.set_backend("numpy", precision="64b")
PyhfManager.pyhf.pdf.log.setLevel(logging.CRITICAL)
PyhfManager.pyhf.workspace.log.setLevel(logging.CRITICAL)
self.shim = importlib.import_module("pyhf.optimize.common").shim
def __repr__(self) -> Text:
return (
f"pyhfManager(pyhf_backend='{self.backend}',"
+ f" available_backends={self.available_backends})"
)
@property
def grad_available(self) -> bool:
"""
Returns true or false depending on the availablility
of the gradient within the current backend.
"""
return self.backend != "numpy"
@property
def available_backends(self) -> List[Text]:
"""
Retreive the names of available backends
Returns:
``List[Text]``:
list of available backends
"""
return [
bd for bd in ["numpy", "tensorflow", "jax"] if importlib.util.find_spec(bd)
]
@property
def backend(self) -> Text:
"""Retreive current backend name"""
return PyhfManager.pyhf.tensorlib.name
@backend.setter
def backend(self, backend: Text) -> None:
"""
Modify pyhf backend.
Args:
backend (`Text`): backend type. see pyhf for details
"""
if backend not in self.available_backends:
raise MethodNotAvailable(f"{backend} backend currently not available.")
PyhfManager.pyhf.set_backend(backend)
@property
def backend_accessor(self):
"""access to current pyhf backend"""
return importlib.import_module(self.backend)
def jit(self, function: Callable) -> Callable:
"""Jit the given function if available"""
if self.backend == "jax":
return self.backend_accessor.jit(function)
return function