AFL.double_agent.BoTorchMixin module

BoTorch machinery shared by the GP regressors, classifiers and acquisition ops.

Everything BoTorch specific lives on BoTorchMixin: the lazy import of the optional dependency, the tensor helpers that move data between xarray and torch, GP construction, and the bound helpers that read a PipelineOp’s own attributes so an op can call self.candidates(dataset) instead of threading the same variable and dimension names through every call. ScoreFunctionAcquisition is the one module-level object: a torch.nn.Module adapter for user score functions.

class AFL.double_agent.BoTorchMixin.BoTorchMixin

Bases: object

All BoTorch machinery an op needs, as methods on one class.

Mix this into a PipelineOp subclass (list it before the op base class) and call _init_botorch() at the end of __init__.

The class has two layers. Stateless helpers (static/class methods such as optimization_bounds_to_tensor(), fit_single_task_gp(), evaluate_acquisition_on_grid()) take explicit arguments and can be called on the class itself, e.g. BoTorchMixin.fit_single_task_gp(train_x, train_y). Bound helpers (candidates(), bounds_tensor(), load_model(), optimize_acqf(), …) read the following attributes, which the op sets itself:

grid_variable, grid_dim

Candidate grid and the dimension indexing its points; the last remaining dimension (or component_dim if the op defines one) indexes components.

feature_input_variable, predictor_input_variable, sample_dim, objective_direction

Only needed by training_data() (regressors).

bounds, is_simplex, model_variable, model_uri, num_restarts, raw_samples

Set by _init_botorch().

The mixin does not define __init__ so it composes with any PipelineOp constructor signature.

BOUNDS_OUTPUT = 'bounds'

dimension and coordinate names used for the bounds output

static best_grid_point(values: Tensor | ndarray, candidate_x: Tensor) → tuple[Tensor, float]

Fallback when optimization fails: the grid point with the largest (finite) value.

Returns (x, value) with x shaped (1, d).

property botorch: dict

The BoTorch symbols, imported on first use so import AFL works without BoTorch.

bounds_tensor(dataset: Dataset, write_output: bool = True) → Tensor

self.bounds as a (2, d) tensor aligned with the grid’s components.

When write_output is true and the grid has component names, the bounds are also written to the bounds output so downstream ops can read them.

classmethod build_builtin_acquisition(acquisition_kind: str, model, best_f: float)

Instantiate one of the built-in BoTorch acquisition functions.

acquisition_kind is "logei" (analytic LogEI) or "qlogei" (Monte Carlo qLogEI for joint batches).

classmethod build_single_task_gp(train_x: Tensor, train_y: Tensor, standardize: bool = True)

Construct an unfitted SingleTaskGP. Used to rebuild saved models.

candidates(dataset: Dataset) → Tensor

The candidate grid grid_variable as an (n, d) tensor.

component_names(dataset: Dataset) → ndarray | None

Component coordinate values of the grid, or None when the grid has no component coords.

constraints(n_dim: int) → list[tuple[Tensor, Tensor, float]] | None

Simplex inequality constraints when is_simplex is set, otherwise None.

classmethod dataset_to_botorch_candidates(dataset: Dataset, grid_variable: str, grid_dim: str) → Tensor

The candidate grid as an (n, d) double tensor.

classmethod dataset_to_botorch_training_data(dataset: Dataset, feature_input_variable: str, predictor_input_variable: str, sample_dim: str, objective_direction: Literal['maximize', 'minimize'] = 'minimize') → tuple[Tensor, Tensor]

(train_x, train_y) double tensors; train_y is negated when minimizing.

static evaluate_acquisition_on_grid(acq_function, candidate_x: Tensor, q: int = 1) → Tensor

Evaluate an acquisition function at every row of candidate_x (n, d).

Points are evaluated individually when q <= 1. For joint acquisitions with q > 1 the grid is scanned with a sliding window of q consecutive points; windows that do not fit are padded with -inf.

classmethod fit_single_task_gp(train_x: Tensor, train_y: Tensor, standardize: bool = True)

Build and fit a SingleTaskGP by maximizing the exact marginal log likelihood.

