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..47a8d5f93 --- /dev/null +++ b/docs/examples/single_objective_bayes_opt/baxus.ipynb @@ -0,0 +1,235 @@ +{ + "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" + ] + }, + { + "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" + ] + }, + { + "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()" + ] + }, + { + "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()" + ] + }, + { + "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(\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)" + ] + }, + { + "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" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b36961c9", + "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]" + ] + } + ], + "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..7d1a749ac --- /dev/null +++ b/xopt/generators/bayesian/baxus.py @@ -0,0 +1,638 @@ +"""BAxUS generator - Bayesian Optimization with Adaptively Expanding Subspaces. + +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 +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, + 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 to unit geometric mean. + + The geometric mean is taken as a product of per-element roots (the reference + 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): + """Trust-region state, in normalized [0, 1] target-space units. + + The generator owns dimensionality (``embedding.target_dim``) and passes its + derived failure tolerance into every ``update``. + """ + + length: 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 + # None until a BO-phase evaluation is folded in + best_value: float | None = None + restart_triggered: bool = False + + 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) + 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 >= 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): + """Sparse random embedding of shape (target_dim, input_dim). + + 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]] + + @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 lift(self, Z: np.ndarray) -> np.ndarray: + """Lift from target subspace to full input space: ``X = Z @ S``.""" + 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(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, 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 = 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) + + 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): + name = "baxus" + supports_batch_generation: bool = False + # 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 instead" + ), + ) + eval_budget: PositiveInt | None = Field( + default=None, + description=( + "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)", + ) + 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 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, " + "which prevents re-ingesting history on resume" + ), + ) + # 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; the low-noise prior " + "matches 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: + """Create the embedding and trust region, and size the Sobol seed quota. + + A dump -> construct round trip supplies these 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: + values["trust_region"] = BAxUSTrustRegion( + length=float(values.get("length_init") or 0.8) + ) + + return values + + def model_post_init(self, context: Any, /) -> None: + """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( + "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.""" + # 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: + 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, 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) + + def _get_torch_bounds(self) -> torch.Tensor: + """The full embedded target space, ``[-1, 1]^target_dim``. + + 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: + """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() + + 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 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) + + 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``). + """ + self.data = data + + 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 region + (matching the reference), and ``tr_observed_rows`` keeps a resumed run + from re-ingesting history. + """ + 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 + + self.trust_region.update( + 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, + # 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", + self.trust_region.length, + previous_dim, + self.embedding.target_dim, + ) + self.n_expansions += 1 + # 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 + + 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.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, 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``. + + 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( + data[self._objective_name].to_numpy(dtype=np.float64), + dtype=torch.double, + ) + best_idx = int(torch.argmax(self._objective_weight() * y)) + + Z = self.embedding.project(self._normalized_inputs(data)) + center = torch.tensor(Z[best_idx], dtype=torch.double) + + 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 _failure_tolerance(self) -> int: + """The trust region's failure tolerance, budget-aware when possible. + + 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 ``target_dim_init``. + """ + target_dim = self.embedding.target_dim + input_dim = self.embedding.input_dim + if target_dim >= input_dim: + return target_dim + if self.eval_budget is None: + 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 + n_splits = math.ceil(math.log(input_dim / d_init, b + 1) - 1e-9) + # 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(self.trust_region.length_min / self.length_init, 0.5)), + ) + # spend that share evenly over the halvings needed to reach length_min + 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 _lift_to_vocs(self, z: np.ndarray) -> np.ndarray: + """Lift through the embedding and scale to 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 = 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: + """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. + + ``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: + # 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: + 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. + + [-1, 1] rather than ``vocs.normalize_inputs``' [0, 1]: these feed a matrix + product with the embedding, which is sign-symmetric. + """ + 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 _draw_sobol_point(self) -> np.ndarray: + """Next quasi-random seed point in target space, ``[-1, 1]^target_dim``. + + 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: + 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. + + 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 + + # 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) + + 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..76d9c512f --- /dev/null +++ b/xopt/tests/generators/bayesian/test_baxus.py @@ -0,0 +1,1096 @@ +import math +import warnings +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.generators import get_generator_dynamic +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 + + +# 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(): + 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: + 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: + 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): + 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: + for _ in range(gen.n_initial_sobol): + point = gen.generate(1)[0] + gen.add_data(pd.DataFrame([{**point, **_sphere(point)}])) + + +def _force_expansion(gen: BAxUSGenerator) -> None: + # 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: + # 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) + if with_data and gen.data is not None: + restored.data = gen.data.copy() + return restored + + +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 + + +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: + gen.embedding = BAxUSEmbedding.create( + input_dim=gen.embedding.input_dim, target_dim=target_dim, seed=0 + ) + gen.trust_region = BAxUSTrustRegion() + + +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 + + # 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() + + @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_registry_bases_and_gp_constructor_default(self) -> None: + assert get_generator_dynamic("baxus") is BAxUSGenerator + assert issubclass(BAxUSGenerator, ExpectedImprovementGenerator) + + 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 + + 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] == [] + + +# 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) + g3 = BAxUSGenerator(vocs=vocs, seed=1) + + assert g1.embedding.matrix == g2.embedding.matrix + assert g1.embedding.matrix != g3.embedding.matrix + + 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) + + 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)]) + + # 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 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) + + assert gen._draw_sobol_point().shape == (1, 8) + + +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={"x0": [0.0, 10.0], "x1": [-165.0, 165.0], "x2": [0.5, 0.6]}, + objectives={"f": "MAXIMIZE"}, + ) + calls = {"n": 0} + + 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) + + 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 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 + 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) == 4 + assert len(gen._finite_data()) == 2 + + # 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 + + +# 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 + + tr.update(torch.tensor([next(improving)]), failure_tolerance=1) + assert tr.length == pytest.approx(0.4, rel=1e-12) # exactly doubled + + # 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_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(100): + if tr.restart_triggered: + break + tr.update(torch.tensor([-999.0]), failure_tolerance=1) + assert tr.restart_triggered + assert tr.length < tr.length_min + + # 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)} + + # 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() + + assert gen.trust_region.best_value is None + assert gen.tr_observed_rows == 0 + + 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) + + tr = gen.trust_region + before = (tr.length, tr.success_counter, tr.failure_counter, tr.best_value) + observed = gen.tr_observed_rows + + 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), + target_dim_init=target_dim_init, + n_initial_sobol=10, + seed=0, + eval_budget=eval_budget, + ) + + # 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: + # 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=_simple_vocs(n_vars=32), target_dim_init=2, seed=0, eval_budget=203 + ) + _seed_sobol(gen) + restored = _round_trip(gen, with_data=True) + + 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 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_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=16), + target_dim_init=2, + seed=0, + length_init=0.5**4, + ) + _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") + + 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()) + + +class TestBAxUSEmbedding: + 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_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 + + 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) + + 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}" + ) + + @pytest.mark.parametrize("seed", [0, 1, 2]) + def test_randomized_expansion_keeps_the_matrix_well_formed(self, seed: int) -> None: + 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 column/row checks above force a signed + # permutation by pigeonhole + assert emb.target_dim == self.INPUT_DIM + + +# 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}) + + 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 + + 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 + + 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 + + # 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]: + 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": 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 = 6) -> 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_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() + 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}" + ) + + 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"] + + # 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", "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) + + +# 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 + + bounds = gen._get_optimization_bounds() + center = (bounds[0] + bounds[1]) / 2.0 + width = bounds[1] - bounds[0] + + 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()}" + + 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 + + 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: + 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 + ) + + 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()}" + + @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_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) + ) + + # 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() + ) + + 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) + # the ordering keeps the normalization from inverting the weights + assert weights.argsort().tolist() == lengthscales.argsort().tolist() + + +# 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()) + + 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) + + # 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]] + + @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 + ) + + 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) + + 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] + + # 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 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]