|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | +# |
| 4 | +# This source code is licensed under the MIT license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Generic BoTorch-compatible model wrappers. |
| 8 | +
|
| 9 | +This module provides lightweight Model subclasses that integrate with BoTorch's |
| 10 | +ModelList, acquisition functions, and optimize_acqf without requiring GP training: |
| 11 | +
|
| 12 | +- ``LinearModel``: A linear model f(x) = c^T x with known coefficient uncertainty |
| 13 | + and optional per-class coefficient indexing. |
| 14 | +- ``FixedFeatureModel``: Wraps any model to fix a subset of inputs to constants. |
| 15 | +""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +import torch |
| 20 | +from botorch.models.model import Model |
| 21 | +from botorch.posteriors import GPyTorchPosterior, Posterior |
| 22 | +from gpytorch.distributions import MultivariateNormal |
| 23 | +from linear_operator.operators import DiagLinearOperator |
| 24 | +from torch import Tensor |
| 25 | + |
| 26 | + |
| 27 | +class LinearModel(Model): |
| 28 | + """A linear model f(x) = c^T x with known coefficient uncertainty. |
| 29 | +
|
| 30 | + Supports two modes: |
| 31 | +
|
| 32 | + 1. **Simple** (``class_dim=None``): A single set of coefficients applied to |
| 33 | + all inputs. ``coefficients`` and ``coefficient_vars`` are ``(d,)``-dim. |
| 34 | +
|
| 35 | + 2. **Class-indexed** (``class_dim`` specified): Per-class coefficients stored |
| 36 | + as a ``(K, d)`` tensor. The input column at ``class_dim`` is used as an |
| 37 | + integer index to select the appropriate coefficient row. The class column |
| 38 | + is excluded from the dot product. This generalizes naturally to any |
| 39 | + number of classes K. |
| 40 | +
|
| 41 | + No training data or fitting required. This is equivalent to Bayesian |
| 42 | + linear regression with a diagonal prior on the coefficients and no |
| 43 | + observations — the "posterior" is the prior itself. |
| 44 | +
|
| 45 | + This model integrates seamlessly with BoTorch's ``ModelList``, |
| 46 | + ``FixedFeatureModel``, acquisition functions, and ``optimize_acqf``. |
| 47 | + """ |
| 48 | + |
| 49 | + def __init__( |
| 50 | + self, |
| 51 | + coefficients: Tensor, |
| 52 | + coefficient_vars: Tensor, |
| 53 | + class_dim: int | None = None, |
| 54 | + ): |
| 55 | + """Constructs a LinearModel with known coefficients and their variances. |
| 56 | +
|
| 57 | + Args: |
| 58 | + coefficients: Coefficient means. Shape ``(d,)`` for simple mode or |
| 59 | + ``(K, d)`` for class-indexed mode (K classes, d features). |
| 60 | + coefficient_vars: Coefficient variances. Same shape as |
| 61 | + ``coefficients``. |
| 62 | + class_dim: If specified, the input column index that contains the |
| 63 | + integer class label. This column is used for indexing into |
| 64 | + the coefficient tensor and excluded from the linear computation. |
| 65 | + """ |
| 66 | + super().__init__() |
| 67 | + if (coefficient_vars < 0).any(): |
| 68 | + raise ValueError("coefficient_vars must be non-negative.") |
| 69 | + self.register_buffer("_coefficients", coefficients) |
| 70 | + self.register_buffer("_coefficient_vars", coefficient_vars) |
| 71 | + self._class_dim = class_dim |
| 72 | + self._num_outputs = 1 |
| 73 | + |
| 74 | + @property |
| 75 | + def num_outputs(self) -> int: |
| 76 | + """The number of outputs (always 1).""" |
| 77 | + return self._num_outputs |
| 78 | + |
| 79 | + def posterior( |
| 80 | + self, X: Tensor, observation_noise: bool = False, **kwargs |
| 81 | + ) -> GPyTorchPosterior: |
| 82 | + """Computes the posterior at input X. |
| 83 | +
|
| 84 | + Args: |
| 85 | + X: Input Tensor. For simple mode: ``(... x d)``-dim. |
| 86 | + For class-indexed mode: ``(... x (d + 1))``-dim where column |
| 87 | + ``class_dim`` contains integer class indices. |
| 88 | + observation_noise: Ignored (included for API compatibility). |
| 89 | +
|
| 90 | + Returns: |
| 91 | + A ``GPyTorchPosterior`` with the computed mean and variance. |
| 92 | + """ |
| 93 | + if self._class_dim is not None: |
| 94 | + # Extract class indices and remove class column from features |
| 95 | + class_indices = X[..., self._class_dim].long() |
| 96 | + feature_dims = [i for i in range(X.shape[-1]) if i != self._class_dim] |
| 97 | + X_features = X[..., feature_dims] |
| 98 | + # Index into per-class coefficients (cast to input dtype) |
| 99 | + coeffs = self._coefficients[class_indices].to(X.dtype) |
| 100 | + coeff_vars = self._coefficient_vars[class_indices].to(X.dtype) |
| 101 | + mean = (X_features * coeffs).sum(-1) |
| 102 | + var = (X_features**2 * coeff_vars).sum(-1) |
| 103 | + else: |
| 104 | + coefficients = self._coefficients.to(X.dtype) |
| 105 | + coefficient_vars = self._coefficient_vars.to(X.dtype) |
| 106 | + mean = X @ coefficients |
| 107 | + var = (X**2) @ coefficient_vars |
| 108 | + |
| 109 | + mvn = MultivariateNormal(mean, DiagLinearOperator(var)) |
| 110 | + return GPyTorchPosterior(mvn) |
| 111 | + |
| 112 | + |
| 113 | +class FixedFeatureModel(Model): |
| 114 | + """Wraps a model to fix a subset of inputs to constant values. |
| 115 | +
|
| 116 | + At evaluation time the fixed features are spliced back into the input |
| 117 | + tensor before delegating to the ``base_model``. This is useful for |
| 118 | + optimization over a subset of input dimensions while keeping others |
| 119 | + constant (e.g., fixing Material Source or Temperature). |
| 120 | + """ |
| 121 | + |
| 122 | + def __init__( |
| 123 | + self, |
| 124 | + base_model: Model, |
| 125 | + dim: int, |
| 126 | + indices: list[int] | Tensor, |
| 127 | + values: list[float] | Tensor, |
| 128 | + ): |
| 129 | + """A wrapper around a model that fixes some inputs to specific values. |
| 130 | +
|
| 131 | + Args: |
| 132 | + base_model: The base model to wrap. |
| 133 | + dim: The total input dimensionality expected by the base model. |
| 134 | + The FixedFeatureModel accepts ``(dim - len(indices))``-dimensional |
| 135 | + inputs and reconstructs the full ``dim``-dimensional input. |
| 136 | + indices: The indices of the inputs to fix. |
| 137 | + values: The values to fix the inputs to. |
| 138 | +
|
| 139 | + Raises: |
| 140 | + ValueError: If indices and values do not have the same length. |
| 141 | + """ |
| 142 | + super().__init__() |
| 143 | + self.base_model = base_model |
| 144 | + if len(indices) != len(values): |
| 145 | + raise ValueError("indices and values do not have the same length.") |
| 146 | + # Sort by index so that the boolean mask assignment in |
| 147 | + # _add_fixed_features places values at the correct positions. |
| 148 | + indices = torch.as_tensor(indices) |
| 149 | + values = torch.as_tensor(values) |
| 150 | + sort_order = indices.argsort() |
| 151 | + indices = indices[sort_order] |
| 152 | + values = values[sort_order] |
| 153 | + self._dim = dim |
| 154 | + self.register_buffer("_indices", indices) |
| 155 | + self.register_buffer( |
| 156 | + "_fixed", |
| 157 | + torch.tensor( |
| 158 | + [i in indices for i in torch.arange(dim, dtype=indices.dtype)] |
| 159 | + ), |
| 160 | + ) |
| 161 | + self.register_buffer("_values", values) |
| 162 | + |
| 163 | + def _add_fixed_features(self, X: Tensor) -> Tensor: |
| 164 | + """Reconstructs the full input tensor by splicing in fixed values. |
| 165 | +
|
| 166 | + Args: |
| 167 | + X: A ``(... x d_free)``-dim Tensor of free features. |
| 168 | +
|
| 169 | + Returns: |
| 170 | + A ``(... x d_full)``-dim Tensor with fixed features inserted. |
| 171 | + """ |
| 172 | + tkwargs = {"dtype": X.dtype, "device": X.device} |
| 173 | + Z = torch.zeros(*X.shape[:-1], X.shape[-1] + len(self._indices), **tkwargs) |
| 174 | + Z[..., self._fixed] = self._values.to(X.dtype) |
| 175 | + Z[..., ~self._fixed] = X |
| 176 | + return Z |
| 177 | + |
| 178 | + def forward(self, X: Tensor, *args, **kwargs) -> Tensor: |
| 179 | + """Forward pass with fixed features added. |
| 180 | +
|
| 181 | + Args: |
| 182 | + X: The ``(batch_shape x d_free)``-dim input Tensor. |
| 183 | +
|
| 184 | + Returns: |
| 185 | + The model output Tensor. |
| 186 | + """ |
| 187 | + return self.base_model.forward(self._add_fixed_features(X), *args, **kwargs) |
| 188 | + |
| 189 | + def posterior(self, X: Tensor, *args, **kwargs) -> Posterior: |
| 190 | + """Computes the posterior with fixed features added. |
| 191 | +
|
| 192 | + Args: |
| 193 | + X: The ``(batch_shape x d_free)``-dim input Tensor. |
| 194 | +
|
| 195 | + Returns: |
| 196 | + The posterior of the base model evaluated at the augmented input. |
| 197 | + """ |
| 198 | + return self.base_model.posterior(self._add_fixed_features(X), *args, **kwargs) |
| 199 | + |
| 200 | + @property |
| 201 | + def num_outputs(self) -> int: |
| 202 | + """The number of outputs of the base model.""" |
| 203 | + return self.base_model.num_outputs |
| 204 | + |
| 205 | + def subset_output(self, idcs: list[int]) -> FixedFeatureModel: |
| 206 | + """Returns a new ``FixedFeatureModel`` whose base model is subset to |
| 207 | + the given output indices. |
| 208 | +
|
| 209 | + Args: |
| 210 | + idcs: Output indices to keep. |
| 211 | +
|
| 212 | + Returns: |
| 213 | + A ``FixedFeatureModel`` wrapping the subset base model. |
| 214 | + """ |
| 215 | + return FixedFeatureModel( |
| 216 | + base_model=self.base_model.subset_output(idcs), |
| 217 | + dim=self._dim, |
| 218 | + indices=self._indices, |
| 219 | + values=self._values, |
| 220 | + ) |
0 commit comments