Surrogates

banditry.surrogates.gp.GP(num_cont, num_enum, num_out, **conf)

Bases: BaseModel

Exact Gaussian-process surrogate built on gpytorch's ExactGP.

Implements the :class:BaseModel fit/predict interface: :meth:fit learns the hyperparameters by maximising the exact marginal log likelihood and :meth:predict returns the posterior (mean, variance). Continuous inputs are min-max scaled to [-1, 1] and targets are standardised internally; Cholesky failures are retried with increasing jitter.

Config keys (passed via **conf): - lr (float, default 3e-2): learning rate for the hyperparameter optimiser. - num_epochs (int, default 100): number of optimisation steps. - verbose (bool, default False): log the loss during fitting. - print_every (int, default 10): logging period (in epochs) when verbose. - pred_likeli (bool, default True): predictive variance includes observation noise when True; when False it is the variance of the latent function f only. - noise_lb (float, default 1e-5): lower bound on the likelihood noise. - noise_guess (float, default 0.01): location of the log-normal noise prior. - ard_kernel (bool, default True): ARD lengthscales in the default kernel. - optimizer (str, default "adam"): one of "adam", "lbfgs", "sgld". - num_uniqs (list[int], required when num_enum > 0): cardinality of each categorical column. - emb_sizes (list[int], optional): embedding sizes for the categorical columns. - kern (gpytorch kernel, optional): covariance override; defaults to :func:default_kern. mean, fe and product_kernel are likewise forwarded to :class:GPyTorchModel.

Source code in src/banditry/surrogates/gp.py

def __init__(self, num_cont, num_enum, num_out, **conf):
    super().__init__(num_cont, num_enum, num_out, **conf)
    self.lr = conf.get("lr", 3e-2)
    self.num_epochs = conf.get("num_epochs", 100)
    self.verbose = conf.get("verbose", False)
    self.print_every = conf.get("print_every", 10)
    self.pred_likeli = conf.get("pred_likeli", True)
    self.noise_lb = conf.get("noise_lb", 1e-5)
    self.noise_guess = conf.get("noise_guess", 0.01)
    self.ard_kernel = conf.get("ard_kernel", True)
    self.xscaler = TorchMinMaxScaler((-1, 1))
    self.yscaler = TorchStandardScaler()
    self.optimizer = conf.get("optimizer", "adam")

banditry.surrogates.svgp.SVGP(num_cont, num_enum, num_out, **conf)

Bases: BaseModel

Sparse variational GP surrogate trained by minibatch ELBO maximisation.

Uses inducing points (one ApproximateGP per output), so it scales to larger datasets than the exact :class:~banditry.surrogates.gp.GP, supports multi-output targets with missing (NaN) entries, and can optionally use natural gradient descent for the variational parameters. Continuous inputs are min-max scaled to [-1, 1] and targets are standardised internally.

Config keys (passed via **conf): - num_inducing (int, default 128): number of inducing points (capped at the number of training rows). - batch_size (int, default 64): minibatch size for the ELBO. - num_epochs (int, default 300): number of training epochs. - lr (float, default 1e-2): learning rate for kernel/likelihood hyperparameters. - lr_vp (float, default 1e-1): learning rate for the variational parameters. - lr_fe (float, default 1e-3): learning rate for the feature extractor. - use_ngd (bool, default False): natural gradient descent for the variational parameters. - ard_kernel (bool, default True): ARD lengthscales in the default kernel. - beta (float, default 1.0): KL weight in the variational ELBO. - noise_lb (float, default 1e-5): lower bound on the likelihood noise. - verbose (bool, default False): log the loss during fitting. - print_every (int, default 10): logging period (in epochs) when verbose. - pred_likeli (bool, default True): predictive variance includes observation noise when True; latent f variance only when False. - num_uniqs (list[int], required when num_enum > 0): cardinality of each categorical column.

Also forwarded to :class:`SVGPModel`: ``emb_sizes``, ``fe``, ``mean``, ``kern``,
``product_kernel``, ``learn_u``.

Source code in src/banditry/surrogates/svgp.py

def __init__(self, num_cont, num_enum, num_out, **conf):
    super().__init__(num_cont, num_enum, num_out, **conf)

    self.use_ngd = conf.get("use_ngd", False)
    self.lr = conf.get("lr", 1e-2)
    self.lr_vp = conf.get("lr_vp", 1e-1)
    self.lr_fe = conf.get("lr_fe", 1e-3)
    self.num_inducing = conf.get("num_inducing", 128)
    self.ard_kernel = conf.get("ard_kernel", True)
    self.pred_likeli = conf.get("pred_likeli", True)
    self.beta = conf.get("beta", 1.0)

    self.batch_size = conf.get("batch_size", 64)
    self.num_epochs = conf.get("num_epochs", 300)
    self.verbose = conf.get("verbose", False)
    self.print_every = conf.get("print_every", 10)
    self.noise_lb = conf.get("noise_lb", 1e-5)
    self.xscaler = TorchMinMaxScaler((-1, 1))
    self.yscaler = TorchStandardScaler()

