Core API

This section documents the core classes and functions of JaxBo.

BOBE Class

The main Bayesian Optimization for Bayesian Evidence class.

class BOBE.bo.BOBE(loglikelihood, param_list=None, param_bounds=None, param_labels=None, likelihood_name=None, confidence_for_unbounded=0.9999995, gp_kwargs={}, n_cobaya_init=4, n_sobol_init=4, init_train_x=None, init_train_y=None, resume=False, resume_file=None, save_dir='.', save=True, save_step=5, optimizer='scipy', use_clf=False, clf_type='svm', clf_nsigma_threshold=20, clf_use_size=10, clf_update_step=1, minus_inf=-10000000000.0, seed=None, verbosity='INFO', dynamic_dispatch=False)[source]

Bases: object

__init__(loglikelihood, param_list=None, param_bounds=None, param_labels=None, likelihood_name=None, confidence_for_unbounded=0.9999995, gp_kwargs={}, n_cobaya_init=4, n_sobol_init=4, init_train_x=None, init_train_y=None, resume=False, resume_file=None, save_dir='.', save=True, save_step=5, optimizer='scipy', use_clf=False, clf_type='svm', clf_nsigma_threshold=20, clf_use_size=10, clf_update_step=1, minus_inf=-10000000000.0, seed=None, verbosity='INFO', dynamic_dispatch=False)[source]

Initialize the BOBE (Bayesian Optimization for Bayesian Evidence) sampler.

Parameters:

  • loglikelihood (Union[Callable, str, Dict[str, Any], Likelihood]) – Log-likelihood specification. Can be: - A callable function (requires param_list and param_bounds) - A string path to Cobaya YAML file (automatically creates CobayaLikelihood) - A dict with Cobaya info (automatically creates CobayaLikelihood) - A Likelihood instance (param_list, param_bounds ignored)

  • param_list (List[str]) – Names of parameters. Required if loglikelihood is a callable. Ignored for Cobaya likelihoods (extracted from YAML/dict).

  • param_bounds (array-like, optional) – Parameter bounds, shape (2, ndim). Required if loglikelihood is a callable. Ignored for Cobaya likelihoods (extracted from priors).

  • param_labels (list of str, optional) – LaTeX labels for parameters. If not provided, uses param_list. Ignored for Cobaya likelihoods (extracted from YAML/dict).

  • likelihood_name (str, optional) – Name for the likelihood (used in output files). If not provided, uses ‘likelihood’ for callables or ‘cobaya_model’ for Cobaya likelihoods.

  • confidence_for_unbounded (float, optional) – Confidence level for unbounded Cobaya priors. Default is 0.9999995. Only used when loglikelihood is a Cobaya YAML file or dict.

  • gp_kwargs (Dict[str, Any]) – Additional keyword arguments to pass to GP constructors. Default is {}.

  • n_cobaya_init (int, optional) – Number of initial points from Cobaya reference distribution. Only used for CobayaLikelihood instances. Default is 4.

  • n_sobol_init (int, optional) – Number of initial Sobol quasi-random points. Default is 4.

  • init_train_x (array-like, optional) – User-provided initial training points in parameter space, shape (n_points, ndim). If provided, these will be added to the initial GP training set. Default is None.

  • init_train_y (array-like, optional) – User-provided initial training values (log-likelihood), shape (n_points, 1) or (n_points,). Must be provided if init_train_x is given. Default is None.

  • resume (bool, optional) – If True, resume from a previous run. Default is False.

  • resume_file (str, optional) – Path to resume from (directory containing GP file). Default is None.

  • save_dir (str, optional) – Directory for saving results. Default is ‘.’.

  • save (bool, optional) – Whether to save results periodically. Default is True.

  • save_step (int, optional) – Save results every save_step iterations. Default is 5.

  • optimizer (str, optional) – Optimizer for GP and acquisition function. Options: ‘scipy’, ‘optax’. Default is ‘scipy’.

  • use_clf (bool, optional) – Whether to use classifier for GP filtering. Default is True.

  • clf_type (str, optional) – Classifier type: ‘svm’, ‘nn’, ‘ellipsoid’. Default is ‘svm’.

  • clf_nsigma_threshold (float, optional) – N-sigma threshold for classifier training. Default is 20.

  • clf_use_size (int, optional) – Minimum dataset size before using classifier. Default is 10.

  • clf_update_step (int, optional) – Update classifier every clf_update_step iterations. Default is 1.

  • minus_inf (float, optional) – Value representing negative infinity for failed evaluations. Default is -1e10.

  • seed (Optional[int]) – Random seed for reproducibility. Default is None.

  • verbosity (str) – Logging verbosity level: ‘DEBUG’, ‘INFO’, ‘WARNING’, ‘ERROR’. Default is ‘INFO’.

  • dynamic_dispatch (bool) – If False (default), use static (round-robin) task distribution across MPI ranks. Static dispatch is fully reproducible across runs at a fixed nprocs and seed because task i always lands on rank i % nprocs, whose RNG is seeded deterministically. If True, use dynamic dispatch (first-available worker). Dynamic dispatch can be faster on heterogeneous tasks but yields non-deterministic task-to-worker assignment, breaking reproducibility for any task that consumes randomness (notably TASK_COBAYA_INIT). Only set to True if reproducibility is not required.

Notes

MPI parallelization is handled automatically and transparently. Users do not need to manage MPI processes explicitly in their scripts. When running with MPI (e.g., mpirun -n 4 python script.py), worker processes automatically participate in parallel likelihood evaluations and GP hyperparameter optimization via the MPI_Pool class, while only the main process (rank 0) runs the optimization loop and manages results. Worker processes enter a waiting loop after initialization and process tasks dispatched by the main process.

