From a2c3c35d971d16c8135c6fc8221d99b654af875f Mon Sep 17 00:00:00 2001 From: Gianluca Martino Date: Fri, 24 Jul 2026 16:27:20 -0700 Subject: [PATCH 1/3] Add BAxUS generator for high-dimensional Bayesian optimization --- README.md | 1 + docs/algorithms.md | 18 +- docs/api/generators/bayesian/baxus.md | 2 + .../single_objective_bayes_opt/baxus.ipynb | 223 +++ docs/index.md | 1 + mkdocs.yml | 2 + xopt/generators/__init__.py | 3 + xopt/generators/bayesian/__init__.py | 2 + xopt/generators/bayesian/baxus.py | 769 ++++++++ xopt/tests/generators/bayesian/test_baxus.py | 1743 +++++++++++++++++ 10 files changed, 2763 insertions(+), 1 deletion(-) create mode 100644 docs/api/generators/bayesian/baxus.md create mode 100644 docs/examples/single_objective_bayes_opt/baxus.ipynb create mode 100644 xopt/generators/bayesian/baxus.py create mode 100644 xopt/tests/generators/bayesian/test_baxus.py diff --git a/README.md b/README.md index 6dc0d7805..a9d036033 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ Currenty **Xopt** provides: optimization - `multi_fidelity` Multi-fidelity single or multi objective optimization - `BAX` Bayesian algorithm execution using virtual measurements + - `baxus` High-dimensional BO in adaptively expanding random subspaces - BO customization: - Trust region BO - Heteroskedastic noise specification diff --git a/docs/algorithms.md b/docs/algorithms.md index dccf210bc..ba650aee6 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -15,11 +15,27 @@ Bayesian Generators All of the generators here use Bayesian optimization (BO) type methods to solve single objective, multi objective and characterization problems. Bayesian generators incorperate unknown constrianing functions into optimization based on what is -specified in `VOCS` +specified in `VOCS` (with the exception of `BAxUSGenerator`, see below). - [`ExpectedImprovementGenerator`](examples/single_objective_bayes_opt/constrained_bo_tutorial.ipynb): implements Expected Improvement single objective BO. Automatically balances trade-offs between exploration and exploitation and is thus useful for general purpose optimization. +- [`BAxUSGenerator`](examples/single_objective_bayes_opt/baxus.ipynb): implements + Bayesian Optimization with Adaptively Expanding Subspaces (BAxUS, + Papenmeier et al., NeurIPS 2022) for high-dimensional single objective + problems. Optimization happens inside a low-dimensional random embedding of + the input space that expands adaptively when a trust region collapses, + making it effective when many input dimensions are irrelevant. + Because the model lives in the embedded subspace rather than in `VOCS` space, + this generator is more restricted than the others: it supports a single + objective over continuous variables only, and rejects constraints, + observables, discrete and contextual variables, `fixed_features`, + `max_travel_distances`, `custom_objective` and `n_interpolate_points`. It + generates one candidate at a time and does not support `visualize_model`. + Set `eval_budget` to the total number of planned evaluations - the reference + expansion schedule is derived from it, and without it a slower fallback + heuristic is used. Note also that the acquisition function is analytic LogEI + over the trust region rather than the Thompson sampling used in the paper. - [`UpperConfidenceBoundGenerator`](examples/single_objective_bayes_opt/upper_confidence_bound.ipynb): implements Upper Confidence Bound single objective BO. Requires a hyperparameter `beta` that explicitly sets the tradeoff between exploration and exploitation. Default value of `beta=2` is a good diff --git a/docs/api/generators/bayesian/baxus.md b/docs/api/generators/bayesian/baxus.md new file mode 100644 index 000000000..f30d6da4d --- /dev/null +++ b/docs/api/generators/bayesian/baxus.md @@ -0,0 +1,2 @@ +# BAxUS Generator +::: xopt.generators.bayesian.baxus diff --git a/docs/examples/single_objective_bayes_opt/baxus.ipynb b/docs/examples/single_objective_bayes_opt/baxus.ipynb new file mode 100644 index 000000000..50cf9da57 --- /dev/null +++ b/docs/examples/single_objective_bayes_opt/baxus.ipynb @@ -0,0 +1,223 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0776d0a0", + "metadata": {}, + "source": [ + "# High-dimensional Bayesian optimization with BAxUS\n", + "\n", + "BAxUS (Bayesian Optimization with Adaptively Expanding Subspaces,\n", + "[Papenmeier et al., NeurIPS 2022](https://arxiv.org/abs/2304.11468)) optimizes\n", + "high-dimensional problems by working in a low-dimensional random embedding of\n", + "the input space. A trust region controls local search in the embedded space;\n", + "when it collapses without progress, the embedding *expands*, adding dimensions\n", + "until (in the limit) it recovers the full input space. This makes BAxUS\n", + "effective when many input dimensions are irrelevant.\n", + "\n", + "This example optimizes a 100-dimensional problem in which only 3 dimensions\n", + "actually affect the objective.\n", + "\n", + "**Scope.** Because the model lives in the embedded subspace rather than in\n", + "`VOCS` space, this generator is more restricted than the other Bayesian\n", + "generators. It handles a single objective over continuous variables and rejects\n", + "constraints, observables, discrete and contextual variables, `fixed_features`,\n", + "`max_travel_distances`, `custom_objective` and `n_interpolate_points`. It\n", + "generates one candidate at a time, and `visualize_model` is not available.\n", + "\n", + "**Differences from the paper.** The acquisition function here is the analytic\n", + "LogEI inherited from `ExpectedImprovementGenerator`, evaluated over the trust\n", + "region, whereas the paper uses Thompson sampling over a masked discrete\n", + "candidate set. The initial subspace dimensionality is the user-set\n", + "`target_dim_init` rather than being derived from the input dimensionality.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e822fcb3", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "from xopt import Xopt, Evaluator\n", + "from xopt.vocs import VOCS\n", + "from xopt.generators.bayesian import BAxUSGenerator\n", + "\n", + "N_DIM = 100\n", + "N_EFFECTIVE = 3\n", + "N_STEPS = 60\n", + "\n", + "\n", + "def sphere_embedded(inputs: dict) -> dict:\n", + " \"\"\"Sphere on the first 3 coordinates; the other 97 are irrelevant.\"\"\"\n", + " x = np.array([inputs[f\"x{i}\"] for i in range(N_EFFECTIVE)])\n", + " return {\"f\": float(np.sum(x**2))}\n", + "\n", + "\n", + "vocs = VOCS(\n", + " variables={f\"x{i}\": [-1.0, 1.0] for i in range(N_DIM)},\n", + " objectives={\"f\": \"MINIMIZE\"},\n", + ")\n", + "vocs\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08e06987", + "metadata": {}, + "outputs": [], + "source": [ + "generator = BAxUSGenerator(\n", + " vocs=vocs,\n", + " target_dim_init=2,\n", + " seed=0,\n", + " # total planned evaluations: the expansion schedule is paced against this,\n", + " # so it should match the number of steps actually run below\n", + " eval_budget=N_STEPS,\n", + ")\n", + "evaluator = Evaluator(function=sphere_embedded)\n", + "X = Xopt(evaluator=evaluator, generator=generator)\n", + "X\n" + ] + }, + { + "cell_type": "markdown", + "id": "6837a540", + "metadata": {}, + "source": [ + "## Running the optimization\n", + "\n", + "The generator draws its own Sobol seed points in the embedded space (no\n", + "`random_evaluate` needed), then switches to trust-region LogEI. The embedding\n", + "expands automatically when the trust region collapses." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59ce976b", + "metadata": {}, + "outputs": [], + "source": [ + "target_dims = []\n", + "for _ in range(N_STEPS):\n", + " X.step()\n", + " target_dims.append(X.generator.embedding.target_dim)\n", + "\n", + "print(f\"best f: {X.data['f'].min():.3e}\")\n", + "print(f\"embedding grew: {sorted(set(target_dims))}\")\n", + "X.data[[\"f\"]].tail()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8c081907", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import pyplot as plt\n", + "\n", + "fig, (ax0, ax1) = plt.subplots(2, 1, sharex=True, figsize=(6, 6))\n", + "best = X.data[\"f\"].cummin()\n", + "ax0.plot(best.values)\n", + "ax0.set_ylabel(\"best f\")\n", + "ax0.set_yscale(\"log\")\n", + "\n", + "ax1.step(range(len(target_dims)), target_dims, where=\"post\")\n", + "ax1.set_ylabel(\"embedding target_dim\")\n", + "ax1.set_xlabel(\"evaluation\")\n", + "\n", + "# mark each embedding expansion on both panels\n", + "for i in np.flatnonzero(np.diff(target_dims)) + 1:\n", + " for ax in (ax0, ax1):\n", + " ax.axvline(i, color=\"grey\", ls=\"--\", lw=0.8)\n", + "ax0.set_title(\"BAxUS: convergence and subspace growth\")\n", + "plt.tight_layout()\n" + ] + }, + { + "cell_type": "markdown", + "id": "b39a3c12", + "metadata": {}, + "source": [ + "## Serialization and resuming\n", + "\n", + "`BAxUSGenerator` owns its own state (the embedding, the trust region, and how\n", + "much history has already been folded into it), so it is a `StateOwner`: saving\n", + "and reloading through `Xopt` reattaches the data without replaying that history.\n", + "Round-tripping a checkpoint therefore leaves the run exactly where it was.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7c0483bd", + "metadata": {}, + "outputs": [], + "source": [ + "checkpoint = X.yaml()\n", + "resumed = Xopt.from_yaml(checkpoint)\n", + "\n", + "print(\"target_dim :\", X.generator.embedding.target_dim,\n", + " \"->\", resumed.generator.embedding.target_dim)\n", + "print(\"tr length :\", X.generator.trust_region.length,\n", + " \"->\", resumed.generator.trust_region.length)\n", + "print(\"rows folded:\", X.generator.tr_observed_rows,\n", + " \"->\", resumed.generator.tr_observed_rows)\n", + "\n", + "resumed.step() # continues the run\n", + "len(resumed.data)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Dumping the generator on its own works too, but a trained generator's dump\n", + "carries the live botorch `model` and a `computation_time` frame, neither of\n", + "which is YAML-safe -- drop both before serializing by hand.\n" + ], + "id": "fbb00bc6" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import yaml\n", + "\n", + "dump = X.generator.model_dump()\n", + "dump.pop(\"model\", None)\n", + "dump.pop(\"computation_time\", None)\n", + "yaml.safe_dump(dump)[:200]\n" + ], + "id": "b36961c9" + } + ], + "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.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/index.md b/docs/index.md index 58fcb9514..25ae10588 100644 --- a/docs/index.md +++ b/docs/index.md @@ -34,6 +34,7 @@ Currenty **Xopt** provides: optimization - [`multi_fidelity`](examples/single_objective_bayes_opt/multi_fidelity_simple.ipynb) Multi-fidelity single or multi objective optimization - [`BAX`](examples/single_objective_bayes_opt/bax_tutorial.ipynb) Bayesian algorithm execution using virtual measurements + - [`baxus`](examples/single_objective_bayes_opt/baxus.ipynb) High-dimensional BO in adaptively expanding random subspaces - BO customization: - [Trust region BO](examples/trust_region_bo/turbo_basics.ipynb) - [Heteroskedastic noise specification](examples/single_objective_bayes_opt/heteroskedastic_noise_tutorial.ipynb) diff --git a/mkdocs.yml b/mkdocs.yml index 91a7af1fe..0060db8e9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -43,6 +43,7 @@ nav: - 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 + - BAxUS high-dimensional optimization: examples/single_objective_bayes_opt/baxus.ipynb - Trust Region Bayesian Optimization: - Basic example: examples/trust_region_bo/turbo_basics.ipynb - Optimization: examples/trust_region_bo/turbo_optimize.ipynb @@ -82,6 +83,7 @@ nav: - Time-Dependent Bayesian Generators: api/generators/bayesian/time_dependent.md - Bayesian Exploration: api/generators/bayesian/bayesian_exploration.md - Expected Improvement: api/generators/bayesian/expected_improvement.md + - BAxUS (high-dimensional): api/generators/bayesian/baxus.md - Multi-Objective BO: api/generators/bayesian/mobo.md - Multi-Generation Gaussian Process Optimization: api/generators/bayesian/mggpo.md - Multi-Fidelity: api/generators/bayesian/multi_fidelity.md diff --git a/xopt/generators/__init__.py b/xopt/generators/__init__.py index 776ed8d32..fd9758a5a 100644 --- a/xopt/generators/__init__.py +++ b/xopt/generators/__init__.py @@ -24,6 +24,7 @@ "expected_improvement", "multi_fidelity", "bax", + "baxus", }, "ga": {"cnsga", "nsga2"}, "es": {"extremum_seeking"}, @@ -88,6 +89,7 @@ def get_generator_dynamic(name: str) -> type[Generator]: UpperConfidenceBoundGenerator, ) from xopt.generators.bayesian.bax_generator import BaxGenerator + from xopt.generators.bayesian.baxus import BAxUSGenerator registered_generators: list[type[Generator]] = [ UpperConfidenceBoundGenerator, @@ -97,6 +99,7 @@ def get_generator_dynamic(name: str) -> type[Generator]: ExpectedImprovementGenerator, MultiFidelityGenerator, BaxGenerator, + BAxUSGenerator, ] for gen in registered_generators: generators[gen.name] = gen diff --git a/xopt/generators/bayesian/__init__.py b/xopt/generators/bayesian/__init__.py index 44a9b7bea..4c60fd58e 100644 --- a/xopt/generators/bayesian/__init__.py +++ b/xopt/generators/bayesian/__init__.py @@ -1,6 +1,7 @@ import torch from xopt.generators.bayesian.bax_generator import BaxGenerator +from xopt.generators.bayesian.baxus import BAxUSGenerator from xopt.generators.bayesian.bayesian_exploration import BayesianExplorationGenerator from xopt.generators.bayesian.bayesian_generator import BayesianGenerator from xopt.generators.bayesian.expected_improvement import ( @@ -35,6 +36,7 @@ "TDExpectedImprovementGenerator", "MGGPOGenerator", "BaxGenerator", + "BAxUSGenerator", "BayesianGenerator", "TimeDependentBayesianGenerator", "TurboController", diff --git a/xopt/generators/bayesian/baxus.py b/xopt/generators/bayesian/baxus.py new file mode 100644 index 000000000..e10a00661 --- /dev/null +++ b/xopt/generators/bayesian/baxus.py @@ -0,0 +1,769 @@ +"""BAxUS generator - Bayesian Optimization with Adaptively Expanding Subspaces. + +Starts in a low-dimensional random subspace and gradually expands it, making +it effective when many input dimensions are irrelevant. + +The inherited ``data`` frame is the single source of truth; all persistent +state lives in serializable pydantic fields, and the generator is a +``StateOwner`` so ``Xopt.from_yaml`` reattaches data without replaying the +trust-region history it already reflects. The trust region advances once per +``generate`` call (as ``TurboController`` does), which keeps the trajectory +independent of how results were batched into ``add_data``. + +Randomness is seeded rather than removed: the embedding's creation and each of +its expansions draw from ``seed`` combined with the expansion count, so a +dump -> construct round trip reproduces the same subspaces. + +The BO phase is ``BayesianGenerator.generate`` re-based into the embedded target +space through four override seams - ``get_training_data`` (finite-objective rows +only), ``train_model`` (fit on projections), ``_get_torch_bounds`` / +``_get_optimization_bounds`` (target-space boxes), and ``_process_candidates`` +(lift back to vocs space). Everything else in that pipeline is inert here, +because the options it serves are rejected at construction. + +Deviations from the reference: + - ``target_dim_init`` is a user knob rather than being derived from the + input dimensionality. + - The acquisition is the inherited analytic LogEI over the trust-region box, + not Thompson sampling over a masked discrete candidate set. + - Data not generated through the current embedding is least-squares + projected into target space. + +Reference: + Papenmeier et al., "Increasing the Scope as You Learn: Adaptive Bayesian + Optimization in Nested Subspaces", NeurIPS 2022. + https://arxiv.org/abs/2304.11468 + Tutorial: https://botorch.org/docs/tutorials/baxus +""" + +import logging +import math +import warnings +from typing import Any + +import numpy as np +import pandas as pd +import torch +from pydantic import ( + Field, + PositiveInt, + PrivateAttr, + SerializeAsAny, + ValidationInfo, + field_validator, + model_validator, +) +from torch.nn import Module +from torch.quasirandom import SobolEngine + +from xopt.errors import GeneratorWarning +from xopt.generator import StateOwner +from xopt.generators.bayesian.base_model import ModelConstructor +from xopt.generators.bayesian.bayesian_generator import formatted_base_docstring +from xopt.generators.bayesian.expected_improvement import ExpectedImprovementGenerator +from xopt.generators.bayesian.models.standard import StandardModelConstructor +from xopt.generators.bayesian.utils import set_botorch_weights +from xopt.pydantic import XoptBaseModel +from xopt.vocs import VOCS, convert_numpy_to_inputs, get_variable_bounds_array + +logger = logging.getLogger(__name__) + + +def _normalize_lengthscale_weights( + lengthscales: torch.Tensor, target_dim: int +) -> torch.Tensor: + """Scale lengthscales to mean 1, then unit geometric mean (volume-preserving). + + The geometric mean is taken as a product of per-element roots (the reference + form). Forming ``weights.prod()`` first underflows to zero once target_dim + reaches a few hundred, which would make every weight ``inf`` and silently + widen the trust region to the whole domain. + """ + weights = lengthscales / lengthscales.mean() + return weights / weights.pow(1.0 / target_dim).prod() + + +class BAxUSTrustRegion(XoptBaseModel): + """Serializable trust-region state, in normalized [0, 1] target-space units. + + ``update`` takes weighted objective values (MAXIMIZE -> +y, MINIMIZE -> -y), + so higher is always better. ``best_value=None`` means no BO-phase + evaluation has been observed yet. + """ + + target_dim: PositiveInt + length: float = 0.8 + length_init: float = 0.8 + length_min: float = 0.5**7 + length_max: float = 1.6 + failure_counter: int = 0 + success_counter: int = 0 + success_tolerance: int = 3 + failure_tolerance: PositiveInt = Field( + default=0, + validate_default=True, + description=( + "Failures before shrinking; 0 derives ceil(target_dim / 2), re-derived " + "on every BO-phase ingest (budget-aware when eval_budget is set)" + ), + ) + best_value: float | None = None + restart_triggered: bool = False + + @field_validator("failure_tolerance", mode="before") + @classmethod + def _default_failure_tolerance(cls, value: Any, info: ValidationInfo) -> Any: + """0 (the field default) means "derive from target_dim".""" + if value: + return value + # absent only if target_dim itself failed validation; PositiveInt then rejects 0 + target_dim = info.data.get("target_dim") + return math.ceil(target_dim / 2) if target_dim else value + + def update(self, y_weighted: torch.Tensor) -> None: + """Adjust the trust region from a batch of weighted objective values.""" + best = self.best_value + y_max = float(y_weighted.max()) + improved = best is None or y_max > best + 1e-3 * abs(best) + if improved: + self.success_counter += 1 + self.failure_counter = 0 + else: + self.success_counter = 0 + self.failure_counter += 1 + + if self.success_counter >= self.success_tolerance: + self.length = min(2.0 * self.length, self.length_max) + self.success_counter = 0 + elif self.failure_counter >= self.failure_tolerance: + self.length /= 2.0 + self.failure_counter = 0 + + self.best_value = y_max if best is None else max(best, y_max) + + if self.length < self.length_min: + self.restart_triggered = True + + +class BAxUSEmbedding(XoptBaseModel): + """Serializable sparse random embedding. + + ``matrix`` has shape (target_dim, input_dim) with exactly one signed unit + entry per column. Both ``create`` and ``expand`` consume randomness; each + takes an explicit seed so a run can be reproduced from a dump. + """ + + matrix: list[list[float]] + + @classmethod + def create( + cls, input_dim: int, target_dim: int, seed: int | None + ) -> "BAxUSEmbedding": + rng = np.random.default_rng(seed) + S = np.zeros((target_dim, input_dim), dtype=np.float64) + perm = rng.permutation(input_dim) + target_rows = np.arange(input_dim) % target_dim + signs = rng.choice([-1.0, 1.0], size=input_dim) + S[target_rows, perm] = signs + return cls(matrix=S.tolist()) + + @property + def target_dim(self) -> int: + return len(self.matrix) + + @property + def input_dim(self) -> int: + return len(self.matrix[0]) + + def _as_array(self) -> np.ndarray: + return np.asarray(self.matrix, dtype=np.float64) + + def lift(self, Z: np.ndarray) -> np.ndarray: + """Lift from target subspace to full input space: ``X = Z @ S``.""" + return Z @ self._as_array() + + def project(self, X: np.ndarray) -> np.ndarray: + """Project from full input space to target subspace via the pseudo-inverse.""" + return X @ np.linalg.pinv(self._as_array()) + + def expand( + self, new_bins_on_split: int, seed: int | None = None + ) -> "BAxUSEmbedding": + """Split each bin into up to ``new_bins_on_split + 1`` sub-bins, preserving signs. + + The dimensions contributing to a bin are randomly permuted before being + split, matching the reference implementation - without it the split is a + fixed function of column index, so dimensions that are adjacent in index + and share a bin are never separated before full dimensionality. Pass a + ``seed`` to keep an expansion reproducible across a dump -> construct + round trip. + """ + rng = np.random.default_rng(seed) + S = self._as_array() + old_target_dim, input_dim = S.shape + n_sub = new_bins_on_split + 1 + new_target_dim = min(old_target_dim * n_sub, input_dim) + + new_S = np.zeros((new_target_dim, input_dim), dtype=np.float64) + new_row = 0 + for row in S: + bin_cols = rng.permutation(np.nonzero(row)[0]) + for sub_cols in np.array_split(bin_cols, min(n_sub, bin_cols.size)): + new_S[new_row, sub_cols] = row[sub_cols] + new_row += 1 + return BAxUSEmbedding(matrix=new_S.tolist()) + + +class BAxUSGenerator(ExpectedImprovementGenerator, StateOwner): + """Bayesian Optimization with Adaptively Expanding Subspaces (BAxUS). + + Operates in a low-dimensional random embedding of the input space and + gradually expands it when the trust region shrinks below a threshold. + The acquisition (analytic LogEI) is inherited from + ``ExpectedImprovementGenerator``; only the model/bounds seams are re-based + into the embedded target space. + + This is a ``StateOwner``: the trust region and embedding are advanced once + per ``generate`` call, so restoring a saved run reattaches data without + replaying (and corrupting) the state that was just deserialized. + """ + + name = "baxus" + supports_batch_generation: bool = False + supports_single_objective: bool = True + # the embedded target-space GP models only the objective over continuous + # variables; constrained, discrete, observable, and contextual vocs are rejected + supports_constraints: bool = False + supports_discrete_variables: bool = False + supports_contextual_variables: bool = False + supports_no_objective: bool = False + + __doc__ = ( + "Bayesian optimization with adaptively expanding subspaces (BAxUS)\n" + + formatted_base_docstring() + ) + + # BAxUS manages its own trust region in the embedded target space + _compatible_turbo_controllers = [] + + target_dim_init: PositiveInt = Field( + default=2, + description="Initial target dimensionality of the embedding", + ) + n_initial_sobol: int = Field( + default=0, + ge=0, + description="Number of initial Sobol points (0 = auto-computed)", + ) + length_init: float = Field( + default=0.8, + description="Initial trust-region side length", + ) + new_bins_on_split: PositiveInt = Field( + default=3, + description="Number of new bins created per existing bin on expansion", + ) + seed: int | None = Field( + default=None, + description=( + "Random seed for the embedding and Sobol initialization. GP fitting and " + "acquisition optimization use torch's global RNG and are not covered by it." + ), + ) + eval_budget: PositiveInt | None = Field( + default=None, + description=( + "Total planned evaluations for the run, seed points included (the seed " + "quota is subtracted internally). Enables the reference budget-aware " + "failure tolerance; None keeps the ceil(target_dim / 2) heuristic." + ), + ) + embedding: BAxUSEmbedding = Field( + description="Sparse random embedding (auto-created from vocs and seed when omitted)", + ) + trust_region: BAxUSTrustRegion = Field( + description="Trust-region state (auto-created when omitted)", + ) + sobol_draws: int = Field( + default=0, + ge=0, + description="Number of Sobol seed points drawn so far (used to resume the sequence)", + ) + n_expansions: int = Field( + default=0, + ge=0, + description=( + "Number of embedding expansions performed so far; combined with seed to " + "keep each expansion's randomization reproducible across a round trip" + ), + ) + tr_observed_rows: int = Field( + default=0, + ge=0, + description=( + "Number of finite-objective rows already folded into the trust region " + "(prevents re-ingesting history on resume)" + ), + ) + # a plain default instance, matching BayesianGenerator: pydantic deep-copies + # it per instance, and get_generator_defaults only understands `default` + gp_constructor: SerializeAsAny[ModelConstructor] = Field( + StandardModelConstructor(use_low_noise_prior=True), + description=( + "Constructor used to generate the target-space model. Defaults to a " + "low-noise prior, matching the near-interpolating GP the reference " + "trust-region logic assumes" + ), + ) + + _sobol: SobolEngine | None = PrivateAttr(default=None) + + @model_validator(mode="before") + @classmethod + def _init_components(cls, values: Any) -> Any: + """Fill in the vocs-derived defaults on fresh construction. + + Creates the embedding and trust region and sizes the Sobol seed quota + from the embedding. A dump -> construct round trip supplies the + components and skips their creation. + """ + if not isinstance(values, dict) or values.get("vocs") is None: + return values + vocs = values["vocs"] + if isinstance(vocs, dict): + # xopt's VOCS validators pop "type" off entry dicts in place, so write + # the parsed object back before the outer "vocs" field validation + values["vocs"] = vocs = VOCS(**vocs) + input_dim = len(vocs.variable_names) + + # this runs before pydantic applies defaults, so the `or` fallbacks below + # must mirror the target_dim_init / length_init field defaults + emb = values.get("embedding") + if emb is None: + target_dim = min(int(values.get("target_dim_init") or 2), input_dim) + emb = BAxUSEmbedding.create( + input_dim=input_dim, target_dim=target_dim, seed=values.get("seed") + ) + values["embedding"] = emb + # the embedding arrives as a raw dict on a dump -> construct round trip + emb_target_dim = len(emb["matrix"]) if isinstance(emb, dict) else emb.target_dim + + if not values.get("n_initial_sobol"): + values["n_initial_sobol"] = max(2, emb_target_dim + 1) + + if values.get("trust_region") is None: + length_init = float(values.get("length_init") or 0.8) + values["trust_region"] = BAxUSTrustRegion( + target_dim=emb_target_dim, + length=length_init, + length_init=length_init, + ) + + return values + + def model_post_init(self, context: Any, /) -> None: + """Warn once that the reference expansion schedule is disabled. + + This hook runs exactly once per construction. A model validator of + either mode would instead re-fire on every field assignment under + ``validate_assignment``, repeating the warning on each trust-region + update and embedding expansion. + """ + super().model_post_init(context) + if self.eval_budget is None: + warnings.warn( + "BAxUS: eval_budget is not set, so the reference budget-aware failure " + "tolerance is replaced by the ceil(target_dim / 2) heuristic. The " + "embedding then expands far more slowly and may never reach full " + "dimensionality within a typical run - set eval_budget to the total " + "number of planned evaluations.", + GeneratorWarning, + stacklevel=2, + ) + + @model_validator(mode="after") + def _validate_budget_covers_seeding(self) -> "BAxUSGenerator": + """Reject an eval_budget smaller than the Sobol seed quota.""" + if self.eval_budget is not None and self.eval_budget < self.n_initial_sobol: + raise ValueError( + f"eval_budget ({self.eval_budget}) must cover the seed quota (n_initial_sobol={self.n_initial_sobol})" + ) + return self + + @model_validator(mode="after") + def _reject_unsupported_options(self) -> "BAxUSGenerator": + """vocs-space point controls cannot be honored in an embedded subspace.""" + for option in ( + "fixed_features", + "max_travel_distances", + "custom_objective", + "n_interpolate_points", + ): + if getattr(self, option) is not None: + raise ValueError(f"BAxUS does not support {option}") + # the target-space model has a single outcome, while the inherited + # acquisition scalarizes over every vocs output + if self.vocs.observables: + raise ValueError("BAxUS does not support observables") + return self + + def generate(self, n_candidates: int = 1) -> list[dict[str, float]]: + """Sobol seed points first, then LogEI within the trust region. + + The BO phase defers to ``BayesianGenerator.generate`` through the + ``get_training_data`` / ``_process_candidates`` / ``_get_torch_bounds`` + seams, so only the seed branch and the trust-region advance are + BAxUS-specific. Only the BO phase records ``computation_time`` rows; + the Sobol seed phase trains no model. + """ + # the guard is repeated from the base call because the seed branch + # below returns without ever reaching it + if n_candidates > 1 and not self.supports_batch_generation: + raise NotImplementedError( + "This Bayesian algorithm does not currently support parallel candidate generation" + ) + + if len(self._finite_data()) < self.n_initial_sobol: + self.n_candidates = n_candidates + z = torch.tensor(self._draw_sobol_point()) + return self._process_candidates(z).to_dict("records") + + # fold newly observed results in before training, so an expansion + # rebuilds the model in the new target space + self._advance_trust_region() + return super().generate(n_candidates) + + def get_training_data(self, data: pd.DataFrame) -> pd.DataFrame: + """Restrict the GP training set to rows with a finite objective.""" + return data[self._finite_objective_mask(data)] + + def _process_candidates(self, candidates: torch.Tensor) -> pd.DataFrame: + """Lift target-space candidates into vocs space. + + Replaces the base implementation wholesale: its discrete snapping and + fixed-feature handling are keyed to vocs-space columns, and BAxUS + rejects both options at construction. + """ + x_raw = self._lift_to_vocs(candidates.detach().cpu().numpy()) + return convert_numpy_to_inputs(self.vocs, x_raw, include_constants=False) + + def _get_torch_bounds(self) -> torch.Tensor: + """The full embedded target space, ``[-1, 1]^target_dim``. + + This deliberately re-points a base-class seam from vocs space into + target space. It is safe because BAxUS overrides or rejects every other + consumer: ``_get_optimization_bounds`` below, ``max_travel_distances``, + and ``visualize_model``. + """ + ones = torch.ones(self.embedding.target_dim, dtype=torch.double) + return torch.stack([-ones, ones]) + + def get_optimum(self) -> pd.DataFrame: + """Select the best point given by the model's posterior mean. + + Optimizes over the full embedded target space and lifts the result to + vocs space. The base implementation's constraint, fixed-feature, + contextual, and discrete handling is inert here because BAxUS rejects + all of those at construction. + """ + # an expansion clears the model, and a freshly restored run has none; + # retrain so the model always matches the current embedding + if self.model is None: + self.train_model() + return super().get_optimum() + + def visualize_model(self, **kwargs): + """Not supported: the model lives in the embedded target space.""" + raise NotImplementedError( + "BAxUS models live in the embedded target space; vocs-space model " + "visualization is not supported" + ) + + def add_data(self, new_data: pd.DataFrame) -> None: + """Ingest evaluation results. + + Pure ingestion - the trust region is advanced in ``generate`` instead, so + the state machine follows optimization iterations rather than how results + were chunked into ``add_data`` calls. + + Non-finite objective rows stay in ``data`` but are excluded from the GP + training set and the trust region. A batch missing the objective column + entirely (what an evaluator returns for a failed evaluation under + ``strict=False``) is ingested the same way. + """ + super().add_data(new_data) + + n_bad = int((~self._finite_objective_mask(new_data)).sum()) + if n_bad: + logger.warning( + "BAxUS: dropping %d row(s) with non-finite objective from training data", + n_bad, + ) + + def set_data(self, data: pd.DataFrame) -> None: + """Reattach a full dataset without replaying trust-region updates. + + ``Xopt`` calls this instead of ``add_data`` when loading a saved run + (see ``StateOwner``), so the deserialized embedding and trust region + survive a checkpoint round trip intact. + """ + self.data = data + + def _expansion_seed(self) -> int | None: + """Seed for the next expansion, derived so a round trip reproduces it.""" + if self.seed is None: + return None + return self.seed + 104729 * (self.n_expansions + 1) + + def _advance_trust_region(self) -> None: + """Fold newly observed BO-phase results into the trust region. + + Called once per ``generate``, mirroring how ``BayesianGenerator`` drives + ``TurboController.update_state``. Sobol seed points never move the trust + region (matches the reference), and ``tr_observed_rows`` records how much + history has already been applied so a resumed run does not re-ingest it. + """ + data = self._finite_data() + n_finite = len(data) + start = max(self.tr_observed_rows, self.n_initial_sobol) + if n_finite <= start: + return + + y = data[self._objective_name].to_numpy(dtype=np.float64)[start:] + self.tr_observed_rows = n_finite + + # refresh lazily so an eval_budget assigned after construction takes effect + self._refresh_failure_tolerance() + self.trust_region.update( + torch.tensor(self._objective_weight() * y, dtype=torch.double) + ) + + if self.trust_region.restart_triggered: + previous_dim = self.embedding.target_dim + self.embedding = self.embedding.expand( + self.new_bins_on_split, seed=self._expansion_seed() + ) + logger.info( + "BAxUS: trust region too small (%.4e), expanded embedding from %d to %d dims", + self.trust_region.length, + previous_dim, + self.embedding.target_dim, + ) + self.n_expansions += 1 + self.trust_region = BAxUSTrustRegion( + target_dim=self.embedding.target_dim, + length=self.length_init, + length_init=self.length_init, + best_value=self.trust_region.best_value, + ) + self._refresh_failure_tolerance() + # the fitted model lives in the old target space + self.model = None + + def train_model( + self, data: pd.DataFrame | None = None, update_internal: bool = True + ) -> Module: + """Fit the GP on target-space projections of the finite-objective data.""" + data = data if data is not None else self._finite_data() + if data.empty: + raise ValueError("no data available to build model") + + Z = self._target_space(data) + input_names = [f"z{i}" for i in range(self.embedding.target_dim)] + objective_name = self._objective_name + frame = pd.DataFrame(Z, columns=input_names) + frame[objective_name] = data[objective_name].to_numpy(dtype=np.float64) + + # xopt's constructor downgrades a failed fit to a warning and returns an + # untrained model (expected occasionally right after an embedding + # expansion); re-emit captured warnings and log the degradation + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + model = self.gp_constructor.build_model( + input_names=input_names, + outcome_names=[objective_name], + data=frame, + input_bounds={name: [-1.0, 1.0] for name in input_names}, + **self.tkwargs, + ) + for caught_warning in caught: + warnings.warn_explicit( + caught_warning.message, + caught_warning.category, + caught_warning.filename, + caught_warning.lineno, + ) + if any("Model fitting failed" in str(w.message) for w in caught): + logger.warning( + "BAxUS: GP fit failed at target_dim=%d - proceeding with an untrained model for this iteration", + self.embedding.target_dim, + ) + if update_internal: + self.model = model + return model + + def _get_optimization_bounds(self) -> torch.Tensor: + """Trust-region bounds in target space [-1, 1]^target_dim. + + Centered on the incumbent, per-dimension widths scaled by lengthscale + weights. Computed in [0, 1] units (BoTorch reference semantics for the + length fields), then mapped to [-1, 1]. + """ + data = self._finite_data() + y = torch.tensor( + data[self._objective_name].to_numpy(dtype=np.float64), + dtype=torch.double, + ) + best_idx = int(torch.argmax(self._objective_weight() * y)) + + Z = self._target_space(data) + center01 = (torch.tensor(Z[best_idx], dtype=torch.double) + 1.0) / 2.0 + + weights = self._lengthscale_weights(self.model) + half_length = self.trust_region.length / 2.0 + lb01 = torch.clamp(center01 - half_length * weights, 0.0, 1.0) + ub01 = torch.clamp(center01 + half_length * weights, 0.0, 1.0) + return torch.stack([2.0 * lb01 - 1.0, 2.0 * ub01 - 1.0]).to(**self.tkwargs) + + def _refresh_failure_tolerance(self) -> None: + """Set the trust region's failure tolerance, budget-aware when possible. + + Implements the reference schedule (Papenmeier et al., Alg. 1), pacing + embedding expansions so the final split lands as the evaluation budget + runs out; falls back to the ``ceil(target_dim / 2)`` heuristic without + ``eval_budget``. At full dimensionality the tolerance widens to + ``target_dim``. The planned split count is generalized to + ``ceil(log_{b+1}(input_dim / d_init))`` to honor a user-chosen + ``target_dim_init``. + """ + region = self.trust_region + # the embedding is the single source of truth for dimensionality + target_dim = self.embedding.target_dim + input_dim = self.embedding.input_dim + if target_dim >= input_dim: + region.failure_tolerance = target_dim + return + if self.eval_budget is None: + # no budget: leave the ceil(target_dim / 2) heuristic the trust region + # already derived for this target_dim + return + d_init = min(self.target_dim_init, input_dim) + b = self.new_bins_on_split + # The epsilon keeps exact powers of b+1 from rounding up on float noise. + n_splits = math.ceil(math.log(input_dim / d_init, b + 1) - 1e-9) + # each split multiplies dimensionality by (b + 1), so the planned splits + # span a total growth factor of (b + 1) ** (n_splits + 1) + growth = (b + 1) ** (n_splits + 1) + # evaluations left for the BO phase, then this target_dim's share of them + bo_budget = max(1, self.eval_budget - self.n_initial_sobol) + split_budget = round(b * bo_budget * target_dim / (d_init * (growth - 1))) + # max(1, ...): length_init == length_min would otherwise divide by zero + halvings = max( + 1, math.floor(math.log(region.length_min / region.length_init, 0.5)) + ) + # spend that share evenly over the halvings needed to reach length_min + region.failure_tolerance = min( + target_dim, max(1, math.floor(split_budget / halvings)) + ) + + @property + def _objective_name(self) -> str: + """The single objective this generator models (multi-objective is rejected).""" + return self.vocs.objective_names[0] + + def _vocs_bounds(self) -> tuple[np.ndarray, np.ndarray]: + """Per-variable lower and upper vocs bounds, ordered as ``variable_names``.""" + lb, ub = get_variable_bounds_array(self.vocs) + return lb, ub + + def _lift_to_vocs(self, z: np.ndarray) -> np.ndarray: + """Lift target-space rows to raw vocs-space values. + + Lifts through the embedding, clips to ``[-1, 1]``, then scales to the + per-variable vocs bounds. + """ + x_norm = np.clip(self.embedding.lift(z), -1.0, 1.0) + lb, ub = self._vocs_bounds() + return lb + (x_norm + 1.0) / 2.0 * (ub - lb) + + def _objective_weight(self) -> float: + """xopt objective weight: +1 for MAXIMIZE, -1 for MINIMIZE.""" + weights = set_botorch_weights(self.vocs) + return float(weights[self.vocs.output_names.index(self._objective_name)]) + + def _finite_objective_mask(self, frame: pd.DataFrame) -> np.ndarray: + """Rows of ``frame`` with a finite objective; all-False if the column is absent. + + Note that ``vocs.extract_data(return_valid=True)`` cannot stand in here: + it filters on feasibility, which is all-True without constraints. + """ + column = frame.get(self._objective_name) + if column is None: + # every evaluation in the frame failed, so the column was never created + return np.zeros(len(frame), dtype=bool) + y = pd.to_numeric(column, errors="coerce").to_numpy(dtype=np.float64) + return np.isfinite(y) + + def _finite_data(self) -> pd.DataFrame: + """Rows of ``data`` with a finite objective - the GP training set.""" + if self.data is None or self.data.empty: + return pd.DataFrame( + columns=[*self.vocs.variable_names, self._objective_name] + ) + return self.data[self._finite_objective_mask(self.data)] + + def _normalized_inputs(self, data: pd.DataFrame) -> np.ndarray: + """Variable columns scaled to [-1, 1] per vocs bounds. + + Deliberately [-1, 1] rather than ``vocs.normalize_inputs``' [0, 1]: these + feed a matrix product with the embedding, which is sign-symmetric. + """ + lb, ub = self._vocs_bounds() + X = data[self.vocs.variable_names].to_numpy(dtype=np.float64) + return 2.0 * (X - lb) / (ub - lb) - 1.0 + + def _target_space(self, data: pd.DataFrame) -> np.ndarray: + """Variable columns of ``data`` projected into the embedded target space.""" + return self.embedding.project(self._normalized_inputs(data)) + + def _draw_sobol_point(self) -> np.ndarray: + """Next quasi-random seed point in target space [-1, 1]^target_dim. + + After a dump -> construct round trip the engine is rebuilt and + fast-forwarded by ``sobol_draws`` (exact continuation needs a fixed + ``seed``). + """ + target_dim = self.embedding.target_dim + if self._sobol is None or self._sobol.dimension != target_dim: + self._sobol = SobolEngine( + dimension=target_dim, scramble=True, seed=self.seed + ) + if self.sobol_draws: + self._sobol.fast_forward(self.sobol_draws) + z = 2.0 * self._sobol.draw(1).to(dtype=torch.double).numpy() - 1.0 + self.sobol_draws += 1 + return z + + def _lengthscale_weights(self, model: Module) -> torch.Tensor: + """Per-dimension trust-region scaling weights from the fitted model. + + Tries a gpytorch kernel lengthscale, then a duck-typed + ``get_lengthscale_weights`` hook, then falls back to uniform weights + (isotropic region). + """ + target_dim = self.embedding.target_dim + single = model.models[0] if hasattr(model, "models") else model + + # each getattr tolerates a None input, so the chain needs no nesting + covar = getattr(single, "covar_module", None) + kernel = getattr(covar, "base_kernel", covar) + lengthscale = getattr(kernel, "lengthscale", None) + if lengthscale is not None: + ls = torch.atleast_1d(lengthscale.detach().squeeze()) + return _normalize_lengthscale_weights(ls, target_dim) + + hook = getattr(single, "get_lengthscale_weights", None) + weights = hook(target_dim) if hook is not None else None + if weights is not None: + return weights + return torch.ones(target_dim, dtype=torch.double) diff --git a/xopt/tests/generators/bayesian/test_baxus.py b/xopt/tests/generators/bayesian/test_baxus.py new file mode 100644 index 000000000..9f93782c5 --- /dev/null +++ b/xopt/tests/generators/bayesian/test_baxus.py @@ -0,0 +1,1743 @@ +"""Tests for the BAxUS generator.""" + +import logging +import math +import warnings +from dataclasses import dataclass +from typing import Any, cast +from unittest.mock import MagicMock + +import numpy as np +import pandas as pd +import pytest +import torch +import yaml +from botorch.acquisition import LogExpectedImprovement +from pydantic import ValidationError +from xopt import VOCS, Xopt +from xopt import Evaluator as XoptEvaluator +from xopt.errors import GeneratorWarning, VOCSError +from xopt.generator import StateOwner +from xopt.generators import get_generator_dynamic, list_available_generators +from xopt.generators.bayesian.bayesian_generator import BayesianGenerator +from xopt.generators.bayesian.expected_improvement import ( + ExpectedImprovementGenerator, +) +from xopt.generators.bayesian.baxus import ( + BAxUSEmbedding, + BAxUSGenerator, + BAxUSTrustRegion, + _normalize_lengthscale_weights, +) +from xopt.generators.bayesian.objectives import CustomXoptObjective +from xopt.vocs import ContextualVariable, DiscreteVariable + + +@pytest.fixture(autouse=True) +def _quiet_missing_eval_budget(): + """Silence the eval_budget advisory across this module. + + Most tests here construct generators without ``eval_budget`` on purpose, + because they exercise something unrelated to the expansion schedule, and the + advisory would otherwise fire on nearly every construction. The contract + itself stays pinned by ``TestBAxUSEvalBudgetWarning``, whose ``pytest.warns`` + and ``catch_warnings`` blocks install their own filters over this one. + """ + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="BAxUS: eval_budget is not set", + category=GeneratorWarning, + ) + yield + + +def _simple_vocs(n_vars: int = 6) -> VOCS: + """Create a VOCS with *n_vars* variables for testing.""" + variables = {f"x{i}": [-1.0, 1.0] for i in range(n_vars)} + return VOCS(variables=variables, objectives={"f": "MAXIMIZE"}) + + +def _sphere(inputs: dict[str, float]) -> dict[str, float]: + """Negative sphere - maximum at the origin.""" + return {"f": -sum(v**2 for v in inputs.values())} + + +def _assert_point_in_bounds(point: dict[str, float], vocs: VOCS) -> None: + """Assert every variable value lies within its VOCS domain.""" + for name in vocs.variable_names: + lo, hi = vocs.variables[name].domain + assert lo <= point[name] <= hi, f"{name}={point[name]} outside [{lo}, {hi}]" + + +class _MinimalCustomObjective(CustomXoptObjective): + """Concrete CustomXoptObjective, just enough to exercise the rejection path + (CustomXoptObjective is abstract, so a bare instance can't reach our validator).""" + + def forward( + self, samples: torch.Tensor, X: torch.Tensor | None = None + ) -> torch.Tensor: + return samples + + +def _seed_sobol(gen: BAxUSGenerator) -> None: + """Run the Sobol init phase so subsequent generate() calls take the BO path.""" + for _ in range(gen.n_initial_sobol): + point = gen.generate(1)[0] + gen.add_data(pd.DataFrame([{**point, **_sphere(point)}])) + + +def _bo_step(gen: BAxUSGenerator) -> dict[str, float]: + """Run one BO generate->evaluate->add_data cycle, asserting the candidate stays in bounds.""" + point = gen.generate(1)[0] + _assert_point_in_bounds(point, gen.vocs) + gen.add_data(pd.DataFrame([{**point, **_sphere(point)}])) + return point + + +def _force_expansion(gen: BAxUSGenerator) -> None: + """Put the trust region below length_min with an unbeatable incumbent, so the + next fold-in trips the restart -> expand branch.""" + gen.trust_region.length = gen.trust_region.length_min / 2.0 + gen.trust_region.best_value = 1e9 # guarantee a failure (no improvement) + + +def _round_trip(gen: BAxUSGenerator, *, with_data: bool = False) -> BAxUSGenerator: + """Dump -> construct, the documented resume path. + + The live botorch model is not serializable and is popped by hand, exactly as + the docs instruct a user to do. + """ + dump = gen.model_dump() + dump.pop("model", None) + restored = BAxUSGenerator(**dump) + if with_data and gen.data is not None: + restored.data = gen.data.copy() + return restored + + +def _fold_in_results(gen: BAxUSGenerator) -> None: + """Advance the trust region over results ingested since the last generate. + + The trust region is advanced at the start of ``generate`` (mirroring how + ``BayesianGenerator`` drives ``TurboController.update_state``), so folding in + the most recent evaluations means asking for the next candidate. + """ + gen.generate(1) + + +def _set_target_dim(gen: BAxUSGenerator, target_dim: int) -> None: + """Move the generator to a given target dimensionality, keeping the + embedding and trust region consistent (the generator enforces that).""" + gen.embedding = BAxUSEmbedding.create( + input_dim=gen.embedding.input_dim, target_dim=target_dim, seed=0 + ) + gen.trust_region = BAxUSTrustRegion(target_dim=target_dim) + + +class TestBAxUSGeneratorCreation: + """Test generator initialisation.""" + + def test_baxus_generator_creation(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs()) + + assert gen.name == "baxus" + assert gen.trust_region.target_dim == 2 + assert gen.embedding.target_dim == 2 + assert gen.embedding.input_dim == 6 + + def test_custom_target_dim(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=10), target_dim_init=4) + assert gen.trust_region.target_dim == 4 + assert gen.embedding.target_dim == 4 + assert gen.embedding.input_dim == 10 + + def test_target_dim_capped_to_input_dim(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=3), target_dim_init=10) + assert gen.trust_region.target_dim == 3 + + def test_missing_vocs_raises_validation_error(self) -> None: + """Without vocs, ``_init_components`` must skip auto-creation and let + pydantic report the missing required fields, not crash internally.""" + with pytest.raises(ValidationError, match="vocs"): + BAxUSGenerator() + + def test_embedding_structure(self) -> None: + """Each column of S should have exactly one non-zero entry.""" + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=8), target_dim_init=3) + S = np.asarray(gen.embedding.matrix) + for col_idx in range(S.shape[1]): + assert np.count_nonzero(S[:, col_idx]) == 1 + assert abs(S[:, col_idx]).max() == 1.0 + + +class TestBAxUSInitialSobol: + """Test that early generate() calls return Sobol points in bounds.""" + + def test_baxus_initial_sobol(self) -> None: + vocs = _simple_vocs() + gen = BAxUSGenerator(vocs=vocs) + + for _ in range(gen.n_initial_sobol): + candidates = gen.generate(1) + assert len(candidates) == 1 + _assert_point_in_bounds(candidates[0], vocs) + gen.add_data(pd.DataFrame([{**candidates[0], **_sphere(candidates[0])}])) + + +class TestBAxUSSeed: + """The ``seed`` field pins the embedding and Sobol initialisation. + + GP fitting/acquisition use torch's global RNG and are out of scope. + """ + + def test_same_seed_gives_identical_embedding(self) -> None: + vocs = _simple_vocs(n_vars=6) + g1 = BAxUSGenerator(vocs=vocs, seed=0) + g2 = BAxUSGenerator(vocs=vocs, seed=0) + assert g1.embedding.matrix == g2.embedding.matrix + + def test_different_seeds_give_different_embedding(self) -> None: + vocs = _simple_vocs(n_vars=6) + g1 = BAxUSGenerator(vocs=vocs, seed=0) + g2 = BAxUSGenerator(vocs=vocs, seed=1) + assert g1.embedding.matrix != g2.embedding.matrix + + def test_seed_none_still_works(self) -> None: + vocs = _simple_vocs(n_vars=4) + gen = BAxUSGenerator(vocs=vocs, seed=None) + point = gen.generate(1)[0] + _assert_point_in_bounds(point, vocs) + + +class TestBAxUSSobolSequence: + """The Sobol phase must be one reproducible low-discrepancy sequence, + not independent scrambles from a fresh engine per call.""" + + def test_same_seed_gives_identical_sobol_points(self) -> None: + vocs = _simple_vocs(n_vars=6) + g1 = BAxUSGenerator(vocs=vocs, seed=7) + g2 = BAxUSGenerator(vocs=vocs, seed=7) + pts1 = [g1.generate(1)[0] for _ in range(g1.n_initial_sobol)] + pts2 = [g2.generate(1)[0] for _ in range(g2.n_initial_sobol)] + assert pts1 == pts2 + + def test_sobol_points_advance(self) -> None: + """A fresh ``draw(1)`` engine per call would repeat the same first point.""" + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=6), seed=3) + pts = [tuple(gen.generate(1)[0].values()) for _ in range(gen.n_initial_sobol)] + assert len(set(pts)) == len(pts) + + +class TestBAxUSEndToEnd: + """End-to-end: Sobol init then BO steps.""" + + def test_baxus_end_to_end(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) + + _seed_sobol(gen) + for _ in range(3): + _bo_step(gen) + + assert gen.data is not None + assert len(gen.data) == gen.n_initial_sobol + 3 + + +def _asymmetric_vocs() -> VOCS: + """VOCS with per-variable bounds at different scales - catches bounds orientation bugs.""" + return VOCS( + variables={"x0": [0.0, 10.0], "x1": [-165.0, 165.0], "x2": [0.5, 0.6]}, + objectives={"f": "MAXIMIZE"}, + ) + + +class TestBAxUSAsymmetricBounds: + """Guard against the v2->v3 `vocs.bounds` orientation flip (was (2, D), now (D, 2)).""" + + def test_points_lie_in_per_variable_bounds(self) -> None: + vocs = _asymmetric_vocs() + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2) + + _seed_sobol(gen) + # A wrong bounds orientation produces out-of-box points in the BO phase. + for _ in range(3): + _bo_step(gen) + + +class TestBAxUSSteering: + """BO-phase points must be better *on average* than the Sobol seed points. + + Guards against a min-form/max-form inversion (``get_objective_data`` + negates MAXIMIZE objectives). Never compare best-over-all-data against + best-over-seeds - the seeds are a subset of all data, so that is + vacuously true and lets the generator anti-optimize undetected. + + The optimum is deliberately placed OFF-CENTER. With an optimum at the origin + - the center of the symmetric [-1, 1] domain - these tests only measure + "did the candidate move toward the middle of the box", which a generator + that ignores the acquisition function entirely (proposing a constant zero in + embedded space) satisfies perfectly. The offset makes them measure + optimization instead. + """ + + N_SOBOL = 6 + N_BO_STEPS = 15 + OPTIMUM = 0.6 + + def _run_phases(self, direction: str) -> tuple[list[float], list[float]]: + """Run Sobol seeding then BO steps; return (sobol_ys, bo_ys). + + The objective peaks at ``x_i = OPTIMUM`` for every i: MAXIMIZE gets + -||x - c||^2 (values <= 0), MINIMIZE gets +||x - c||^2 (values >= 0). + """ + torch.manual_seed(0) + vocs = VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(6)}, + objectives={"f": direction}, + ) + gen = BAxUSGenerator( + vocs=vocs, + target_dim_init=2, + n_initial_sobol=self.N_SOBOL, + seed=0, + eval_budget=self.N_SOBOL + self.N_BO_STEPS, + ) + sign = -1.0 if direction == "MAXIMIZE" else 1.0 + + ys: list[float] = [] + for _ in range(self.N_SOBOL + self.N_BO_STEPS): + point = gen.generate(1)[0] + _assert_point_in_bounds(point, vocs) + y = sign * sum((v - self.OPTIMUM) ** 2 for v in point.values()) + ys.append(y) + gen.add_data(pd.DataFrame([{**point, "f": y}])) + return ys[: self.N_SOBOL], ys[self.N_SOBOL :] + + def test_maximize_bo_phase_beats_sobol_phase(self) -> None: + sobol_ys, bo_ys = self._run_phases("MAXIMIZE") + assert np.mean(bo_ys) > np.mean(sobol_ys), ( + f"MAXIMIZE BO phase is worse than random seeding (direction inverted?): " + f"sobol mean={np.mean(sobol_ys):.4f} bo mean={np.mean(bo_ys):.4f}" + ) + + def test_minimize_bo_phase_beats_sobol_phase(self) -> None: + sobol_ys, bo_ys = self._run_phases("MINIMIZE") + assert np.mean(bo_ys) < np.mean(sobol_ys), ( + f"MINIMIZE BO phase is worse than random seeding (direction inverted?): " + f"sobol mean={np.mean(sobol_ys):.4f} bo mean={np.mean(bo_ys):.4f}" + ) + + +class TestBAxUSNaNIngestion: + """NaN objectives (failed evaluations) must not poison the training data + or the trust-region update.""" + + def test_nan_row_dropped_from_training_data(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) + + point = gen.generate(1)[0] + gen.add_data(pd.DataFrame([{**point, "f": float("nan")}])) + + assert len(gen.data) == 1 # row is kept in the canonical history... + assert gen._finite_data().empty # ...but excluded from GP training + + def test_mixed_batch_keeps_only_finite_rows(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) + + rows = [] + for y in (1.0, float("nan"), 2.0): + point = gen.generate(1)[0] + rows.append({**point, "f": y}) + gen.add_data(pd.DataFrame(rows)) + + assert len(gen.data) == 3 + assert len(gen._finite_data()) == 2 + + def test_bo_step_survives_nan_ingestion(self) -> None: + """After a NaN row mid-run, the next BO candidate is still finite and in bounds.""" + vocs = _simple_vocs(n_vars=4) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2) + + _seed_sobol(gen) + point = gen.generate(1)[0] + gen.add_data(pd.DataFrame([{**point, "f": float("nan")}])) + + candidate = gen.generate(1)[0] + assert all(math.isfinite(v) for v in candidate.values()) + _assert_point_in_bounds(candidate, vocs) + + +class TestBAxUSTrustRegionSeedPhase: + """The trust region must track BO-phase evaluations only: the BoTorch + reference starts BO with a pristine trust region, so Sobol seeding must + not move the counters or ``best_value``.""" + + def test_seed_phase_leaves_state_pristine(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + + _seed_sobol(gen) + + assert gen.trust_region.length == gen.trust_region.length_init + assert gen.trust_region.success_counter == 0 + assert gen.trust_region.failure_counter == 0 + assert gen.trust_region.best_value is None + + def test_first_bo_result_updates_state(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + + _seed_sobol(gen) + _bo_step(gen) + assert gen.trust_region.best_value is None # not folded in yet + _fold_in_results(gen) + + assert gen.trust_region.best_value is not None + assert gen.tr_observed_rows == len(gen.data) + + def test_straddling_batch_counts_only_post_quota_rows(self) -> None: + """A batch that crosses the seed quota updates the trust region with the + BO-phase rows only.""" + vocs = _simple_vocs(n_vars=4) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0) + + pts = [gen.generate(1)[0] for _ in range(3)] + rows = [ + {**pts[0], "f": 100.0}, + {**pts[1], "f": 100.0}, + {**pts[2], "f": 1.0}, + ] + gen.add_data(pd.DataFrame(rows)) + _fold_in_results(gen) + + # MAXIMIZE ⇒ weighted value is +f; only the BO row (f=1.0) should count. + assert gen.trust_region.best_value == 1.0 + + def test_trust_region_is_independent_of_ingestion_chunking(self) -> None: + """Same observations, different add_data batching ⇒ identical state. + + The state machine is driven by generate calls, so feeding results one at + a time or all at once must not change the trajectory (a warm start via + ``Xopt(..., data=...)`` arrives as one big frame). + """ + vocs = _simple_vocs(n_vars=4) + rng = np.random.default_rng(0) + pts = [ + {f"x{i}": float(v) for i, v in enumerate(rng.uniform(-1, 1, 4))} + for _ in range(12) + ] + rows = [{**p, **_sphere(p)} for p in pts] + + def state(batched: bool) -> tuple[float, int, int, float]: + gen = BAxUSGenerator( + vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0 + ) + if batched: + gen.add_data(pd.DataFrame(rows)) + else: + for row in rows: + gen.add_data(pd.DataFrame([row])) + _fold_in_results(gen) + tr = gen.trust_region + return (tr.length, tr.success_counter, tr.failure_counter, tr.best_value) + + assert state(batched=True) == state(batched=False) + + def test_baxus_embedding_expansion(self) -> None: + """Expansion grows target_dim, resets the region, and preserves best_value.""" + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=8), target_dim_init=2, seed=0) + _seed_sobol(gen) + + # Force the next fold-in to trip the restart branch. + _force_expansion(gen) + _bo_step(gen) + _fold_in_results(gen) + + assert gen.trust_region.target_dim > 2 + assert gen.embedding.target_dim == gen.trust_region.target_dim + assert gen.trust_region.length == gen.trust_region.length_init + assert gen.trust_region.best_value == 1e9 + assert not gen.trust_region.restart_triggered + assert gen.n_expansions == 1 + + +class TestBAxUSBatchGeneration: + """``generate(2)`` must raise (matching xopt's guard), not silently + under-deliver one point.""" + + def test_generate_multiple_candidates_raises(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) + with pytest.raises(NotImplementedError, match="parallel candidate generation"): + gen.generate(2) + + +class TestBAxUSTrainModel: + """train_model must refuse to fit a GP with no finite-objective data.""" + + def test_raises_when_no_data_available(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) + with pytest.raises(ValueError, match="no data available to build model"): + gen.train_model() + + +class TestBAxUSViaXopt: + """Full integration through the Xopt runner.""" + + def test_baxus_via_xopt(self) -> None: + vocs = _simple_vocs(n_vars=4) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2) + evaluator = XoptEvaluator(function=_sphere) + X = Xopt(evaluator=evaluator, generator=gen) + + # Seed with random points (like the real phases.py flow), then BO. + X.random_evaluate(gen.n_initial_sobol) + for _ in range(3): + X.step() + + assert len(X.data) == gen.n_initial_sobol + 3 + + for name in vocs.variable_names: + lo, hi = vocs.variables[name].domain + values = X.data[name].to_numpy() + assert np.all(values >= lo - 1e-9), f"{name} below lower bound" + assert np.all(values <= hi + 1e-9), f"{name} above upper bound" + + +class TestBAxUSLengthscaleWeights: + """_lengthscale_weights: gpytorch kernel -> duck-typed hook -> uniform fallback.""" + + def _gen(self, target_dim: int = 2) -> BAxUSGenerator: + return BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=target_dim) + + @staticmethod + def _kernel_model(lengthscale: torch.Tensor) -> MagicMock: + """A mock shaped like a gpytorch model: covar_module.base_kernel.lengthscale.""" + kernel = MagicMock() + kernel.lengthscale = lengthscale + covar = MagicMock() + covar.base_kernel = kernel + model = MagicMock(spec=["covar_module"]) + model.covar_module = covar + return model + + def test_gpytorch_kernel_lengthscale_used(self) -> None: + gen = self._gen() + model = self._kernel_model(torch.tensor([[0.5, 2.0]], dtype=torch.double)) + + weights = gen._lengthscale_weights(model) + assert weights.shape == (2,) + assert not torch.allclose(weights, torch.ones(2, dtype=torch.double)) + + def test_scalar_lengthscale_unsqueezed(self) -> None: + gen = self._gen(target_dim=1) + model = self._kernel_model(torch.tensor(0.5, dtype=torch.double)) + + assert gen._lengthscale_weights(model).shape == (1,) + + def test_duck_typed_hook_used_when_no_kernel(self) -> None: + gen = self._gen() + model = MagicMock(spec=["get_lengthscale_weights"]) + model.get_lengthscale_weights.return_value = torch.tensor( + [0.5, 1.5], dtype=torch.double + ) + + weights = gen._lengthscale_weights(model) + assert torch.allclose(weights, torch.tensor([0.5, 1.5], dtype=torch.double)) + model.get_lengthscale_weights.assert_called_once_with(2) + + def test_uniform_fallback_when_hook_returns_none(self) -> None: + gen = self._gen() + model = MagicMock(spec=["get_lengthscale_weights"]) + model.get_lengthscale_weights.return_value = None + + assert torch.allclose( + gen._lengthscale_weights(model), torch.ones(2, dtype=torch.double) + ) + + def test_uniform_fallback_for_bare_model(self) -> None: + gen = self._gen() + model = MagicMock(spec=["posterior", "num_outputs"]) + + assert torch.allclose( + gen._lengthscale_weights(model), torch.ones(2, dtype=torch.double) + ) + + +class TestBAxUSConfig: + def test_explicit_n_initial_sobol_kept(self) -> None: + gen = BAxUSGenerator( + vocs=_simple_vocs(n_vars=4), target_dim_init=2, n_initial_sobol=5 + ) + assert gen.n_initial_sobol == 5 + + def test_auto_n_initial_sobol(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=6), target_dim_init=3) + assert gen.n_initial_sobol == 4 # max(2, target_dim + 1) + + +class TestBAxUSTrustRegion: + """Trust-region state bookkeeping (weighted convention: higher is better).""" + + def test_success_expands_length(self) -> None: + tr = BAxUSTrustRegion(target_dim=2, length=0.8) + tr.best_value = 0.0 + for i in range(tr.success_tolerance): + tr.update(torch.tensor([10.0 * (i + 1)])) + assert tr.length > 0.8 + + def test_failure_shrinks_length(self) -> None: + tr = BAxUSTrustRegion(target_dim=2, length=0.8) + tr.best_value = 999.0 + for _ in range(tr.failure_tolerance): + tr.update(torch.tensor([-999.0])) + assert tr.length < 0.8 + + def test_restart_triggered_on_tiny_length(self) -> None: + tr = BAxUSTrustRegion(target_dim=2, length=0.008) + tr.best_value = 999.0 + for _ in range(100): + tr.update(torch.tensor([-999.0])) + if tr.restart_triggered: + break + assert tr.restart_triggered + + def test_first_update_counts_as_success(self) -> None: + """best_value=None (pristine) means the first observed batch always improves.""" + tr = BAxUSTrustRegion(target_dim=2) + tr.update(torch.tensor([-5.0])) + assert tr.best_value == -5.0 + assert tr.failure_counter == 0 + + def test_failure_tolerance_scales_with_target_dim(self) -> None: + assert BAxUSTrustRegion(target_dim=2).failure_tolerance == 1 + assert BAxUSTrustRegion(target_dim=10).failure_tolerance == 5 + + def test_round_trips_through_dump(self) -> None: + tr = BAxUSTrustRegion( + target_dim=3, length=0.4, success_counter=2, best_value=1.5 + ) + assert BAxUSTrustRegion(**tr.model_dump()) == tr + + +@dataclass +class _TutorialBaxusState: + """Oracle: the reference ``BaxusState`` from the BoTorch BAxUS tutorial, + trimmed to the fields the failure-tolerance formula needs.""" + + dim: int + eval_budget: int + new_bins_on_split: int = 3 + d_init: int = 0 + target_dim: int = 0 + n_splits: int = 0 + length_init: float = 0.8 + length_min: float = 0.5**7 + + def __post_init__(self) -> None: + n_splits = round(math.log(self.dim, self.new_bins_on_split + 1)) + self.d_init = int( + 1 + + np.argmin( + np.abs( + (1 + np.arange(self.new_bins_on_split)) + * (1 + self.new_bins_on_split) ** n_splits + - self.dim + ) + ) + ) + self.target_dim = self.d_init + self.n_splits = n_splits + + @property + def split_budget(self) -> int: + growth: int = (self.new_bins_on_split + 1) ** (self.n_splits + 1) + return round( + -1 + * (self.new_bins_on_split * self.eval_budget * self.target_dim) + / (self.d_init * (1 - growth)) + ) + + @property + def failure_tolerance(self) -> int: + if self.target_dim == self.dim: + return self.target_dim + k = math.floor(math.log(self.length_min / self.length_init, 0.5)) + return min(self.target_dim, max(1, math.floor(self.split_budget / k))) + + +class TestBAxUSFailureTolerance: + """Budget-aware failure tolerance (paper Alg. 1) against the tutorial oracle. + + input_dim=16 with b=3 is an exact power of b+1, where the oracle's derived + d_init (=1) and split count (=2) coincide with ours - the regime in which a + verbatim comparison is meaningful. + """ + + def _gen(self, eval_budget: int | None, target_dim_init: int = 1) -> BAxUSGenerator: + return BAxUSGenerator( + vocs=_simple_vocs(n_vars=16), + target_dim_init=target_dim_init, + n_initial_sobol=10, + seed=0, + eval_budget=eval_budget, + ) + + # bo_budget=93 with target_dim=8 is the discriminating pair: computing the + # budget without subtracting the 10 seed points gives 6 instead of 5, so + # this pins the total-including-seeds semantics of eval_budget. + @pytest.mark.parametrize("bo_budget", [30, 93, 200]) + def test_matches_tutorial_oracle_across_dims(self, bo_budget: int) -> None: + oracle = _TutorialBaxusState(dim=16, eval_budget=bo_budget) + assert oracle.d_init == 1 # sanity: oracle derives the d_init we configure + gen = self._gen( + eval_budget=bo_budget + 10 + ) # our field is total incl. the 10 seeds + + for target_dim in (1, 4, 8, 16): + oracle.target_dim = target_dim + _set_target_dim(gen, target_dim) + gen._refresh_failure_tolerance() + assert gen.trust_region.failure_tolerance == oracle.failure_tolerance + + def test_differs_from_heuristic(self) -> None: + """The paper schedule is not just ceil(d/2) in disguise.""" + gen = self._gen(eval_budget=210) + _set_target_dim(gen, 4) + gen._refresh_failure_tolerance() + assert gen.trust_region.failure_tolerance == 4 # heuristic would say 2 + + def test_heuristic_kept_without_budget(self) -> None: + gen = self._gen(eval_budget=None, target_dim_init=4) + assert gen.trust_region.failure_tolerance == 2 # ceil(4/2) + gen._refresh_failure_tolerance() + assert gen.trust_region.failure_tolerance == 2 + + def test_full_dim_widens_to_target_dim(self) -> None: + """At full dimensionality the tolerance is target_dim, budget or not.""" + gen = self._gen(eval_budget=None, target_dim_init=16) + gen._refresh_failure_tolerance() + assert gen.trust_region.failure_tolerance == 16 + + def test_bo_phase_applies_budget_tolerance(self) -> None: + """Folding in a BO result refreshes the tolerance (also covers a d_init the + reference would not derive itself: n_splits generalizes).""" + vocs = _simple_vocs(n_vars=16) + gen = BAxUSGenerator( + vocs=vocs, target_dim_init=4, n_initial_sobol=2, seed=0, eval_budget=210 + ) + _seed_sobol(gen) + _bo_step(gen) + _fold_in_results(gen) + assert ( + gen.trust_region.failure_tolerance == 4 + ) # paper value; heuristic would say 2 + + def test_expansion_recomputes_tolerance(self) -> None: + vocs = _simple_vocs(n_vars=8) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0, eval_budget=50) + _seed_sobol(gen) + + _force_expansion(gen) + _bo_step(gen) + _fold_in_results(gen) + + assert gen.trust_region.target_dim == 8 # expanded to full dim + assert gen.trust_region.failure_tolerance == 8 + + def test_expansion_recomputes_budget_tolerance_below_full_dim(self) -> None: + """An expansion landing below input_dim exercises the budget formula, + not just the full-dim widening.""" + vocs = _simple_vocs(n_vars=32) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0, eval_budget=203) + _seed_sobol(gen) + + _force_expansion(gen) + _bo_step(gen) + _fold_in_results(gen) + + assert gen.trust_region.target_dim == 8 # 2 * (b + 1), still < 32 + # d_init=2, n_splits=2, budget=200: round(4800/126)=38 -> floor(38/6)=6 + assert gen.trust_region.failure_tolerance == 6 + + def test_explicit_failure_tolerance_survives_construction(self) -> None: + """Static-model semantics only: failure_tolerance is derived state and + the generator re-derives it on every BO-phase ingest.""" + assert ( + BAxUSTrustRegion(target_dim=8, failure_tolerance=7).failure_tolerance == 7 + ) + + def test_budget_below_seed_quota_rejected(self) -> None: + with pytest.raises(ValidationError, match="must cover the seed quota"): + self._gen(eval_budget=5) # quota is 10 + + def test_dump_round_trips_budget_and_tolerance(self) -> None: + gen = self._gen(eval_budget=210) + _set_target_dim(gen, 4) + gen._refresh_failure_tolerance() + + restored = _round_trip(gen) + + assert restored.eval_budget == 210 + assert ( + restored.trust_region.failure_tolerance + == gen.trust_region.failure_tolerance + ) + + +class TestBAxUSPostExpansionFit: + """Right after an embedding expansion the projected training data has + duplicated coordinates, which can make the covariance indefinite; the run + must survive the degraded fit and log it.""" + + def test_generate_survives_post_expansion_training(self) -> None: + vocs = _simple_vocs(n_vars=8) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) + _seed_sobol(gen) + + _force_expansion(gen) + _bo_step(gen) + + # this generate expands first, then trains and proposes in the new space + candidate = gen.generate(1)[0] + assert gen.embedding.target_dim == 8 + _assert_point_in_bounds(candidate, vocs) + assert all(math.isfinite(v) for v in candidate.values()) + + def test_fit_failure_warning_surfaces_in_logs( + self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch + ) -> None: + """train_model must mirror xopt's fit-failure warning into the logger + and still re-emit the original warning unchanged.""" + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + _seed_sobol(gen) + + constructor_cls = type(gen.gp_constructor) + real_build = constructor_cls.build_model + + def failing_build(self: Any, **kwargs: Any) -> Any: + warnings.warn( + "Model fitting failed. Returning untrained model.", stacklevel=1 + ) + return real_build(self, **kwargs) + + monkeypatch.setattr(constructor_cls, "build_model", failing_build) + + logger_name = "xopt.generators.bayesian.baxus" + with ( + caplog.at_level(logging.WARNING, logger=logger_name), + pytest.warns(UserWarning, match="Model fitting failed"), + ): + gen.train_model() + + assert any( + "GP fit failed at target_dim=2" in record.message + for record in caplog.records + ) + + +class TestBAxUSEmbedding: + """Sparse random embedding: structure, projection round-trip, expansion.""" + + def test_create_structure(self) -> None: + """Each input column has exactly one signed unit entry.""" + emb = BAxUSEmbedding.create(input_dim=8, target_dim=3, seed=0) + S = np.asarray(emb.matrix) + assert S.shape == (3, 8) + for col in range(8): + assert np.count_nonzero(S[:, col]) == 1 + assert abs(S[:, col]).max() == 1.0 + + def test_create_assigns_both_signs(self) -> None: + """Sign randomization is what makes the embedding an unbiased projection; + an all-positive matrix is still structurally valid, so check the signs.""" + emb = BAxUSEmbedding.create(input_dim=200, target_dim=4, seed=0) + matrix = np.asarray(emb.matrix) + + positive = int((matrix == 1.0).sum()) + negative = int((matrix == -1.0).sum()) + assert positive and negative, "embedding signs are not randomized" + assert 0.3 < positive / (positive + negative) < 0.7 + + def test_create_is_seed_deterministic(self) -> None: + e1 = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=0) + e2 = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=0) + e3 = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=1) + assert e1.matrix == e2.matrix + assert e1.matrix != e3.matrix + + def test_lift_project_round_trip(self) -> None: + """Points lifted from target space project back exactly (S rows are orthogonal).""" + emb = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=0) + Z = np.array([[0.5, -0.25], [1.0, 1.0]]) + assert np.allclose(emb.project(emb.lift(Z)), Z) + + def test_expand_grows_target_dim_and_preserves_signs(self) -> None: + emb = BAxUSEmbedding.create(input_dim=12, target_dim=2, seed=0) + grown = emb.expand(new_bins_on_split=3) + S_old, S_new = np.asarray(emb.matrix), np.asarray(grown.matrix) + assert grown.target_dim == min(2 * 4, 12) + assert grown.input_dim == 12 + # Every column keeps its sign and moves to exactly one (new) row. + for col in range(12): + assert np.count_nonzero(S_new[:, col]) == 1 + assert S_new[:, col].sum() == S_old[:, col].sum() + + def test_expand_caps_at_input_dim(self) -> None: + emb = BAxUSEmbedding.create(input_dim=2, target_dim=1, seed=0) + grown = emb.expand(new_bins_on_split=5) + assert grown.target_dim == 2 + + def test_round_trips_through_dump(self) -> None: + emb = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=3) + assert BAxUSEmbedding(**emb.model_dump()) == emb + + +class TestBAxUSUnsupportedOptions: + """vocs-space point controls cannot be honored in an embedded subspace.""" + + @pytest.mark.parametrize( + "kwargs", + [ + {"fixed_features": {"x0": 0.5}}, + {"max_travel_distances": [0.1] * 6}, + {"n_interpolate_points": 3}, + ], + ) + def test_rejected_at_construction(self, kwargs: dict[str, Any]) -> None: + with pytest.raises(ValidationError, match="BAxUS does not support"): + BAxUSGenerator(vocs=_simple_vocs(), **kwargs) + + def test_custom_objective_rejected_at_construction(self) -> None: + """custom_objective must pass its own CustomXoptObjective type validation + before ours can fire, so it needs a real (minimal) subclass instance + rather than a plain kwargs value.""" + vocs = _simple_vocs() + objective = _MinimalCustomObjective(vocs=vocs) + with pytest.raises( + ValidationError, match="BAxUS does not support custom_objective" + ): + BAxUSGenerator(vocs=vocs, custom_objective=objective) + + def test_turbo_controller_rejected(self) -> None: + with pytest.raises( + ValidationError, match="no turbo controllers are compatible" + ): + BAxUSGenerator(vocs=_simple_vocs(), turbo_controller="optimize") + + @pytest.mark.parametrize( + ("extra_variables", "vocs_kwargs", "match"), + [ + pytest.param( + {}, + {"constraints": {"c": ["LESS_THAN", 0.5]}}, + "does not support constraints", + id="constraints", + ), + pytest.param( + {"d": DiscreteVariable(values=[0.0, 1.0, 2.0])}, + {}, + "does not support discrete variables", + id="discrete", + ), + pytest.param( + {"c": ContextualVariable()}, + {}, + "does not support contextual variables", + id="contextual", + ), + ], + ) + def test_unsupported_vocs_rejected( + self, extra_variables: dict[str, Any], vocs_kwargs: dict[str, Any], match: str + ) -> None: + """The target-space GP models only the objective over continuous + vocs-space variables, so each of these must be rejected up front. Each + case guards a supports_* = False override against the True default + inherited from BayesianGenerator / ExpectedImprovementGenerator.""" + vocs = VOCS( + variables={ + **{f"x{i}": [-1.0, 1.0] for i in range(3)}, + **extra_variables, + }, + objectives={"f": "MAXIMIZE"}, + **vocs_kwargs, + ) + with pytest.raises(VOCSError, match=match): + BAxUSGenerator(vocs=vocs) + + +class TestBAxUSAcquisitionPinning: + """Pin the acquisition contract: analytic LogEI on the target-space model + with best_f = best weighted finite objective. Inherited from + ExpectedImprovementGenerator, so these fail loudly if an xopt upgrade + changes EI's acquisition semantics.""" + + def test_analytic_logei_best_f_ignores_nan(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + _seed_sobol(gen) + point = gen.generate(1)[0] + gen.add_data(pd.DataFrame([{**point, "f": float("nan")}])) + + acq = gen.get_acquisition(gen.train_model()) + + assert isinstance(acq, LogExpectedImprovement) + expected = np.nanmax(gen.data["f"].to_numpy(dtype=np.float64)) + assert float(cast(torch.Tensor, acq.best_f)) == pytest.approx(expected) + + def test_minimize_negates_best_f(self) -> None: + vocs = VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, + objectives={"f": "MINIMIZE"}, + ) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) + _seed_sobol(gen) + + acq = gen.get_acquisition(gen.train_model()) + + assert isinstance(acq, LogExpectedImprovement) + expected = -np.min(gen.data["f"].to_numpy(dtype=np.float64)) + assert float(cast(torch.Tensor, acq.best_f)) == pytest.approx(expected) + + +class TestBAxUSIsBayesianGenerator: + def test_subclasses_bayesian_generator(self) -> None: + assert isinstance(BAxUSGenerator(vocs=_simple_vocs()), BayesianGenerator) + + def test_subclasses_expected_improvement_generator(self) -> None: + assert issubclass(BAxUSGenerator, ExpectedImprovementGenerator) + + def test_default_gp_constructor_is_standard(self) -> None: + assert BAxUSGenerator(vocs=_simple_vocs()).gp_constructor.name == "standard" + + +class TestBAxUSSerialization: + """Dump -> construct -> reattach data must resume the run. + + Resume contract: ``gen.model_dump()`` then pop ``"model"`` - xopt ignores + the ``exclude`` argument, so a trained generator's dump carries the live + botorch model (not yaml-safe) unless popped explicitly. ``computation_time`` + needs the same treatment for raw yaml: it is a live ``pd.DataFrame`` once a + BO-phase ``generate()`` has run. + + Data is reattached with ``set_data`` (equivalently, by assigning + ``gen.data``). Note this is what ``Xopt`` itself does for a ``StateOwner``, + so the ordinary ``Xopt.from_yaml`` path resumes correctly too - see + ``TestBAxUSCheckpointIntegrity``. + """ + + def test_dump_excludes_data_and_round_trips_components(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(), seed=7) + _seed_sobol(gen) + + dump = gen.model_dump() + dump.pop("model", None) + assert "data" not in dump + + gen2 = BAxUSGenerator(**dump) + assert gen2.embedding == gen.embedding + assert gen2.trust_region == gen.trust_region + assert gen2.sobol_draws == gen.sobol_draws + + def test_dump_is_yaml_safe(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(), seed=7) + dump = gen.model_dump() + dump.pop("model", None) + mid_dump = BAxUSGenerator(**dump).model_dump() + mid_dump.pop("model", None) + reloaded = yaml.safe_load(yaml.safe_dump(mid_dump)) + assert BAxUSGenerator(**reloaded).embedding == gen.embedding + + # Post-BO a live botorch model and a populated computation_time frame + # ride along; without the pops, yaml.safe_dump raises RepresenterError + # on the trained ModelListGP / the pd.DataFrame. + _seed_sobol(gen) + _bo_step(gen) + bo_dump = gen.model_dump() + bo_dump.pop("model", None) + bo_dump.pop("computation_time", None) + yaml.safe_dump(bo_dump) # must not raise + + gen2 = BAxUSGenerator(**bo_dump) + assert gen2.embedding == gen.embedding + assert gen2.trust_region == gen.trust_region + + def test_sobol_sequence_resumes_exactly(self) -> None: + """Seed phase is bit-identical after a round trip (fast-forwarded engine).""" + g1 = BAxUSGenerator(vocs=_simple_vocs(), seed=7) + for _ in range(2): + point = g1.generate(1)[0] + g1.add_data(pd.DataFrame([{**point, **_sphere(point)}])) + + g2 = _round_trip(g1, with_data=True) + + assert g1.generate(1)[0] == g2.generate(1)[0] + + def test_bo_phase_state_identical_after_round_trip(self) -> None: + """With the same torch seed, the resumed generator proposes the same point.""" + g1 = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + _seed_sobol(g1) + _bo_step(g1) + + g2 = _round_trip(g1, with_data=True) + + torch.manual_seed(0) + p1 = g1.generate(1)[0] + torch.manual_seed(0) + p2 = g2.generate(1)[0] + assert p1 == pytest.approx(p2) + + +class TestBAxUSRegistry: + """The generator must be reachable through xopt's dynamic registry.""" + + def test_get_generator_dynamic(self): + assert get_generator_dynamic("baxus") is BAxUSGenerator + + def test_listed_in_available_generators(self): + assert "baxus" in list_available_generators() + + +class TestBAxUSGetOptimum: + """get_optimum must search the full embedded target space (not vocs-space + bounds against the target-space model) and lift the result to vocs space.""" + + def test_get_optimum_returns_single_row_in_bounds(self) -> None: + vocs = _simple_vocs(n_vars=4) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) + + _seed_sobol(gen) + for _ in range(3): + _bo_step(gen) + + opt = gen.get_optimum() + + assert isinstance(opt, pd.DataFrame) + assert len(opt) == 1 + assert list(opt.columns) == vocs.variable_names + _assert_point_in_bounds(opt.iloc[0].to_dict(), vocs) + + +class TestBAxUSVisualizeModel: + """The inherited visualization is keyed to vocs-space variables and is + meaningless (or crashes) against the embedded target-space model.""" + + def test_visualize_model_raises_not_implemented(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + + with pytest.raises(NotImplementedError, match="target space"): + gen.visualize_model() + + +class TestBAxUSFailureToleranceZeroDivision: + """length_init == length_min makes the reference halvings formula 0, + which must not divide by zero.""" + + def test_equal_length_init_and_min_does_not_raise(self) -> None: + gen = BAxUSGenerator( + vocs=_simple_vocs(n_vars=8), + target_dim_init=2, + n_initial_sobol=2, + eval_budget=50, + length_init=0.5**7, + ) + gen._refresh_failure_tolerance() + assert gen.trust_region.failure_tolerance >= 1 + + +class TestBAxUSOptimizationBounds: + """The trust-region box itself. + + No end-to-end test can pin this down: replacing _get_optimization_bounds with + the full domain, or centering it on the worst point, leaves every behavioral + test passing. It needs direct assertions on the returned box. + """ + + @staticmethod + def _gen_with_known_incumbent() -> tuple[BAxUSGenerator, int]: + """Four points near the origin; row 2 is the unique best (MAXIMIZE).""" + vocs = _simple_vocs(n_vars=4) + gen = BAxUSGenerator( + vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=50 + ) + rows = [ + {"x0": 0.10, "x1": -0.05, "x2": 0.00, "x3": 0.10, "f": -1.0}, + {"x0": -0.10, "x1": 0.05, "x2": 0.10, "x3": -0.05, "f": -2.0}, + {"x0": 0.05, "x1": 0.10, "x2": -0.05, "x3": 0.00, "f": 5.0}, # best + {"x0": 0.00, "x1": 0.00, "x2": 0.05, "x3": 0.05, "f": -3.0}, + ] + gen.add_data(pd.DataFrame(rows)) + return gen, 2 + + def _incumbent_in_target_space(self, gen: BAxUSGenerator, row: int) -> torch.Tensor: + Z = gen.embedding.project(gen._normalized_inputs(gen.data)) + return torch.tensor(Z[row], dtype=torch.double) + + def test_box_is_centered_on_the_incumbent(self) -> None: + """argmin instead of argmax here would center on the worst point.""" + gen, best_row = self._gen_with_known_incumbent() + gen.trust_region.length = 0.4 + + bounds = gen._get_optimization_bounds() + center = (bounds[0] + bounds[1]) / 2.0 + + assert torch.allclose( + center, self._incumbent_in_target_space(gen, best_row), atol=1e-12 + ) + + def test_minimize_centers_on_the_smallest_objective(self) -> None: + """The incumbent is direction-aware, not just the numeric maximum.""" + vocs = VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, + objectives={"f": "MINIMIZE"}, + ) + gen = BAxUSGenerator( + vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=50 + ) + rows = [ + {"x0": 0.10, "x1": -0.05, "x2": 0.00, "x3": 0.10, "f": 4.0}, + {"x0": 0.05, "x1": 0.10, "x2": -0.05, "x3": 0.00, "f": -7.0}, # best + {"x0": 0.00, "x1": 0.00, "x2": 0.05, "x3": 0.05, "f": 2.0}, + ] + gen.add_data(pd.DataFrame(rows)) + gen.trust_region.length = 0.4 + + bounds = gen._get_optimization_bounds() + center = (bounds[0] + bounds[1]) / 2.0 + Z = gen.embedding.project(gen._normalized_inputs(gen.data)) + + assert torch.allclose( + center, torch.tensor(Z[1], dtype=torch.double), atol=1e-12 + ) + + @pytest.mark.parametrize("length", [0.1, 0.25, 0.4]) + def test_box_width_is_the_trust_region_length(self, length: float) -> None: + """In [-1, 1] target units the unclamped width is 2 * length (the length + fields are in [0, 1] units, per the BoTorch reference).""" + gen, _ = self._gen_with_known_incumbent() + gen.trust_region.length = length + + bounds = gen._get_optimization_bounds() + width = bounds[1] - bounds[0] + + assert torch.allclose( + width, torch.full_like(width, 2.0 * length), atol=1e-12 + ), f"expected width {2.0 * length}, got {width.tolist()}" + + def test_box_stays_inside_the_domain(self) -> None: + """A wide region around an incumbent at the edge must clamp, not overflow.""" + vocs = _simple_vocs(n_vars=4) + gen = BAxUSGenerator( + vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=50 + ) + gen.add_data( + pd.DataFrame( + [ + {"x0": 1.0, "x1": 1.0, "x2": 1.0, "x3": 1.0, "f": 5.0}, + {"x0": -1.0, "x1": -1.0, "x2": -1.0, "x3": -1.0, "f": -5.0}, + ] + ) + ) + gen.trust_region.length = gen.trust_region.length_max + + bounds = gen._get_optimization_bounds() + + assert bool((bounds[0] >= -1.0).all()) + assert bool((bounds[1] <= 1.0).all()) + assert bool((bounds[1] > bounds[0]).all()) + + def test_lengthscale_weights_stretch_the_box_per_dimension(self) -> None: + """A dimension with a longer lengthscale gets a wider side.""" + gen, _ = self._gen_with_known_incumbent() + gen.trust_region.length = 0.2 + + uniform = gen._get_optimization_bounds() + gen._lengthscale_weights = lambda model: torch.tensor( # type: ignore[method-assign] + [0.5, 2.0], dtype=torch.double + ) + weighted = gen._get_optimization_bounds() + + uniform_width = uniform[1] - uniform[0] + weighted_width = weighted[1] - weighted[0] + assert weighted_width[0] < uniform_width[0] + assert weighted_width[1] > uniform_width[1] + + +class TestBAxUSCandidateComesFromTheAcquisition: + """The BO candidate must be the acquisition optimum inside the trust region. + + No objective-value comparison can establish this. On a quadratic, the center + of a symmetric domain beats the average random point by exactly the variance + of a random draw - for *any* placement of the optimum - so a generator that + discarded the acquisition result and always proposed the center would still + "beat Sobol on average". The candidate has to be checked against the box the + acquisition was optimized over. + """ + + @staticmethod + def _gen_with_offset_incumbent( + z_star: np.ndarray | None = None, + ) -> tuple[BAxUSGenerator, VOCS]: + """Incumbent placed at a known, off-origin point of the target space.""" + vocs = _simple_vocs(n_vars=6) + gen = BAxUSGenerator( + vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=40 + ) + # lift an explicit target-space point so the incumbent projects back to it + if z_star is None: + z_star = np.array([[0.8, -0.7]]) + x_star = gen.embedding.lift(z_star)[0] + best = {name: float(v) for name, v in zip(vocs.variable_names, x_star)} + + rng = np.random.default_rng(0) + rows = [{**best, "f": 10.0}] + for _ in range(5): + other = { + name: float(v) + for name, v in zip(vocs.variable_names, rng.uniform(-1, 1, 6)) + } + rows.append({**other, "f": -5.0}) + gen.add_data(pd.DataFrame(rows)) + return gen, vocs + + def test_candidate_lies_inside_the_trust_region(self) -> None: + gen, vocs = self._gen_with_offset_incumbent() + + gen.generate(1) # fold in results and fit the model + gen.trust_region.length = 0.1 # narrow box, far from the origin + point = gen.generate(1)[0] + + # no data was added, so this reproduces the box the candidate came from + bounds = gen._get_optimization_bounds() + lb, ub = bounds[0].numpy(), bounds[1].numpy() + + assert not ((lb <= 0.0).all() and (ub >= 0.0).all()), ( + "box contains the origin, so this test could not detect a constant proposal" + ) + + z = gen.embedding.project(gen._normalized_inputs(pd.DataFrame([point])))[0] + assert (z >= lb - 1e-9).all(), f"candidate {z} below trust region {lb}" + assert (z <= ub + 1e-9).all(), f"candidate {z} above trust region {ub}" + + def test_candidate_tracks_the_incumbent(self) -> None: + """Move the incumbent, and the proposal must follow it. + + Lengthscale weights are pinned to 1 so the box is exactly ``length`` wide + around the incumbent. Left to a real fit, a handful of points produces + extreme ARD ratios that widen one side to the whole domain, and the + proposal is then free to sit anywhere along it. + """ + points = [] + for z_star in (np.array([[0.7, 0.7]]), np.array([[-0.7, -0.7]])): + gen, _ = self._gen_with_offset_incumbent(z_star) + + gen.generate(1) + gen.trust_region.length = 0.2 + gen._lengthscale_weights = lambda model: torch.ones( # type: ignore[method-assign] + 2, dtype=torch.double + ) + candidate = gen.generate(1)[0] + points.append( + gen.embedding.project( + gen._normalized_inputs(pd.DataFrame([candidate])) + )[0] + ) + + assert points[0][0] > points[1][0], ( + f"proposal did not follow the incumbent: {points[0]} vs {points[1]}" + ) + assert points[0][1] > points[1][1] + + +class TestBAxUSLengthscaleWeightNormalization: + """The volume-preserving normalization must survive high target_dim. + + Forming the raw product before taking the root underflows to zero once + target_dim reaches a few hundred, which makes every weight inf and silently + widens the trust region to the entire domain. + """ + + @pytest.mark.parametrize("target_dim", [8, 128, 600, 1200]) + def test_weights_are_finite_and_volume_preserving(self, target_dim: int) -> None: + torch.manual_seed(0) + lengthscales = ( + torch.distributions.LogNormal(0.0, 2.0).sample((target_dim,)).double() + ) + + weights = _normalize_lengthscale_weights(lengthscales, target_dim) + + assert bool(torch.isfinite(weights).all()), "weights underflowed to inf" + assert bool((weights > 0).all()) + # unit geometric mean + assert float(weights.log().mean().exp()) == pytest.approx(1.0, rel=1e-9) + + def test_ordering_is_preserved(self) -> None: + lengthscales = torch.tensor([0.5, 4.0, 1.0], dtype=torch.double) + weights = _normalize_lengthscale_weights(lengthscales, 3) + assert weights.argsort().tolist() == lengthscales.argsort().tolist() + + +class TestBAxUSExactScheduleConstants: + """The trust region halves and doubles by exactly a factor of two. + + Asserting only the direction (got smaller / got larger) leaves the schedule + free to drift, which changes how many failures it takes to trigger an + embedding expansion. + """ + + def test_failure_halves_the_length_exactly(self) -> None: + tr = BAxUSTrustRegion(target_dim=4, failure_tolerance=1, length=0.8) + tr.best_value = 10.0 + tr.update(torch.tensor([1.0])) # no improvement + assert tr.length == pytest.approx(0.4, rel=1e-12) + + def test_success_doubles_the_length_exactly(self) -> None: + tr = BAxUSTrustRegion(target_dim=4, success_tolerance=1, length=0.4) + tr.best_value = 0.0 + tr.update(torch.tensor([5.0])) # improvement + assert tr.length == pytest.approx(0.8, rel=1e-12) + + def test_growth_is_capped_at_length_max(self) -> None: + tr = BAxUSTrustRegion( + target_dim=4, success_tolerance=1, length=1.0, length_max=1.6 + ) + tr.best_value = 0.0 + for step in range(3): + tr.update(torch.tensor([10.0 * (step + 1)])) + assert tr.length == pytest.approx(1.6, rel=1e-12) + + def test_improvement_needs_to_clear_the_tolerance(self) -> None: + """A gain below 0.1% of the incumbent counts as a failure.""" + tr = BAxUSTrustRegion(target_dim=4, failure_tolerance=1, length=0.8) + tr.best_value = 100.0 + tr.update(torch.tensor([100.05])) # +0.05% -> not an improvement + assert tr.failure_counter == 0 # reset after the shrink fired + assert tr.length == pytest.approx(0.4, rel=1e-12) + + tr2 = BAxUSTrustRegion(target_dim=4, failure_tolerance=1, length=0.8) + tr2.best_value = 100.0 + tr2.update(torch.tensor([100.5])) # +0.5% -> an improvement + assert tr2.length == pytest.approx(0.8, rel=1e-12) + assert tr2.success_counter == 1 + + +class TestBAxUSSobolDomain: + """Seed points must cover the whole target cube, not just one orthant.""" + + def test_sobol_draws_span_negative_and_positive(self) -> None: + gen = BAxUSGenerator( + vocs=_simple_vocs(n_vars=6), + target_dim_init=2, + n_initial_sobol=32, + seed=0, + eval_budget=60, + ) + zs = np.vstack([gen._draw_sobol_point() for _ in range(32)]) + + assert zs.min() < -0.2, f"no negative seed coordinates (min {zs.min():.3f})" + assert zs.max() > 0.2, f"no positive seed coordinates (max {zs.max():.3f})" + assert (zs >= -1.0).all() and (zs <= 1.0).all() + assert abs(float(zs.mean())) < 0.25, "seed points are not centered" + + +class TestBAxUSExpansionRandomization: + """Splitting must randomize which dimensions land in which sub-bin. + + The reference permutes the contributing dimensions before splitting. Without + it the split is a fixed function of column index, so two dimensions that are + adjacent in index and share a bin are never separated until the embedding + reaches full dimensionality - which is exactly where BAxUS has no advantage + left. That failure is invisible to structural checks: the matrix is still a + valid partition either way. + """ + + INPUT_DIM = 64 + B = 3 + + @staticmethod + def _bin_of(emb: BAxUSEmbedding, col: int) -> int: + matrix = np.asarray(emb.matrix) + return int(np.nonzero(matrix[:, col])[0][0]) + + def _collision_rate(self, target_dim: int, trials: int = 200) -> float: + """How often input dims 0 and 1 (adjacent in index) share a bin.""" + collisions = 0 + for trial in range(trials): + emb = BAxUSEmbedding.create( + input_dim=self.INPUT_DIM, target_dim=2, seed=trial + ) + while emb.target_dim < target_dim: + emb = emb.expand(self.B, seed=trial * 7919 + emb.target_dim) + collisions += self._bin_of(emb, 0) == self._bin_of(emb, 1) + return collisions / trials + + def test_expansion_separates_adjacent_dimensions(self) -> None: + """Collision rate must fall as the embedding grows. + + A deterministic split leaves it pinned at P(same initial bin) = 1/2. + """ + rate_8 = self._collision_rate(target_dim=8) + rate_32 = self._collision_rate(target_dim=32) + + assert rate_8 < 0.30, f"dims 0/1 still collide {rate_8:.0%} of the time at d=8" + assert rate_32 < rate_8, ( + f"expanding did not decorrelate the split: {rate_8:.3f} -> {rate_32:.3f}" + ) + + def test_expand_is_reproducible_for_a_given_seed(self) -> None: + emb = BAxUSEmbedding.create(input_dim=self.INPUT_DIM, target_dim=2, seed=0) + assert emb.expand(self.B, seed=7) == emb.expand(self.B, seed=7) + + def test_different_seeds_give_different_partitions(self) -> None: + emb = BAxUSEmbedding.create(input_dim=self.INPUT_DIM, target_dim=2, seed=0) + assert emb.expand(self.B, seed=1) != emb.expand(self.B, seed=2) + + @pytest.mark.parametrize("seed", [0, 1, 2]) + def test_randomized_expansion_keeps_the_matrix_well_formed(self, seed: int) -> None: + """Every input dimension stays assigned to exactly one bin, with its sign.""" + emb = BAxUSEmbedding.create(input_dim=self.INPUT_DIM, target_dim=2, seed=seed) + original = np.asarray(emb.matrix) + original_sign = { + col: original[:, col][np.nonzero(original[:, col])[0][0]] + for col in range(self.INPUT_DIM) + } + + while emb.target_dim < self.INPUT_DIM: + emb = emb.expand(self.B, seed=seed * 31 + emb.target_dim) + matrix = np.asarray(emb.matrix) + + assert (np.count_nonzero(matrix, axis=0) == 1).all(), ( + "column lost/duplicated" + ) + assert (np.count_nonzero(matrix, axis=1) > 0).all(), "empty bin" + for col in range(self.INPUT_DIM): + nz = np.nonzero(matrix[:, col])[0][0] + assert matrix[nz, col] == original_sign[col], "sign changed on split" + + # at full dimensionality the embedding is a signed permutation + assert emb.target_dim == self.INPUT_DIM + assert (np.count_nonzero(np.asarray(emb.matrix), axis=1) == 1).all() + + def test_generator_expansion_survives_a_round_trip(self) -> None: + """The expansion seed is derived from (seed, n_expansions), so a run + restored mid-flight expands to the same embedding.""" + vocs = _simple_vocs(n_vars=16) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=3, eval_budget=60) + _seed_sobol(gen) + + restored = _round_trip(gen, with_data=True) + + for generator in (gen, restored): + _force_expansion(generator) + generator.add_data( + pd.DataFrame([{**{f"x{i}": 0.1 for i in range(16)}, "f": -1.0}]) + ) + generator.generate(1) + + assert gen.n_expansions == 1 + assert restored.embedding == gen.embedding + + +class TestBAxUSCheckpointIntegrity: + """Saving and reloading through Xopt must not advance the algorithm. + + Xopt's data validator hands a restored dataframe to the generator; for a + generator that owns state, that must not replay the history it already + reflects. BAxUS is a StateOwner precisely so this is a no-op. + """ + + @staticmethod + def _state(gen: BAxUSGenerator) -> dict[str, Any]: + tr = gen.trust_region + return { + "target_dim": gen.embedding.target_dim, + "matrix": gen.embedding.matrix, + "length": tr.length, + "success_counter": tr.success_counter, + "failure_counter": tr.failure_counter, + "failure_tolerance": tr.failure_tolerance, + "best_value": tr.best_value, + "sobol_draws": gen.sobol_draws, + "n_expansions": gen.n_expansions, + "tr_observed_rows": gen.tr_observed_rows, + } + + def _run(self, steps: int = 12) -> Xopt: + vocs = _simple_vocs(n_vars=6) + X = Xopt( + evaluator=XoptEvaluator(function=_sphere), + generator=BAxUSGenerator( + vocs=vocs, target_dim_init=2, seed=0, eval_budget=40 + ), + ) + for _ in range(steps): + X.step() + return X + + def test_is_a_state_owner(self) -> None: + assert isinstance(BAxUSGenerator(vocs=_simple_vocs(), seed=0), StateOwner) + + def test_repeated_checkpoint_cycles_do_not_move_the_trust_region(self) -> None: + """No evaluations happen between cycles, so nothing may change.""" + X = self._run() + expected = self._state(X.generator) + assert len(X.data) == 12 + + for cycle in range(6): + X = Xopt.from_yaml(X.yaml()) + assert len(X.data) == 12, f"data changed on cycle {cycle}" + assert self._state(X.generator) == expected, ( + f"generator state drifted on checkpoint cycle {cycle}" + ) + + def test_resume_folds_in_only_the_new_results(self) -> None: + """The pre-restore history must not be re-ingested. + + A replay would re-consume everything from the seed quota onward; a + correct resume advances the trust region by exactly the one row that was + evaluated after the checkpoint. + """ + X = self._run() + before = self._state(X.generator) + + resumed = Xopt.from_yaml(X.yaml()) + resumed.step() + after = self._state(resumed.generator) + + assert after["tr_observed_rows"] == before["tr_observed_rows"] + 1 + assert after["n_expansions"] == before["n_expansions"] + # one tick only: at most one counter moved, by one + assert abs(after["failure_counter"] - before["failure_counter"]) <= 1 + assert abs(after["success_counter"] - before["success_counter"]) <= 1 + + def test_resumed_run_continues_the_same_trajectory(self) -> None: + """One more step on the original and on the restored run agree. + + The proposed point is not compared exactly: xopt's yaml writer keeps 10 + significant digits, and that ~1e-10 perturbation of the training data is + amplified by GP hyperparameter fitting and the L-BFGS acquisition + optimization. The algorithm state is insensitive to it. + """ + X = self._run() + resumed = Xopt.from_yaml(X.yaml()) + + for runner in (X, resumed): + runner.step() + + original, restored = self._state(X.generator), self._state(resumed.generator) + for key in ( + "target_dim", + "matrix", + "length", + "success_counter", + "failure_counter", + "failure_tolerance", + "sobol_draws", + "n_expansions", + "tr_observed_rows", + ): + assert original[key] == restored[key], f"{key} diverged after resume" + assert original["best_value"] == pytest.approx(restored["best_value"], rel=1e-6) + + +class TestBAxUSFailedEvaluations: + """An evaluator that raises returns a row with no objective column at all.""" + + def test_add_data_tolerates_a_missing_objective_column(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + + failed = {f"x{i}": 0.1 for i in range(4)} + failed["xopt_error"] = True + gen.add_data(pd.DataFrame([failed])) + + assert len(gen.data) == 1 + assert gen._finite_data().empty + + def test_xopt_run_survives_a_failing_evaluator(self) -> None: + """The generator must not be the reason a non-strict run dies.""" + vocs = _simple_vocs(n_vars=4) + calls = {"n": 0} + + def flaky(inputs: dict[str, float]) -> dict[str, float]: + calls["n"] += 1 + if calls["n"] == 4: + raise RuntimeError("evaluation blew up") + return _sphere(inputs) + + X = Xopt( + evaluator=XoptEvaluator(function=flaky), + generator=BAxUSGenerator( + vocs=vocs, target_dim_init=2, seed=0, eval_budget=30 + ), + strict=False, + ) + for _ in range(8): + X.step() + + assert len(X.data) == 8 + assert len(X.generator._finite_data()) == 7 # the failed row is excluded + + +class TestBAxUSObservables: + """The target-space model has a single outcome while the inherited + acquisition scalarizes over every vocs output, so observables cannot work.""" + + def test_observables_rejected_at_construction(self) -> None: + vocs = VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, + objectives={"f": "MAXIMIZE"}, + observables=["g"], + ) + with pytest.raises(ValidationError, match="does not support observables"): + BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) + + +class TestBAxUSEvalBudgetWarning: + """Without eval_budget the reference expansion schedule is unavailable, so + construction warns. It must warn exactly once: the generator assigns to its + own fields on every BO step, and a model validator of either mode re-fires + on assignment under validate_assignment, which would bury the run in + duplicate warnings.""" + + def test_warns_once_on_construction(self) -> None: + with pytest.warns(GeneratorWarning, match="eval_budget is not set") as caught: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + assert len(caught) == 1 + + with warnings.catch_warnings(record=True) as on_assignment: + warnings.simplefilter("always") + gen.sobol_draws = 3 + gen.n_expansions = 1 + assert [w for w in on_assignment if w.category is GeneratorWarning] == [] + + def test_no_warning_when_budget_is_set(self) -> None: + vocs = _simple_vocs(n_vars=4) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0, eval_budget=50) + assert [w for w in caught if w.category is GeneratorWarning] == [] + + +class TestBAxUSComputationTime: + """generate() must record training/acquisition-optimization timing during + the BO phase, matching the base class's bookkeeping.""" + + def test_bo_generate_records_one_row(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + _seed_sobol(gen) + + assert gen.computation_time is None + _bo_step(gen) + + assert isinstance(gen.computation_time, pd.DataFrame) + assert len(gen.computation_time) == 1 + assert list(gen.computation_time.columns) == [ + "training", + "acquisition_optimization", + ] + + def test_second_bo_generate_appends_row(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + _seed_sobol(gen) + + _bo_step(gen) + _bo_step(gen) + + assert len(gen.computation_time) == 2 From f889fb5fd89f5d267e893b0d4d0cd3907f3bcf30 Mon Sep 17 00:00:00 2001 From: Gianluca Martino Date: Fri, 24 Jul 2026 16:42:06 -0700 Subject: [PATCH 2/3] Clear notebook outputs and formatting in the BAxUS example --- .../single_objective_bayes_opt/baxus.ipynb | 44 ++++++++++++------- 1 file changed, 28 insertions(+), 16 deletions(-) diff --git a/docs/examples/single_objective_bayes_opt/baxus.ipynb b/docs/examples/single_objective_bayes_opt/baxus.ipynb index 50cf9da57..47a8d5f93 100644 --- a/docs/examples/single_objective_bayes_opt/baxus.ipynb +++ b/docs/examples/single_objective_bayes_opt/baxus.ipynb @@ -60,7 +60,7 @@ " variables={f\"x{i}\": [-1.0, 1.0] for i in range(N_DIM)},\n", " objectives={\"f\": \"MINIMIZE\"},\n", ")\n", - "vocs\n" + "vocs" ] }, { @@ -80,7 +80,7 @@ ")\n", "evaluator = Evaluator(function=sphere_embedded)\n", "X = Xopt(evaluator=evaluator, generator=generator)\n", - "X\n" + "X" ] }, { @@ -109,7 +109,7 @@ "\n", "print(f\"best f: {X.data['f'].min():.3e}\")\n", "print(f\"embedding grew: {sorted(set(target_dims))}\")\n", - "X.data[[\"f\"]].tail()\n" + "X.data[[\"f\"]].tail()" ] }, { @@ -136,7 +136,7 @@ " for ax in (ax0, ax1):\n", " ax.axvline(i, color=\"grey\", ls=\"--\", lw=0.8)\n", "ax0.set_title(\"BAxUS: convergence and subspace growth\")\n", - "plt.tight_layout()\n" + "plt.tight_layout()" ] }, { @@ -162,30 +162,43 @@ "checkpoint = X.yaml()\n", "resumed = Xopt.from_yaml(checkpoint)\n", "\n", - "print(\"target_dim :\", X.generator.embedding.target_dim,\n", - " \"->\", resumed.generator.embedding.target_dim)\n", - "print(\"tr length :\", X.generator.trust_region.length,\n", - " \"->\", resumed.generator.trust_region.length)\n", - "print(\"rows folded:\", X.generator.tr_observed_rows,\n", - " \"->\", resumed.generator.tr_observed_rows)\n", + "print(\n", + " \"target_dim :\",\n", + " X.generator.embedding.target_dim,\n", + " \"->\",\n", + " resumed.generator.embedding.target_dim,\n", + ")\n", + "print(\n", + " \"tr length :\",\n", + " X.generator.trust_region.length,\n", + " \"->\",\n", + " resumed.generator.trust_region.length,\n", + ")\n", + "print(\n", + " \"rows folded:\",\n", + " X.generator.tr_observed_rows,\n", + " \"->\",\n", + " resumed.generator.tr_observed_rows,\n", + ")\n", "\n", "resumed.step() # continues the run\n", - "len(resumed.data)\n" + "len(resumed.data)" ] }, { "cell_type": "markdown", + "id": "fbb00bc6", "metadata": {}, "source": [ "Dumping the generator on its own works too, but a trained generator's dump\n", "carries the live botorch `model` and a `computation_time` frame, neither of\n", "which is YAML-safe -- drop both before serializing by hand.\n" - ], - "id": "fbb00bc6" + ] }, { "cell_type": "code", "execution_count": null, + "id": "b36961c9", "metadata": {}, "outputs": [], "source": [ @@ -194,9 +207,8 @@ "dump = X.generator.model_dump()\n", "dump.pop(\"model\", None)\n", "dump.pop(\"computation_time\", None)\n", - "yaml.safe_dump(dump)[:200]\n" - ], - "id": "b36961c9" + "yaml.safe_dump(dump)[:200]" + ] } ], "metadata": { From 012044748e8dd3130b58ae21eaefe9981a276293 Mon Sep 17 00:00:00 2001 From: Gianluca Martino Date: Mon, 27 Jul 2026 11:45:06 -0700 Subject: [PATCH 3/3] Simplify the BAxUS generator and consolidate its tests without losing detection --- xopt/generators/bayesian/baxus.py | 439 ++-- xopt/tests/generators/bayesian/test_baxus.py | 2237 +++++++----------- 2 files changed, 949 insertions(+), 1727 deletions(-) diff --git a/xopt/generators/bayesian/baxus.py b/xopt/generators/bayesian/baxus.py index e10a00661..7d1a749ac 100644 --- a/xopt/generators/bayesian/baxus.py +++ b/xopt/generators/bayesian/baxus.py @@ -1,39 +1,26 @@ """BAxUS generator - Bayesian Optimization with Adaptively Expanding Subspaces. -Starts in a low-dimensional random subspace and gradually expands it, making -it effective when many input dimensions are irrelevant. - -The inherited ``data`` frame is the single source of truth; all persistent -state lives in serializable pydantic fields, and the generator is a -``StateOwner`` so ``Xopt.from_yaml`` reattaches data without replaying the -trust-region history it already reflects. The trust region advances once per -``generate`` call (as ``TurboController`` does), which keeps the trajectory -independent of how results were batched into ``add_data``. - -Randomness is seeded rather than removed: the embedding's creation and each of -its expansions draw from ``seed`` combined with the expansion count, so a -dump -> construct round trip reproduces the same subspaces. - -The BO phase is ``BayesianGenerator.generate`` re-based into the embedded target -space through four override seams - ``get_training_data`` (finite-objective rows -only), ``train_model`` (fit on projections), ``_get_torch_bounds`` / -``_get_optimization_bounds`` (target-space boxes), and ``_process_candidates`` -(lift back to vocs space). Everything else in that pipeline is inert here, -because the options it serves are rejected at construction. - -Deviations from the reference: - - ``target_dim_init`` is a user knob rather than being derived from the - input dimensionality. - - The acquisition is the inherited analytic LogEI over the trust-region box, - not Thompson sampling over a masked discrete candidate set. - - Data not generated through the current embedding is least-squares - projected into target space. - -Reference: - Papenmeier et al., "Increasing the Scope as You Learn: Adaptive Bayesian - Optimization in Nested Subspaces", NeurIPS 2022. - https://arxiv.org/abs/2304.11468 - Tutorial: https://botorch.org/docs/tutorials/baxus +Optimizes in a low-dimensional random embedding of the input space and expands it +whenever the trust region shrinks below a threshold, which makes it effective when +many input dimensions are irrelevant. + +The analytic LogEI acquisition is inherited from ``ExpectedImprovementGenerator`` +and re-based into the embedded target space through ``get_training_data``, +``train_model``, ``_get_torch_bounds`` / ``_get_optimization_bounds`` and +``_process_candidates``. All persistent state lives in serializable pydantic +fields and the generator is a ``StateOwner``, so a dump -> construct round trip +reproduces the run. + +Deviations from the reference: ``target_dim_init`` is a user knob rather than +derived from the input dimensionality, the acquisition is the inherited analytic +LogEI over the trust-region box rather than Thompson sampling over a masked +candidate set, and data not generated through the current embedding is +least-squares projected into target space. + +Papenmeier et al., "Increasing the Scope as You Learn: Adaptive Bayesian +Optimization in Nested Subspaces", NeurIPS 2022 +https://arxiv.org/abs/2304.11468 +Tutorial: https://botorch.org/docs/tutorials/baxus """ import logging @@ -49,8 +36,6 @@ PositiveInt, PrivateAttr, SerializeAsAny, - ValidationInfo, - field_validator, model_validator, ) from torch.nn import Module @@ -72,56 +57,35 @@ def _normalize_lengthscale_weights( lengthscales: torch.Tensor, target_dim: int ) -> torch.Tensor: - """Scale lengthscales to mean 1, then unit geometric mean (volume-preserving). + """Scale lengthscales to mean 1, then to unit geometric mean. The geometric mean is taken as a product of per-element roots (the reference - form). Forming ``weights.prod()`` first underflows to zero once target_dim - reaches a few hundred, which would make every weight ``inf`` and silently - widen the trust region to the whole domain. + form); forming the product first underflows to zero past a few hundred + dimensions, making every weight inf. """ weights = lengthscales / lengthscales.mean() return weights / weights.pow(1.0 / target_dim).prod() class BAxUSTrustRegion(XoptBaseModel): - """Serializable trust-region state, in normalized [0, 1] target-space units. + """Trust-region state, in normalized [0, 1] target-space units. - ``update`` takes weighted objective values (MAXIMIZE -> +y, MINIMIZE -> -y), - so higher is always better. ``best_value=None`` means no BO-phase - evaluation has been observed yet. + The generator owns dimensionality (``embedding.target_dim``) and passes its + derived failure tolerance into every ``update``. """ - target_dim: PositiveInt length: float = 0.8 - length_init: float = 0.8 length_min: float = 0.5**7 length_max: float = 1.6 failure_counter: int = 0 success_counter: int = 0 success_tolerance: int = 3 - failure_tolerance: PositiveInt = Field( - default=0, - validate_default=True, - description=( - "Failures before shrinking; 0 derives ceil(target_dim / 2), re-derived " - "on every BO-phase ingest (budget-aware when eval_budget is set)" - ), - ) + # None until a BO-phase evaluation is folded in best_value: float | None = None restart_triggered: bool = False - @field_validator("failure_tolerance", mode="before") - @classmethod - def _default_failure_tolerance(cls, value: Any, info: ValidationInfo) -> Any: - """0 (the field default) means "derive from target_dim".""" - if value: - return value - # absent only if target_dim itself failed validation; PositiveInt then rejects 0 - target_dim = info.data.get("target_dim") - return math.ceil(target_dim / 2) if target_dim else value - - def update(self, y_weighted: torch.Tensor) -> None: - """Adjust the trust region from a batch of weighted objective values.""" + def update(self, y_weighted: torch.Tensor, failure_tolerance: int) -> None: + """Adjust the trust region from weighted objective values (higher is better).""" best = self.best_value y_max = float(y_weighted.max()) improved = best is None or y_max > best + 1e-3 * abs(best) @@ -135,7 +99,7 @@ def update(self, y_weighted: torch.Tensor) -> None: if self.success_counter >= self.success_tolerance: self.length = min(2.0 * self.length, self.length_max) self.success_counter = 0 - elif self.failure_counter >= self.failure_tolerance: + elif self.failure_counter >= failure_tolerance: self.length /= 2.0 self.failure_counter = 0 @@ -146,11 +110,11 @@ def update(self, y_weighted: torch.Tensor) -> None: class BAxUSEmbedding(XoptBaseModel): - """Serializable sparse random embedding. + """Sparse random embedding of shape (target_dim, input_dim). - ``matrix`` has shape (target_dim, input_dim) with exactly one signed unit - entry per column. Both ``create`` and ``expand`` consume randomness; each - takes an explicit seed so a run can be reproduced from a dump. + Exactly one signed unit entry per column. ``create`` and ``expand`` both + consume randomness and take an explicit seed, so a run can be reproduced from + a dump. """ matrix: list[list[float]] @@ -175,31 +139,26 @@ def target_dim(self) -> int: def input_dim(self) -> int: return len(self.matrix[0]) - def _as_array(self) -> np.ndarray: - return np.asarray(self.matrix, dtype=np.float64) - def lift(self, Z: np.ndarray) -> np.ndarray: """Lift from target subspace to full input space: ``X = Z @ S``.""" - return Z @ self._as_array() + return Z @ np.asarray(self.matrix, dtype=np.float64) def project(self, X: np.ndarray) -> np.ndarray: """Project from full input space to target subspace via the pseudo-inverse.""" - return X @ np.linalg.pinv(self._as_array()) + return X @ np.linalg.pinv(np.asarray(self.matrix, dtype=np.float64)) def expand( self, new_bins_on_split: int, seed: int | None = None ) -> "BAxUSEmbedding": - """Split each bin into up to ``new_bins_on_split + 1`` sub-bins, preserving signs. - - The dimensions contributing to a bin are randomly permuted before being - split, matching the reference implementation - without it the split is a - fixed function of column index, so dimensions that are adjacent in index - and share a bin are never separated before full dimensionality. Pass a - ``seed`` to keep an expansion reproducible across a dump -> construct - round trip. + """Split each bin into up to ``new_bins_on_split + 1`` sub-bins, keeping signs. + + Contributing dimensions are permuted before the split, matching the + reference - otherwise the split is a fixed function of column index, so + dimensions adjacent in index that share a bin are never separated before + full dimensionality. """ rng = np.random.default_rng(seed) - S = self._as_array() + S = np.asarray(self.matrix, dtype=np.float64) old_target_dim, input_dim = S.shape n_sub = new_bins_on_split + 1 new_target_dim = min(old_target_dim * n_sub, input_dim) @@ -215,22 +174,8 @@ def expand( class BAxUSGenerator(ExpectedImprovementGenerator, StateOwner): - """Bayesian Optimization with Adaptively Expanding Subspaces (BAxUS). - - Operates in a low-dimensional random embedding of the input space and - gradually expands it when the trust region shrinks below a threshold. - The acquisition (analytic LogEI) is inherited from - ``ExpectedImprovementGenerator``; only the model/bounds seams are re-based - into the embedded target space. - - This is a ``StateOwner``: the trust region and embedding are advanced once - per ``generate`` call, so restoring a saved run reattaches data without - replaying (and corrupting) the state that was just deserialized. - """ - name = "baxus" supports_batch_generation: bool = False - supports_single_objective: bool = True # the embedded target-space GP models only the objective over continuous # variables; constrained, discrete, observable, and contextual vocs are rejected supports_constraints: bool = False @@ -248,71 +193,69 @@ class BAxUSGenerator(ExpectedImprovementGenerator, StateOwner): target_dim_init: PositiveInt = Field( default=2, - description="Initial target dimensionality of the embedding", + description="initial target dimensionality of the embedding", ) n_initial_sobol: int = Field( default=0, ge=0, - description="Number of initial Sobol points (0 = auto-computed)", + description="number of initial Sobol points (0 = auto-computed)", ) length_init: float = Field( default=0.8, - description="Initial trust-region side length", + description="initial trust-region side length", ) new_bins_on_split: PositiveInt = Field( default=3, - description="Number of new bins created per existing bin on expansion", + description="number of new bins created per existing bin on expansion", ) seed: int | None = Field( default=None, description=( - "Random seed for the embedding and Sobol initialization. GP fitting and " - "acquisition optimization use torch's global RNG and are not covered by it." + "random seed for the embedding and Sobol initialization; GP fitting and " + "acquisition optimization use torch's global RNG instead" ), ) eval_budget: PositiveInt | None = Field( default=None, description=( - "Total planned evaluations for the run, seed points included (the seed " - "quota is subtracted internally). Enables the reference budget-aware " - "failure tolerance; None keeps the ceil(target_dim / 2) heuristic." + "total planned evaluations for the run, seed points included; enables the " + "reference budget-aware failure tolerance, None keeps ceil(target_dim / 2)" ), ) embedding: BAxUSEmbedding = Field( - description="Sparse random embedding (auto-created from vocs and seed when omitted)", + description="sparse random embedding (auto-created from vocs and seed when omitted)", ) trust_region: BAxUSTrustRegion = Field( - description="Trust-region state (auto-created when omitted)", + description="trust-region state (auto-created when omitted)", ) sobol_draws: int = Field( default=0, ge=0, - description="Number of Sobol seed points drawn so far (used to resume the sequence)", + description="number of Sobol seed points drawn so far, used to resume the sequence", ) n_expansions: int = Field( default=0, ge=0, description=( - "Number of embedding expansions performed so far; combined with seed to " - "keep each expansion's randomization reproducible across a round trip" + "number of embedding expansions so far, combined with seed to keep each " + "expansion reproducible across a round trip" ), ) tr_observed_rows: int = Field( default=0, ge=0, description=( - "Number of finite-objective rows already folded into the trust region " - "(prevents re-ingesting history on resume)" + "number of finite-objective rows already folded into the trust region, " + "which prevents re-ingesting history on resume" ), ) - # a plain default instance, matching BayesianGenerator: pydantic deep-copies - # it per instance, and get_generator_defaults only understands `default` + # plain default instance, matching BayesianGenerator: pydantic deep-copies it + # per instance, and get_generator_defaults only understands `default` gp_constructor: SerializeAsAny[ModelConstructor] = Field( StandardModelConstructor(use_low_noise_prior=True), description=( - "Constructor used to generate the target-space model. Defaults to a " - "low-noise prior, matching the near-interpolating GP the reference " - "trust-region logic assumes" + "constructor used to generate the target-space model; the low-noise prior " + "matches the near-interpolating GP the reference trust-region logic assumes" ), ) @@ -321,11 +264,9 @@ class BAxUSGenerator(ExpectedImprovementGenerator, StateOwner): @model_validator(mode="before") @classmethod def _init_components(cls, values: Any) -> Any: - """Fill in the vocs-derived defaults on fresh construction. + """Create the embedding and trust region, and size the Sobol seed quota. - Creates the embedding and trust region and sizes the Sobol seed quota - from the embedding. A dump -> construct round trip supplies the - components and skips their creation. + A dump -> construct round trip supplies these and skips their creation. """ if not isinstance(values, dict) or values.get("vocs") is None: return values @@ -352,23 +293,16 @@ def _init_components(cls, values: Any) -> Any: values["n_initial_sobol"] = max(2, emb_target_dim + 1) if values.get("trust_region") is None: - length_init = float(values.get("length_init") or 0.8) values["trust_region"] = BAxUSTrustRegion( - target_dim=emb_target_dim, - length=length_init, - length_init=length_init, + length=float(values.get("length_init") or 0.8) ) return values def model_post_init(self, context: Any, /) -> None: - """Warn once that the reference expansion schedule is disabled. - - This hook runs exactly once per construction. A model validator of - either mode would instead re-fire on every field assignment under - ``validate_assignment``, repeating the warning on each trust-region - update and embedding expansion. - """ + """Warn once if the budget is missing.""" + # runs exactly once per construction; a model validator of either mode + # would re-fire on every field assignment under validate_assignment super().model_post_init(context) if self.eval_budget is None: warnings.warn( @@ -408,23 +342,14 @@ def _reject_unsupported_options(self) -> "BAxUSGenerator": return self def generate(self, n_candidates: int = 1) -> list[dict[str, float]]: - """Sobol seed points first, then LogEI within the trust region. - - The BO phase defers to ``BayesianGenerator.generate`` through the - ``get_training_data`` / ``_process_candidates`` / ``_get_torch_bounds`` - seams, so only the seed branch and the trust-region advance are - BAxUS-specific. Only the BO phase records ``computation_time`` rows; - the Sobol seed phase trains no model. - """ - # the guard is repeated from the base call because the seed branch - # below returns without ever reaching it + """Sobol seed points first, then LogEI within the trust region.""" + # repeated from the base call, which the seed branch below never reaches if n_candidates > 1 and not self.supports_batch_generation: raise NotImplementedError( "This Bayesian algorithm does not currently support parallel candidate generation" ) if len(self._finite_data()) < self.n_initial_sobol: - self.n_candidates = n_candidates z = torch.tensor(self._draw_sobol_point()) return self._process_candidates(z).to_dict("records") @@ -440,9 +365,8 @@ def get_training_data(self, data: pd.DataFrame) -> pd.DataFrame: def _process_candidates(self, candidates: torch.Tensor) -> pd.DataFrame: """Lift target-space candidates into vocs space. - Replaces the base implementation wholesale: its discrete snapping and - fixed-feature handling are keyed to vocs-space columns, and BAxUS - rejects both options at construction. + Replaces the base implementation, whose discrete snapping and + fixed-feature handling are keyed to vocs-space columns. """ x_raw = self._lift_to_vocs(candidates.detach().cpu().numpy()) return convert_numpy_to_inputs(self.vocs, x_raw, include_constants=False) @@ -450,24 +374,15 @@ def _process_candidates(self, candidates: torch.Tensor) -> pd.DataFrame: def _get_torch_bounds(self) -> torch.Tensor: """The full embedded target space, ``[-1, 1]^target_dim``. - This deliberately re-points a base-class seam from vocs space into - target space. It is safe because BAxUS overrides or rejects every other - consumer: ``_get_optimization_bounds`` below, ``max_travel_distances``, - and ``visualize_model``. + Re-points a base-class seam from vocs space into target space; safe + because every other consumer is overridden or rejected. """ ones = torch.ones(self.embedding.target_dim, dtype=torch.double) return torch.stack([-ones, ones]) def get_optimum(self) -> pd.DataFrame: - """Select the best point given by the model's posterior mean. - - Optimizes over the full embedded target space and lifts the result to - vocs space. The base implementation's constraint, fixed-feature, - contextual, and discrete handling is inert here because BAxUS rejects - all of those at construction. - """ - # an expansion clears the model, and a freshly restored run has none; - # retrain so the model always matches the current embedding + """Best point given by the model's posterior mean, lifted to vocs space.""" + # an expansion clears the model, and a freshly restored run has none if self.model is None: self.train_model() return super().get_optimum() @@ -482,14 +397,10 @@ def visualize_model(self, **kwargs): def add_data(self, new_data: pd.DataFrame) -> None: """Ingest evaluation results. - Pure ingestion - the trust region is advanced in ``generate`` instead, so - the state machine follows optimization iterations rather than how results - were chunked into ``add_data`` calls. - - Non-finite objective rows stay in ``data`` but are excluded from the GP - training set and the trust region. A batch missing the objective column - entirely (what an evaluator returns for a failed evaluation under - ``strict=False``) is ingested the same way. + Pure ingestion - the trust region advances in ``generate`` instead, so the + state machine follows optimization iterations rather than ``add_data`` + chunking. Rows without a finite objective stay in ``data`` but are kept + out of the GP training set and the trust region. """ super().add_data(new_data) @@ -503,25 +414,18 @@ def add_data(self, new_data: pd.DataFrame) -> None: def set_data(self, data: pd.DataFrame) -> None: """Reattach a full dataset without replaying trust-region updates. - ``Xopt`` calls this instead of ``add_data`` when loading a saved run - (see ``StateOwner``), so the deserialized embedding and trust region - survive a checkpoint round trip intact. + ``Xopt`` calls this instead of ``add_data`` when loading a saved run (see + ``StateOwner``). """ self.data = data - def _expansion_seed(self) -> int | None: - """Seed for the next expansion, derived so a round trip reproduces it.""" - if self.seed is None: - return None - return self.seed + 104729 * (self.n_expansions + 1) - def _advance_trust_region(self) -> None: """Fold newly observed BO-phase results into the trust region. Called once per ``generate``, mirroring how ``BayesianGenerator`` drives - ``TurboController.update_state``. Sobol seed points never move the trust - region (matches the reference), and ``tr_observed_rows`` records how much - history has already been applied so a resumed run does not re-ingest it. + ``TurboController.update_state``. Sobol seed points never move the region + (matching the reference), and ``tr_observed_rows`` keeps a resumed run + from re-ingesting history. """ data = self._finite_data() n_finite = len(data) @@ -532,16 +436,19 @@ def _advance_trust_region(self) -> None: y = data[self._objective_name].to_numpy(dtype=np.float64)[start:] self.tr_observed_rows = n_finite - # refresh lazily so an eval_budget assigned after construction takes effect - self._refresh_failure_tolerance() self.trust_region.update( - torch.tensor(self._objective_weight() * y, dtype=torch.double) + torch.tensor(self._objective_weight() * y, dtype=torch.double), + self._failure_tolerance(), ) if self.trust_region.restart_triggered: previous_dim = self.embedding.target_dim self.embedding = self.embedding.expand( - self.new_bins_on_split, seed=self._expansion_seed() + self.new_bins_on_split, + # derived seed, so a round trip reproduces the expansion + seed=None + if self.seed is None + else self.seed + 104729 * (self.n_expansions + 1), ) logger.info( "BAxUS: trust region too small (%.4e), expanded embedding from %d to %d dims", @@ -550,13 +457,11 @@ def _advance_trust_region(self) -> None: self.embedding.target_dim, ) self.n_expansions += 1 - self.trust_region = BAxUSTrustRegion( - target_dim=self.embedding.target_dim, - length=self.length_init, - length_init=self.length_init, - best_value=self.trust_region.best_value, - ) - self._refresh_failure_tolerance() + # the reference reset, minus its counter zeroing: the shrink that + # trips a restart has just reset both counters, so only length and + # the flag can differ - and best_value deliberately survives + self.trust_region.length = self.length_init + self.trust_region.restart_triggered = False # the fitted model lives in the old target space self.model = None @@ -568,46 +473,33 @@ def train_model( if data.empty: raise ValueError("no data available to build model") - Z = self._target_space(data) + Z = self.embedding.project(self._normalized_inputs(data)) input_names = [f"z{i}" for i in range(self.embedding.target_dim)] objective_name = self._objective_name frame = pd.DataFrame(Z, columns=input_names) frame[objective_name] = data[objective_name].to_numpy(dtype=np.float64) # xopt's constructor downgrades a failed fit to a warning and returns an - # untrained model (expected occasionally right after an embedding - # expansion); re-emit captured warnings and log the degradation - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - model = self.gp_constructor.build_model( - input_names=input_names, - outcome_names=[objective_name], - data=frame, - input_bounds={name: [-1.0, 1.0] for name in input_names}, - **self.tkwargs, - ) - for caught_warning in caught: - warnings.warn_explicit( - caught_warning.message, - caught_warning.category, - caught_warning.filename, - caught_warning.lineno, - ) - if any("Model fitting failed" in str(w.message) for w in caught): - logger.warning( - "BAxUS: GP fit failed at target_dim=%d - proceeding with an untrained model for this iteration", - self.embedding.target_dim, - ) + # untrained model, which is survivable and happens occasionally right + # after an embedding expansion + model = self.gp_constructor.build_model( + input_names=input_names, + outcome_names=[objective_name], + data=frame, + input_bounds={name: [-1.0, 1.0] for name in input_names}, + **self.tkwargs, + ) if update_internal: self.model = model return model def _get_optimization_bounds(self) -> torch.Tensor: - """Trust-region bounds in target space [-1, 1]^target_dim. + """Trust-region bounds in target space, ``[-1, 1]^target_dim``. - Centered on the incumbent, per-dimension widths scaled by lengthscale - weights. Computed in [0, 1] units (BoTorch reference semantics for the - length fields), then mapped to [-1, 1]. + The reference ``create_candidate`` box: centered on the incumbent with + half-width ``length * weights``, clamped to the domain. A half-width of + ``length`` in this side-2 domain covers the same fraction of each side as + TuRBO's ``length / 2`` does in its unit domain. """ data = self._finite_data() y = torch.tensor( @@ -616,74 +508,61 @@ def _get_optimization_bounds(self) -> torch.Tensor: ) best_idx = int(torch.argmax(self._objective_weight() * y)) - Z = self._target_space(data) - center01 = (torch.tensor(Z[best_idx], dtype=torch.double) + 1.0) / 2.0 + Z = self.embedding.project(self._normalized_inputs(data)) + center = torch.tensor(Z[best_idx], dtype=torch.double) - weights = self._lengthscale_weights(self.model) - half_length = self.trust_region.length / 2.0 - lb01 = torch.clamp(center01 - half_length * weights, 0.0, 1.0) - ub01 = torch.clamp(center01 + half_length * weights, 0.0, 1.0) - return torch.stack([2.0 * lb01 - 1.0, 2.0 * ub01 - 1.0]).to(**self.tkwargs) + half_width = self.trust_region.length * self._lengthscale_weights(self.model) + lb = torch.clamp(center - half_width, -1.0, 1.0) + ub = torch.clamp(center + half_width, -1.0, 1.0) + return torch.stack([lb, ub]).to(**self.tkwargs) - def _refresh_failure_tolerance(self) -> None: - """Set the trust region's failure tolerance, budget-aware when possible. + def _failure_tolerance(self) -> int: + """The trust region's failure tolerance, budget-aware when possible. - Implements the reference schedule (Papenmeier et al., Alg. 1), pacing - embedding expansions so the final split lands as the evaluation budget - runs out; falls back to the ``ceil(target_dim / 2)`` heuristic without - ``eval_budget``. At full dimensionality the tolerance widens to + Derived on demand, as in the reference, so a changed ``eval_budget`` or an + expanded embedding is honored automatically. The reference schedule + (Papenmeier et al., Alg. 1) paces expansions so the final split lands as + the budget runs out; without ``eval_budget`` it falls back to + ``ceil(target_dim / 2)``, and at full dimensionality it widens to ``target_dim``. The planned split count is generalized to - ``ceil(log_{b+1}(input_dim / d_init))`` to honor a user-chosen - ``target_dim_init``. + ``ceil(log_{b+1}(input_dim / d_init))`` to honor ``target_dim_init``. """ - region = self.trust_region - # the embedding is the single source of truth for dimensionality target_dim = self.embedding.target_dim input_dim = self.embedding.input_dim if target_dim >= input_dim: - region.failure_tolerance = target_dim - return + return target_dim if self.eval_budget is None: - # no budget: leave the ceil(target_dim / 2) heuristic the trust region - # already derived for this target_dim - return + return math.ceil(target_dim / 2) d_init = min(self.target_dim_init, input_dim) b = self.new_bins_on_split - # The epsilon keeps exact powers of b+1 from rounding up on float noise. + # the epsilon keeps exact powers of b + 1 from rounding up on float noise n_splits = math.ceil(math.log(input_dim / d_init, b + 1) - 1e-9) - # each split multiplies dimensionality by (b + 1), so the planned splits - # span a total growth factor of (b + 1) ** (n_splits + 1) + # each split multiplies dimensionality by (b + 1) growth = (b + 1) ** (n_splits + 1) # evaluations left for the BO phase, then this target_dim's share of them bo_budget = max(1, self.eval_budget - self.n_initial_sobol) split_budget = round(b * bo_budget * target_dim / (d_init * (growth - 1))) # max(1, ...): length_init == length_min would otherwise divide by zero halvings = max( - 1, math.floor(math.log(region.length_min / region.length_init, 0.5)) + 1, + math.floor(math.log(self.trust_region.length_min / self.length_init, 0.5)), ) # spend that share evenly over the halvings needed to reach length_min - region.failure_tolerance = min( - target_dim, max(1, math.floor(split_budget / halvings)) - ) + return min(target_dim, max(1, math.floor(split_budget / halvings))) @property def _objective_name(self) -> str: """The single objective this generator models (multi-objective is rejected).""" return self.vocs.objective_names[0] - def _vocs_bounds(self) -> tuple[np.ndarray, np.ndarray]: - """Per-variable lower and upper vocs bounds, ordered as ``variable_names``.""" - lb, ub = get_variable_bounds_array(self.vocs) - return lb, ub - def _lift_to_vocs(self, z: np.ndarray) -> np.ndarray: - """Lift target-space rows to raw vocs-space values. + """Lift through the embedding and scale to vocs bounds. - Lifts through the embedding, clips to ``[-1, 1]``, then scales to the - per-variable vocs bounds. + No clamp is needed: every caller optimizes or draws within [-1, 1], and the + embedding maps each coordinate to a signed copy of one input. """ - x_norm = np.clip(self.embedding.lift(z), -1.0, 1.0) - lb, ub = self._vocs_bounds() + x_norm = self.embedding.lift(z) + lb, ub = get_variable_bounds_array(self.vocs) return lb + (x_norm + 1.0) / 2.0 * (ub - lb) def _objective_weight(self) -> float: @@ -694,8 +573,8 @@ def _objective_weight(self) -> float: def _finite_objective_mask(self, frame: pd.DataFrame) -> np.ndarray: """Rows of ``frame`` with a finite objective; all-False if the column is absent. - Note that ``vocs.extract_data(return_valid=True)`` cannot stand in here: - it filters on feasibility, which is all-True without constraints. + ``vocs.extract_data(return_valid=True)`` cannot stand in - it filters on + feasibility, which is all-True without constraints. """ column = frame.get(self._objective_name) if column is None: @@ -706,32 +585,27 @@ def _finite_objective_mask(self, frame: pd.DataFrame) -> np.ndarray: def _finite_data(self) -> pd.DataFrame: """Rows of ``data`` with a finite objective - the GP training set.""" - if self.data is None or self.data.empty: - return pd.DataFrame( - columns=[*self.vocs.variable_names, self._objective_name] - ) - return self.data[self._finite_objective_mask(self.data)] + if self.data is None: + return pd.DataFrame() + return self.get_training_data(self.data) def _normalized_inputs(self, data: pd.DataFrame) -> np.ndarray: """Variable columns scaled to [-1, 1] per vocs bounds. - Deliberately [-1, 1] rather than ``vocs.normalize_inputs``' [0, 1]: these - feed a matrix product with the embedding, which is sign-symmetric. + [-1, 1] rather than ``vocs.normalize_inputs``' [0, 1]: these feed a matrix + product with the embedding, which is sign-symmetric. """ - lb, ub = self._vocs_bounds() + lb, ub = get_variable_bounds_array(self.vocs) X = data[self.vocs.variable_names].to_numpy(dtype=np.float64) return 2.0 * (X - lb) / (ub - lb) - 1.0 - def _target_space(self, data: pd.DataFrame) -> np.ndarray: - """Variable columns of ``data`` projected into the embedded target space.""" - return self.embedding.project(self._normalized_inputs(data)) - def _draw_sobol_point(self) -> np.ndarray: - """Next quasi-random seed point in target space [-1, 1]^target_dim. + """Next quasi-random seed point in target space, ``[-1, 1]^target_dim``. - After a dump -> construct round trip the engine is rebuilt and - fast-forwarded by ``sobol_draws`` (exact continuation needs a fixed - ``seed``). + The cached engine is rebuilt and fast-forwarded by ``sobol_draws`` after a + round trip (exact continuation needs a fixed ``seed``), and whenever the + embedding has expanded under it - which ``set_data`` can expose by dropping + back below the seed quota. """ target_dim = self.embedding.target_dim if self._sobol is None or self._sobol.dimension != target_dim: @@ -747,9 +621,8 @@ def _draw_sobol_point(self) -> np.ndarray: def _lengthscale_weights(self, model: Module) -> torch.Tensor: """Per-dimension trust-region scaling weights from the fitted model. - Tries a gpytorch kernel lengthscale, then a duck-typed - ``get_lengthscale_weights`` hook, then falls back to uniform weights - (isotropic region). + Falls back to uniform weights (an isotropic region) for kernels that have + no lengthscale. """ target_dim = self.embedding.target_dim single = model.models[0] if hasattr(model, "models") else model @@ -762,8 +635,4 @@ def _lengthscale_weights(self, model: Module) -> torch.Tensor: ls = torch.atleast_1d(lengthscale.detach().squeeze()) return _normalize_lengthscale_weights(ls, target_dim) - hook = getattr(single, "get_lengthscale_weights", None) - weights = hook(target_dim) if hook is not None else None - if weights is not None: - return weights return torch.ones(target_dim, dtype=torch.double) diff --git a/xopt/tests/generators/bayesian/test_baxus.py b/xopt/tests/generators/bayesian/test_baxus.py index 9f93782c5..76d9c512f 100644 --- a/xopt/tests/generators/bayesian/test_baxus.py +++ b/xopt/tests/generators/bayesian/test_baxus.py @@ -1,9 +1,5 @@ -"""Tests for the BAxUS generator.""" - -import logging import math import warnings -from dataclasses import dataclass from typing import Any, cast from unittest.mock import MagicMock @@ -17,9 +13,7 @@ from xopt import VOCS, Xopt from xopt import Evaluator as XoptEvaluator from xopt.errors import GeneratorWarning, VOCSError -from xopt.generator import StateOwner -from xopt.generators import get_generator_dynamic, list_available_generators -from xopt.generators.bayesian.bayesian_generator import BayesianGenerator +from xopt.generators import get_generator_dynamic from xopt.generators.bayesian.expected_improvement import ( ExpectedImprovementGenerator, ) @@ -30,19 +24,14 @@ _normalize_lengthscale_weights, ) from xopt.generators.bayesian.objectives import CustomXoptObjective -from xopt.vocs import ContextualVariable, DiscreteVariable +# most tests here construct generators without eval_budget on purpose, so the +# advisory would otherwise fire on nearly every construction; the contract itself +# stays pinned by test_missing_eval_budget_warns_exactly_once, whose own filters +# override this one @pytest.fixture(autouse=True) def _quiet_missing_eval_budget(): - """Silence the eval_budget advisory across this module. - - Most tests here construct generators without ``eval_budget`` on purpose, - because they exercise something unrelated to the expansion schedule, and the - advisory would otherwise fire on nearly every construction. The contract - itself stays pinned by ``TestBAxUSEvalBudgetWarning``, whose ``pytest.warns`` - and ``catch_warnings`` blocks install their own filters over this one. - """ with warnings.catch_warnings(): warnings.filterwarnings( "ignore", @@ -53,61 +42,60 @@ def _quiet_missing_eval_budget(): def _simple_vocs(n_vars: int = 6) -> VOCS: - """Create a VOCS with *n_vars* variables for testing.""" variables = {f"x{i}": [-1.0, 1.0] for i in range(n_vars)} return VOCS(variables=variables, objectives={"f": "MAXIMIZE"}) def _sphere(inputs: dict[str, float]) -> dict[str, float]: - """Negative sphere - maximum at the origin.""" + # negative sphere - maximum at the origin return {"f": -sum(v**2 for v in inputs.values())} def _assert_point_in_bounds(point: dict[str, float], vocs: VOCS) -> None: - """Assert every variable value lies within its VOCS domain.""" for name in vocs.variable_names: lo, hi = vocs.variables[name].domain assert lo <= point[name] <= hi, f"{name}={point[name]} outside [{lo}, {hi}]" +# CustomXoptObjective is abstract, so the rejection path needs a real subclass class _MinimalCustomObjective(CustomXoptObjective): - """Concrete CustomXoptObjective, just enough to exercise the rejection path - (CustomXoptObjective is abstract, so a bare instance can't reach our validator).""" - def forward( self, samples: torch.Tensor, X: torch.Tensor | None = None ) -> torch.Tensor: return samples +def _vocs_with(**vocs_kwargs: Any) -> dict: + # ``vocs=`` kwargs for a minimal VOCS carrying one unsupported feature + return { + "vocs": VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(3)}, + objectives={"f": "MAXIMIZE"}, + **vocs_kwargs, + ) + } + + +def _bare_gen() -> BAxUSGenerator: + return BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + + def _seed_sobol(gen: BAxUSGenerator) -> None: - """Run the Sobol init phase so subsequent generate() calls take the BO path.""" for _ in range(gen.n_initial_sobol): point = gen.generate(1)[0] gen.add_data(pd.DataFrame([{**point, **_sphere(point)}])) -def _bo_step(gen: BAxUSGenerator) -> dict[str, float]: - """Run one BO generate->evaluate->add_data cycle, asserting the candidate stays in bounds.""" - point = gen.generate(1)[0] - _assert_point_in_bounds(point, gen.vocs) - gen.add_data(pd.DataFrame([{**point, **_sphere(point)}])) - return point - - def _force_expansion(gen: BAxUSGenerator) -> None: - """Put the trust region below length_min with an unbeatable incumbent, so the - next fold-in trips the restart -> expand branch.""" + # below length_min with an unbeatable incumbent, so the next fold-in trips + # the restart -> expand branch gen.trust_region.length = gen.trust_region.length_min / 2.0 gen.trust_region.best_value = 1e9 # guarantee a failure (no improvement) def _round_trip(gen: BAxUSGenerator, *, with_data: bool = False) -> BAxUSGenerator: - """Dump -> construct, the documented resume path. - - The live botorch model is not serializable and is popped by hand, exactly as - the docs instruct a user to do. - """ + # the documented resume path: the live botorch model is not serializable and + # is popped by hand, exactly as the docs instruct dump = gen.model_dump() dump.pop("model", None) restored = BAxUSGenerator(**dump) @@ -116,561 +104,384 @@ def _round_trip(gen: BAxUSGenerator, *, with_data: bool = False) -> BAxUSGenerat return restored -def _fold_in_results(gen: BAxUSGenerator) -> None: - """Advance the trust region over results ingested since the last generate. +def _gen_with_target_space_data( + z_rows: list[list[float]], + ys: list[float], + *, + n_vars: int = 4, + direction: str = "MAXIMIZE", +) -> BAxUSGenerator: + # rows are given in target space and lifted into vocs space, so each one + # projects back to exactly the z_rows entry it came from - which is what makes + # "the box is centered on *this* point" assertable. ys is in weighted form + # (higher is better) and is mirrored for MINIMIZE, so the same row stays the + # incumbent in either direction. + vocs = VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(n_vars)}, + objectives={"f": direction}, + ) + gen = BAxUSGenerator( + vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=50 + ) + sign = 1.0 if direction == "MAXIMIZE" else -1.0 + X = gen.embedding.lift(np.array(z_rows)) + gen.add_data( + pd.DataFrame( + [ + { + **{name: float(v) for name, v in zip(vocs.variable_names, x)}, + "f": sign * y, + } + for x, y in zip(X, ys) + ] + ) + ) + return gen + - The trust region is advanced at the start of ``generate`` (mirroring how - ``BayesianGenerator`` drives ``TurboController.update_state``), so folding in - the most recent evaluations means asking for the next candidate. - """ - gen.generate(1) +def _in_target_space(gen: BAxUSGenerator, data: pd.DataFrame) -> np.ndarray: + return gen.embedding.project(gen._normalized_inputs(data)) def _set_target_dim(gen: BAxUSGenerator, target_dim: int) -> None: - """Move the generator to a given target dimensionality, keeping the - embedding and trust region consistent (the generator enforces that).""" gen.embedding = BAxUSEmbedding.create( input_dim=gen.embedding.input_dim, target_dim=target_dim, seed=0 ) - gen.trust_region = BAxUSTrustRegion(target_dim=target_dim) - - -class TestBAxUSGeneratorCreation: - """Test generator initialisation.""" - - def test_baxus_generator_creation(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs()) - - assert gen.name == "baxus" - assert gen.trust_region.target_dim == 2 - assert gen.embedding.target_dim == 2 - assert gen.embedding.input_dim == 6 - - def test_custom_target_dim(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=10), target_dim_init=4) - assert gen.trust_region.target_dim == 4 - assert gen.embedding.target_dim == 4 - assert gen.embedding.input_dim == 10 - - def test_target_dim_capped_to_input_dim(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=3), target_dim_init=10) - assert gen.trust_region.target_dim == 3 + gen.trust_region = BAxUSTrustRegion() - def test_missing_vocs_raises_validation_error(self) -> None: - """Without vocs, ``_init_components`` must skip auto-creation and let - pydantic report the missing required fields, not crash internally.""" - with pytest.raises(ValidationError, match="vocs"): - BAxUSGenerator() - def test_embedding_structure(self) -> None: - """Each column of S should have exactly one non-zero entry.""" - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=8), target_dim_init=3) - S = np.asarray(gen.embedding.matrix) - for col_idx in range(S.shape[1]): - assert np.count_nonzero(S[:, col_idx]) == 1 - assert abs(S[:, col_idx]).max() == 1.0 - - -class TestBAxUSInitialSobol: - """Test that early generate() calls return Sobol points in bounds.""" - - def test_baxus_initial_sobol(self) -> None: - vocs = _simple_vocs() - gen = BAxUSGenerator(vocs=vocs) - - for _ in range(gen.n_initial_sobol): - candidates = gen.generate(1) - assert len(candidates) == 1 - _assert_point_in_bounds(candidates[0], vocs) - gen.add_data(pd.DataFrame([{**candidates[0], **_sphere(candidates[0])}])) +class TestBAxUSConstruction: + @pytest.mark.parametrize( + ("n_vars", "kwargs", "expected"), + [(6, {}, 2), (10, {"target_dim_init": 4}, 4), (3, {"target_dim_init": 10}, 3)], + ids=["defaults-to-2", "honors-the-knob", "capped-at-input-dim"], + ) + def test_embedding_sized_from_vocs( + self, n_vars: int, kwargs: dict[str, Any], expected: int + ) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=n_vars), **kwargs) + assert gen.embedding.target_dim == expected + assert gen.embedding.input_dim == n_vars -class TestBAxUSSeed: - """The ``seed`` field pins the embedding and Sobol initialisation. + # without vocs, _init_components must skip auto-creation and let pydantic + # report the missing fields rather than crashing internally; a budget below + # the seed quota would leave no BO phase at all; generate(2) must raise + # rather than silently under-deliver one point; and the inherited + # visualization is keyed to vocs-space variables + @pytest.mark.parametrize( + ("action", "error", "match"), + [ + (BAxUSGenerator, ValidationError, "vocs"), + ( + lambda: BAxUSGenerator( + vocs=_simple_vocs(), n_initial_sobol=10, eval_budget=5 + ), + ValidationError, + "must cover the seed quota", + ), + ( + lambda: _bare_gen().generate(2), + NotImplementedError, + "parallel candidate", + ), + ( + lambda: _bare_gen().visualize_model(), + NotImplementedError, + "target space", + ), + ], + ids=[ + "missing-vocs", + "budget-below-seed-quota", + "batch-generation", + "visualize-model", + ], + ) + def test_hard_refusals( + self, action: Any, error: type[Exception], match: str + ) -> None: + with pytest.raises(error, match=match): + action() - GP fitting/acquisition use torch's global RNG and are out of scope. - """ + @pytest.mark.parametrize( + ("kwargs", "expected"), + # an explicit quota wins; otherwise it is max(2, target_dim + 1) + [ + ({"target_dim_init": 2, "n_initial_sobol": 5}, 5), + ({"target_dim_init": 3}, 4), + ], + ids=["explicit", "derived"], + ) + def test_n_initial_sobol_is_explicit_or_derived( + self, kwargs: dict[str, Any], expected: int + ) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=6), **kwargs) + assert gen.n_initial_sobol == expected - def test_same_seed_gives_identical_embedding(self) -> None: - vocs = _simple_vocs(n_vars=6) - g1 = BAxUSGenerator(vocs=vocs, seed=0) - g2 = BAxUSGenerator(vocs=vocs, seed=0) - assert g1.embedding.matrix == g2.embedding.matrix + def test_registry_bases_and_gp_constructor_default(self) -> None: + assert get_generator_dynamic("baxus") is BAxUSGenerator + assert issubclass(BAxUSGenerator, ExpectedImprovementGenerator) - def test_different_seeds_give_different_embedding(self) -> None: - vocs = _simple_vocs(n_vars=6) - g1 = BAxUSGenerator(vocs=vocs, seed=0) - g2 = BAxUSGenerator(vocs=vocs, seed=1) - assert g1.embedding.matrix != g2.embedding.matrix + gen = BAxUSGenerator(vocs=_simple_vocs()) + assert gen.gp_constructor.name == "standard" + # the near-interpolating GP the reference trust-region logic assumes + assert gen.gp_constructor.use_low_noise_prior is True + + def test_missing_eval_budget_warns_exactly_once(self) -> None: + # the generator assigns to its own fields on every BO step, and a model + # validator of either mode re-fires on assignment under + # validate_assignment, which would bury the run in duplicate warnings + with pytest.warns(GeneratorWarning, match="eval_budget is not set") as caught: + gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + assert len(caught) == 1 - def test_seed_none_still_works(self) -> None: - vocs = _simple_vocs(n_vars=4) - gen = BAxUSGenerator(vocs=vocs, seed=None) - point = gen.generate(1)[0] - _assert_point_in_bounds(point, vocs) + with warnings.catch_warnings(record=True) as on_assignment: + warnings.simplefilter("always") + gen.sobol_draws = 3 + gen.n_expansions = 1 + assert [w for w in on_assignment if w.category is GeneratorWarning] == [] + with warnings.catch_warnings(record=True) as with_budget: + warnings.simplefilter("always") + BAxUSGenerator( + vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0, eval_budget=50 + ) + assert [w for w in with_budget if w.category is GeneratorWarning] == [] -class TestBAxUSSobolSequence: - """The Sobol phase must be one reproducible low-discrepancy sequence, - not independent scrambles from a fresh engine per call.""" - def test_same_seed_gives_identical_sobol_points(self) -> None: +# seed pins the embedding and the Sobol initialisation; GP fitting and +# acquisition use torch's global RNG and are out of scope +class TestBAxUSSeeding: + def test_seed_pins_the_embedding_and_the_sobol_sequence(self) -> None: + # the Sobol phase must be one reproducible sequence - a fresh draw(1) + # engine per call would instead repeat the same first point every time vocs = _simple_vocs(n_vars=6) g1 = BAxUSGenerator(vocs=vocs, seed=7) g2 = BAxUSGenerator(vocs=vocs, seed=7) - pts1 = [g1.generate(1)[0] for _ in range(g1.n_initial_sobol)] - pts2 = [g2.generate(1)[0] for _ in range(g2.n_initial_sobol)] - assert pts1 == pts2 - - def test_sobol_points_advance(self) -> None: - """A fresh ``draw(1)`` engine per call would repeat the same first point.""" - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=6), seed=3) - pts = [tuple(gen.generate(1)[0].values()) for _ in range(gen.n_initial_sobol)] - assert len(set(pts)) == len(pts) - - -class TestBAxUSEndToEnd: - """End-to-end: Sobol init then BO steps.""" - - def test_baxus_end_to_end(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) - - _seed_sobol(gen) - for _ in range(3): - _bo_step(gen) - - assert gen.data is not None - assert len(gen.data) == gen.n_initial_sobol + 3 - - -def _asymmetric_vocs() -> VOCS: - """VOCS with per-variable bounds at different scales - catches bounds orientation bugs.""" - return VOCS( - variables={"x0": [0.0, 10.0], "x1": [-165.0, 165.0], "x2": [0.5, 0.6]}, - objectives={"f": "MAXIMIZE"}, - ) - + g3 = BAxUSGenerator(vocs=vocs, seed=1) -class TestBAxUSAsymmetricBounds: - """Guard against the v2->v3 `vocs.bounds` orientation flip (was (2, D), now (D, 2)).""" - - def test_points_lie_in_per_variable_bounds(self) -> None: - vocs = _asymmetric_vocs() - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2) + assert g1.embedding.matrix == g2.embedding.matrix + assert g1.embedding.matrix != g3.embedding.matrix - _seed_sobol(gen) - # A wrong bounds orientation produces out-of-box points in the BO phase. - for _ in range(3): - _bo_step(gen) + pts1 = [g1.generate(1)[0] for _ in range(g1.n_initial_sobol)] + pts2 = [g2.generate(1)[0] for _ in range(g2.n_initial_sobol)] + assert pts1 == pts2 + distinct = {tuple(p.values()) for p in pts1} + assert len(distinct) == len(pts1) -class TestBAxUSSteering: - """BO-phase points must be better *on average* than the Sobol seed points. + def test_raw_draws_span_the_cube_and_follow_the_target_dim(self) -> None: + gen = BAxUSGenerator( + vocs=_simple_vocs(n_vars=8), + target_dim_init=2, + n_initial_sobol=32, + seed=0, + eval_budget=60, + ) + zs = np.vstack([gen._draw_sobol_point() for _ in range(32)]) - Guards against a min-form/max-form inversion (``get_objective_data`` - negates MAXIMIZE objectives). Never compare best-over-all-data against - best-over-seeds - the seeds are a subset of all data, so that is - vacuously true and lets the generator anti-optimize undetected. + # the draws must cover the whole cube, not just one orthant + assert zs.shape == (32, 2) + assert zs.min() < -0.2, f"no negative seed coordinates (min {zs.min():.3f})" + assert zs.max() > 0.2, f"no positive seed coordinates (max {zs.max():.3f})" + assert (zs >= -1.0).all() and (zs <= 1.0).all() + assert abs(float(zs.mean())) < 0.25, "seed points are not centered" - The optimum is deliberately placed OFF-CENTER. With an optimum at the origin - - the center of the symmetric [-1, 1] domain - these tests only measure - "did the candidate move toward the middle of the box", which a generator - that ignores the acquisition function entirely (proposing a constant zero in - embedded space) satisfies perfectly. The offset makes them measure - optimization instead. - """ + # the engine is cached, so it has to be rebuilt when the embedding grows + # under it; without the rebuild the stale engine feeds a wrongly-sized + # vector into the lift and numpy raises on the matmul + _set_target_dim(gen, 8) - N_SOBOL = 6 - N_BO_STEPS = 15 - OPTIMUM = 0.6 + assert gen._draw_sobol_point().shape == (1, 8) - def _run_phases(self, direction: str) -> tuple[list[float], list[float]]: - """Run Sobol seeding then BO steps; return (sobol_ys, bo_ys). - The objective peaks at ``x_i = OPTIMUM`` for every i: MAXIMIZE gets - -||x - c||^2 (values <= 0), MINIMIZE gets +||x - c||^2 (values >= 0). - """ - torch.manual_seed(0) +class TestBAxUSEndToEnd: + def test_xopt_run_survives_failures_and_stays_in_bounds(self) -> None: + # asymmetric per-variable bounds guard against the v2->v3 vocs.bounds + # orientation flip (was (2, D), now (D, 2)): a wrong orientation produces + # out-of-box points in the BO phase. get_optimum is checked for the same + # reason - it optimizes in target space and lifts, so bounds at different + # scales are what make a broken lift visible. The run is warm-started + # with random points (like the real phases.py flow) rather than the + # generator's own Sobol phase, so those rows reach the model by + # least-squares projection, and the evaluator fails once mid-run: the + # failed row stays in the history, is excluded from training, and later + # BO steps still propose in-bounds points. vocs = VOCS( - variables={f"x{i}": [-1.0, 1.0] for i in range(6)}, - objectives={"f": direction}, - ) - gen = BAxUSGenerator( - vocs=vocs, - target_dim_init=2, - n_initial_sobol=self.N_SOBOL, - seed=0, - eval_budget=self.N_SOBOL + self.N_BO_STEPS, + variables={"x0": [0.0, 10.0], "x1": [-165.0, 165.0], "x2": [0.5, 0.6]}, + objectives={"f": "MAXIMIZE"}, ) - sign = -1.0 if direction == "MAXIMIZE" else 1.0 + calls = {"n": 0} - ys: list[float] = [] - for _ in range(self.N_SOBOL + self.N_BO_STEPS): - point = gen.generate(1)[0] - _assert_point_in_bounds(point, vocs) - y = sign * sum((v - self.OPTIMUM) ** 2 for v in point.values()) - ys.append(y) - gen.add_data(pd.DataFrame([{**point, "f": y}])) - return ys[: self.N_SOBOL], ys[self.N_SOBOL :] - - def test_maximize_bo_phase_beats_sobol_phase(self) -> None: - sobol_ys, bo_ys = self._run_phases("MAXIMIZE") - assert np.mean(bo_ys) > np.mean(sobol_ys), ( - f"MAXIMIZE BO phase is worse than random seeding (direction inverted?): " - f"sobol mean={np.mean(sobol_ys):.4f} bo mean={np.mean(bo_ys):.4f}" - ) + def flaky(inputs: dict[str, float]) -> dict[str, float]: + calls["n"] += 1 + if calls["n"] == 6: # the third BO step, after the 3 warm-start rows + raise RuntimeError("evaluation blew up") + return _sphere(inputs) - def test_minimize_bo_phase_beats_sobol_phase(self) -> None: - sobol_ys, bo_ys = self._run_phases("MINIMIZE") - assert np.mean(bo_ys) < np.mean(sobol_ys), ( - f"MINIMIZE BO phase is worse than random seeding (direction inverted?): " - f"sobol mean={np.mean(sobol_ys):.4f} bo mean={np.mean(bo_ys):.4f}" - ) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0, eval_budget=30) + X = Xopt(evaluator=XoptEvaluator(function=flaky), generator=gen, strict=False) + X.random_evaluate(gen.n_initial_sobol) + assert X.generator.computation_time is None # no model trained yet + for _ in range(8): + X.step() + + assert len(X.data) == gen.n_initial_sobol + 8 + # the failed row is kept but excluded from training + assert len(X.generator._finite_data()) == len(X.data) - 1 + for _, row in X.data.iterrows(): + _assert_point_in_bounds(row.to_dict(), vocs) + timings = X.generator.computation_time + assert len(timings) == 8 + assert list(timings.columns) == ["training", "acquisition_optimization"] + + opt = X.generator.get_optimum() + assert len(opt) == 1 + assert list(opt.columns) == vocs.variable_names + _assert_point_in_bounds(opt.iloc[0].to_dict(), vocs) -class TestBAxUSNaNIngestion: - """NaN objectives (failed evaluations) must not poison the training data - or the trust-region update.""" - def test_nan_row_dropped_from_training_data(self) -> None: +class TestBAxUSTrainingData: + def test_unusable_rows_are_kept_but_excluded_from_training(self) -> None: gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) + with pytest.raises(ValueError, match="no data available to build model"): + gen.train_model() + point = gen.generate(1)[0] gen.add_data(pd.DataFrame([{**point, "f": float("nan")}])) - - assert len(gen.data) == 1 # row is kept in the canonical history... - assert gen._finite_data().empty # ...but excluded from GP training - - def test_mixed_batch_keeps_only_finite_rows(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) + assert len(gen.data) == 1 + assert gen._finite_data().empty rows = [] for y in (1.0, float("nan"), 2.0): point = gen.generate(1)[0] rows.append({**point, "f": y}) gen.add_data(pd.DataFrame(rows)) - - assert len(gen.data) == 3 + assert len(gen.data) == 4 assert len(gen._finite_data()) == 2 - def test_bo_step_survives_nan_ingestion(self) -> None: - """After a NaN row mid-run, the next BO candidate is still finite and in bounds.""" - vocs = _simple_vocs(n_vars=4) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2) - - _seed_sobol(gen) - point = gen.generate(1)[0] - gen.add_data(pd.DataFrame([{**point, "f": float("nan")}])) - - candidate = gen.generate(1)[0] - assert all(math.isfinite(v) for v in candidate.values()) - _assert_point_in_bounds(candidate, vocs) - - -class TestBAxUSTrustRegionSeedPhase: - """The trust region must track BO-phase evaluations only: the BoTorch - reference starts BO with a pristine trust region, so Sobol seeding must - not move the counters or ``best_value``.""" - - def test_seed_phase_leaves_state_pristine(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - - _seed_sobol(gen) - - assert gen.trust_region.length == gen.trust_region.length_init - assert gen.trust_region.success_counter == 0 - assert gen.trust_region.failure_counter == 0 - assert gen.trust_region.best_value is None - - def test_first_bo_result_updates_state(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - - _seed_sobol(gen) - _bo_step(gen) - assert gen.trust_region.best_value is None # not folded in yet - _fold_in_results(gen) - - assert gen.trust_region.best_value is not None - assert gen.tr_observed_rows == len(gen.data) - - def test_straddling_batch_counts_only_post_quota_rows(self) -> None: - """A batch that crosses the seed quota updates the trust region with the - BO-phase rows only.""" - vocs = _simple_vocs(n_vars=4) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0) - - pts = [gen.generate(1)[0] for _ in range(3)] - rows = [ - {**pts[0], "f": 100.0}, - {**pts[1], "f": 100.0}, - {**pts[2], "f": 1.0}, - ] - gen.add_data(pd.DataFrame(rows)) - _fold_in_results(gen) - - # MAXIMIZE ⇒ weighted value is +f; only the BO row (f=1.0) should count. - assert gen.trust_region.best_value == 1.0 - - def test_trust_region_is_independent_of_ingestion_chunking(self) -> None: - """Same observations, different add_data batching ⇒ identical state. - - The state machine is driven by generate calls, so feeding results one at - a time or all at once must not change the trajectory (a warm start via - ``Xopt(..., data=...)`` arrives as one big frame). - """ - vocs = _simple_vocs(n_vars=4) - rng = np.random.default_rng(0) - pts = [ - {f"x{i}": float(v) for i, v in enumerate(rng.uniform(-1, 1, 4))} - for _ in range(12) - ] - rows = [{**p, **_sphere(p)} for p in pts] - - def state(batched: bool) -> tuple[float, int, int, float]: - gen = BAxUSGenerator( - vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0 - ) - if batched: - gen.add_data(pd.DataFrame(rows)) - else: - for row in rows: - gen.add_data(pd.DataFrame([row])) - _fold_in_results(gen) - tr = gen.trust_region - return (tr.length, tr.success_counter, tr.failure_counter, tr.best_value) - - assert state(batched=True) == state(batched=False) - - def test_baxus_embedding_expansion(self) -> None: - """Expansion grows target_dim, resets the region, and preserves best_value.""" - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=8), target_dim_init=2, seed=0) - _seed_sobol(gen) - - # Force the next fold-in to trip the restart branch. - _force_expansion(gen) - _bo_step(gen) - _fold_in_results(gen) - - assert gen.trust_region.target_dim > 2 - assert gen.embedding.target_dim == gen.trust_region.target_dim - assert gen.trust_region.length == gen.trust_region.length_init - assert gen.trust_region.best_value == 1e9 - assert not gen.trust_region.restart_triggered - assert gen.n_expansions == 1 - - -class TestBAxUSBatchGeneration: - """``generate(2)`` must raise (matching xopt's guard), not silently - under-deliver one point.""" - - def test_generate_multiple_candidates_raises(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) - with pytest.raises(NotImplementedError, match="parallel candidate generation"): - gen.generate(2) - - -class TestBAxUSTrainModel: - """train_model must refuse to fit a GP with no finite-objective data.""" - - def test_raises_when_no_data_available(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2) - with pytest.raises(ValueError, match="no data available to build model"): - gen.train_model() - - -class TestBAxUSViaXopt: - """Full integration through the Xopt runner.""" - - def test_baxus_via_xopt(self) -> None: - vocs = _simple_vocs(n_vars=4) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2) - evaluator = XoptEvaluator(function=_sphere) - X = Xopt(evaluator=evaluator, generator=gen) - - # Seed with random points (like the real phases.py flow), then BO. - X.random_evaluate(gen.n_initial_sobol) - for _ in range(3): - X.step() - - assert len(X.data) == gen.n_initial_sobol + 3 - - for name in vocs.variable_names: - lo, hi = vocs.variables[name].domain - values = X.data[name].to_numpy() - assert np.all(values >= lo - 1e-9), f"{name} below lower bound" - assert np.all(values <= hi + 1e-9), f"{name} above upper bound" - - -class TestBAxUSLengthscaleWeights: - """_lengthscale_weights: gpytorch kernel -> duck-typed hook -> uniform fallback.""" - - def _gen(self, target_dim: int = 2) -> BAxUSGenerator: - return BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=target_dim) - - @staticmethod - def _kernel_model(lengthscale: torch.Tensor) -> MagicMock: - """A mock shaped like a gpytorch model: covar_module.base_kernel.lengthscale.""" - kernel = MagicMock() - kernel.lengthscale = lengthscale - covar = MagicMock() - covar.base_kernel = kernel - model = MagicMock(spec=["covar_module"]) - model.covar_module = covar - return model - - def test_gpytorch_kernel_lengthscale_used(self) -> None: - gen = self._gen() - model = self._kernel_model(torch.tensor([[0.5, 2.0]], dtype=torch.double)) - - weights = gen._lengthscale_weights(model) - assert weights.shape == (2,) - assert not torch.allclose(weights, torch.ones(2, dtype=torch.double)) - - def test_scalar_lengthscale_unsqueezed(self) -> None: - gen = self._gen(target_dim=1) - model = self._kernel_model(torch.tensor(0.5, dtype=torch.double)) - - assert gen._lengthscale_weights(model).shape == (1,) - - def test_duck_typed_hook_used_when_no_kernel(self) -> None: - gen = self._gen() - model = MagicMock(spec=["get_lengthscale_weights"]) - model.get_lengthscale_weights.return_value = torch.tensor( - [0.5, 1.5], dtype=torch.double - ) - - weights = gen._lengthscale_weights(model) - assert torch.allclose(weights, torch.tensor([0.5, 1.5], dtype=torch.double)) - model.get_lengthscale_weights.assert_called_once_with(2) - - def test_uniform_fallback_when_hook_returns_none(self) -> None: - gen = self._gen() - model = MagicMock(spec=["get_lengthscale_weights"]) - model.get_lengthscale_weights.return_value = None - - assert torch.allclose( - gen._lengthscale_weights(model), torch.ones(2, dtype=torch.double) - ) - - def test_uniform_fallback_for_bare_model(self) -> None: - gen = self._gen() - model = MagicMock(spec=["posterior", "num_outputs"]) - - assert torch.allclose( - gen._lengthscale_weights(model), torch.ones(2, dtype=torch.double) + # an evaluator that raises returns a row with no objective column at all + gen.add_data( + pd.DataFrame([{**{f"x{i}": 0.1 for i in range(4)}, "xopt_error": True}]) ) + assert len(gen.data) == 5 + assert len(gen._finite_data()) == 2 -class TestBAxUSConfig: - def test_explicit_n_initial_sobol_kept(self) -> None: - gen = BAxUSGenerator( - vocs=_simple_vocs(n_vars=4), target_dim_init=2, n_initial_sobol=5 - ) - assert gen.n_initial_sobol == 5 - - def test_auto_n_initial_sobol(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=6), target_dim_init=3) - assert gen.n_initial_sobol == 4 # max(2, target_dim + 1) +# trust-region state bookkeeping (weighted convention: higher is better), then +# which evaluations the generator folds into it - BO-phase only, since the +# reference starts BO with a pristine region +class TestBAxUSTrustRegion: + def test_success_path_doubles_the_length_then_caps_it(self) -> None: + # asserting only the direction (got larger) would leave the schedule free + # to drift, which changes how many failures trigger an expansion + tr = BAxUSTrustRegion(length=0.2) + tr.best_value = 0.0 + improving = iter([100.0 * i for i in range(1, 100)]) + for _ in range(tr.success_tolerance - 1): + tr.update(torch.tensor([next(improving)]), failure_tolerance=1) + assert tr.length == pytest.approx(0.2, rel=1e-12) # tolerance not reached -class TestBAxUSTrustRegion: - """Trust-region state bookkeeping (weighted convention: higher is better).""" + tr.update(torch.tensor([next(improving)]), failure_tolerance=1) + assert tr.length == pytest.approx(0.4, rel=1e-12) # exactly doubled - def test_success_expands_length(self) -> None: - tr = BAxUSTrustRegion(target_dim=2, length=0.8) - tr.best_value = 0.0 - for i in range(tr.success_tolerance): - tr.update(torch.tensor([10.0 * (i + 1)])) - assert tr.length > 0.8 + # 0.4 -> 0.8 -> 1.6, then pinned: a third block would reach 3.2 uncapped + for _ in range(3 * tr.success_tolerance): + tr.update(torch.tensor([next(improving)]), failure_tolerance=1) + assert tr.length == pytest.approx(tr.length_max, rel=1e-12) - def test_failure_shrinks_length(self) -> None: - tr = BAxUSTrustRegion(target_dim=2, length=0.8) + def test_failure_path_halves_the_length_then_restarts(self) -> None: + # the exact halving schedule is pinned by the gain-0.05%-fails case below + tr = BAxUSTrustRegion(length=0.8) tr.best_value = 999.0 - for _ in range(tr.failure_tolerance): - tr.update(torch.tensor([-999.0])) - assert tr.length < 0.8 - def test_restart_triggered_on_tiny_length(self) -> None: - tr = BAxUSTrustRegion(target_dim=2, length=0.008) - tr.best_value = 999.0 for _ in range(100): - tr.update(torch.tensor([-999.0])) if tr.restart_triggered: break + tr.update(torch.tensor([-999.0]), failure_tolerance=1) assert tr.restart_triggered + assert tr.length < tr.length_min - def test_first_update_counts_as_success(self) -> None: - """best_value=None (pristine) means the first observed batch always improves.""" - tr = BAxUSTrustRegion(target_dim=2) - tr.update(torch.tensor([-5.0])) - assert tr.best_value == -5.0 - assert tr.failure_counter == 0 + # a gain below 0.1% of the incumbent counts as a failure, and a pristine + # region always counts the first batch as a success + @pytest.mark.parametrize( + ("best_value", "y", "success_counter", "length"), + [(None, -5.0, 1, 0.8), (100.0, 100.05, 0, 0.4), (100.0, 100.5, 1, 0.8)], + ids=["pristine-always-improves", "gain-0.05%-fails", "gain-0.5%-improves"], + ) + def test_improvement_needs_to_clear_the_tolerance( + self, best_value: float | None, y: float, success_counter: int, length: float + ) -> None: + # success_counter and length are what tell the branches apart: on the + # failure branch best_value is still assigned (best is None short-circuits + # the max) and failure_counter is reset by the shrink it triggers + tr = BAxUSTrustRegion(length=0.8) + tr.best_value = best_value + + tr.update(torch.tensor([y]), failure_tolerance=1) + + assert tr.success_counter == success_counter + assert tr.failure_counter == 0 # a shrink resets it either way + assert tr.length == pytest.approx(length, rel=1e-12) + assert tr.best_value == max(y, best_value if best_value is not None else y) + + @pytest.mark.parametrize("direction", ["MAXIMIZE", "MINIMIZE"]) + def test_seed_rows_never_fold_in_and_bo_results_fold_in_exactly_once( + self, direction: str + ) -> None: + # the seed quota is what protects the pristine state, not the early + # return, so the advance has to be driven directly. tr_observed_rows is + # what makes the fold-in once-only: without it every advance re-folds the + # whole BO history, which is silent while results keep arriving but ticks + # the counters on an advance with no new results. + vocs = VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, + objectives={"f": direction}, + ) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0) + sign = 1.0 if direction == "MAXIMIZE" else -1.0 + row = {f"x{i}": 0.1 for i in range(4)} - def test_failure_tolerance_scales_with_target_dim(self) -> None: - assert BAxUSTrustRegion(target_dim=2).failure_tolerance == 1 - assert BAxUSTrustRegion(target_dim=10).failure_tolerance == 5 + # seed rows are given the better weighted value (+100), so a fold-in of + # the seed quota would surface below as best_value == 100.0 + gen.add_data(pd.DataFrame([{**row, "f": sign * 100.0}] * 2)) + gen._advance_trust_region() - def test_round_trips_through_dump(self) -> None: - tr = BAxUSTrustRegion( - target_dim=3, length=0.4, success_counter=2, best_value=1.5 - ) - assert BAxUSTrustRegion(**tr.model_dump()) == tr - - -@dataclass -class _TutorialBaxusState: - """Oracle: the reference ``BaxusState`` from the BoTorch BAxUS tutorial, - trimmed to the fields the failure-tolerance formula needs.""" - - dim: int - eval_budget: int - new_bins_on_split: int = 3 - d_init: int = 0 - target_dim: int = 0 - n_splits: int = 0 - length_init: float = 0.8 - length_min: float = 0.5**7 - - def __post_init__(self) -> None: - n_splits = round(math.log(self.dim, self.new_bins_on_split + 1)) - self.d_init = int( - 1 - + np.argmin( - np.abs( - (1 + np.arange(self.new_bins_on_split)) - * (1 + self.new_bins_on_split) ** n_splits - - self.dim - ) - ) - ) - self.target_dim = self.d_init - self.n_splits = n_splits - - @property - def split_budget(self) -> int: - growth: int = (self.new_bins_on_split + 1) ** (self.n_splits + 1) - return round( - -1 - * (self.new_bins_on_split * self.eval_budget * self.target_dim) - / (self.d_init * (1 - growth)) - ) + assert gen.trust_region.best_value is None + assert gen.tr_observed_rows == 0 - @property - def failure_tolerance(self) -> int: - if self.target_dim == self.dim: - return self.target_dim - k = math.floor(math.log(self.length_min / self.length_init, 0.5)) - return min(self.target_dim, max(1, math.floor(self.split_budget / k))) + gen.add_data(pd.DataFrame([{**row, "f": sign * 1.0}])) + gen._advance_trust_region() + assert gen.trust_region.best_value == 1.0 # only the post-quota row counts + assert gen.tr_observed_rows == len(gen.data) -class TestBAxUSFailureTolerance: - """Budget-aware failure tolerance (paper Alg. 1) against the tutorial oracle. + tr = gen.trust_region + before = (tr.length, tr.success_counter, tr.failure_counter, tr.best_value) + observed = gen.tr_observed_rows - input_dim=16 with b=3 is an exact power of b+1, where the oracle's derived - d_init (=1) and split count (=2) coincide with ours - the regime in which a - verbatim comparison is meaningful. - """ + gen._advance_trust_region() + gen._advance_trust_region() + tr = gen.trust_region + assert ( + tr.length, + tr.success_counter, + tr.failure_counter, + tr.best_value, + ) == before + assert gen.tr_observed_rows == observed + + +# budget-aware failure tolerance (paper Alg. 1). input_dim=16 with b=3 is an +# exact power of b+1, where the BoTorch tutorial's BaxusState derives the same +# d_init (=1) and split count (=2) we configure - the regime in which a verbatim +# comparison against the reference is meaningful. +class TestBAxUSFailureTolerance: def _gen(self, eval_budget: int | None, target_dim_init: int = 1) -> BAxUSGenerator: return BAxUSGenerator( vocs=_simple_vocs(n_vars=16), @@ -680,774 +491,135 @@ def _gen(self, eval_budget: int | None, target_dim_init: int = 1) -> BAxUSGenera eval_budget=eval_budget, ) - # bo_budget=93 with target_dim=8 is the discriminating pair: computing the - # budget without subtracting the 10 seed points gives 6 instead of 5, so - # this pins the total-including-seeds semantics of eval_budget. - @pytest.mark.parametrize("bo_budget", [30, 93, 200]) - def test_matches_tutorial_oracle_across_dims(self, bo_budget: int) -> None: - oracle = _TutorialBaxusState(dim=16, eval_budget=bo_budget) - assert oracle.d_init == 1 # sanity: oracle derives the d_init we configure - gen = self._gen( - eval_budget=bo_budget + 10 - ) # our field is total incl. the 10 seeds - - for target_dim in (1, 4, 8, 16): - oracle.target_dim = target_dim - _set_target_dim(gen, target_dim) - gen._refresh_failure_tolerance() - assert gen.trust_region.failure_tolerance == oracle.failure_tolerance - - def test_differs_from_heuristic(self) -> None: - """The paper schedule is not just ceil(d/2) in disguise.""" - gen = self._gen(eval_budget=210) - _set_target_dim(gen, 4) - gen._refresh_failure_tolerance() - assert gen.trust_region.failure_tolerance == 4 # heuristic would say 2 - - def test_heuristic_kept_without_budget(self) -> None: - gen = self._gen(eval_budget=None, target_dim_init=4) - assert gen.trust_region.failure_tolerance == 2 # ceil(4/2) - gen._refresh_failure_tolerance() - assert gen.trust_region.failure_tolerance == 2 - - def test_full_dim_widens_to_target_dim(self) -> None: - """At full dimensionality the tolerance is target_dim, budget or not.""" - gen = self._gen(eval_budget=None, target_dim_init=16) - gen._refresh_failure_tolerance() - assert gen.trust_region.failure_tolerance == 16 + # expected values are the tutorial's BaxusState.failure_tolerance at dim=16, + # b=3, d_init=1, with eval_budget as the total including the 10 seed points. + # (103, 8) is the discriminating pair: computing the budget without + # subtracting the seeds gives 6 instead of 5. + @pytest.mark.parametrize( + ("eval_budget", "target_dim", "expected"), + [(40, 4, 1), (40, 8, 1), (103, 4, 3), (103, 8, 5), (210, 4, 4), (210, 8, 8)], + ) + def test_matches_the_reference_schedule( + self, eval_budget: int, target_dim: int, expected: int + ) -> None: + gen = self._gen(eval_budget=eval_budget) + _set_target_dim(gen, target_dim) + assert gen._failure_tolerance() == expected + + def test_branch_selection(self) -> None: + # without a budget: the ceil(target_dim / 2) heuristic + heuristic = self._gen(eval_budget=None, target_dim_init=4) + assert heuristic._failure_tolerance() == 2 # ceil(4/2) + + # at full dimensionality: target_dim, budget or not + full_dim = self._gen(eval_budget=None, target_dim_init=16) + assert full_dim._failure_tolerance() == 16 + + # length_init == length_min makes the reference halvings formula 0, which + # must not divide by zero: it is floored to 1, so the budget share is + # spent in one go - round(3 * 48 * 2 / (2 * 15)) = 10, capped at target_dim + degenerate = BAxUSGenerator( + vocs=_simple_vocs(n_vars=8), + target_dim_init=2, + n_initial_sobol=2, + eval_budget=50, + length_init=0.5**7, + ) + assert degenerate._failure_tolerance() == 2 def test_bo_phase_applies_budget_tolerance(self) -> None: - """Folding in a BO result refreshes the tolerance (also covers a d_init the - reference would not derive itself: n_splits generalizes).""" - vocs = _simple_vocs(n_vars=16) + # also covers a d_init the reference would not derive itself: n_splits + # generalizes + gen = BAxUSGenerator( + vocs=_simple_vocs(n_vars=16), + target_dim_init=4, + n_initial_sobol=2, + seed=0, + eval_budget=210, + ) + assert gen._failure_tolerance() == 4 # paper value; heuristic would say 2 + + # ...and the fold-in honors it: one failure against a tolerance of 4 is + # counted but must not shrink the region, which a hardwired tolerance of + # 1 would + row = {f"x{i}": 0.1 for i in range(16)} + gen.add_data(pd.DataFrame([{**row, "f": 0.0}] * 3)) # 2 seeds + 1 BO row + gen.trust_region.best_value = 1e9 # the BO result cannot improve on it + gen._advance_trust_region() + assert gen.trust_region.failure_counter == 1 + assert gen.trust_region.length == gen.length_init + + +# generator-level embedding expansion; the matrix-level split is +# TestBAxUSEmbedding's job. Right after an expansion the projected training data +# has duplicated coordinates, which can make the covariance indefinite: xopt +# downgrades the resulting fit failure to a warning and returns an untrained +# model, which the run must survive. +class TestBAxUSExpansion: + def test_expansion_resets_the_region_and_survives_a_round_trip(self) -> None: + # the generate that folds the results in expands first, then trains and + # proposes in the new space, so its candidate doubles as the + # post-expansion survival check. The tolerance is recomputed for the new + # target_dim: d_init=2, n_splits=2, budget=200 -> round(4800/126)=38, + # floor(38/6)=6. The expansion seed is derived from (seed, n_expansions), + # so a run restored mid-flight expands to the same embedding. gen = BAxUSGenerator( - vocs=vocs, target_dim_init=4, n_initial_sobol=2, seed=0, eval_budget=210 + vocs=_simple_vocs(n_vars=32), target_dim_init=2, seed=0, eval_budget=203 ) _seed_sobol(gen) - _bo_step(gen) - _fold_in_results(gen) - assert ( - gen.trust_region.failure_tolerance == 4 - ) # paper value; heuristic would say 2 - - def test_expansion_recomputes_tolerance(self) -> None: - vocs = _simple_vocs(n_vars=8) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0, eval_budget=50) - _seed_sobol(gen) - - _force_expansion(gen) - _bo_step(gen) - _fold_in_results(gen) - - assert gen.trust_region.target_dim == 8 # expanded to full dim - assert gen.trust_region.failure_tolerance == 8 - - def test_expansion_recomputes_budget_tolerance_below_full_dim(self) -> None: - """An expansion landing below input_dim exercises the budget formula, - not just the full-dim widening.""" - vocs = _simple_vocs(n_vars=32) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0, eval_budget=203) - _seed_sobol(gen) - - _force_expansion(gen) - _bo_step(gen) - _fold_in_results(gen) - - assert gen.trust_region.target_dim == 8 # 2 * (b + 1), still < 32 - # d_init=2, n_splits=2, budget=200: round(4800/126)=38 -> floor(38/6)=6 - assert gen.trust_region.failure_tolerance == 6 - - def test_explicit_failure_tolerance_survives_construction(self) -> None: - """Static-model semantics only: failure_tolerance is derived state and - the generator re-derives it on every BO-phase ingest.""" - assert ( - BAxUSTrustRegion(target_dim=8, failure_tolerance=7).failure_tolerance == 7 - ) - - def test_budget_below_seed_quota_rejected(self) -> None: - with pytest.raises(ValidationError, match="must cover the seed quota"): - self._gen(eval_budget=5) # quota is 10 - - def test_dump_round_trips_budget_and_tolerance(self) -> None: - gen = self._gen(eval_budget=210) - _set_target_dim(gen, 4) - gen._refresh_failure_tolerance() - - restored = _round_trip(gen) - - assert restored.eval_budget == 210 - assert ( - restored.trust_region.failure_tolerance - == gen.trust_region.failure_tolerance - ) - - -class TestBAxUSPostExpansionFit: - """Right after an embedding expansion the projected training data has - duplicated coordinates, which can make the covariance indefinite; the run - must survive the degraded fit and log it.""" - - def test_generate_survives_post_expansion_training(self) -> None: - vocs = _simple_vocs(n_vars=8) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) - _seed_sobol(gen) - - _force_expansion(gen) - _bo_step(gen) + restored = _round_trip(gen, with_data=True) - # this generate expands first, then trains and proposes in the new space + row = {f"x{i}": 0.1 for i in range(32)} + for generator in (gen, restored): + _force_expansion(generator) # best_value = 1e9 trips the restart + generator.add_data(pd.DataFrame([{**row, "f": -1.0}])) candidate = gen.generate(1)[0] + restored._advance_trust_region() # the expansion, without the GP fit + assert gen.embedding.target_dim == 8 - _assert_point_in_bounds(candidate, vocs) + assert gen.trust_region.length == gen.length_init + assert gen.trust_region.best_value == 1e9 # carried across the reset + assert not gen.trust_region.restart_triggered + assert gen.n_expansions == 1 + assert gen._failure_tolerance() == 6 + _assert_point_in_bounds(candidate, gen.vocs) assert all(math.isfinite(v) for v in candidate.values()) + assert restored.n_expansions == 1 + assert restored.embedding == gen.embedding - def test_fit_failure_warning_surfaces_in_logs( - self, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch - ) -> None: - """train_model must mirror xopt's fit-failure warning into the logger - and still re-emit the original warning unchanged.""" - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - _seed_sobol(gen) - - constructor_cls = type(gen.gp_constructor) - real_build = constructor_cls.build_model - - def failing_build(self: Any, **kwargs: Any) -> Any: - warnings.warn( - "Model fitting failed. Returning untrained model.", stacklevel=1 - ) - return real_build(self, **kwargs) - - monkeypatch.setattr(constructor_cls, "build_model", failing_build) - - logger_name = "xopt.generators.bayesian.baxus" - with ( - caplog.at_level(logging.WARNING, logger=logger_name), - pytest.warns(UserWarning, match="Model fitting failed"), - ): - gen.train_model() - - assert any( - "GP fit failed at target_dim=2" in record.message - for record in caplog.records - ) - - -class TestBAxUSEmbedding: - """Sparse random embedding: structure, projection round-trip, expansion.""" - - def test_create_structure(self) -> None: - """Each input column has exactly one signed unit entry.""" - emb = BAxUSEmbedding.create(input_dim=8, target_dim=3, seed=0) - S = np.asarray(emb.matrix) - assert S.shape == (3, 8) - for col in range(8): - assert np.count_nonzero(S[:, col]) == 1 - assert abs(S[:, col]).max() == 1.0 - - def test_create_assigns_both_signs(self) -> None: - """Sign randomization is what makes the embedding an unbiased projection; - an all-positive matrix is still structurally valid, so check the signs.""" - emb = BAxUSEmbedding.create(input_dim=200, target_dim=4, seed=0) - matrix = np.asarray(emb.matrix) - - positive = int((matrix == 1.0).sum()) - negative = int((matrix == -1.0).sum()) - assert positive and negative, "embedding signs are not randomized" - assert 0.3 < positive / (positive + negative) < 0.7 - - def test_create_is_seed_deterministic(self) -> None: - e1 = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=0) - e2 = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=0) - e3 = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=1) - assert e1.matrix == e2.matrix - assert e1.matrix != e3.matrix - - def test_lift_project_round_trip(self) -> None: - """Points lifted from target space project back exactly (S rows are orthogonal).""" - emb = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=0) - Z = np.array([[0.5, -0.25], [1.0, 1.0]]) - assert np.allclose(emb.project(emb.lift(Z)), Z) - - def test_expand_grows_target_dim_and_preserves_signs(self) -> None: - emb = BAxUSEmbedding.create(input_dim=12, target_dim=2, seed=0) - grown = emb.expand(new_bins_on_split=3) - S_old, S_new = np.asarray(emb.matrix), np.asarray(grown.matrix) - assert grown.target_dim == min(2 * 4, 12) - assert grown.input_dim == 12 - # Every column keeps its sign and moves to exactly one (new) row. - for col in range(12): - assert np.count_nonzero(S_new[:, col]) == 1 - assert S_new[:, col].sum() == S_old[:, col].sum() - - def test_expand_caps_at_input_dim(self) -> None: - emb = BAxUSEmbedding.create(input_dim=2, target_dim=1, seed=0) - grown = emb.expand(new_bins_on_split=5) - assert grown.target_dim == 2 - - def test_round_trips_through_dump(self) -> None: - emb = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=3) - assert BAxUSEmbedding(**emb.model_dump()) == emb - - -class TestBAxUSUnsupportedOptions: - """vocs-space point controls cannot be honored in an embedded subspace.""" - - @pytest.mark.parametrize( - "kwargs", - [ - {"fixed_features": {"x0": 0.5}}, - {"max_travel_distances": [0.1] * 6}, - {"n_interpolate_points": 3}, - ], - ) - def test_rejected_at_construction(self, kwargs: dict[str, Any]) -> None: - with pytest.raises(ValidationError, match="BAxUS does not support"): - BAxUSGenerator(vocs=_simple_vocs(), **kwargs) - - def test_custom_objective_rejected_at_construction(self) -> None: - """custom_objective must pass its own CustomXoptObjective type validation - before ours can fire, so it needs a real (minimal) subclass instance - rather than a plain kwargs value.""" - vocs = _simple_vocs() - objective = _MinimalCustomObjective(vocs=vocs) - with pytest.raises( - ValidationError, match="BAxUS does not support custom_objective" - ): - BAxUSGenerator(vocs=vocs, custom_objective=objective) - - def test_turbo_controller_rejected(self) -> None: - with pytest.raises( - ValidationError, match="no turbo controllers are compatible" - ): - BAxUSGenerator(vocs=_simple_vocs(), turbo_controller="optimize") - - @pytest.mark.parametrize( - ("extra_variables", "vocs_kwargs", "match"), - [ - pytest.param( - {}, - {"constraints": {"c": ["LESS_THAN", 0.5]}}, - "does not support constraints", - id="constraints", - ), - pytest.param( - {"d": DiscreteVariable(values=[0.0, 1.0, 2.0])}, - {}, - "does not support discrete variables", - id="discrete", - ), - pytest.param( - {"c": ContextualVariable()}, - {}, - "does not support contextual variables", - id="contextual", - ), - ], - ) - def test_unsupported_vocs_rejected( - self, extra_variables: dict[str, Any], vocs_kwargs: dict[str, Any], match: str - ) -> None: - """The target-space GP models only the objective over continuous - vocs-space variables, so each of these must be rejected up front. Each - case guards a supports_* = False override against the True default - inherited from BayesianGenerator / ExpectedImprovementGenerator.""" - vocs = VOCS( - variables={ - **{f"x{i}": [-1.0, 1.0] for i in range(3)}, - **extra_variables, - }, - objectives={"f": "MAXIMIZE"}, - **vocs_kwargs, - ) - with pytest.raises(VOCSError, match=match): - BAxUSGenerator(vocs=vocs) - - -class TestBAxUSAcquisitionPinning: - """Pin the acquisition contract: analytic LogEI on the target-space model - with best_f = best weighted finite objective. Inherited from - ExpectedImprovementGenerator, so these fail loudly if an xopt upgrade - changes EI's acquisition semantics.""" - - def test_analytic_logei_best_f_ignores_nan(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - _seed_sobol(gen) - point = gen.generate(1)[0] - gen.add_data(pd.DataFrame([{**point, "f": float("nan")}])) - - acq = gen.get_acquisition(gen.train_model()) - - assert isinstance(acq, LogExpectedImprovement) - expected = np.nanmax(gen.data["f"].to_numpy(dtype=np.float64)) - assert float(cast(torch.Tensor, acq.best_f)) == pytest.approx(expected) - - def test_minimize_negates_best_f(self) -> None: - vocs = VOCS( - variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, - objectives={"f": "MINIMIZE"}, - ) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) - _seed_sobol(gen) - - acq = gen.get_acquisition(gen.train_model()) - - assert isinstance(acq, LogExpectedImprovement) - expected = -np.min(gen.data["f"].to_numpy(dtype=np.float64)) - assert float(cast(torch.Tensor, acq.best_f)) == pytest.approx(expected) - - -class TestBAxUSIsBayesianGenerator: - def test_subclasses_bayesian_generator(self) -> None: - assert isinstance(BAxUSGenerator(vocs=_simple_vocs()), BayesianGenerator) - - def test_subclasses_expected_improvement_generator(self) -> None: - assert issubclass(BAxUSGenerator, ExpectedImprovementGenerator) - - def test_default_gp_constructor_is_standard(self) -> None: - assert BAxUSGenerator(vocs=_simple_vocs()).gp_constructor.name == "standard" - - -class TestBAxUSSerialization: - """Dump -> construct -> reattach data must resume the run. - - Resume contract: ``gen.model_dump()`` then pop ``"model"`` - xopt ignores - the ``exclude`` argument, so a trained generator's dump carries the live - botorch model (not yaml-safe) unless popped explicitly. ``computation_time`` - needs the same treatment for raw yaml: it is a live ``pd.DataFrame`` once a - BO-phase ``generate()`` has run. - - Data is reattached with ``set_data`` (equivalently, by assigning - ``gen.data``). Note this is what ``Xopt`` itself does for a ``StateOwner``, - so the ordinary ``Xopt.from_yaml`` path resumes correctly too - see - ``TestBAxUSCheckpointIntegrity``. - """ - - def test_dump_excludes_data_and_round_trips_components(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(), seed=7) - _seed_sobol(gen) - - dump = gen.model_dump() - dump.pop("model", None) - assert "data" not in dump - - gen2 = BAxUSGenerator(**dump) - assert gen2.embedding == gen.embedding - assert gen2.trust_region == gen.trust_region - assert gen2.sobol_draws == gen.sobol_draws - - def test_dump_is_yaml_safe(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(), seed=7) - dump = gen.model_dump() - dump.pop("model", None) - mid_dump = BAxUSGenerator(**dump).model_dump() - mid_dump.pop("model", None) - reloaded = yaml.safe_load(yaml.safe_dump(mid_dump)) - assert BAxUSGenerator(**reloaded).embedding == gen.embedding - - # Post-BO a live botorch model and a populated computation_time frame - # ride along; without the pops, yaml.safe_dump raises RepresenterError - # on the trained ModelListGP / the pd.DataFrame. - _seed_sobol(gen) - _bo_step(gen) - bo_dump = gen.model_dump() - bo_dump.pop("model", None) - bo_dump.pop("computation_time", None) - yaml.safe_dump(bo_dump) # must not raise - - gen2 = BAxUSGenerator(**bo_dump) - assert gen2.embedding == gen.embedding - assert gen2.trust_region == gen.trust_region - - def test_sobol_sequence_resumes_exactly(self) -> None: - """Seed phase is bit-identical after a round trip (fast-forwarded engine).""" - g1 = BAxUSGenerator(vocs=_simple_vocs(), seed=7) - for _ in range(2): - point = g1.generate(1)[0] - g1.add_data(pd.DataFrame([{**point, **_sphere(point)}])) - - g2 = _round_trip(g1, with_data=True) - - assert g1.generate(1)[0] == g2.generate(1)[0] - - def test_bo_phase_state_identical_after_round_trip(self) -> None: - """With the same torch seed, the resumed generator proposes the same point.""" - g1 = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - _seed_sobol(g1) - _bo_step(g1) - - g2 = _round_trip(g1, with_data=True) - - torch.manual_seed(0) - p1 = g1.generate(1)[0] - torch.manual_seed(0) - p2 = g2.generate(1)[0] - assert p1 == pytest.approx(p2) - - -class TestBAxUSRegistry: - """The generator must be reachable through xopt's dynamic registry.""" - - def test_get_generator_dynamic(self): - assert get_generator_dynamic("baxus") is BAxUSGenerator - - def test_listed_in_available_generators(self): - assert "baxus" in list_available_generators() - - -class TestBAxUSGetOptimum: - """get_optimum must search the full embedded target space (not vocs-space - bounds against the target-space model) and lift the result to vocs space.""" - - def test_get_optimum_returns_single_row_in_bounds(self) -> None: - vocs = _simple_vocs(n_vars=4) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) - - _seed_sobol(gen) - for _ in range(3): - _bo_step(gen) - - opt = gen.get_optimum() - - assert isinstance(opt, pd.DataFrame) - assert len(opt) == 1 - assert list(opt.columns) == vocs.variable_names - _assert_point_in_bounds(opt.iloc[0].to_dict(), vocs) - - -class TestBAxUSVisualizeModel: - """The inherited visualization is keyed to vocs-space variables and is - meaningless (or crashes) against the embedded target-space model.""" - - def test_visualize_model_raises_not_implemented(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - - with pytest.raises(NotImplementedError, match="target space"): - gen.visualize_model() - - -class TestBAxUSFailureToleranceZeroDivision: - """length_init == length_min makes the reference halvings formula 0, - which must not divide by zero.""" - - def test_equal_length_init_and_min_does_not_raise(self) -> None: - gen = BAxUSGenerator( - vocs=_simple_vocs(n_vars=8), - target_dim_init=2, - n_initial_sobol=2, - eval_budget=50, - length_init=0.5**7, - ) - gen._refresh_failure_tolerance() - assert gen.trust_region.failure_tolerance >= 1 - - -class TestBAxUSOptimizationBounds: - """The trust-region box itself. - - No end-to-end test can pin this down: replacing _get_optimization_bounds with - the full domain, or centering it on the worst point, leaves every behavioral - test passing. It needs direct assertions on the returned box. - """ - - @staticmethod - def _gen_with_known_incumbent() -> tuple[BAxUSGenerator, int]: - """Four points near the origin; row 2 is the unique best (MAXIMIZE).""" - vocs = _simple_vocs(n_vars=4) - gen = BAxUSGenerator( - vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=50 - ) - rows = [ - {"x0": 0.10, "x1": -0.05, "x2": 0.00, "x3": 0.10, "f": -1.0}, - {"x0": -0.10, "x1": 0.05, "x2": 0.10, "x3": -0.05, "f": -2.0}, - {"x0": 0.05, "x1": 0.10, "x2": -0.05, "x3": 0.00, "f": 5.0}, # best - {"x0": 0.00, "x1": 0.00, "x2": 0.05, "x3": 0.05, "f": -3.0}, - ] - gen.add_data(pd.DataFrame(rows)) - return gen, 2 - - def _incumbent_in_target_space(self, gen: BAxUSGenerator, row: int) -> torch.Tensor: - Z = gen.embedding.project(gen._normalized_inputs(gen.data)) - return torch.tensor(Z[row], dtype=torch.double) - - def test_box_is_centered_on_the_incumbent(self) -> None: - """argmin instead of argmax here would center on the worst point.""" - gen, best_row = self._gen_with_known_incumbent() - gen.trust_region.length = 0.4 - - bounds = gen._get_optimization_bounds() - center = (bounds[0] + bounds[1]) / 2.0 - - assert torch.allclose( - center, self._incumbent_in_target_space(gen, best_row), atol=1e-12 - ) - - def test_minimize_centers_on_the_smallest_objective(self) -> None: - """The incumbent is direction-aware, not just the numeric maximum.""" - vocs = VOCS( - variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, - objectives={"f": "MINIMIZE"}, - ) - gen = BAxUSGenerator( - vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=50 - ) - rows = [ - {"x0": 0.10, "x1": -0.05, "x2": 0.00, "x3": 0.10, "f": 4.0}, - {"x0": 0.05, "x1": 0.10, "x2": -0.05, "x3": 0.00, "f": -7.0}, # best - {"x0": 0.00, "x1": 0.00, "x2": 0.05, "x3": 0.05, "f": 2.0}, - ] - gen.add_data(pd.DataFrame(rows)) - gen.trust_region.length = 0.4 - - bounds = gen._get_optimization_bounds() - center = (bounds[0] + bounds[1]) / 2.0 - Z = gen.embedding.project(gen._normalized_inputs(gen.data)) - - assert torch.allclose( - center, torch.tensor(Z[1], dtype=torch.double), atol=1e-12 - ) - - @pytest.mark.parametrize("length", [0.1, 0.25, 0.4]) - def test_box_width_is_the_trust_region_length(self, length: float) -> None: - """In [-1, 1] target units the unclamped width is 2 * length (the length - fields are in [0, 1] units, per the BoTorch reference).""" - gen, _ = self._gen_with_known_incumbent() - gen.trust_region.length = length - - bounds = gen._get_optimization_bounds() - width = bounds[1] - bounds[0] - - assert torch.allclose( - width, torch.full_like(width, 2.0 * length), atol=1e-12 - ), f"expected width {2.0 * length}, got {width.tolist()}" - - def test_box_stays_inside_the_domain(self) -> None: - """A wide region around an incumbent at the edge must clamp, not overflow.""" - vocs = _simple_vocs(n_vars=4) - gen = BAxUSGenerator( - vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=50 - ) - gen.add_data( - pd.DataFrame( - [ - {"x0": 1.0, "x1": 1.0, "x2": 1.0, "x3": 1.0, "f": 5.0}, - {"x0": -1.0, "x1": -1.0, "x2": -1.0, "x3": -1.0, "f": -5.0}, - ] - ) - ) - gen.trust_region.length = gen.trust_region.length_max - - bounds = gen._get_optimization_bounds() - - assert bool((bounds[0] >= -1.0).all()) - assert bool((bounds[1] <= 1.0).all()) - assert bool((bounds[1] > bounds[0]).all()) - - def test_lengthscale_weights_stretch_the_box_per_dimension(self) -> None: - """A dimension with a longer lengthscale gets a wider side.""" - gen, _ = self._gen_with_known_incumbent() - gen.trust_region.length = 0.2 - - uniform = gen._get_optimization_bounds() - gen._lengthscale_weights = lambda model: torch.tensor( # type: ignore[method-assign] - [0.5, 2.0], dtype=torch.double - ) - weighted = gen._get_optimization_bounds() - - uniform_width = uniform[1] - uniform[0] - weighted_width = weighted[1] - weighted[0] - assert weighted_width[0] < uniform_width[0] - assert weighted_width[1] > uniform_width[1] - - -class TestBAxUSCandidateComesFromTheAcquisition: - """The BO candidate must be the acquisition optimum inside the trust region. - - No objective-value comparison can establish this. On a quadratic, the center - of a symmetric domain beats the average random point by exactly the variance - of a random draw - for *any* placement of the optimum - so a generator that - discarded the acquisition result and always proposed the center would still - "beat Sobol on average". The candidate has to be checked against the box the - acquisition was optimized over. - """ - - @staticmethod - def _gen_with_offset_incumbent( - z_star: np.ndarray | None = None, - ) -> tuple[BAxUSGenerator, VOCS]: - """Incumbent placed at a known, off-origin point of the target space.""" - vocs = _simple_vocs(n_vars=6) - gen = BAxUSGenerator( - vocs=vocs, target_dim_init=2, n_initial_sobol=2, seed=0, eval_budget=40 - ) - # lift an explicit target-space point so the incumbent projects back to it - if z_star is None: - z_star = np.array([[0.8, -0.7]]) - x_star = gen.embedding.lift(z_star)[0] - best = {name: float(v) for name, v in zip(vocs.variable_names, x_star)} - - rng = np.random.default_rng(0) - rows = [{**best, "f": 10.0}] - for _ in range(5): - other = { - name: float(v) - for name, v in zip(vocs.variable_names, rng.uniform(-1, 1, 6)) - } - rows.append({**other, "f": -5.0}) - gen.add_data(pd.DataFrame(rows)) - return gen, vocs - - def test_candidate_lies_inside_the_trust_region(self) -> None: - gen, vocs = self._gen_with_offset_incumbent() - - gen.generate(1) # fold in results and fit the model - gen.trust_region.length = 0.1 # narrow box, far from the origin - point = gen.generate(1)[0] - - # no data was added, so this reproduces the box the candidate came from - bounds = gen._get_optimization_bounds() - lb, ub = bounds[0].numpy(), bounds[1].numpy() - - assert not ((lb <= 0.0).all() and (ub >= 0.0).all()), ( - "box contains the origin, so this test could not detect a constant proposal" - ) - - z = gen.embedding.project(gen._normalized_inputs(pd.DataFrame([point])))[0] - assert (z >= lb - 1e-9).all(), f"candidate {z} below trust region {lb}" - assert (z <= ub + 1e-9).all(), f"candidate {z} above trust region {ub}" - - def test_candidate_tracks_the_incumbent(self) -> None: - """Move the incumbent, and the proposal must follow it. - - Lengthscale weights are pinned to 1 so the box is exactly ``length`` wide - around the incumbent. Left to a real fit, a handful of points produces - extreme ARD ratios that widen one side to the whole domain, and the - proposal is then free to sit anywhere along it. - """ - points = [] - for z_star in (np.array([[0.7, 0.7]]), np.array([[-0.7, -0.7]])): - gen, _ = self._gen_with_offset_incumbent(z_star) - - gen.generate(1) - gen.trust_region.length = 0.2 - gen._lengthscale_weights = lambda model: torch.ones( # type: ignore[method-assign] - 2, dtype=torch.double - ) - candidate = gen.generate(1)[0] - points.append( - gen.embedding.project( - gen._normalized_inputs(pd.DataFrame([candidate])) - )[0] - ) - - assert points[0][0] > points[1][0], ( - f"proposal did not follow the incumbent: {points[0]} vs {points[1]}" - ) - assert points[0][1] > points[1][1] - - -class TestBAxUSLengthscaleWeightNormalization: - """The volume-preserving normalization must survive high target_dim. - - Forming the raw product before taking the root underflows to zero once - target_dim reaches a few hundred, which makes every weight inf and silently - widens the trust region to the entire domain. - """ - - @pytest.mark.parametrize("target_dim", [8, 128, 600, 1200]) - def test_weights_are_finite_and_volume_preserving(self, target_dim: int) -> None: - torch.manual_seed(0) - lengthscales = ( - torch.distributions.LogNormal(0.0, 2.0).sample((target_dim,)).double() - ) - - weights = _normalize_lengthscale_weights(lengthscales, target_dim) - - assert bool(torch.isfinite(weights).all()), "weights underflowed to inf" - assert bool((weights > 0).all()) - # unit geometric mean - assert float(weights.log().mean().exp()) == pytest.approx(1.0, rel=1e-9) - - def test_ordering_is_preserved(self) -> None: - lengthscales = torch.tensor([0.5, 4.0, 1.0], dtype=torch.double) - weights = _normalize_lengthscale_weights(lengthscales, 3) - assert weights.argsort().tolist() == lengthscales.argsort().tolist() - - -class TestBAxUSExactScheduleConstants: - """The trust region halves and doubles by exactly a factor of two. - - Asserting only the direction (got smaller / got larger) leaves the schedule - free to drift, which changes how many failures it takes to trigger an - embedding expansion. - """ - - def test_failure_halves_the_length_exactly(self) -> None: - tr = BAxUSTrustRegion(target_dim=4, failure_tolerance=1, length=0.8) - tr.best_value = 10.0 - tr.update(torch.tensor([1.0])) # no improvement - assert tr.length == pytest.approx(0.4, rel=1e-12) - - def test_success_doubles_the_length_exactly(self) -> None: - tr = BAxUSTrustRegion(target_dim=4, success_tolerance=1, length=0.4) - tr.best_value = 0.0 - tr.update(torch.tensor([5.0])) # improvement - assert tr.length == pytest.approx(0.8, rel=1e-12) - - def test_growth_is_capped_at_length_max(self) -> None: - tr = BAxUSTrustRegion( - target_dim=4, success_tolerance=1, length=1.0, length_max=1.6 - ) - tr.best_value = 0.0 - for step in range(3): - tr.update(torch.tensor([10.0 * (step + 1)])) - assert tr.length == pytest.approx(1.6, rel=1e-12) - - def test_improvement_needs_to_clear_the_tolerance(self) -> None: - """A gain below 0.1% of the incumbent counts as a failure.""" - tr = BAxUSTrustRegion(target_dim=4, failure_tolerance=1, length=0.8) - tr.best_value = 100.0 - tr.update(torch.tensor([100.05])) # +0.05% -> not an improvement - assert tr.failure_counter == 0 # reset after the shrink fired - assert tr.length == pytest.approx(0.4, rel=1e-12) - - tr2 = BAxUSTrustRegion(target_dim=4, failure_tolerance=1, length=0.8) - tr2.best_value = 100.0 - tr2.update(torch.tensor([100.5])) # +0.5% -> an improvement - assert tr2.length == pytest.approx(0.8, rel=1e-12) - assert tr2.success_counter == 1 - - -class TestBAxUSSobolDomain: - """Seed points must cover the whole target cube, not just one orthant.""" - - def test_sobol_draws_span_negative_and_positive(self) -> None: + def test_accumulated_failures_expand_organically(self) -> None: + # the composition _force_expansion skips: real BO steps whose results + # never improve, counted by generate's own advance until the region + # collapses below length_min, restarts, and expands. Without a budget + # the tolerance is ceil(2 / 2) = 1, so every miss halves; length_init + # is shortened so the collapse takes 3 halvings instead of 7. gen = BAxUSGenerator( - vocs=_simple_vocs(n_vars=6), + vocs=_simple_vocs(n_vars=16), target_dim_init=2, - n_initial_sobol=32, seed=0, - eval_budget=60, + length_init=0.5**4, ) - zs = np.vstack([gen._draw_sobol_point() for _ in range(32)]) - - assert zs.min() < -0.2, f"no negative seed coordinates (min {zs.min():.3f})" - assert zs.max() > 0.2, f"no positive seed coordinates (max {zs.max():.3f})" - assert (zs >= -1.0).all() and (zs <= 1.0).all() - assert abs(float(zs.mean())) < 0.25, "seed points are not centered" + _seed_sobol(gen) + ys = iter(range(0, -10, -1)) # the first fold-in succeeds, the rest fail + for _ in range(10): + point = gen.generate(1)[0] + if gen.n_expansions: + break + gen.add_data(pd.DataFrame([{**point, "f": float(next(ys))}])) + else: + pytest.fail("failures never accumulated into an expansion") -class TestBAxUSExpansionRandomization: - """Splitting must randomize which dimensions land in which sub-bin. + assert gen.embedding.target_dim == 8 + assert gen.trust_region.length == gen.length_init + assert gen.trust_region.best_value == 0.0 # carried across the reset + assert not gen.trust_region.restart_triggered + assert gen._failure_tolerance() == 4 # recomputed: ceil(8 / 2) + _assert_point_in_bounds(point, gen.vocs) + assert all(math.isfinite(v) for v in point.values()) - The reference permutes the contributing dimensions before splitting. Without - it the split is a fixed function of column index, so two dimensions that are - adjacent in index and share a bin are never separated until the embedding - reaches full dimensionality - which is exactly where BAxUS has no advantage - left. That failure is invisible to structural checks: the matrix is still a - valid partition either way. - """ +class TestBAxUSEmbedding: INPUT_DIM = 64 B = 3 @@ -1457,7 +629,7 @@ def _bin_of(emb: BAxUSEmbedding, col: int) -> int: return int(np.nonzero(matrix[:, col])[0][0]) def _collision_rate(self, target_dim: int, trials: int = 200) -> float: - """How often input dims 0 and 1 (adjacent in index) share a bin.""" + # how often input dims 0 and 1 (adjacent in index) share a bin collisions = 0 for trial in range(trials): emb = BAxUSEmbedding.create( @@ -1468,11 +640,31 @@ def _collision_rate(self, target_dim: int, trials: int = 200) -> float: collisions += self._bin_of(emb, 0) == self._bin_of(emb, 1) return collisions / trials - def test_expansion_separates_adjacent_dimensions(self) -> None: - """Collision rate must fall as the embedding grows. + def test_create_structure_and_signs(self) -> None: + # an all-positive matrix is structurally valid but biased, so the + # structural check alone would not catch it + emb = BAxUSEmbedding.create(input_dim=200, target_dim=4, seed=0) + S = np.asarray(emb.matrix) + assert S.shape == (4, 200) + for col in range(200): + assert np.count_nonzero(S[:, col]) == 1 + assert abs(S[:, col]).max() == 1.0 - A deterministic split leaves it pinned at P(same initial bin) = 1/2. - """ + positive = int((S == 1.0).sum()) + negative = int((S == -1.0).sum()) + assert positive and negative, "embedding signs are not randomized" + assert 0.3 < positive / (positive + negative) < 0.7 + + def test_lift_project_round_trip(self) -> None: + # points lifted from target space project back exactly (S rows are orthogonal) + emb = BAxUSEmbedding.create(input_dim=6, target_dim=2, seed=0) + Z = np.array([[0.5, -0.25], [1.0, 1.0]]) + assert np.allclose(emb.project(emb.lift(Z)), Z) + + def test_expansion_separates_adjacent_dimensions(self) -> None: + # a deterministic split leaves the collision rate pinned at + # P(same initial bin) = 1/2, which structural checks cannot see: the + # matrix is a valid partition either way rate_8 = self._collision_rate(target_dim=8) rate_32 = self._collision_rate(target_dim=32) @@ -1481,17 +673,8 @@ def test_expansion_separates_adjacent_dimensions(self) -> None: f"expanding did not decorrelate the split: {rate_8:.3f} -> {rate_32:.3f}" ) - def test_expand_is_reproducible_for_a_given_seed(self) -> None: - emb = BAxUSEmbedding.create(input_dim=self.INPUT_DIM, target_dim=2, seed=0) - assert emb.expand(self.B, seed=7) == emb.expand(self.B, seed=7) - - def test_different_seeds_give_different_partitions(self) -> None: - emb = BAxUSEmbedding.create(input_dim=self.INPUT_DIM, target_dim=2, seed=0) - assert emb.expand(self.B, seed=1) != emb.expand(self.B, seed=2) - @pytest.mark.parametrize("seed", [0, 1, 2]) def test_randomized_expansion_keeps_the_matrix_well_formed(self, seed: int) -> None: - """Every input dimension stays assigned to exactly one bin, with its sign.""" emb = BAxUSEmbedding.create(input_dim=self.INPUT_DIM, target_dim=2, seed=seed) original = np.asarray(emb.matrix) original_sign = { @@ -1511,37 +694,133 @@ def test_randomized_expansion_keeps_the_matrix_well_formed(self, seed: int) -> N nz = np.nonzero(matrix[:, col])[0][0] assert matrix[nz, col] == original_sign[col], "sign changed on split" - # at full dimensionality the embedding is a signed permutation + # at full dimensionality the column/row checks above force a signed + # permutation by pigeonhole assert emb.target_dim == self.INPUT_DIM - assert (np.count_nonzero(np.asarray(emb.matrix), axis=1) == 1).all() - def test_generator_expansion_survives_a_round_trip(self) -> None: - """The expansion seed is derived from (seed, n_expansions), so a run - restored mid-flight expands to the same embedding.""" - vocs = _simple_vocs(n_vars=16) - gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=3, eval_budget=60) - _seed_sobol(gen) - restored = _round_trip(gen, with_data=True) +# everything BAxUS refuses at construction. The constraints case guards a +# supports_constraints = False override against the True inherited from +# ExpectedImprovementGenerator, so it surfaces as VOCSError; everything else is +# BAxUS' own validator raising ValidationError. Discrete/contextual variables are +# base-class defaults BAxUS merely re-declares, covered by test_generator.py. +class TestBAxUSUnsupportedOptions: + @pytest.mark.parametrize( + ("kwargs", "error", "match"), + [ + ( + {"fixed_features": {"x0": 0.5}}, + ValidationError, + "support fixed_features", + ), + ( + {"max_travel_distances": [0.1] * 6}, + ValidationError, + "max_travel_distances", + ), + ({"n_interpolate_points": 3}, ValidationError, "n_interpolate_points"), + # custom_objective must clear its own CustomXoptObjective type validation + # before ours can fire, so it needs a real (minimal) subclass instance + ( + {"custom_objective": _MinimalCustomObjective(vocs=_simple_vocs())}, + ValidationError, + "support custom_objective", + ), + ({"turbo_controller": "optimize"}, ValidationError, "no turbo controllers"), + ( + _vocs_with(constraints={"c": ["LESS_THAN", 0.5]}), + VOCSError, + "constraints", + ), + (_vocs_with(observables=["g"]), ValidationError, "support observables"), + ], + ids=[ + "fixed_features", + "max_travel_distances", + "n_interpolate_points", + "custom_objective", + "turbo_controller", + "constraints", + "observables", + ], + ) + def test_rejected_at_construction( + self, kwargs: dict[str, Any], error: type[Exception], match: str + ) -> None: + with pytest.raises(error, match=match): + BAxUSGenerator(**{"vocs": _simple_vocs(), **kwargs}) - for generator in (gen, restored): - _force_expansion(generator) - generator.add_data( - pd.DataFrame([{**{f"x{i}": 0.1 for i in range(16)}, "f": -1.0}]) - ) - generator.generate(1) + def test_declares_its_unsupported_features(self) -> None: + gen = BAxUSGenerator(vocs=_simple_vocs()) + assert not gen.supports_batch_generation + assert not gen.supports_constraints + assert not gen.supports_discrete_variables + assert not gen.supports_contextual_variables + assert not gen.supports_no_objective + + +# saving and reloading must continue a run, not restart or advance it. By hand +# the resume contract is model_dump() then pop "model" - xopt ignores the exclude +# argument, so a trained generator's dump carries the live botorch model. +# Through Xopt.from_yaml the same set_data is what Xopt calls for a StateOwner, +# which must not replay the trust-region history the frame already reflects. +class TestBAxUSResume: + def test_dump_is_yaml_safe_and_round_trips_components(self) -> None: + # eval_budget rides along because the failure tolerance is derived from it + # at every trust-region update: losing it silently downgrades a resumed + # run to the ceil(target_dim / 2) heuristic. + gen = BAxUSGenerator(vocs=_simple_vocs(), seed=7, eval_budget=50) + _seed_sobol(gen) # so sobol_draws is non-zero and has to survive the trip - assert gen.n_expansions == 1 - assert restored.embedding == gen.embedding + dump = gen.model_dump() + dump.pop("model", None) + mid = BAxUSGenerator(**dump) + assert mid.embedding == gen.embedding + assert mid.trust_region == gen.trust_region + assert mid.eval_budget == gen.eval_budget + assert mid.sobol_draws == gen.sobol_draws > 0 -class TestBAxUSCheckpointIntegrity: - """Saving and reloading through Xopt must not advance the algorithm. + mid_dump = mid.model_dump() + mid_dump.pop("model", None) + reloaded = yaml.safe_load(yaml.safe_dump(mid_dump)) + assert BAxUSGenerator(**reloaded).embedding == gen.embedding - Xopt's data validator hands a restored dataframe to the generator; for a - generator that owns state, that must not replay the history it already - reflects. BAxUS is a StateOwner precisely so this is a no-op. - """ + # Post-BO a live botorch model and a populated computation_time frame + # ride along; without the pops, yaml.safe_dump raises RepresenterError + # on the trained ModelListGP / the pd.DataFrame. + point = gen.generate(1)[0] + gen.add_data(pd.DataFrame([{**point, **_sphere(point)}])) + bo_dump = gen.model_dump() + bo_dump.pop("model", None) + bo_dump.pop("computation_time", None) + yaml.safe_dump(bo_dump) # must not raise + assert BAxUSGenerator(**bo_dump).trust_region == gen.trust_region + + # a restored run continues the sequence rather than restarting it: in the seed + # phase bit-identically, in the BO phase only up to GP fitting and + # acquisition-optimization noise, so both are driven from the same torch seed + @pytest.mark.parametrize("phase", ["sobol", "bo"]) + def test_round_trip_proposes_the_same_next_point(self, phase: str) -> None: + g1 = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=7) + if phase == "sobol": + # stop one point short of the quota, so the next call still seeds + for _ in range(g1.n_initial_sobol - 1): + point = g1.generate(1)[0] + g1.add_data(pd.DataFrame([{**point, **_sphere(point)}])) + else: + _seed_sobol(g1) + # a post-quota row by hand, so the round trip starts in the BO phase + # without paying for a fit here; the compared generates fit anyway + g1.add_data(pd.DataFrame([{**{f"x{i}": 0.1 for i in range(4)}, "f": -0.5}])) + + g2 = _round_trip(g1, with_data=True) + + torch.manual_seed(0) + p1 = g1.generate(1)[0] + torch.manual_seed(0) + p2 = g2.generate(1)[0] + assert p1 == pytest.approx(p2) @staticmethod def _state(gen: BAxUSGenerator) -> dict[str, Any]: @@ -1552,14 +831,14 @@ def _state(gen: BAxUSGenerator) -> dict[str, Any]: "length": tr.length, "success_counter": tr.success_counter, "failure_counter": tr.failure_counter, - "failure_tolerance": tr.failure_tolerance, + "failure_tolerance": gen._failure_tolerance(), "best_value": tr.best_value, "sobol_draws": gen.sobol_draws, "n_expansions": gen.n_expansions, "tr_observed_rows": gen.tr_observed_rows, } - def _run(self, steps: int = 12) -> Xopt: + def _run(self, steps: int = 6) -> Xopt: vocs = _simple_vocs(n_vars=6) X = Xopt( evaluator=XoptEvaluator(function=_sphere), @@ -1571,173 +850,247 @@ def _run(self, steps: int = 12) -> Xopt: X.step() return X - def test_is_a_state_owner(self) -> None: - assert isinstance(BAxUSGenerator(vocs=_simple_vocs(), seed=0), StateOwner) - - def test_repeated_checkpoint_cycles_do_not_move_the_trust_region(self) -> None: - """No evaluations happen between cycles, so nothing may change.""" + def test_checkpoint_cycles_hold_state_and_resume_folds_in_only_new_results( + self, + ) -> None: + # no evaluations happen between cycles, so nothing may change across them. + # The proposed point is not compared exactly: xopt's yaml writer keeps 10 + # significant digits, and that ~1e-10 perturbation is amplified by GP + # hyperparameter fitting and the L-BFGS acquisition optimization. X = self._run() - expected = self._state(X.generator) - assert len(X.data) == 12 - - for cycle in range(6): - X = Xopt.from_yaml(X.yaml()) - assert len(X.data) == 12, f"data changed on cycle {cycle}" - assert self._state(X.generator) == expected, ( + before = self._state(X.generator) + assert len(X.data) == 6 + # the row the resumed run must fold in: the last one evaluated pre-restore + new_y = float(X.data["f"].to_numpy(dtype=np.float64)[-1]) + oracle = X.generator.trust_region.model_copy(deep=True) + + resumed = X + for cycle in range(2): + resumed = Xopt.from_yaml(resumed.yaml()) + assert len(resumed.data) == 6, f"data changed on cycle {cycle}" + assert self._state(resumed.generator) == before, ( f"generator state drifted on checkpoint cycle {cycle}" ) - def test_resume_folds_in_only_the_new_results(self) -> None: - """The pre-restore history must not be re-ingested. - - A replay would re-consume everything from the seed quota onward; a - correct resume advances the trust region by exactly the one row that was - evaluated after the checkpoint. - """ - X = self._run() - before = self._state(X.generator) - - resumed = Xopt.from_yaml(X.yaml()) - resumed.step() + for runner in (X, resumed): + runner.step() after = self._state(resumed.generator) assert after["tr_observed_rows"] == before["tr_observed_rows"] + 1 assert after["n_expansions"] == before["n_expansions"] - # one tick only: at most one counter moved, by one - assert abs(after["failure_counter"] - before["failure_counter"]) <= 1 - assert abs(after["success_counter"] - before["success_counter"]) <= 1 - - def test_resumed_run_continues_the_same_trajectory(self) -> None: - """One more step on the original and on the restored run agree. - - The proposed point is not compared exactly: xopt's yaml writer keeps 10 - significant digits, and that ~1e-10 perturbation of the training data is - amplified by GP hyperparameter fitting and the L-BFGS acquisition - optimization. The algorithm state is insensitive to it. - """ - X = self._run() - resumed = Xopt.from_yaml(X.yaml()) - for runner in (X, resumed): - runner.step() + # exactly one tick, not "at most one": replaying that single observation + # through a copy of the pre-restore region must reproduce the state. A + # replay of the history would move the counters further (or not at all). + oracle.update( + torch.tensor([X.generator._objective_weight() * new_y], dtype=torch.double), + X.generator._failure_tolerance(), + ) + assert (after["success_counter"], after["failure_counter"]) == ( + oracle.success_counter, + oracle.failure_counter, + ) + assert after["length"] == pytest.approx(oracle.length, rel=1e-12) + # ...and the same step on the un-checkpointed run lands in the same place; + # only the components the oracle replay above does not already pin original, restored = self._state(X.generator), self._state(resumed.generator) - for key in ( - "target_dim", - "matrix", - "length", - "success_counter", - "failure_counter", - "failure_tolerance", - "sobol_draws", - "n_expansions", - "tr_observed_rows", - ): + for key in ("target_dim", "matrix", "failure_tolerance", "sobol_draws"): assert original[key] == restored[key], f"{key} diverged after resume" assert original["best_value"] == pytest.approx(restored["best_value"], rel=1e-6) -class TestBAxUSFailedEvaluations: - """An evaluator that raises returns a row with no objective column at all.""" +# the trust-region box, and the lengthscale weights that shape it. No end-to-end +# test can pin the box down: replacing _get_optimization_bounds with the full +# domain, or centering it on the worst point, leaves every behavioral test +# passing - it needs direct assertions on the returned box. +class TestBAxUSOptimizationBounds: + Z_ROWS = [[0.1, 0.7], [-0.6, 0.3], [0.5, -0.4], [-0.2, -0.8]] + YS = [-1.0, -2.0, 5.0, -3.0] # row 2 best; row 3 is what argmin would pick + BEST_ROW = 2 + + @classmethod + def _gen(cls, direction: str = "MAXIMIZE") -> BAxUSGenerator: + return _gen_with_target_space_data(cls.Z_ROWS, cls.YS, direction=direction) + + # argmin instead of argmax here would center on the worst point, and the + # incumbent is direction-aware rather than just the numeric maximum. In + # [-1, 1] target units the unclamped width is 2 * length (the length fields + # are in [0, 1] units, per the BoTorch reference). + @pytest.mark.parametrize( + ("direction", "length"), + [("MAXIMIZE", 0.1), ("MAXIMIZE", 0.25), ("MAXIMIZE", 0.4), ("MINIMIZE", 0.4)], + ) + def test_box_is_centered_on_the_incumbent_and_2_length_wide( + self, direction: str, length: float + ) -> None: + gen = self._gen(direction) + gen.trust_region.length = length - def test_add_data_tolerates_a_missing_objective_column(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) + bounds = gen._get_optimization_bounds() + center = (bounds[0] + bounds[1]) / 2.0 + width = bounds[1] - bounds[0] - failed = {f"x{i}": 0.1 for i in range(4)} - failed["xopt_error"] = True - gen.add_data(pd.DataFrame([failed])) + incumbent = _in_target_space(gen, gen.data)[self.BEST_ROW] + assert torch.allclose( + center, torch.tensor(incumbent, dtype=torch.double), atol=1e-12 + ) + assert torch.allclose( + width, torch.full_like(width, 2.0 * length), atol=1e-12 + ), f"expected width {2.0 * length}, got {width.tolist()}" - assert len(gen.data) == 1 - assert gen._finite_data().empty + def test_box_stays_inside_the_domain(self) -> None: + # at length_max the box is wider than the domain on both sides of the + # incumbent, so it must clamp rather than overflow + gen = self._gen() + gen.trust_region.length = gen.trust_region.length_max - def test_xopt_run_survives_a_failing_evaluator(self) -> None: - """The generator must not be the reason a non-strict run dies.""" - vocs = _simple_vocs(n_vars=4) - calls = {"n": 0} + bounds = gen._get_optimization_bounds() - def flaky(inputs: dict[str, float]) -> dict[str, float]: - calls["n"] += 1 - if calls["n"] == 4: - raise RuntimeError("evaluation blew up") - return _sphere(inputs) + assert bool((bounds[0] >= -1.0).all()) + assert bool((bounds[1] <= 1.0).all()) + assert bool((bounds[1] > bounds[0]).all()) - X = Xopt( - evaluator=XoptEvaluator(function=flaky), - generator=BAxUSGenerator( - vocs=vocs, target_dim_init=2, seed=0, eval_budget=30 - ), - strict=False, + def test_lengthscale_weights_stretch_the_box_per_dimension(self) -> None: + gen = self._gen() + gen.trust_region.length = 0.2 + gen._lengthscale_weights = lambda model: torch.tensor( # type: ignore[method-assign] + [0.5, 2.0], dtype=torch.double ) - for _ in range(8): - X.step() - assert len(X.data) == 8 - assert len(X.generator._finite_data()) == 7 # the failed row is excluded + weighted = gen._get_optimization_bounds() + # 2 * length * weight per dimension, unclamped at this incumbent + width = weighted[1] - weighted[0] + assert torch.allclose( + width, torch.tensor([0.2, 0.8], dtype=torch.double), atol=1e-12 + ), f"expected widths [0.2, 0.8], got {width.tolist()}" -class TestBAxUSObservables: - """The target-space model has a single outcome while the inherited - acquisition scalarizes over every vocs output, so observables cannot work.""" + @staticmethod + def _kernel_model(lengthscale: torch.Tensor) -> MagicMock: + # a mock shaped like a gpytorch model: covar_module.base_kernel.lengthscale + kernel = MagicMock() + kernel.lengthscale = lengthscale + covar = MagicMock() + covar.base_kernel = kernel + model = MagicMock(spec=["covar_module"]) + model.covar_module = covar + return model - def test_observables_rejected_at_construction(self) -> None: - vocs = VOCS( - variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, - objectives={"f": "MAXIMIZE"}, - observables=["g"], + def test_kernel_lengthscale_is_used_with_a_uniform_fallback(self) -> None: + # a vector lengthscale gives non-uniform weights, a scalar one is + # unsqueezed rather than collapsing to a 0-dim tensor, and a model with no + # kernel at all falls back to an isotropic region + gen = _bare_gen() + model = self._kernel_model(torch.tensor([[0.5, 2.0]], dtype=torch.double)) + + weights = gen._lengthscale_weights(model) + assert weights.shape == (2,) + assert not torch.allclose(weights, torch.ones(2, dtype=torch.double)) + + scalar_gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=1) + scalar_model = self._kernel_model(torch.tensor(0.5, dtype=torch.double)) + assert scalar_gen._lengthscale_weights(scalar_model).shape == (1,) + + bare = MagicMock(spec=["posterior", "num_outputs"]) + assert torch.allclose( + gen._lengthscale_weights(bare), torch.ones(2, dtype=torch.double) ) - with pytest.raises(ValidationError, match="does not support observables"): - BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) + # forming the raw product before taking the root underflows to zero once + # target_dim reaches a few hundred, making every weight inf and silently + # widening the trust region to the entire domain + @pytest.mark.parametrize("target_dim", [8, 1200]) + def test_weights_are_finite_volume_preserving_and_ordered( + self, target_dim: int + ) -> None: + torch.manual_seed(0) + lengthscales = ( + torch.distributions.LogNormal(0.0, 2.0).sample((target_dim,)).double() + ) -class TestBAxUSEvalBudgetWarning: - """Without eval_budget the reference expansion schedule is unavailable, so - construction warns. It must warn exactly once: the generator assigns to its - own fields on every BO step, and a model validator of either mode re-fires - on assignment under validate_assignment, which would bury the run in - duplicate warnings.""" + weights = _normalize_lengthscale_weights(lengthscales, target_dim) - def test_warns_once_on_construction(self) -> None: - with pytest.warns(GeneratorWarning, match="eval_budget is not set") as caught: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - assert len(caught) == 1 + assert bool(torch.isfinite(weights).all()), "weights underflowed to inf" + assert bool((weights > 0).all()) + # unit geometric mean + assert float(weights.log().mean().exp()) == pytest.approx(1.0, rel=1e-9) + # the ordering keeps the normalization from inverting the weights + assert weights.argsort().tolist() == lengthscales.argsort().tolist() - with warnings.catch_warnings(record=True) as on_assignment: - warnings.simplefilter("always") - gen.sobol_draws = 3 - gen.n_expansions = 1 - assert [w for w in on_assignment if w.category is GeneratorWarning] == [] - def test_no_warning_when_budget_is_set(self) -> None: - vocs = _simple_vocs(n_vars=4) - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0, eval_budget=50) - assert [w for w in caught if w.category is GeneratorWarning] == [] +# analytic LogEI on the target-space model with best_f = the best weighted finite +# objective, inherited from ExpectedImprovementGenerator. No objective-value +# comparison can establish that the candidate comes from it: on a quadratic the +# center of a symmetric domain beats the average random point for *any* placement +# of the optimum, so a generator that always proposed the center would still +# "beat Sobol on average" - the candidate has to be checked against the box. +class TestBAxUSAcquisition: + @pytest.mark.parametrize("direction", ["MAXIMIZE", "MINIMIZE"]) + def test_best_f_is_the_best_weighted_finite_objective(self, direction: str) -> None: + # MINIMIZE negates; a NaN row never reaches best_f in either direction + vocs = VOCS( + variables={f"x{i}": [-1.0, 1.0] for i in range(4)}, + objectives={"f": direction}, + ) + gen = BAxUSGenerator(vocs=vocs, target_dim_init=2, seed=0) + gen.add_data( + pd.DataFrame( + [ + {**{f"x{i}": x for i in range(4)}, "f": f} + for x, f in [ + (-0.5, -3.0), + (0.1, 2.0), + (0.6, 0.5), + (0.9, float("nan")), + ] + ] + ) + ) + acq = gen.get_acquisition(gen.train_model()) -class TestBAxUSComputationTime: - """generate() must record training/acquisition-optimization timing during - the BO phase, matching the base class's bookkeeping.""" + ys = gen.data["f"].to_numpy(dtype=np.float64) + expected = np.nanmax(ys) if direction == "MAXIMIZE" else -np.nanmin(ys) + assert isinstance(acq, LogExpectedImprovement) + assert float(cast(torch.Tensor, acq.best_f)) == pytest.approx(expected) - def test_bo_generate_records_one_row(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - _seed_sobol(gen) + # five clearly worse points to surround the incumbent with + Z_OTHERS = [[-0.9, 0.9], [0.2, 0.4], [-0.3, -0.6], [0.6, 0.1], [-0.5, -0.2]] - assert gen.computation_time is None - _bo_step(gen) + @classmethod + def _gen_with_incumbent_at(cls, z_star: list[float]) -> BAxUSGenerator: + return _gen_with_target_space_data( + [z_star, *cls.Z_OTHERS], [10.0, *[-5.0] * 5], n_vars=6 + ) - assert isinstance(gen.computation_time, pd.DataFrame) - assert len(gen.computation_time) == 1 - assert list(gen.computation_time.columns) == [ - "training", - "acquisition_optimization", - ] + def test_candidate_stays_in_the_box_and_tracks_the_incumbent(self) -> None: + # lengthscale weights are pinned to 1 so the box is exactly ``length`` wide + # around the incumbent; left to a real fit, a handful of points produces + # extreme ARD ratios that widen one side to the whole domain + points = [] + for z_star in ([0.7, 0.7], [-0.7, -0.7]): + gen = self._gen_with_incumbent_at(z_star) - def test_second_bo_generate_appends_row(self) -> None: - gen = BAxUSGenerator(vocs=_simple_vocs(n_vars=4), target_dim_init=2, seed=0) - _seed_sobol(gen) + gen._advance_trust_region() # fold the results in without a fit + gen.trust_region.length = 0.2 + gen._lengthscale_weights = lambda model: torch.ones( # type: ignore[method-assign] + 2, dtype=torch.double + ) + candidate = gen.generate(1)[0] + z = _in_target_space(gen, pd.DataFrame([candidate]))[0] - _bo_step(gen) - _bo_step(gen) + # no data was added, so this reproduces the box the candidate came from + bounds = gen._get_optimization_bounds() + lb, ub = bounds[0].numpy(), bounds[1].numpy() + assert not ((lb <= 0.0).all() and (ub >= 0.0).all()), ( + "box contains the origin: cannot detect a constant proposal" + ) + assert (z >= lb - 1e-9).all(), f"candidate {z} below trust region {lb}" + assert (z <= ub + 1e-9).all(), f"candidate {z} above trust region {ub}" + points.append(z) - assert len(gen.computation_time) == 2 + assert points[0][0] > points[1][0], ( + f"proposal did not follow the incumbent: {points[0]} vs {points[1]}" + ) + assert points[0][1] > points[1][1]