Optimisation

Acquisition optimisers:

banditry.optimisation_oracles.gen_alg.EvolutionOpt(design_space, **conf)

Evolutionary acquisition optimiser wrapping pymoo's mixed-variable GAs.

Minimises an acquisition Problem (typically a ContextualProblem) over the design space with a genetic algorithm chosen via :class:GAEnum: unless a method is passed explicitly, GAEnum.determine(problem.n_obj, compute_budget_high=True) picks U-NSGA-III for one or two objectives and NSGA-III for three or more. All algorithms run with mixed-variable mating and duplicate elimination, so continuous, integer and categorical variables are handled natively.

Parameters:

Name	Type	Description	Default
design_space	DesignSpace	The search space to optimise over.	required
**conf		Optional settings — pop (population size, default 100), max_iters (generation cap, default 500), verbose (pymoo verbosity, default False), repair (optional pymoo repair operator, default None) and sobol_init (default True; stored but currently unused).	{}

Source code in src/banditry/optimisation_oracles/gen_alg.py

def __init__(self, design_space: DesignSpace, **conf):
    self.space = design_space
    self.pop = conf.get("pop", 100)
    self.max_iters = conf.get("max_iters", 500)
    self.verbose = conf.get("verbose", False)
    self.repair = conf.get("repair", None)
    self.sobol_init = conf.get("sobol_init", True)

termnation_condition(problem)

Combine pymoo's default termination with an n_gen cap of max_iters.

Source code in src/banditry/optimisation_oracles/gen_alg.py

def termnation_condition(self, problem: Problem):
    """Combine pymoo's default termination with an ``n_gen`` cap of ``max_iters``."""
    def_term = default_termination(problem)
    max_iter_term = get_termination("n_gen", self.max_iters)
    return TerminationCollection(def_term, max_iter_term)

optimise(problem, initial_suggest=None, return_pop=False, method=None, seed=None)

Run the evolutionary search and decode the results to raw design-space values.

Reads the fixed context values (problem.fix) and free variables (problem.vars) off the problem; when every variable is fixed the fixed values are returned directly without running the GA. Otherwise an initial population is built via :func:get_init_pop, the algorithm selected by method (or :meth:GAEnum.determine) is run under the combined default/n_gen termination, and the resulting candidates are inverse-transformed back to raw parameter values with fixed context columns re-attached.

Parameters:

Name	Type	Description	Default
problem	Problem	pymoo Problem to minimise; typically a ContextualProblem.	required
initial_suggest	DataFrame	Optional raw-space DataFrame used to seed the initial population (and to fill free values when everything is fixed).	None
return_pop		If True, return the whole final population (sorted by the first objective) instead of the optimiser's solution set.	False
method	GAEnum	Explicit :class:GAEnum member overriding the automatic selection.	None
seed	int | None	Seed passed to pymoo.minimize; when None one is drawn from the globally seeded numpy RNG so the run stays reproducible.	None

Returns:

Type	Description
DataFrame	DataFrame of recommended configurations in raw design-space values, one row
DataFrame	per candidate, with columns ordered as space.para_names. The raw pymoo
DataFrame	result is stored on self.res.

Source code in src/banditry/optimisation_oracles/gen_alg.py