update_gp(new_pts_u, new_vals, step=0, verbose=True)[source]

Update the GP with new points and values, and track hyperparameters.

Uses pool for parallel GP fitting when refitting is needed. Refits based on number of points added to GP since last fit.

get_next_batch(acq_kwargs, n_batch, n_restarts, maxiter, early_stop_patience, step, verbose=True)[source]

Get the next batch of points using the acquisition function, and track acquisition values.

evaluate_likelihood(new_pts_u, step, verbose=True)[source]

Evaluate the likelihood for new points using pool.

Parameters:

  • new_pts_u (array-like) – Points in unit cube space to evaluate, shape (n_points, ndim).

  • step (int) – Current iteration number.

  • verbose (bool, optional) – Whether to log detailed information.

Returns:

new_vals – Evaluated likelihood values, shape (n_points, 1).

Return type:

jax.numpy.ndarray

check_max_evals_and_gpsize(current_evals)[source]

Check if the maximum evaluations or GP size has been reached.

Parameters:

current_evals – Current number of objective evaluations.

finalise_results()[source]
check_convergence_ei(step, acq_val)[source]

Check convergence for EI/LogEI based on the acquisition function value.

Parameters:

  • step – Current iteration number.

  • acq_val – Current acquisition function value.

Returns:

Whether convergence is achieved based on acquisition value.

Return type:

bool

check_convergence_logz(step, logz_dict, equal_samples, equal_logl, verbose=True, save_checkpoint=True)[source]

Check if the nested sampling has converged and compute KL divergence metrics.

Parameters:

  • step – Current iteration number

  • logz_dict – Dictionary with logz bounds and mean

  • ns_samples – Nested sampling samples with x, weights, logl

  • threshold – LogZ convergence threshold

Returns:

Whether convergence is achieved based on logz only

Return type:

bool

run(acq='wipstd', min_evals=None, max_evals=None, max_gp_size=None, logz_threshold=None, convergence_n_iters=1, ei_goal=1e-10, do_final_ns=False, fit_n_points=None, batch_size=None, ns_n_points=None, num_hmc_warmup=None, num_hmc_samples=None, mc_points_size=None, thinning=4, num_chains=None, mc_points_method='NUTS', zeta_ei=0.0)[source]

Run the Bayesian Optimization loop.

Parameters:

  • acq (Union[str, Tuple[str]]) – Acquisition function(s) to use: ‘WIPV’, ‘EI’, ‘LogEI’, ‘WIPStd’.

  • min_evals (Optional[int]) – Minimum number of likelihood evaluations before checking convergence. If None, uses dimension-based default from _get_dimension_based_defaults().

  • max_evals (Optional[int]) – Maximum number of likelihood evaluations. If None, uses dimension-based default from _get_dimension_based_defaults().

  • max_gp_size (Optional[int]) – Maximum number of points used to train the GP. If None, uses dimension-based default from _get_dimension_based_defaults().

  • logz_threshold (Optional[float]) – Convergence threshold for log evidence change (WIPV/WIPStd). If None, uses dimension-based default from _get_dimension_based_defaults().

  • convergence_n_iters (int) – Number of successive iterations meeting threshold for convergence. Default is 1.

  • ei_goal (float) – Goal value for EI/LogEI acquisition convergence. Default is 1e-10.

  • do_final_ns (bool) – Whether to run final nested sampling at convergence (WIPV/WIPStd). Default is False.

  • fit_n_points (Optional[int]) – Refit GP hyperparameters after adding this many new points to the GP. If None, uses dimension-based default from _get_dimension_based_defaults().

  • batch_size (Optional[int]) – Batch size for WIPV/WIPStd acquisition. If None, uses dimension-based default from _get_dimension_based_defaults().

  • ns_n_points (Optional[int]) – Run nested sampling after adding this many new points to the GP (for WIPV/WIPStd). If None, uses dimension-based default from _get_dimension_based_defaults().

  • num_hmc_warmup (Optional[int]) – Number of HMC warmup steps. If None, uses dimension-based default from _get_dimension_based_defaults().

  • num_hmc_samples (Optional[int]) – Number of HMC samples to draw. If None, uses dimension-based default from _get_dimension_based_defaults().

  • mc_points_size (Optional[int]) – Number of MC points for WIPV acquisition. If None, uses dimension-based default from _get_dimension_based_defaults().

  • thinning (int) – Thinning factor for MC samples. Default is 4.

  • num_chains (Optional[int]) – Number of parallel HMC chains. If None, uses dimension-based default from _get_dimension_based_defaults().

  • mc_points_method (str) – Method for generating MC points: ‘NUTS’, ‘NS’, or ‘uniform’. Default is ‘NUTS’.

  • zeta_ei (float) – Exploration parameter for EI acquisition. Default is 0.0.

Returns:

Results dictionary containing samples, GP, likelihood, and convergence information. Keys include:

Return type:

dict

run_EI(ii=0)[source]

Run the optimization loop for EI/LogEI acquisition functions.

run_weighted_integrated_posterior(acq_func_class, ii=0)[source]

Run the optimization loop for Weighted Integrated Posterior acquisition functions (WIPV or WIPStd).

Parameters:

  • acq_func_class (class) – The acquisition function class to use (WIPV or WIPStd).

  • ii (int, optional) – Starting iteration number. Default is 0.

run_WIPStd(ii=0)[source]

Run optimization loop for WIPStd acquisition function.

run_WIPV(ii=0)[source]

Run optimization loop for WIPV acquisition function.