bonni.optimize_bonni

Contents

bonni.optimize_bonni#

bonni.optimize_bonni(fn, bounds, num_bonni_iterations, num_random_samples=None, direction='minimize', seed=None, save_path=None, xs=None, ys=None, gs=None, num_iter_until_recompile=50, ensemble_size=20, training_plot_directory=None, surrogate_plot_directory=None, num_acq_optim_samples=100, custom_ei_config=None, custom_optim_config=None, custom_base_model_config=None, num_embedding_channels=1, non_diff_params=None, num_acq_optim_runs=5, num_initial_acq_samples=10)[source]#

Optimize any black-box function with BONNI. Executes the Bayesian Optimization (BO) loop using an MLP Ensemble surrogate model.

Parameters:

  • fn (Callable[[np.ndarray], tuple[np.ndarray, np.ndarray]]) – The black-box objective function to optimize. It must accept an input array of shape (D,) and return a tuple (y, g), where y is the scalar objective and g represents gradients or auxiliary outputs of shape (D,). D represents the number of dimensions.

  • bounds (jax.Array | np.ndarray) – A 2D array of shape (D, 2) specifying the (lower, upper) search space boundaries for each of the D input dimensions.

  • num_bonni_iterations (int) – The number of Bayesian Optimization steps (active learning iterations) to perform after initialization. Must be > 0.

  • num_random_samples (int | None, optional) – The number of initial random samples to evaluate before starting the BO loop. If xs/ys/gs are not provided, this must be specified and non-zero. Defaults to None.

  • direction (Literal["maximize", "minimize"], optional) – The optimization goal. If “minimize”, the wrapper internally negates the objective values. Defaults to “minimize”.

  • seed (int | jax.Array | None, optional) – Random seed or JAX PRNGKey for reproducibility. If None, a random seed is generated. Defaults to None.

  • save_path (Path | str | None, optional) – Directory path to save the results of function evaluation as npz-file. Defaults to None.

  • xs (np.ndarray | jax.Array | None, optional) – Existing history of input points. Must be provided if ys and gs are provided. Defaults to None.

  • ys (np.ndarray | jax.Array | None, optional) – Existing history of objective values. Must be provided if xs and gs are provided. Defaults to None.

  • gs (np.ndarray | jax.Array | None, optional) – Existing history of gradients. Must be provided if xxs and ys are provided. Defaults to None.

  • num_iter_until_recompile (int, optional) – The number of optimization iterations, where jax does not recompile functions. Recompilation makes the execution faster, but also takes a lot of time. Therefore, the number of recompilations is a hyperparameter that can be tuned for speed (depending on hardware, other parameters and JAX version). Defaults to 50.

  • ensemble_size (int, optional) – The number of individual MLP models to train within the ensemble surrogate. Defaults to 20.

  • training_plot_directory (Path | None, optional) – Directory to save plots related to model training (losses, etc.). Defaults to None.

  • surrogate_plot_directory (Path | None, optional) – Directory to save visualizations of the surrogate landscape. This is only available if the dimensionality of the function is 1. Defaults to None.

  • num_acq_optim_samples (int, optional) – The number of samples used to optimize the acquisition function during each step. Defaults to 100.

  • custom_ei_config (EIConfig | None, optional) – Custom configuration for the Expected Improvement acquisition function. Defaults to None.

  • custom_optim_config (OptimConfig | None, optional) – Custom configuration for the optimizer (e.g., learning rate, steps). Defaults to None.

  • custom_base_model_config (MLPModelConfig | None, optional) – Custom architecture configuration for the individual MLP models. Defaults to None.

  • num_embedding_channels (int, optional) – The number of embedding channels used in the model input layer. Defaults to 1.

  • non_diff_params (np.ndarray | None, optional) – A boolean mask of shape (D,) indicating which parameters are non-differentiable (True) or differentiable (False). Defaults to None (all assumed differentiable). Note that the function fn still needs to return an array of shape (D,) for the gradients, but the values for non-diff. parameters can be arbitrary.

  • num_acq_optim_runs (int, optional) – The number of indepent acquisition function optimization runs that are performed. Since optimizing the acq. fn. is a difficult problem, restarts can increase sample quality greatly. Defaults to 5.

  • num_initial_acq_samples (int, optional) – Number of random samples, where the maximum is selected as startpoint for optimizing the acquisition function. Defaults to 10.

Returns:

A tuple containing the full history of:

  • xs: Input parameters (shape (N, D)).

  • ys: Objective values (shape (N,)).

  • gs: Gradients/Auxiliary outputs (shape (N, D)).

Here D is the number of dimensions and N the sample count.

Return type:

tuple[np.ndarray, np.ndarray, np.ndarray]

Contents