def optimise(
    self,
    problem: Problem,
    initial_suggest: pd.DataFrame = None,
    return_pop=False,
    method: GAEnum = None,
    seed: int | None = None,
) -> pd.DataFrame:
    """Run the evolutionary search and decode the results to raw design-space values.

    Reads the fixed context values (``problem.fix``) and free variables
    (``problem.vars``) off the problem; when every variable is fixed the fixed values
    are returned directly without running the GA. Otherwise an initial population is
    built via :func:`get_init_pop`, the algorithm selected by ``method`` (or
    :meth:`GAEnum.determine`) is run under the combined default/``n_gen`` termination,
    and the resulting candidates are inverse-transformed back to raw parameter values
    with fixed context columns re-attached.

    Args:
        problem: pymoo ``Problem`` to minimise; typically a ``ContextualProblem``.
        initial_suggest: Optional raw-space DataFrame used to seed the initial
            population (and to fill free values when everything is fixed).
        return_pop: If ``True``, return the whole final population (sorted by the
            first objective) instead of the optimiser's solution set.
        method: Explicit :class:`GAEnum` member overriding the automatic selection.
        seed: Seed passed to ``pymoo.minimize``; when ``None`` one is drawn from the
            globally seeded numpy RNG so the run stays reproducible.

    Returns:
        DataFrame of recommended configurations in raw design-space values, one row
        per candidate, with columns ordered as ``space.para_names``. The raw pymoo
        result is stored on ``self.res``.
    """
    fix_input = getattr(problem, "fix", None) or {}

    problem_vars = getattr(problem, "vars", None)
    active_para_names = list(self.space.para_names) if problem_vars is None else list(problem_vars.keys())

    if len(active_para_names) == 0:
        self.res = None
        return self._fixed_only_recommendation(fix_input, initial_suggest)

    init_pop = get_init_pop(self.space, self.pop, initial_suggest, active_para_names)

    if method is None:
        method = GAEnum.determine(problem.n_obj, compute_budget_high=True)

    algo = method.auto(
        n_dim=problem.n_obj, n_points=self.pop, pop_size=self.pop, repair=self.repair, sampling=init_pop
    )

    # pymoo.minimize(seed=None) draws from a non-deterministic source
    # rather than the globally-seeded numpy RNG, so reproducibility
    # requires an explicit seed. We derive one from np.random — which
    # IS seeded by Experiment.run's seed_everything — so the chain
    # stays deterministic without any per-call seed plumbing upstream.
    if seed is None:
        seed = int(np.random.randint(0, 2**31 - 1))

    res = minimize(problem, algo, termination=self.termnation_condition(problem), seed=seed, verbose=self.verbose)

    if res.X is None:
        import banditry.logging_utils as log

        log.debug("Optimisation terminated with no solutions found")

    if res.X is not None and not return_pop:
        candidates = self._as_candidate_list(res.X)
    else:
        candidates = [p for p in res.pop]
        if problem.n_obj == 1 and not return_pop and len(candidates) > 0:
            candidates = [candidates[np.random.choice(len(candidates))]]
        candidates = sorted(candidates, key=lambda x: x.F[0])
        candidates = [c.X for c in candidates]

    self.res = res
    return self._decode_candidates(candidates, active_para_names, fix_input)

banditry.optimisation_oracles.sgld.SGLD(params, lr=0.001, prior_precision=0.0001, temperature=1.0, precondition=False, precond_alpha=0.99, precond_eps=1e-08, generator=None)

Bases: Optimizer

Stochastic gradient Langevin dynamics optimizer (Welling & Teh, 2011).

A torch.optim.Optimizer whose update is a gradient step on the negative log posterior — the loss gradient plus prior_precision * param from an isotropic zero-mean Gaussian prior — followed by injected Gaussian noise with standard deviation sqrt(2 * lr * temperature), so iterates approximately sample from the posterior rather than converge to a mode. With precondition=True an RMSProp-style diagonal preconditioner (an EMA of squared gradients) rescales both the gradient step and the noise, as in pSGLD (Li et al., 2016). All hyperparameters are per parameter group, so the prior can be disabled for selected parameters (e.g. the learned observation noise). Used by LangevinSampler to draw approximate posterior samples of ValueFunction weights.

Parameters:

Name	Type	Description	Default
params		Iterable of parameters or parameter-group dicts; per-group overrides of the defaults below are honoured.	required
lr	float	Learning rate (the SGLD step size eps_t).	0.001
prior_precision	float	Precision of the Gaussian prior, applied as weight decay; 0 disables the prior term.	0.0001
temperature	float	Scales the injected-noise variance; 0 reduces the update to plain (optionally preconditioned) SGD.	1.0
precondition	bool	Enable RMSProp-style diagonal preconditioning.	False
precond_alpha	float	EMA decay rate of the squared-gradient accumulator.	0.99
precond_eps	float	Small constant added to the preconditioner denominator for stability.	1e-08
generator	Generator | None	Optional torch.Generator used for the injected noise.	None

