.github/workflows/python-publish.yml

@@ -26,11 +26,22 @@ jobs:
       uses: actions/setup-python@v4
       with:
         python-version: '3.x'
-    - name: Install dependencies
+    - name: Install poetry
+      uses: abatilo/actions-poetry@v4
+      with:
+        poetry-version: '2.0.1'
+    - name: Setup a local virtual environment
       run: |
-        python -m pip install --upgrade pip
-        pip install setuptools build wheel twine
+        poetry config virtualenvs.create true --local
+        poetry config virtualenvs.in-project true --local
+    - name: Define a cache for the virtual environment based on the dependencies lock file
+      uses: actions/cache@v3
+      with:
+        path: ./.venv
+        key: venv-${{ hashFiles('poetry.lock') }}
+    - name: Install the project dependencies
+      run: poetry install --only main
     - name: Build package
-      run: python -m build
+      run: poetry build
     - name: Publish package
-      run: twine upload -u __token__ -p ${{ secrets.PYPI_API_TOKEN }} dist/*
+      run: poetry publish -u __token__ -p ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/tests.yml

@@ -21,18 +21,37 @@ permissions:
 
 jobs:
   test:
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+        poetry-version: ["2.0.1"]
+        os: [ubuntu-latest]
 
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
 
     steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.x'
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install -e .[tests]
-    - name: Test with pytest
-      run: pytest tests --doctest-modules --junitxml=junit/test-results.xml --cov=boax --cov-report=xml --cov-report=html
+      - uses: actions/checkout@v4
+      - name: Setup python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install poetry
+        uses: abatilo/actions-poetry@v4
+        with:
+          poetry-version: ${{ matrix.poetry-version }}
+      - name: Setup a local virtual environment
+        run: |
+          poetry config virtualenvs.create true --local
+          poetry config virtualenvs.in-project true --local
+      - name: Define a cache for the virtual environment based on the dependencies lock file
+        uses: actions/cache@v3
+        with:
+          path: ./.venv
+          key: venv-${{ hashFiles('poetry.lock') }}
+      - name: Install the project dependencies
+        run: poetry install --with dev,tests
+      - name: Static code analysis with ruff
+        run: poetry run ruff check boax tests
+      - name: Test with pytest
+        run: poetry run pytest tests --doctest-modules --junitxml=junit/test-results.xml --cov=boax --cov-report=xml --cov-report=html

.readthedocs.yml

@@ -4,17 +4,13 @@ build:
   os: "ubuntu-22.04"
   tools:
     python: "3.11"
-
-python:
-  install:
-    # Equivalent to 'pip install .'
-    - method: pip
-      path: .
-    # Equivalent to 'pip install .[docs]'
-    - method: pip
-      path: .
-      extra_requirements:
-        - docs
+  jobs:
+    install:
+      - pip install poetry
+      - poetry install --with docs
+    build:
+      html:
+        - poetry run python -m sphinx -T -b html -d _build/doctrees -D language=en docs/source $READTHEDOCS_OUTPUT/html
 
 sphinx:
   configuration: docs/source/conf.py

README.md

@@ -12,17 +12,28 @@
 
 ## Overview
 
-Boax is a composable library of core components for Bayesian Optimization that is **designed for flexibility**. It comes with low-level interfaces for:
+Boax is a composable library of core components for Bayesian Optimization
+that is **designed for flexibility**.
 
+It comes with high-level interfaces for:
+* **Experiments** (`boax.experiments`):
+  * Bayesian Optimization Setups
+  * Bandit Optimization Setups
+  * Search Spaces
+
+And with low-level interfaces for:
+* **Constructing acquisition functions** (`boax.acquisition`):
+  * Acquisition Functions
+  * Surrogate Models
+* **Constructing policy functions** (`boax.policies`):
+  * Policy Functions
+  * Believes
 * **Core capabilities** (`boax.core`):
   * Common Distributions
-  * Monte-Carlo Samplers
-* **Fitting a surrogate model to data** (`boax.prediction`):
-  * Model Functions
+  * Gaussian Process Models
   * Objective Functions
-* **Constructing and optimizing acquisition functions** (`boax.optimization`):
-  * Acquisition Functions
-  * Optimizer Functions
+  * Quasi-Newton Optimizers
+  * Monte-Carlo Samplers
 
 ## Installation
 
@@ -40,63 +51,76 @@ pip install git+https://github.com/Lando-L/boax.git
 
 ## Basic Usage
 
-Here is a basic example of using the Boax API for defining a Gaussian Process model, constructing an Acquisition function, and generating the next batch of data points to query. For more details check out the [docs](https://boax.readthedocs.io/en/latest/).
+Here is a basic example of using the Boax for hyperparamter tuning.
+For more details check out the [docs](https://boax.readthedocs.io/en/latest/).
 
-1. Defining a Gaussian Process model:
+1. Setting up classification task:
 
 ```python
-from boax.prediction import models
-
-model = models.gaussian_process.exact(
-  models.means.zero(),
-  models.kernels.scaled(
-    models.kernels.rbf(1.0), 0.5
-  ),
-  models.likelihoods.gaussian(1e-4),
-  x_train,
-  y_train,
-)
+  from sklearn.model_selection import train_test_split
+  from sklearn.preprocessing import StandardScaler
+  from sklearn.svm import SVC
+
+  iris = load_iris()
+  X = iris.data
+  y = iris.target
+
+  X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)
+  scaler = StandardScaler()
+  X_train = scaler.fit_transform(X_train)
+  X_test = scaler.transform(X_test)
+
+  def evaluate(C, gamma):
+    svc = SVC(C=C, gamma=gamma, kernel='rbf')
+    svc.fit(X_train, y_train)
+    return svc.score(X_test, y_test)
 ```
 
-2. Constructing an Acquisition function.
+2. Setting up a bayesian optimization experiment.
 
 ```python
-from jax import vmap
-from boax.optimization import acquisitions
-
-acqf = models.outcome_transformed(
-  vmap(model),
-  acquisitions.upper_confidence_bound(2.0)
-)
+  from jax import config
+  config.update("jax_enable_x64", True)
+  from boax.experiments import optimization
+
+  experiment = optimization(
+    parameters=[
+      {
+        'name': 'C',
+        'type': 'log_range',
+        'bounds': [1, 1_000],
+      },
+      {
+        'name': 'gamma',
+        'type': 'log_range',
+        'bounds': [1e-4, 1e-3],
+      },
+    ],
+    batch_size=4,
+  )
 ```
 
-3. Generating the next batch of data points to query.
+3. Running the trial for N = 25 steps.
 
 ```python
-from jax import numpy as jnp
-from jax import random
-from boax.core import distributions, samplers
-from boax.optimization import optimizers
-
-key = random.key(0)
-
-batch_size, num_results, num_restarts = 1, 100, 10
-bounds = jnp.array([[-1.0, 1.0]])
-
-sampler = samplers.halton_uniform(
-  distributions.uniform.uniform(bounds[:, 0], bounds[:, 1])
-)
-
-optimizer = optimizers.batch(
-  optimizers.initializers.q_batch(
-    acqf, sampler, batch_size, num_results, num_restarts,
-  ),
-  optimizers.solvers.scipy(
-    acqf, bounds,  
-  ),
-)
-
-next_x, value = optimizer(key)
+  step, results = None, []
+
+  for _ in range(25):
+    # Retrieve next parameterizations to evaluate
+    step, parameterizations = experiment.next(step, results)
+
+    # Evaluate parameterizations
+    evaluations = [
+      evaluate(**parameterization)
+      for parameterization in parameterizations
+    ]
+
+    results = list(
+        zip(parameterizations, evaluations)
+    )
+
+  # Predicted best
+  experiment.best(step)
 ```
 
 ## Citing Boax
@@ -108,10 +132,9 @@ To cite Boax please use the citation:
   author = {Lando L{\"o}per},
   title = {{B}oax: A Bayesian Optimization library for {JAX}},
   url = {https://github.com/Lando-L/boax},
-  version = {0.1.2},
+  version = {0.2.0},
   year = {2023},
 }
 ```
 
-In the above bibtex entry, the version number
-is intended to be that from [boax/version.py](https://github.com/Lando-L/boax/blob/main/boax/version.py), and the year corresponds to the project's open-source release.
+In the above bibtex entry, the version number is intended to be the latest, and the year corresponds to the project's open-source release.

boax/__init__.py

@@ -14,7 +14,8 @@
 
 """The boax API."""
 
+from . import acquisitions as acquisitions
 from . import core as core
-from . import optimization as optimization
-from . import prediction as prediction
+from . import experiments as experiments
+from . import policies as policies
 from . import utils as utils

boax/optimization/acquisitions/__init__.py → boax/acquisitions/__init__.py

@@ -14,23 +14,16 @@
 
 """The acquisitions sub-package."""
 
-from . import constraints as constraints
+from . import surrogates as surrogates
 from .alias import expected_improvement as expected_improvement
 from .alias import log_expected_improvement as log_expected_improvement
 from .alias import (
   log_probability_of_improvement as log_probability_of_improvement,
 )
-from .alias import posterior_mean as posterior_mean
-from .alias import posterior_scale as posterior_scale
 from .alias import probability_of_improvement as probability_of_improvement
 from .alias import q_expected_improvement as q_expected_improvement
-from .alias import q_knowledge_gradient as q_knowledge_gradient
-from .alias import (
-  q_multi_fidelity_knowledge_gradient as q_multi_fidelity_knowledge_gradient,
-)
+from .alias import q_log_expected_improvement as q_log_expected_improvement
 from .alias import q_probability_of_improvement as q_probability_of_improvement
 from .alias import q_upper_confidence_bound as q_upper_confidence_bound
 from .alias import upper_confidence_bound as upper_confidence_bound
 from .base import Acquisition as Acquisition
-from .transformed import constrained as constrained
-from .transformed import log_constrained as log_constrained