static get_observed_best_f(train_y: Tensor, objective_direction: Literal['maximize', 'minimize'] = 'maximize') → float

Best observed target in original objective units (train_y is in maximization space).

grid_component_dim(dataset: Dataset) → str

The component dimension of the grid (self.component_dim if the op defines one).

static import_botorch()

Import the BoTorch symbols the ops use, with a single helpful error message.

load_model(dataset: Dataset)

Load the model whose ref an upstream op wrote to model_variable.

static make_simplex_constraints(n_dim: int) → list[tuple[Tensor, Tensor, float]]

Return BoTorch-compatible simplex constraints.

BoTorch’s optimize_acqf accepts linear inequality constraints of the form

\[\sum_i a_i x_i \ge b.\]

A simplex equality \(\sum_i x_i = 1\) is therefore encoded as the pair

\[\sum_i x_i \ge 1, \qquad -\sum_i x_i \ge -1,\]

which together are exactly equivalent to the equality constraint.

static optimization_bounds_to_tensor(bounds: Mapping[str, Mapping[str, float] | Sequence[float]] | Sequence[Sequence[float]], component_names: Sequence[Any] | None = None) → Tensor

Convert explicit constructor bounds into a (2, d) BoTorch bounds tensor.

Parameters:

• bounds – Either a mapping from component name to {"min": lower, "max": upper} dictionaries, a mapping from component name to (lower, upper) pairs, or an ordered sequence of (lower, upper) pairs.

• component_names – Optional ordered component names used to align mapping-based bounds with model inputs.

static optimization_bounds_to_xarray(bounds: Tensor, component_names: Sequence[Any], variable_name: str) → DataArray

Represent optimization bounds as an xarray DataArray.

optimize_acqf(acq_function, bounds: Tensor, q: int = 1, *, constraints: Sequence[tuple[Tensor, Tensor, float]] | None = None, with_grad: bool = True, batch_initial_conditions: Tensor | None = None) → tuple[Tensor, Tensor]

Maximize acq_function over bounds with BoTorch’s optimize_acqf.

Uses self.num_restarts / self.raw_samples. with_grad=False makes the L-BFGS-B optimizer use finite-difference gradients, which is required for score functions that are not differentiable by autograd. Returns detached (candidates, values).

static posterior_to_xarray(posterior, grid_index: DataArray, output_prefix: str, objective_direction: Literal['maximize', 'minimize'] = 'maximize') → Dataset

{prefix}_mean / {prefix}_variance over the grid, mean in original objective units.

save_model(model, args: dict | None = None, builder=None) → str | None

Save model to model_uri and write the ref to {prefix}_model; no-op without model_uri.

training_data(dataset: Dataset) → tuple[Tensor, Tensor]

(train_x, train_y) for this op’s feature/predictor variables, in maximization space.

AFL.double_agent.BoTorchMixin.Bounds

{component: {"min": lo, "max": hi}}, {component: (lo, hi)} or an ordered sequence of (lo, hi) pairs.

Type:

Optimization bounds for BoTorch ops

alias of Mapping[str, Mapping[str, float] | Sequence[float]] | Sequence[Sequence[float]]

class AFL.double_agent.BoTorchMixin.ScoreFunctionAcquisition(function, kwargs: Mapping[str, Any] | None = None)

Bases: Module

Adapt a user score function to BoTorch’s acquisition interface.

The wrapped function scores one candidate set: it receives X with shape (q, d) (q jointly proposed design points, no batch dimension) and returns a single scalar, higher being better. This module adds BoTorch’s batch handling: a (b, q, d) input is scored one candidate set at a time and returned as (b,).

X is a torch.Tensor so that torch-based functions are optimized with exact gradients. Functions that work on NumPy arrays also work: if the function cannot handle a tensor that requires gradients, or returns something without autograd history, it is called on a detached tensor (which NumPy accepts as an array) and differentiable is set to False so that the optimizer switches to finite differences.

kwargs are the static args of the acquisition spec, with any Variable placeholders already resolved; the function is called as function(X, **kwargs).

forward(X: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

score(X: Tensor) → Tensor

Score one (q, d) candidate set, returning a scalar tensor.