Raises:

Type	Description
ValueError	If lr or temperature is negative.

Source code in src/banditry/optimisation_oracles/sgld.py

def __init__(
    self,
    params,
    lr: float = 1e-3,
    prior_precision: float = 1e-4,
    temperature: float = 1.0,
    precondition: bool = False,
    precond_alpha: float = 0.99,
    precond_eps: float = 1e-8,
    generator: torch.Generator | None = None,
):
    if lr < 0.0:
        raise ValueError(f"Invalid learning rate: {lr}")
    if temperature < 0.0:
        raise ValueError(f"Invalid temperature: {temperature}")

    defaults = dict(
        lr=lr,
        prior_precision=prior_precision,
        temperature=temperature,
        precondition=precondition,
        precond_alpha=precond_alpha,
        precond_eps=precond_eps,
    )
    super().__init__(params, defaults)
    self.generator = generator

step(noise=True, closure=None)

Perform one SGLD update over all parameter groups.

Parameters:

Name	Type	Description	Default
noise	bool	If False, skip the Langevin noise injection (plain gradient step).	True
closure		Optional callable that re-evaluates the model and returns the loss.	None

Returns:

Type	Description
	The loss returned by closure, or None when no closure is given.

Source code in src/banditry/optimisation_oracles/sgld.py

@torch.no_grad()
def step(self, noise: bool = True, closure=None):
    """Perform one SGLD update over all parameter groups.

    Args:
        noise: If ``False``, skip the Langevin noise injection (plain gradient step).
        closure: Optional callable that re-evaluates the model and returns the loss.

    Returns:
        The loss returned by ``closure``, or ``None`` when no closure is given.
    """
    loss = None
    if closure is not None:
        with torch.enable_grad():
            loss = closure()

    for group in self.param_groups:
        lr = group["lr"]
        prior_prec = group["prior_precision"]
        temperature = group["temperature"]
        use_precond = group["precondition"]
        alpha = group["precond_alpha"]
        eps = group["precond_eps"]

        for p in group["params"]:
            if p.grad is None:
                continue

            g = p.grad
            if prior_prec != 0:
                g = g.add(p, alpha=prior_prec)

            if use_precond:
                state = self.state[p]
                if "v" not in state:
                    state["v"] = torch.zeros_like(p)
                v = state["v"]
                v.mul_(alpha).addcmul_(g, g, value=1.0 - alpha)
                precond = 1.0 / (v.sqrt() + eps)
                p.addcmul_(precond, g, value=-lr)
                if noise:
                    p.add_(
                        precond.sqrt() * self._randn_like(p),
                        alpha=math.sqrt(2.0 * lr * temperature),
                    )
            else:
                p.add_(g, alpha=-lr)
                if noise:
                    p.add_(
                        self._randn_like(p),
                        alpha=math.sqrt(2.0 * lr * temperature),
                    )

    return loss

Acquisition objectives:

banditry.optimisation_subroutines.objectives

Objective(model, **conf)

Bases: ABC

Base class for acquisition objectives evaluated by the evolutionary optimiser.

An objective wraps a surrogate model and maps transformed design points to one or more columns that the optimiser MINIMISES. Subclasses declare how many objective columns (num_obj) and constraint columns (num_constr) their :meth:eval output contains; ContextualProblem splits the returned tensor accordingly (the first num_obj columns become objectives F, the rest constraints G).

Parameters:

Name	Type	Description	Default
model		Surrogate model the objective queries.	required
**conf		Subclass-specific options.	{}

Source code in src/banditry/optimisation_subroutines/objectives.py

def __init__(self, model, **conf):
    self.model = model

num_obj abstractmethod property

Number of objective columns produced by :meth:eval.

num_constr abstractmethod property

Number of constraint columns produced by :meth:eval.

eval(x, xe) abstractmethod

Evaluate the objective at a batch of transformed design points.