banditry.surrogates.svgp.BaseModel(num_cont, num_enum, num_out, **conf)

Bases: ABC

Abstract fit/predict contract shared by all surrogate models.

A surrogate is constructed with the input layout -- num_cont continuous columns, num_enum categorical index columns, num_out outputs -- plus model-specific **conf options. It is trained with fit(Xc, Xe, y) and queried with predict(Xc, Xe), which returns a (mean, variance) tuple, each of shape (n, num_out). sample_y draws samples from the predictive distribution (independent Gaussians by default) and noise reports the observation-noise variance. The class flags support_ts, support_grad, support_multi_output and support_warm_start advertise optional capabilities. When num_enum > 0, conf["num_uniqs"] (list of per-column cardinalities) is required.

Source code in src/banditry/surrogates/svgp.py

def __init__(self, num_cont: int, num_enum: int, num_out: int, **conf):
    self.num_cont = num_cont
    self.num_enum = num_enum
    self.num_out = num_out
    self.conf = conf
    assert self.num_cont >= 0
    assert self.num_enum >= 0
    assert self.num_out > 0
    assert self.num_cont + self.num_enum > 0
    if self.num_enum > 0:
        assert "num_uniqs" in self.conf
        assert isinstance(self.conf["num_uniqs"], list)
        assert len(self.conf["num_uniqs"]) == self.num_enum
    if not self.support_multi_output:
        assert self.num_out == 1, "Model only support single-output"

fit(Xc, Xe, y) abstractmethod

Train the surrogate on continuous inputs Xc, categorical inputs Xe and targets y.

Source code in src/banditry/surrogates/svgp.py

@abstractmethod
def fit(self, Xc: FloatTensor, Xe: LongTensor, y: FloatTensor):
    """Train the surrogate on continuous inputs ``Xc``, categorical inputs ``Xe`` and targets ``y``."""
    pass

predict(Xc, Xe) abstractmethod

Return the predictive (mean, variance), each of shape (n, num_out).

Source code in src/banditry/surrogates/svgp.py

@abstractmethod
def predict(self, Xc: FloatTensor, Xe: LongTensor) -> tuple[FloatTensor, FloatTensor]:
    """Return the predictive ``(mean, variance)``, each of shape ``(n, num_out)``."""
    pass

banditry.surrogates.tsmodel.ValueFunction(num_cont, num_enum, num_out=1, yscaler=None, network_builder=default_network_builder, **conf)

Bases: Module

Neural value model whose weights Thompson sampling draws from the posterior.

Wraps a feature extractor and a value network into a single module mapping (Xc, Xe) to value predictions of shape (n, num_out). Thompson sampling treats the network weights as random: a posterior draw over the weights yields a sampled value function that is then optimised over the design space. Continuous inputs are min-max scaled to [-1, 1] (call :meth:fit_x_scaler first) and predictions are un-scaled through an attached target scaler when one is set (:meth:set_y_scaler).

Config keys (read in __init__ via **conf): - num_uniqs (list[int], required when num_enum > 0): cardinality of each categorical column. - emb_sizes (list[int], optional): embedding sizes for the categorical columns. - fe (nn.Module, optional): feature-extractor override; defaults to :class:DummyFeatureExtractor.

Parameters:

Name	Type	Description	Default
num_cont	int	Number of continuous input columns.	required
num_enum	int	Number of categorical (index) input columns.	required
num_out	int	Number of outputs.	1
yscaler	TorchStandardScaler | None	Optional fitted target scaler used to un-scale predictions.	None
network_builder	Callable[[int, int], Module]	Callable (in_dim, out_dim) -> nn.Module building the value network; defaults to :func:default_network_builder.	default_network_builder
**conf		See "Config keys" above.	{}

Raises:

Type	Description
TypeError	If network_builder does not return an nn.Module.

Source code in src/banditry/surrogates/tsmodel.py

def __init__(
    self,
    num_cont: int,
    num_enum: int,
    num_out: int = 1,
    yscaler: TorchStandardScaler | None = None,
    network_builder: Callable[[int, int], nn.Module] = default_network_builder,
    **conf,
):
    super().__init__()
    assert num_cont >= 0
    assert num_enum >= 0
    assert num_out > 0
    assert num_cont + num_enum > 0
    if num_enum > 0:
        assert "num_uniqs" in conf
        assert isinstance(conf["num_uniqs"], list)
        assert len(conf["num_uniqs"]) == num_enum

    self.num_cont = num_cont
    self.num_enum = num_enum
    self.num_out = num_out
    self.conf = conf
    self.yscaler = deepcopy(yscaler) if yscaler is not None else TorchStandardScaler()

    self.xscaler = TorchMinMaxScaler((-1, 1))
    self.fe = deepcopy(
        conf.get(
            "fe",
            DummyFeatureExtractor(
                num_cont,
                num_enum,
                conf.get("num_uniqs"),
                conf.get("emb_sizes"),
            ),
        )
    )

    self.network_builder = network_builder
    self.net = self.network_builder(self.fe.total_dim, self.num_out)
    if not isinstance(self.net, nn.Module):
        raise TypeError("network_builder must return an nn.Module.")