Contextual bo generator

docs/examples/single_objective_bayes_opt/contextual_bo.ipynb

@@ -0,0 +1,380 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "26334b35",
+   "metadata": {},
+   "source": [
+    "# Contextual Bayesian Optimization\n",
+    "This example compares standard Bayesian optimization against contextual Bayesian optimization on a drifting objective.\n",
+    "\n",
+    "The optimization variable is `x1`, while `context` is observed after each evaluation but cannot be controlled. In the contextual setup, define `context` as a `ContextualVariable` in `VOCS`.\n",
+    "\n",
+    "> `ContextualVariable()` defaults to an unbounded domain `(-inf, inf)`. When finite bounds are needed (for example, model input bounds or visualization meshes), Xopt infers them from current data. If you provide an explicit finite domain, that explicit domain takes precedence over inferred bounds.\n",
+    "\n",
+    "The Bayesian generator conditions on the latest observed context value when optimizing the acquisition function."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d711b0e6",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-01T19:58:42.324890Z",
+     "iopub.status.busy": "2026-06-01T19:58:42.324706Z",
+     "iopub.status.idle": "2026-06-01T19:58:45.435049Z",
+     "shell.execute_reply": "2026-06-01T19:58:45.434387Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import warnings\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from xopt import Evaluator\n",
+    "from xopt.vocs import VOCS, ContextualVariable\n",
+    "\n",
+    "warnings.filterwarnings(\n",
+    "    \"ignore\", message=\"Data \\(outcome observations\\) is not standardized.*\"\n",
+    ")\n",
+    "\n",
+    "SMOKE_TEST = os.environ.get(\"SMOKE_TEST\", \"\").lower() in [\"1\", \"true\", \"yes\"]\n",
+    "N_INITIAL = 3 if SMOKE_TEST else 8\n",
+    "N_STEPS = 8 if SMOKE_TEST else 50"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3f117ad0",
+   "metadata": {},
+   "source": [
+    "## Drifting Objective\n",
+    "We construct a simple non-stationary objective where the optimum moves over time. The evaluator reports both the objective value `y` and the observed context value `context` at each iteration."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f1416bd3",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-01T19:58:45.437428Z",
+     "iopub.status.busy": "2026-06-01T19:58:45.437117Z",
+     "iopub.status.idle": "2026-06-01T19:58:45.441415Z",
+     "shell.execute_reply": "2026-06-01T19:58:45.440827Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "class DriftingObjective:\n",
+    "    def __init__(self, period: int = 20):\n",
+    "        self.period = period\n",
+    "        self.iteration = 0\n",
+    "\n",
+    "    def __call__(self, inputs):\n",
+    "        x = float(inputs[\"x1\"])\n",
+    "        context = 0.5 * (1.0 + np.sin(2.0 * np.pi * self.iteration / self.period))\n",
+    "\n",
+    "        # Peak value is 1.0 when x1 equals the current context.\n",
+    "        y = -((x - context) ** 2)\n",
+    "\n",
+    "        self.iteration += 1\n",
+    "        return {\n",
+    "            \"y\": float(y),\n",
+    "            \"context\": float(context),\n",
+    "            \"tracking_error\": float(abs(x - context)),\n",
+    "        }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f74d152b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-01T19:58:45.442900Z",
+     "iopub.status.busy": "2026-06-01T19:58:45.442733Z",
+     "iopub.status.idle": "2026-06-01T19:58:45.447191Z",
+     "shell.execute_reply": "2026-06-01T19:58:45.446612Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from xopt.base import Xopt\n",
+    "from xopt.generators.bayesian.upper_confidence_bound import (\n",
+    "    UpperConfidenceBoundGenerator,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def run_optimization(use_context: bool):\n",
+    "    objective = DriftingObjective(25)\n",
+    "    evaluator = Evaluator(function=objective)\n",
+    "\n",
+    "    if use_context:\n",
+    "        vocs = VOCS(\n",
+    "            # Default ContextualVariable() uses unbounded domain and infers finite\n",
+    "            # bounds from observed data when needed. To force fixed bounds, use:\n",
+    "            # ContextualVariable(domain=[low, high]).\n",
+    "            variables={\"x1\": [0.0, 1.0], \"context\": ContextualVariable()},\n",
+    "            objectives={\"y\": \"MAXIMIZE\"},\n",
+    "            observables=[\"context\", \"tracking_error\"],\n",
+    "        )\n",
+    "    else:\n",
+    "        vocs = VOCS(\n",
+    "            variables={\"x1\": [0.0, 1.0]},\n",
+    "            objectives={\"y\": \"MAXIMIZE\"},\n",
+    "            observables=[\"context\", \"tracking_error\"],\n",
+    "        )\n",
+    "\n",
+    "    generator = UpperConfidenceBoundGenerator(vocs=vocs)\n",
+    "    generator.gp_constructor.use_low_noise_prior = True\n",
+    "\n",
+    "    X = Xopt(generator=generator, evaluator=evaluator)\n",
+    "\n",
+    "    # Build initial data manually since contextual variables are observed, not sampled.\n",
+    "    initial_inputs = pd.DataFrame({\"x1\": np.linspace(0.05, 0.95, N_INITIAL)})\n",
+    "    X.evaluate_data(initial_inputs)\n",
+    "\n",
+    "    for _ in range(N_STEPS):\n",
+    "        X.step()\n",
+    "\n",
+    "    data = X.data\n",
+    "    data[\"iteration\"] = np.arange(len(data))\n",
+    "\n",
+    "    return X.generator, data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "149e9351",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-01T19:58:45.448624Z",
+     "iopub.status.busy": "2026-06-01T19:58:45.448461Z",
+     "iopub.status.idle": "2026-06-01T19:58:54.598399Z",
+     "shell.execute_reply": "2026-06-01T19:58:54.597607Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "np.random.seed(1)\n",
+    "\n",
+    "plain_generator, plain_data = run_optimization(use_context=False)\n",
+    "contextual_generator, contextual_data = run_optimization(use_context=True)\n",
+    "\n",
+    "print(f\"Standard BO evaluations: {len(plain_data)}\")\n",
+    "print(f\"Contextual BO evaluations: {len(contextual_data)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f6438da2",
+   "metadata": {},
+   "source": [
+    "## Compare Standard BO and Contextual BO\n",
+    "In the first panel, we compare selected `x1` values against the drifting context. In the second panel, we compare rolling objective values. In the third panel, we compare rolling tracking error (lower is better)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "69ab860b",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-01T19:58:54.600257Z",
+     "iopub.status.busy": "2026-06-01T19:58:54.600011Z",
+     "iopub.status.idle": "2026-06-01T19:58:55.021526Z",
+     "shell.execute_reply": "2026-06-01T19:58:55.020720Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "window = 1 if not SMOKE_TEST else 3\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 3, figsize=(16, 4))\n",
+    "\n",
+    "axes[0].plot(plain_data[\"iteration\"], plain_data[\"x1\"], label=\"standard BO\", lw=1.8)\n",
+    "axes[0].plot(\n",
+    "    contextual_data[\"iteration\"],\n",
+    "    contextual_data[\"x1\"],\n",
+    "    label=\"contextual BO\",\n",
+    "    lw=1.8,\n",
+    ")\n",
+    "axes[0].plot(\n",
+    "    plain_data[\"iteration\"],\n",
+    "    plain_data[\"context\"],\n",
+    "    \"k--\",\n",
+    "    label=\"observed context\",\n",
+    "    alpha=0.8,\n",
+    ")\n",
+    "axes[0].set_title(\"Chosen x1 vs Context\")\n",
+    "axes[0].set_xlabel(\"iteration\")\n",
+    "axes[0].set_ylabel(\"value\")\n",
+    "axes[0].legend(loc=\"best\")\n",
+    "\n",
+    "axes[1].plot(\n",
+    "    plain_data[\"iteration\"],\n",
+    "    plain_data[\"y\"].rolling(window, min_periods=1).mean(),\n",
+    "    label=\"standard BO\",\n",
+    ")\n",
+    "axes[1].plot(\n",
+    "    contextual_data[\"iteration\"],\n",
+    "    contextual_data[\"y\"].rolling(window, min_periods=1).mean(),\n",
+    "    label=\"contextual BO\",\n",
+    ")\n",
+    "axes[1].set_title(f\"Rolling mean objective (window={window})\")\n",
+    "axes[1].set_xlabel(\"iteration\")\n",
+    "axes[1].set_ylabel(\"y\")\n",
+    "axes[1].legend(loc=\"best\")\n",
+    "\n",
+    "axes[2].plot(\n",
+    "    plain_data[\"iteration\"],\n",
+    "    plain_data[\"tracking_error\"].rolling(window, min_periods=1).mean(),\n",
+    "    label=\"standard BO\",\n",
+    ")\n",
+    "axes[2].plot(\n",
+    "    contextual_data[\"iteration\"],\n",
+    "    contextual_data[\"tracking_error\"].rolling(window, min_periods=1).mean(),\n",
+    "    label=\"contextual BO\",\n",
+    ")\n",
+    "axes[2].set_title(f\"Rolling mean tracking error (window={window})\")\n",
+    "axes[2].set_xlabel(\"iteration\")\n",
+    "axes[2].set_ylabel(\"|x1 - context|\")\n",
+    "axes[2].legend(loc=\"best\")\n",
+    "\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84bb2d01",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-01T19:58:55.023303Z",
+     "iopub.status.busy": "2026-06-01T19:58:55.023130Z",
+     "iopub.status.idle": "2026-06-01T19:58:55.031514Z",
+     "shell.execute_reply": "2026-06-01T19:58:55.030886Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "comparison = pd.DataFrame(\n",
+    "    {\n",
+    "        \"mean_y\": [plain_data[\"y\"].mean(), contextual_data[\"y\"].mean()],\n",
+    "        \"mean_tracking_error\": [\n",
+    "            plain_data[\"tracking_error\"].mean(),\n",
+    "            contextual_data[\"tracking_error\"].mean(),\n",
+    "        ],\n",
+    "    },\n",
+    "    index=[\"standard_bo\", \"contextual_bo\"],\n",
+    ")\n",
+    "\n",
+    "comparison"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "df537bc9",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-01T19:58:55.033064Z",
+     "iopub.status.busy": "2026-06-01T19:58:55.032880Z",
+     "iopub.status.idle": "2026-06-01T19:58:55.036313Z",
+     "shell.execute_reply": "2026-06-01T19:58:55.035675Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "if (\n",
+    "    comparison.loc[\"contextual_bo\", \"mean_tracking_error\"]\n",
+    "    < comparison.loc[\"standard_bo\", \"mean_tracking_error\"]\n",
+    "):\n",
+    "    print(\"Contextual BO tracks the drifting optimum better in this run.\")\n",
+    "else:\n",
+    "    print(\n",
+    "        \"In this run, contextual BO did not improve tracking error. Try changing the period or number of steps.\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7bc449af",
+   "metadata": {},
+   "source": [
+    "## Visualizing contextual models and acquisition functions\n",
+    "These plots inspect the surrogate models learned by the contextual and standard generators. For the contextual generator, the model depends on both x1 and the observed context value, so visualization can show either the full model or a slice over selected variables.\n",
+    "\n",
+    "Note that we cannot visualize the acquisition function in the 2D plots shown below because it is conditioned on the context value, however we can visualize it in 1D because the context value is given by the reference point."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e346ef02",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "contextual_generator.visualize_model()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a878756e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "contextual_generator.visualize_model(\n",
+    "    variable_names=[\"x1\"],\n",
+    "    show_samples=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c6a47d8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plain_generator.visualize_model()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e6942e0d",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "xopt-dev",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

mkdocs.yml

@@ -42,6 +42,7 @@ nav:
                 - Noisy Bayesian optimization: examples/single_objective_bayes_opt/noisy_bo_tutorial.ipynb
                 - Fast model evaluation: examples/single_objective_bayes_opt/fast_model_eval.ipynb
                 - Hessian kernel: examples/single_objective_bayes_opt/hessian_kernel.ipynb
+                - Contextual Bayesian optimization: examples/single_objective_bayes_opt/contextual_bo.ipynb
             - Trust Region Bayesian Optimization:
                 - Basic example: examples/trust_region_bo/turbo_basics.ipynb
                 - Optimization: examples/trust_region_bo/turbo_optimize.ipynb

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
   "tqdm",
   "orjson",
   "matplotlib",
-  "gest-api"
+  "gest-api>=0.2"
 ]
 description = "Flexible optimization of arbitrary problems in Python."
 dynamic = [ "version" ]