Parameters:

Name	Type	Description	Default
x	Tensor	Continuous features of shape (n, num_cont).	required
xe	Tensor	Categorical (enum) features of shape (n, num_enum).	required

Returns:

Type	Description
Tensor	Tensor whose first num_obj columns are objective values to MINIMISE and
Tensor	whose remaining num_constr columns are constraint values.

Source code in src/banditry/optimisation_subroutines/objectives.py

@abstractmethod
def eval(self, x: Tensor, xe: Tensor) -> Tensor:
    """Evaluate the objective at a batch of transformed design points.

    Args:
        x: Continuous features of shape ``(n, num_cont)``.
        xe: Categorical (enum) features of shape ``(n, num_enum)``.

    Returns:
        Tensor whose first ``num_obj`` columns are objective values to MINIMISE and
        whose remaining ``num_constr`` columns are constraint values.
    """

__call__(x, xe)

Alias for :meth:eval.

Source code in src/banditry/optimisation_subroutines/objectives.py

def __call__(self, x: Tensor, xe: Tensor):
    """Alias for :meth:`eval`."""
    return self.eval(x, xe)

SingleObjective(model, **conf)

Bases: Objective

Convenience base for unconstrained, single-column objectives.

Fixes num_obj = 1 and num_constr = 0; subclasses only implement eval.

Source code in src/banditry/optimisation_subroutines/objectives.py

def __init__(self, model: BaseModel, **conf):
    super().__init__(model, **conf)

Mean(model, **conf)

Bases: SingleObjective

Posterior predictive mean of a single-output surrogate.

Minimising it drives the search towards the lowest predicted value — pure exploitation with no uncertainty bonus.

Source code in src/banditry/optimisation_subroutines/objectives.py

def __init__(self, model: BaseModel, **conf):
    super().__init__(model, **conf)
    assert model.num_out == 1

Sigma(model, **conf)

Bases: SingleObjective

Posterior predictive standard deviation of a single-output surrogate.

Note that objectives are minimised, so this column favours well-explored (low-sigma) points; use its negation (minimising -sigma maximises uncertainty) to reward exploration.

Source code in src/banditry/optimisation_subroutines/objectives.py

def __init__(self, model: BaseModel, **conf):
    super().__init__(model, **conf)
    assert model.num_out == 1

LCB(model, best_y=None, **conf)

Bases: SingleObjective

Lower confidence bound acquisition: py - kappa * ps.

Minimising the LCB balances exploitation (low predicted mean py) with optimism about uncertain points (high predicted standard deviation ps).

Parameters:

Name	Type	Description	Default
model	BaseModel	Surrogate with a predict method returning mean and variance.	required
best_y		Accepted for interface compatibility; unused.	None
**conf		kappa — exploration weight on the standard deviation (default 2.0).	{}

Source code in src/banditry/optimisation_subroutines/objectives.py

def __init__(self, model: BaseModel, best_y=None, **conf):
    super().__init__(model, **conf)
    self.kappa = conf.get("kappa", 2.0)

MACE(model, best_y, **conf)

Bases: Objective

Multi-objective ACquisition function Ensemble (Lyu et al., ICML 2018), as in HEBO.

Exposes three acquisition functions as one simultaneous multi-objective problem — column 0: LCB py - kappa * ps; column 1: negated log expected improvement; column 2: negated log probability of improvement — so the evolutionary optimiser returns a Pareto front trading them off rather than committing to a single acquisition. EI and PI are computed in log space against the jittered incumbent best_y - eps (with observation noise added to the improvement), switching to Mills-ratio asymptotic approximations where the exact values are numerically unstable.

Parameters:

Name	Type	Description	Default
model		Surrogate with a predict method returning mean and variance and a noise attribute.	required
best_y		Incumbent (best observed) value the improvement is measured against.	required
**conf		kappa — LCB exploration weight (default 2.0); eps — improvement jitter subtracted from the incumbent (default 1e-4).	{}

Source code in src/banditry/optimisation_subroutines/objectives.py

