Boax is a Bayesian Optimization library for JAX.
Project description
Boax: A Bayesian Optimization library for JAX.
Overview | Installation | Getting Started | Documentation
Boax is currently in early alpha and under active development!
Overview
Boax is a composable library of core components for Bayesian Optimization that is designed for flexibility. It comes with low-level interfaces for:
- Core capabilities (
boax.core
):- Common Distributions
- Monte-Carlo Samplers
- Fitting a surrogate model to data (
boax.prediction
):- Kernels Functions
- Likelihood Functions
- Mean Functions
- Model Functions
- Objective Functions
- Constructing and optimizing acquisition functions (
boax.optimization
):- Acquisition Functions
- Constraint Functions
- Optimizer Functions
Installation
You can install the latest released version of Boax from PyPI via:
pip install boax
or you can install the latest development version from GitHub:
pip install git+https://github.com/Lando-L/boax.git
Basic Usage
Here is a basic example of using the Boax API for Bayesian Optimization. For more details check out the docs.
- Creation of a prediction model.
def model(params):
return models.outcome_transformed(
models.gaussian_process_regression(
means.zero(),
kernels.rbf(params['length_scale']),
)(
x_train,
y_train,
),
likelihoods.gaussian(params['noise']),
)
- Construction of an acquisition function.
def construct(model):
return optimization.construct(
models.outcome_transformed(
model,
distributions.multivariate_normal.as_normal,
),
acquisitions.upper_confidence_bound(
beta=2.0
),
)
- Selection of the next candidate to query.
def select(key, acqf):
bfgs = optimizers.bfgs(acqf, bounds, x0, 10)
candidates = bfgs.init(key)
next_candidates, values = bfgs.update(candidates)
next_x = next_candidates[jnp.argmax(values)]
next_y = objective(next_x)
return next_x, next_y
Citing Boax
To cite Boax please use the citation:
@software{boax2023github,
author = {Lando L{\"o}per},
title = {{B}oax: A Bayesian Optimization library for {JAX}},
url = {https://github.com/Lando-L/boax},
version = {0.1.0},
year = {2023},
}
In the above bibtex entry, the version number is intended to be that from boax/version.py, and the year corresponds to the project's open-source release.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.