def __init__(self, model, best_y, **conf):
    super().__init__(model, **conf)
    self.kappa = conf.get("kappa", 2.0)
    self.eps = conf.get("eps", 1e-4)
    self.tau = best_y

Mills_ratio_approximations(ps, normalised_improvement) staticmethod

Mills-ratio asymptotic approximations of (log PI, log EI) for large negative improvement.

Source code in src/banditry/optimisation_subroutines/objectives.py

@staticmethod
def Mills_ratio_approximations(ps, normalised_improvement) -> tuple[torch.FloatTensor, torch.FloatTensor]:
    """Mills-ratio asymptotic approximations of ``(log PI, log EI)`` for large negative improvement."""
    logEIapp = ps.log() - 0.5 * normalised_improvement**2 - (normalised_improvement**2 - 1).log()
    logPIapp = -0.5 * normalised_improvement**2 - torch.log(-1 * normalised_improvement) - HALF_LOG_2PI
    return logPIapp, logEIapp

exact_value(ps, normalised_improvement) staticmethod

Exact (log PI, log EI) computed from the standard normal pdf/cdf.

Source code in src/banditry/optimisation_subroutines/objectives.py

@staticmethod
def exact_value(ps, normalised_improvement) -> tuple[torch.FloatTensor, torch.FloatTensor]:
    """Exact ``(log PI, log EI)`` computed from the standard normal pdf/cdf."""
    dist = Normal(0.0, 1.0)
    log_phi = dist.log_prob(normalised_improvement)
    cdf_at_imp = dist.cdf(normalised_improvement)
    return cdf_at_imp.log(), (ps * (cdf_at_imp * normalised_improvement + log_phi.exp())).log()

ThompsonObjective(model)

Bases: Objective

Adapter: exposes a sampled ValueFunction as an Acquisition.

Thompson-sampling acquisition — the objective is simply the sampled model's prediction at the candidate point, reshaped to one column per output. Minimising it selects the action the current posterior draw believes best; num_obj equals the model's num_out and there are no constraints.

Parameters:

Name	Type	Description	Default
model	ValueFunction	A ValueFunction whose weights are a posterior draw (see Sampler).	required

Source code in src/banditry/optimisation_subroutines/objectives.py

def __init__(self, model: ValueFunction):
    self.model = model

banditry.optimisation_subroutines.contextual_problem.ContextualProblem(objective, space, fix=None)

Bases: Problem

pymoo Problem that optimises an acquisition's free variables under a fixed context.

Wraps an :class:~banditry.optimisation_subroutines.objectives.Objective as a mixed-variable pymoo problem. Variables named in fix are pinned: their raw values are transformed once into the optimiser's internal space and broadcast into every candidate at evaluation time, while only the remaining (free) variables are exposed to the optimiser. The per-variable dict handed to pymoo is built by taking DesignSpace.to_pymoo_vars() — Real/Integer variables bounded by the space's transformed opt_lb/opt_ub and Choice variables enumerating category codes — and restricting it to the free names, so bounds and variable types always match the space's transform. Objective and constraint counts come from objective.num_obj and objective.num_constr.

Parameters:

Name	Type	Description	Default
objective	Objective	Acquisition to minimise; also determines n_obj and n_constr.	required
space	DesignSpace	Design space defining all variables and their transforms.	required
fix	dict	Mapping of context-variable names to raw (untransformed) values to hold fixed; None or empty means every variable is optimised.	None

Source code in src/banditry/optimisation_subroutines/contextual_problem.py

def __init__(self, objective: Objective, space: DesignSpace, fix: dict = None):
    self.objective = objective
    self.space = space
    self.fix = fix or {}
    self.fixed_transformed = self._transform_fix_values(self.space, self.fix)

    self.free_para_names = [name for name in self.space.para_names if name not in self.fix]
    all_vars = self.space.to_pymoo_vars()
    vars = {name: all_vars[name] for name in self.free_para_names}
    super().__init__(vars=vars, n_obj=self.objective.num_obj, n_constr=self.objective.num_constr)