BOBE package

Submodules

acquisition

class BOBE.acquisition.AcquisitionFunction(optimizer='scipy', optimizer_options={})[source]

Bases: object

Base class for acquisition functions.

Acquisition functions guide the selection of new points to evaluate by balancing exploration and exploitation. Subclasses must implement the fun and get_next_point methods.

name

Name of the acquisition function.

Type:

str

optimizer

Optimizer to use (‘scipy’ or ‘optax’).

Type:

str

optimizer_options

Additional options for the optimizer.

Type:

dict

acq_optimize

Optimization function (optimize_scipy or optimize_optax).

Type:

callable

name: str = 'BaseAcquisitionFunction'
__init__(optimizer='scipy', optimizer_options={})[source]
fun(x)[source]
get_next_point(gp, acq_kwargs={}, maxiter=500, n_restarts=8, verbose=True, early_stop_patience=25, rng=None)[source]

Optimize the acquisition function to obtain the next point to sample.

Parameters:
  • gp (GP) – Gaussian process model.

  • acq_kwargs (Dict[str, Any]) – Additional arguments for the acquisition function. Default is {}.

  • maxiter (int) – Maximum number of optimization iterations. Default is 500.

  • n_restarts (int) – Number of random restarts for optimization. Default is 8.

  • verbose (bool) – Whether to print optimization progress. Default is True.

  • early_stop_patience (int) – Patience for early stopping. Default is 25.

  • rng (np.random.Generator, optional) – Random number generator. Default is None.

Returns:

(best_point, best_value) where best_point is shape (ndim,) and best_value is the acquisition function value.

Return type:

Tuple[ndarray, float]

get_next_batch(gp, n_batch=1, acq_kwargs={}, maxiter=500, n_restarts=8, verbose=True, early_stop_patience=25, rng=None)[source]

Get the next batch of points to sample.

Return type:

Tuple[ndarray, float]

class BOBE.acquisition.EI(optimizer='scipy', optimizer_options={})[source]

Bases: AcquisitionFunction

Expected Improvement acquisition function.

EI measures the expected improvement over the current best observed value. It balances exploitation (high mean) and exploration (high uncertainty).

The EI criterion is defined as:

EI(x) = E[max(f(x) - f_best - zeta, 0)]

where f_best is the best observed value and zeta is an exploration bonus.

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'EI'
__init__(optimizer='scipy', optimizer_options={})[source]
fun(x, gp, best_y, zeta)[source]

Compute Expected Improvement at point x.

Parameters:
  • x (jnp.ndarray) – Point at which to evaluate EI, shape (ndim,).

  • gp (GP) – Gaussian process model.

  • best_y (float) – Best observed function value.

  • zeta (float) – Exploration bonus parameter.

Returns:

Negative expected improvement (for minimization).

Return type:

float

get_next_point(gp, acq_kwargs, maxiter=250, n_restarts=20, verbose=True, early_stop_patience=25, rng=None)[source]

Optimize the acquisition function to obtain the next point to sample.

Parameters:
  • gp (GP) – Gaussian process model.

  • acq_kwargs (dict, optional) – Additional arguments for the acquisition function. Default is {}.

  • maxiter (int) – Maximum number of optimization iterations. Default is 500.

  • n_restarts (int) – Number of random restarts for optimization. Default is 8.

  • verbose (bool) – Whether to print optimization progress. Default is True.

  • early_stop_patience (int) – Patience for early stopping. Default is 25.

  • rng (np.random.Generator, optional) – Random number generator. Default is None.

Returns:

(best_point, best_value) where best_point is shape (ndim,) and best_value is the acquisition function value.

Return type:

tuple

class BOBE.acquisition.LogEI(optimizer='scipy', optimizer_options={})[source]

Bases: EI

Log Expected Improvement acquisition function.

LogEI computes the logarithm of the Expected Improvement, providing better numerical stability compared to EI, especially when EI values are very small. Uses advanced numerical techniques for accurate computation in extreme cases.

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

References

[1] Ament, S., et al. (2023). “Unexpected Improvements to Expected Improvement

for Bayesian Optimization.” arXiv:2310.20708.

name: str = 'LogEI'
__init__(optimizer='scipy', optimizer_options={})[source]
fun(x, gp, best_y, zeta)[source]

Log Expected Improvement in pure JAX. Returns positive log-EI, so you can maximize directly.

class BOBE.acquisition.WeightedIntegratedPosteriorBase(optimizer='scipy', optimizer_options={})[source]

Bases: AcquisitionFunction

Base class for Weighted Integrated Posterior acquisition functions.

This base class provides common functionality for acquisition functions that integrate over MC samples from the GP posterior, such as WIPV and WIPStd.

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

__init__(optimizer='scipy', optimizer_options={})[source]
get_next_point(gp, acq_kwargs, maxiter=100, n_restarts=1, verbose=True, early_stop_patience=25, rng=None)[source]

Optimize the acquisition function to obtain the next point.

This method is shared between WIPV and WIPStd as they follow the same optimization procedure but differ only in their objective function.

Parameters:
  • gp (GP) – Gaussian process model.

  • acq_kwargs (dict) – Additional arguments containing ‘mc_samples’ and optionally ‘mc_points_size’.

  • maxiter (int) – Maximum optimization iterations. Default is 100.

  • n_restarts (int) – Number of optimization restarts. Default is 1.

  • verbose (bool) – Whether to print progress. Default is True.

  • early_stop_patience (int) – Early stopping patience. Default is 25.

  • rng (np.random.Generator, optional) – Random number generator. Default is None.

Returns:

(best_point, best_value) where best_point is shape (ndim,).

Return type:

tuple

class BOBE.acquisition.WIPV(optimizer='scipy', optimizer_options={})[source]

Bases: WeightedIntegratedPosteriorBase

Weighted Integrated Posterior Variance acquisition function.

WIPV focuses on reducing uncertainty in regions weighted by posterior probability. It integrates the posterior variance over MC samples drawn from the GP posterior, making it particularly effective for Bayesian evidence estimation.

The criterion is defined as:

WIPV(x) = E_{x’ ~ p(x’ | D)}[Var[f(x) | D]]

where the expectation is over MC samples x’ from the posterior.

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'WIPV'
fun(x, gp, mc_points=None, k_train_mc=None)[source]
class BOBE.acquisition.WIPStd(optimizer='scipy', optimizer_options={})[source]

Bases: WeightedIntegratedPosteriorBase

Weighted Integrated Posterior Standard Deviation acquisition function.

WIPStd is similar to WIPV but uses standard deviation instead of variance, which can provide different exploration characteristics. It integrates the posterior standard deviation over MC samples from the GP posterior.

The criterion is defined as:

WIPStd(x) = E_{x’ ~ p(x’ | D)}[Std[f(x) | D]]

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'WIPStd'
fun(x, gp, mc_points=None, k_train_mc=None)[source]
BOBE.acquisition.get_mc_samples(gp, warmup_steps=512, num_samples=1024, thinning=4, method='NUTS', num_chains=4, np_rng=None, rng_key=None)[source]
BOBE.acquisition.get_mc_points(mc_samples, mc_points_size=128, rng=None)[source]

bo

BOBE.bo.load_gp_file(filename, clf)[source]

Load a GP or GPwithClassifier object from a file.

Parameters:

filename (str) – The path to the file from which to load the GP object.

Returns:

The loaded GP or GPwithClassifier object.

Return type:

Union[GP, GPwithClassifier]

BOBE.bo.load_gp_statedict(state_dict, clf)[source]

Load a GP or GPwithClassifier object from a state dictionary.

Parameters:
  • state_dict (Dict[str, Any]) – The state dictionary containing the GP parameters.

  • clf (bool) – Whether to load a GPwithClassifier (True) or a standard GP (False).

Returns:

The loaded GP or GPwithClassifier object.

Return type:

Union[GP, GPwithClassifier]

BOBE.bo.get_dimension_based_defaults(ndim)[source]

Compute reasonable default values for run() parameters based on problem dimension.

This method provides dimension-scaled defaults for parameters that should adapt to the complexity of the problem. Users can override these by providing explicit values to the run() method.

Returns:

Dictionary of default parameter values keyed by parameter name.

Return type:

dict

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.

clf

BOBE.clf.train_svm_classifier(X, Y, settings={}, init_params=None, **kwargs)[source]

Train SVM classifier and return parameters, metrics, and predict function.

BOBE.clf.get_svm_predict_proba_fn(params)[source]

Get prediction function for SVM classifier from parameters (for loading from file).

BOBE.clf.train_nn_classifier(X, Y, settings={}, init_params=None, **kwargs)[source]

Train neural network classifier and return parameters, metrics, and predict function.

BOBE.clf.get_nn_predict_proba_fn(params, settings={}, **kwargs)[source]

Get prediction function for NN classifier from parameters (for loading from file).

BOBE.clf.train_ellipsoid_classifier(X, Y, settings={}, init_params=None, **kwargs)[source]

Train ellipsoid classifier and return parameters, metrics, and predict function.

BOBE.clf.get_ellipsoid_predict_proba_fn(params, settings, d, **kwargs)[source]

Get prediction function for ellipsoid classifier from parameters (for loading from file).

BOBE.clf.svm_predict(x, support_vectors, dual_coef, intercept, gamma)[source]

Compute the decision function for SVM with RBF kernel.

Parameters:
  • x (Array) – Input data point, shape (n_features,)

  • support_vectors (Array) – JAX array of support vectors, shape (n_sv, n_features)

  • dual_coef (Array) – JAX array of dual coefficients, shape (n_sv,)

  • intercept (float) – Scalar bias term.

  • gamma (float) – RBF kernel gamma parameter.

Returns:

Decision function value (scalar). Sign of this value gives the predicted class.

BOBE.clf.svm_predict_proba(x, support_vectors, dual_coef, intercept, gamma)[source]
BOBE.clf.train_with_restarts(train_fn, x, y, n_restarts=2, init_params=None, **train_kwargs)[source]

Train model with multiple restarts using the entire dataset.

Parameters:
  • train_fn (Callable) – Training function that returns (params, metrics)

  • x (Array) – (N, d) features

  • y (Array) – (N,) labels

  • n_restarts (int) – number of random restarts

  • init_params – initial parameters for first restart

  • **train_kwargs – passed to train_fn

Return type:

Tuple[Dict, Dict]

BOBE.clf.train_nn(model, x_train, y_train, init_params=None, **kwargs)[source]

Simplified NN training using entire dataset

BOBE.clf.train_nn_multiple_restarts(model, x, y, **kwargs)[source]

Wrapper for NN training with restarts

BOBE.clf.train_ellipsoid(model, x_train, y_train, init_params=None, **kwargs)[source]

Simplified ellipsoid training using entire dataset

BOBE.clf.train_ellipsoid_multiple_restarts(model, x, y, **kwargs)[source]

Wrapper for ellipsoid training with restarts

clf_gp

class BOBE.clf_gp.GPwithClassifier(train_x=None, train_y=None, clf_type='svm', clf_settings={}, clf_use_size=10, clf_update_step=1, probability_threshold=0.5, minus_inf=-100000.0, clf_threshold=250.0, gp_threshold=500.0, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5.0], tausq=None, tausq_bounds=[0.0001, 10000.0], kernel_variance_prior=None, lengthscale_prior=None, lengthscales=None, kernel_variance=1.0, param_names=None, train_clf_on_init=True)[source]

Bases: GP

__init__(train_x=None, train_y=None, clf_type='svm', clf_settings={}, clf_use_size=10, clf_update_step=1, probability_threshold=0.5, minus_inf=-100000.0, clf_threshold=250.0, gp_threshold=500.0, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5.0], tausq=None, tausq_bounds=[0.0001, 10000.0], kernel_variance_prior=None, lengthscale_prior=None, lengthscales=None, kernel_variance=1.0, param_names=None, train_clf_on_init=True)[source]

Generic Classifier-GP class combining a GP with a classifier. The GP is trained on the data points that are within the GP threshold of the maximum value of the GP.

Parameters:
  • train_x (array-like, shape (n_samples, n_dim)) – Initial training points.

  • train_y (array-like, shape (n_samples,)) – Initial training values.

  • clf_type (str, optional) – Type of classifier (‘svm’, ‘nn’, ‘ellipsoid’, etc.). Default is ‘svm’.

  • clf_params (dict, optional) – Parameters specific to the chosen classifier. Default is None.

  • clf_use_size (int, optional) – Minimum number of points to start using the classifier. Default is 300.

  • clf_update_step (int, optional) – Update classifier every clf_update_step points after clf_use_size is reached. Default is 5.

  • probability_threshold (float, optional) – Threshold for classifier probability/score to consider a point feasible (important for nn, ellipsoid). Default is 0.5.

  • minus_inf (float, optional) – Value used for infeasible predictions. Default is -1e5.

  • clf_threshold (float, optional) – Threshold for initial classifier training labels (if used). If None, gp_threshold might be used or a default calculated.

  • gp_threshold (float, optional) – Threshold for adding points to the GP training set. Default is 5000.

  • noise – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • kernel – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • optimizer – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • kernel_variance_bounds – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • lengthscale_bounds – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • lengthscale_priors – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • lengthscales – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • kernel_variance – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

train_classifier()[source]

Public method to train/retrain the classifier.

predict_mean_single(x)[source]

Single point prediction of mean

predict_var_single(x)[source]
predict_mean_batched(x)[source]
predict_var_batched(x)[source]
predict_single(x)[source]

Predicts the mean and variance of the GP at x but does not unstandardize it. To use with EI and the like.

fantasy_var(new_x, mc_points, k_train_mc)[source]

Computes the fantasy variance, see gp.py for more details. Classifier logic could potentially be added here if needed.

update(new_x, new_y)[source]

Updates the classifier and GP training sets. Retrains classifier/GP based on thresholds and steps.

kernel(x1, x2, lengthscales, kernel_variance, noise, include_noise=True)[source]

Returns the kernel function used by the GP.

get_random_point(rng=None, nstd=None)[source]

Returns a random point in the unit cube.

state_dict()[source]

Returns a dictionary containing the complete state of the GPwithClassifier. This can be used for saving, loading, or copying the GPwithClassifier.

Returns:

state – Dictionary containing all necessary information to reconstruct the GPwithClassifier

Return type:

dict

classmethod from_state_dict(state)[source]

Creates a GPwithClassifier instance from a state dictionary.

Parameters:

state (dict) – State dictionary returned by state_dict()

Returns:

gp_clf – The reconstructed GPwithClassifier object

Return type:

GPwithClassifier

save(filename='gp')[source]

Save the GPwithClassifier state to a file using state_dict.

Parameters:

filename (str) – The filename to save to (with or without .npz extension). Default is ‘gp’.

classmethod load(filename, **kwargs)[source]

Loads a GPwithClassifier from a file

Parameters:
  • filename (str) – The name of the file to load the GPwithClassifier from (with or without .npz extension)

  • **kwargs – Additional keyword arguments to pass to the GPwithClassifier constructor

Returns:

gp_clf – The loaded GPwithClassifier object

Return type:

GPwithClassifier

copy()[source]

Creates a deep copy of the GPwithClassifier using state_dict.

Returns:

gp_clf_copy – A deep copy of the current GPwithClassifier

Return type:

GPwithClassifier

property clf_data_size

Size of the classifier’s training inputs.

property npoints

gp

class BOBE.gp.DummyDistribution[source]

Bases: object

A dummy distribution that always returns log_prob = 0.0

log_prob(x)[source]
BOBE.gp.make_distribution(spec)[source]

Turn a dictionary specification into a NumPyro distribution.

Parameters:

spec (dict) – Dictionary with ‘name’ key for distribution type and additional keyword arguments for the distribution parameters.

Returns:

NumPyro distribution object.

Return type:

Distribution

Examples

>>> spec = {"name": "Normal", "loc": 0.0, "scale": 1.0}
>>> dist = make_distribution(spec)
BOBE.gp.saas_prior_logprob(lengthscales, kernel_variance, tausq)[source]

Compute SAAS prior log probability.

Parameters:
  • lengthscales (jnp.ndarray) – Lengthscale parameters.

  • kernel_variance (float) – Kernel variance parameter.

  • tausq (float) – SAAS tausq parameter.

Returns:

Log probability under SAAS priors.

Return type:

float

BOBE.gp.gp_mll(k, train_y, num_points)[source]

Computes the negative marginal log likelihood of the GP

BOBE.gp.fast_update_cholesky(L, k, k_self)[source]
class BOBE.gp.GP(train_x, train_y, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5], lengthscales=None, kernel_variance=None, kernel_variance_prior=None, lengthscale_prior=None, tausq=None, tausq_bounds=[0.0001, 10000.0], param_names=None)[source]

Bases: object

__init__(train_x, train_y, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5], lengthscales=None, kernel_variance=None, kernel_variance_prior=None, lengthscale_prior=None, tausq=None, tausq_bounds=[0.0001, 10000.0], param_names=None)[source]

Initialize the Gaussian Process model.

Parameters:
  • train_x (jnp.ndarray) – Training inputs, shape (N, D).

  • train_y (jnp.ndarray) – Objective function values at training points, shape (N, 1).

  • noise (float, optional) – Noise parameter added to the diagonal of the kernel. Default is 1e-8.

  • kernel (str, optional) – Kernel to use, either “rbf” or “matern”. Default is “rbf”.

  • optimizer (str, optional) – Optimizer to use for hyperparameter tuning. Default is “scipy”.

  • optimizer_options (dict, optional) – Keyword arguments for the optimizer. Default is {}.

  • kernel_variance_bounds (list, optional) – Bounds for the kernel variance. Default is [1e-4, 1e8].

  • lengthscale_bounds (list, optional) – Bounds for the lengthscales. Default is [0.01, 10].

  • lengthscales (jnp.ndarray, optional) – Initial lengthscale values. If None, defaults to ones. Default is None.

  • kernel_variance (float, optional) – Initial kernel variance. If None, defaults to 1.0. Default is None.

  • kernel_variance_prior (dict or str, optional) – Specification for the kernel variance prior. If None, defaults to {‘name’: ‘LogNormal’, ‘loc’: 0.0, ‘scale’: 1.0}. If ‘fixed’, the kernel variance will be fixed to the initial value and not optimized. Defaults to None.

  • lengthscale_prior (str or dict, optional) – Specification for the lengthscale prior. If ‘DSLP’ or None, uses the DSLP prior. If ‘SAAS’, uses the SAAS prior with tausq parameter. Otherwise, uses the provided distribution spec. Defaults to None.

  • tausq (float, optional) – Initial tausq parameter for SAAS prior. Only used when lengthscale_prior=’SAAS’. If None, defaults to 1.0. Defaults to None.

  • tausq_bounds (list, optional) – Bounds for the tausq parameter (in log space). Only used when lengthscale_prior=’SAAS’. Defaults to [-4, 4].

neg_mll(log_params)[source]

Computes the negative log marginal likelihood for the GP with given hyperparameters.

fit(x0=None, maxiter=500)[source]

Performs a serial fit for a given batch of starting points (x0). This method is called by each MPI process on its assigned chunk.

Parameters:
  • x0 (ndarray) – Array of shape (n_restarts_chunk, n_params) containing starting points for optimization (in log space).

  • maxiter (int) – Maximum number of iterations for the optimizer. Defaults to 500.

Returns:

result – Dictionary containing the best ‘mll’ and corresponding ‘params’ (log space) found.

Return type:

dict

update_hyperparams(hyperparams)[source]

Update the GP hyperparameters and recompute the Cholesky and alphas.

predict_mean_single(x)[source]

Single point prediction of mean

predict_var_single(x)[source]
predict_mean_batched(x)[source]
predict_var_batched(x)[source]
predict_single(x)[source]

Predicts the mean and variance of the GP at x but does not unstandardize it. To use with EI and the like.

predict_batched(x)[source]
update(new_x, new_y)[source]

Updates the GP with new training points and refits the GP if refit is True.

Parameters:
  • refit (bool) – Whether to refit the GP hyperparameters. Default is True.

  • maxiter (int) – The maximum number of iterations for the optax optimizer. Default is 200.

  • n_restarts (int) – The number of restarts for the optax optimizer. Default is 4.

recompute_cholesky()[source]

Recomputes the Cholesky decomposition and alphas. Useful if hyperparameters are changed manually.

fantasy_var(new_x, mc_points, k_train_mc)[source]

Computes the variance of the GP at the mc_points assuming a single point new_x is added to the training set

get_random_point(rng=None, nstd=None)[source]

Returns a random point in the unit cube.

state_dict()[source]

Returns a dictionary containing the complete state of the GP. This can be used for saving, loading, or copying the GP.

Returns:

state – Dictionary containing all necessary information to reconstruct the GP

Return type:

dict

classmethod from_state_dict(state)[source]

Creates a GP instance from a state dictionary.

Parameters:

state (dict) – State dictionary returned by state_dict()

Returns:

gp – The reconstructed GP object

Return type:

GP

classmethod load(filename, **kwargs)[source]

Loads a GP from a file

Parameters:
  • filename (str) – The name of the file to load the GP from (with or without .npz extension)

  • **kwargs – Additional keyword arguments to pass to the GP constructor

Returns:

gp – The loaded GP object

Return type:

GP

save(filename='gp')[source]

Save the GP state to a file using state_dict.

Parameters:

filename (str) – The filename to save to (with or without .npz extension). Default is ‘gp’.

copy()[source]

Creates a deep copy of the GP using state_dict.

Returns:

gp_copy – A deep copy of the current GP

Return type:

GP

property npoints
get_hyperparams()[source]
hyperparams_dict()[source]

kernels

Kernel implementations for Gaussian Process models.

All kernels inherit from the base Kernel class and implement the covariance() method. JAX JIT compilation is handled at higher levels (acquisition functions, optimization).

class BOBE.kernels.Kernel(lengthscales, kernel_variance, noise=1e-08)[source]

Bases: ABC

Abstract base class for all kernels in BOBE.

lengthscales

Lengthscale parameters for each dimension, shape (D,)

Type:

jnp.ndarray

kernel_variance

Overall variance/amplitude of the kernel

Type:

float

noise

Observation noise level

Type:

float

__init__(lengthscales, kernel_variance, noise=1e-08)[source]

Initialize kernel with hyperparameters.

Parameters:
  • lengthscales (jnp.ndarray) – Lengthscale for each input dimension

  • kernel_variance (float) – Kernel variance/amplitude parameter

  • noise (float, optional) – Noise level added to diagonal. Default is 1e-8.

sq_dist(xa, xb)[source]

Compute squared Euclidean distance between two sets of points.

This utility method is used by many kernel implementations.

Parameters:
  • xa (jnp.ndarray) – First set of points, shape (n1, D)

  • xb (jnp.ndarray) – Second set of points, shape (n2, D)

Returns:

sq_dist – Squared distances, shape (n1, n2)

Return type:

jnp.ndarray

abstractmethod covariance(xa, xb, include_noise=True)[source]

Compute covariance matrix between two sets of points.

Parameters:
  • xa (jnp.ndarray) – First set of points, shape (n1, D)

  • xb (jnp.ndarray) – Second set of points, shape (n2, D)

  • include_noise (bool, optional) – Whether to add noise to diagonal (only when xa is xb). Default is True.

Returns:

K – Covariance matrix of shape (n1, n2)

Return type:

jnp.ndarray

diagonal(x, include_noise=True)[source]

Compute only the diagonal of the kernel matrix K(x,x).

For stationary kernels, the diagonal is constant: kernel_variance (+ noise). Override this method if your kernel has a non-constant diagonal.

Parameters:
  • x (jnp.ndarray) – Points at which to compute diagonal, shape (n, D)

  • include_noise (bool, optional) – Whether to include noise in diagonal. Default is True.

Returns:

diag – Diagonal values, shape (n,)

Return type:

jnp.ndarray

update_hyperparams(lengthscales=None, kernel_variance=None, noise=None)[source]

Update kernel hyperparameters.

Parameters:
  • lengthscales (jnp.ndarray, optional) – New lengthscale values

  • kernel_variance (float, optional) – New kernel variance

  • noise (float, optional) – New noise level

__call__(xa, xb, include_noise=True)[source]

Convenience method - same as covariance()

class BOBE.kernels.RBFKernel(lengthscales, kernel_variance, noise=1e-08)[source]

Bases: Kernel

Radial Basis Function (RBF) / Squared Exponential kernel.

k(x, x’) = σ² * exp(-0.5 * ||x - x’||²/ℓ²)

where σ² is kernel_variance and ℓ is lengthscale.

covariance(xa, xb, include_noise=True)[source]

Compute RBF covariance matrix.

Parameters:
  • xa (jnp.ndarray) – First set of input points, shape (n1, d).

  • xb (jnp.ndarray) – Second set of input points, shape (n2, d).

  • include_noise (bool, optional) – Whether to include noise on diagonal. Default is True.

Returns:

Kernel matrix of shape (n1, n2).

Return type:

jnp.ndarray

class BOBE.kernels.MaternKernel(lengthscales, kernel_variance, noise=1e-08)[source]

Bases: Kernel

Matérn-5/2 kernel.

k(x, x’) = σ² * (1 + √5*d + 5*d²/3) * exp(-√5*d)

where d = ||x - x’||/ℓ, σ² is kernel_variance, and ℓ is lengthscale.

covariance(xa, xb, include_noise=True)[source]

Compute Matérn-5/2 covariance matrix.

Parameters:
  • xa (jnp.ndarray) – First set of input points, shape (n1, d).

  • xb (jnp.ndarray) – Second set of input points, shape (n2, d).

  • include_noise (bool, optional) – Whether to include noise on diagonal. Default is True.

Returns:

Kernel matrix of shape (n1, n2).

Return type:

jnp.ndarray

likelihood

class BOBE.likelihood.Likelihood(loglikelihood, param_list, param_labels=None, param_bounds=None, name=None, minus_inf=-10000000000.0)[source]

Bases: object

Base class for log-likelihoods with common evaluation logic.

Parameters:
  • loglikelihood (Callable) – Log-likelihood function that takes parameter array and returns float.

  • param_list (Optional[List[str]]) – List of parameter names.

  • param_labels (Optional[List[str]]) – LaTeX labels for parameters. Default is None (uses param_list).

  • param_bounds (Union[List, ndarray, None]) – Parameter bounds, shape (2, ndim). Default is None (unit cube).

  • name (Optional[str]) – Name for this likelihood. Default is “loglikelihood”.

  • minus_inf (float) – Value to return for failed evaluations. Default is -1e5.

__init__(loglikelihood, param_list, param_labels=None, param_bounds=None, name=None, minus_inf=-10000000000.0)[source]
__call__(X)[source]

Evaluate the likelihood function at a single point.

This method is designed to be called by pool.run_map_objective, which handles distributing tasks across MPI processes in parallel mode or iterating in serial mode.

Parameters:

X (Union[ndarray, List[float]]) – Single input parameter vector.

Returns:

val – Log-likelihood value at the input point.

Return type:

float

class BOBE.likelihood.CobayaLikelihood(input_file_dict, confidence_for_unbounded=0.9999995, minus_inf=-10000000000.0, name='CobayaLikelihood')[source]

Bases: Likelihood

Likelihood wrapper for Cobaya models.

Parameters:
  • input_file_dict (Union[str, Dict[str, Any]]) – Cobaya input YAML file path or input dictionary.

  • confidence_for_unbounded (float) – Confidence level for unbounded priors. Default is 0.9999995.

  • minus_inf (float) – Value to return for failed/minus infinity evaluations. Default is -1e10.

  • name (str) – Name for this likelihood. Default is “CobayaLikelihood”.

__init__(input_file_dict, confidence_for_unbounded=0.9999995, minus_inf=-10000000000.0, name='CobayaLikelihood')[source]
__call__(X)[source]

Evaluate Cobaya likelihood with logprior volume correction. This is added to match Cobaya behaviour.

Return type:

float

samplers

BOBE.samplers.compute_integrals(logl=None, logvol=None, reweight=None, squared=False)[source]
BOBE.samplers.prior_transform(x)[source]
BOBE.samplers.nested_sampling_Dy(gp, mode='acq', ndim=1, dlogz=0.1, dynamic=False, maxcall=5000000, print_progress=True, equal_weights=False, sample_method='rwalk', rng=None)[source]

Nested Sampling using Dynesty

Parameters:
  • gp (GP) – Gaussian Process model

  • ndim (int) – Number of dimensions

  • dlogz (float) – Log evidence goal

  • dynamic (bool) – Use dynamic nested sampling, see Dynesty documentation for more details

  • logz_std (bool) – Compute the upper and lower bounds on logZ using the GP uncertainty

  • maxcall (Optional[int]) – Maximum number of function calls

  • boost_maxcall (int) – Boost the maximum number of function calls

  • print_progress (Optional[bool]) – Print progress of the nested sampling run. If None, automatically disables progress printing in cluster environments and enables it otherwise.

  • equal_weights (bool) – Resample to obtain equal weights

  • sample_method (str) – Sampling method for dynesty

  • rng (random number generator) – Random number generator

Return type:

tuple[ndarray, Dict, bool]

Returns:

  • samples (ndarray) – Equally weighted samples from the nested sampler

  • logz_dict (dict) – Dictionary containing the mean, upper and lower bounds on logZ and the logZ error from the nested sampler

  • success (bool) – Whether the nested sampling run was successful

BOBE.samplers.sample_GP_NUTS(gp, np_rng=None, rng_key=None, num_chains=4, temp=1.0, **kwargs)[source]

Obtain samples from the posterior represented by the GP mean as the logprob. This is a unified function that works for both GP and GPwithClassifier.

Parameters:
  • gp (Union[GP, GPwithClassifier]) – The Gaussian Process model to sample from.

  • np_rng (np.random.Generator, optional) – NumPy random number generator. Default is None.

  • rng_key (jax.random.PRNGKey, optional) – JAX random key. Default is None.

  • num_chains (int, optional) – Number of parallel HMC chains. Default is 4.

  • temp (float, optional) – Temperature parameter for tempering. Default is 1.0.

  • **kwargs (dict) –

    Additional keyword arguments. Can include: - warmup_steps : int, optional

    Number of warmup steps for HMC. If not provided, defaults based on dimensionality.

    • num_samplesint, optional

      Number of samples to draw from each chain. If not provided, defaults based on dimensionality.

    • thinningint, optional

      Thinning factor for samples. If not provided, defaults to 4.

    • dense_massbool, optional

      Whether to use dense mass matrix in NUTS. Default is True.

    • max_tree_depthint, optional

      Maximum tree depth for NUTS. Default is 6.

Returns:

samples_dict – Dictionary containing: - ‘x’: samples array of shape (num_chains * num_samples / thinning, ndim) - ‘logp’: log probabilities for each sample - ‘best’: best sample found - ‘method’: ‘MCMC’

Return type:

dict

optim

BOBE.optim.optimize_optax(fun, fun_args=(), fun_kwargs={}, num_params=1, bounds=None, x0=None, optimizer_options={'early_stop_patience': 25, 'lr': 0.001, 'name': 'adam'}, maxiter=200, n_restarts=1, verbose=False)[source]

SEQUENTIAL OPTIMIZER: Minimize a function using JAX + optax, iterating through restarts with a Python for-loop.

Requires optax to be installed. Install with: pip install ‘BOBE[nn]’ or pip install optax

Return type:

Tuple[Array, float]

BOBE.optim.optimize_optax_vmap(fun, fun_args=(), fun_kwargs={}, num_params=1, bounds=None, x0=None, optimizer_options={'early_stop_patience': 25, 'lr': 0.001, 'name': 'adam'}, maxiter=200, n_restarts=1, verbose=False)[source]

VECTORIZED OPTIMIZER: Minimize a function using JAX + optax, vectorizing over restarts with jax.vmap for parallel execution. Only to be used with EI.

Requires optax to be installed. Install with: pip install ‘BOBE[nn]’ or pip install optax

Return type:

Tuple[Array, float]

BOBE.optim.optimize_scipy(fun, fun_args=(), fun_kwargs={}, num_params=1, bounds=None, x0=None, optimizer_options={'ftol': 1e-06, 'gtol': 1e-06, 'method': 'L-BFGS-B'}, maxiter=200, n_restarts=4, verbose=False)[source]

Standalone method to minimize a function using scipy.optimize.minimize.

Parameters:
  • fun (Callable) – The objective function to minimize

  • fun_args (Optional[Tuple]) – Additional arguments to pass to the function

  • fun_kwargs (Optional[dict]) – Additional keyword arguments to pass to the function

  • num_params (int) – Number of parameters.

  • bounds (Union[List, Tuple, Array, None]) – Parameter bounds in shape (2, num_params).

  • x0 (Optional[Array]) – Initial guess for the parameters, rescaled to unit space.

  • optimizer_kwargs (Optional[dict]) – Additional keyword arguments to pass to the optimizer.

  • maxiter (int) – Maximum number of iterations.

  • n_restarts (int) – Number of restarts for the optimization.

  • verbose (bool) – Whether to print progress messages.

Returns:

The best parameters found and the corresponding function value.

Return type:

Tuple[Array, float]

Module Contents

BOBE - Bayesian Optimization for Expensive Likelihoods using JAX

BOBE is a package for performing Bayesian model comparison for expensive likelihood functions, developed for applications to cosmology. It uses Bayesian Optimization to train Gaussian process surrogates and runs nested sampling/MCMC on the surrogate instead of the underlying expensive likelihood.

Main Components: - BOBE: Main Bayesian Optimization class (accepts raw callable or Likelihood) - GP: Gaussian Process implementation - GPwithClassifier: GP with SVM/NN/Ellipsoid classifier for filtering - Likelihood classes: Likelihood, CobayaLikelihood (optional, wrapping handled automatically) - Acquisition functions: EI, LogEI, WIPV, WIPStd

Quick Start:
>>> from BOBE import BOBE
>>> import numpy as np
>>>
>>> def my_loglike(x):
>>>     return -np.sum(x**2)
>>>
>>> bobe = BOBE(
>>>     loglikelihood=my_loglike,
>>>     param_list=['x', 'y'],
>>>     param_bounds=np.array([[-5, 5], [-5, 5]]).T,
>>>     max_evals=100,
>>> )
>>> results = bobe.run(['wipv'])

MPI parallelization and logging are handled transparently - no explicit setup needed.

class BOBE.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.

class BOBE.GP(train_x, train_y, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5], lengthscales=None, kernel_variance=None, kernel_variance_prior=None, lengthscale_prior=None, tausq=None, tausq_bounds=[0.0001, 10000.0], param_names=None)[source]

Bases: object

__init__(train_x, train_y, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5], lengthscales=None, kernel_variance=None, kernel_variance_prior=None, lengthscale_prior=None, tausq=None, tausq_bounds=[0.0001, 10000.0], param_names=None)[source]

Initialize the Gaussian Process model.

Parameters:
  • train_x (jnp.ndarray) – Training inputs, shape (N, D).

  • train_y (jnp.ndarray) – Objective function values at training points, shape (N, 1).

  • noise (float, optional) – Noise parameter added to the diagonal of the kernel. Default is 1e-8.

  • kernel (str, optional) – Kernel to use, either “rbf” or “matern”. Default is “rbf”.

  • optimizer (str, optional) – Optimizer to use for hyperparameter tuning. Default is “scipy”.

  • optimizer_options (dict, optional) – Keyword arguments for the optimizer. Default is {}.

  • kernel_variance_bounds (list, optional) – Bounds for the kernel variance. Default is [1e-4, 1e8].

  • lengthscale_bounds (list, optional) – Bounds for the lengthscales. Default is [0.01, 10].

  • lengthscales (jnp.ndarray, optional) – Initial lengthscale values. If None, defaults to ones. Default is None.

  • kernel_variance (float, optional) – Initial kernel variance. If None, defaults to 1.0. Default is None.

  • kernel_variance_prior (dict or str, optional) – Specification for the kernel variance prior. If None, defaults to {‘name’: ‘LogNormal’, ‘loc’: 0.0, ‘scale’: 1.0}. If ‘fixed’, the kernel variance will be fixed to the initial value and not optimized. Defaults to None.

  • lengthscale_prior (str or dict, optional) – Specification for the lengthscale prior. If ‘DSLP’ or None, uses the DSLP prior. If ‘SAAS’, uses the SAAS prior with tausq parameter. Otherwise, uses the provided distribution spec. Defaults to None.

  • tausq (float, optional) – Initial tausq parameter for SAAS prior. Only used when lengthscale_prior=’SAAS’. If None, defaults to 1.0. Defaults to None.

  • tausq_bounds (list, optional) – Bounds for the tausq parameter (in log space). Only used when lengthscale_prior=’SAAS’. Defaults to [-4, 4].

neg_mll(log_params)[source]

Computes the negative log marginal likelihood for the GP with given hyperparameters.

fit(x0=None, maxiter=500)[source]

Performs a serial fit for a given batch of starting points (x0). This method is called by each MPI process on its assigned chunk.

Parameters:
  • x0 (ndarray) – Array of shape (n_restarts_chunk, n_params) containing starting points for optimization (in log space).

  • maxiter (int) – Maximum number of iterations for the optimizer. Defaults to 500.

Returns:

result – Dictionary containing the best ‘mll’ and corresponding ‘params’ (log space) found.

Return type:

dict

update_hyperparams(hyperparams)[source]

Update the GP hyperparameters and recompute the Cholesky and alphas.

predict_mean_single(x)[source]

Single point prediction of mean

predict_var_single(x)[source]
predict_mean_batched(x)[source]
predict_var_batched(x)[source]
predict_single(x)[source]

Predicts the mean and variance of the GP at x but does not unstandardize it. To use with EI and the like.

predict_batched(x)[source]
update(new_x, new_y)[source]

Updates the GP with new training points and refits the GP if refit is True.

Parameters:
  • refit (bool) – Whether to refit the GP hyperparameters. Default is True.

  • maxiter (int) – The maximum number of iterations for the optax optimizer. Default is 200.

  • n_restarts (int) – The number of restarts for the optax optimizer. Default is 4.

recompute_cholesky()[source]

Recomputes the Cholesky decomposition and alphas. Useful if hyperparameters are changed manually.

fantasy_var(new_x, mc_points, k_train_mc)[source]

Computes the variance of the GP at the mc_points assuming a single point new_x is added to the training set

get_random_point(rng=None, nstd=None)[source]

Returns a random point in the unit cube.

state_dict()[source]

Returns a dictionary containing the complete state of the GP. This can be used for saving, loading, or copying the GP.

Returns:

state – Dictionary containing all necessary information to reconstruct the GP

Return type:

dict

classmethod from_state_dict(state)[source]

Creates a GP instance from a state dictionary.

Parameters:

state (dict) – State dictionary returned by state_dict()

Returns:

gp – The reconstructed GP object

Return type:

GP

classmethod load(filename, **kwargs)[source]

Loads a GP from a file

Parameters:
  • filename (str) – The name of the file to load the GP from (with or without .npz extension)

  • **kwargs – Additional keyword arguments to pass to the GP constructor

Returns:

gp – The loaded GP object

Return type:

GP

save(filename='gp')[source]

Save the GP state to a file using state_dict.

Parameters:

filename (str) – The filename to save to (with or without .npz extension). Default is ‘gp’.

copy()[source]

Creates a deep copy of the GP using state_dict.

Returns:

gp_copy – A deep copy of the current GP

Return type:

GP

property npoints
get_hyperparams()[source]
hyperparams_dict()[source]
class BOBE.GPwithClassifier(train_x=None, train_y=None, clf_type='svm', clf_settings={}, clf_use_size=10, clf_update_step=1, probability_threshold=0.5, minus_inf=-100000.0, clf_threshold=250.0, gp_threshold=500.0, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5.0], tausq=None, tausq_bounds=[0.0001, 10000.0], kernel_variance_prior=None, lengthscale_prior=None, lengthscales=None, kernel_variance=1.0, param_names=None, train_clf_on_init=True)[source]

Bases: GP

__init__(train_x=None, train_y=None, clf_type='svm', clf_settings={}, clf_use_size=10, clf_update_step=1, probability_threshold=0.5, minus_inf=-100000.0, clf_threshold=250.0, gp_threshold=500.0, noise=1e-08, kernel='rbf', optimizer='scipy', optimizer_options={}, kernel_variance_bounds=[0.0001, 100000000.0], lengthscale_bounds=[0.01, 5.0], tausq=None, tausq_bounds=[0.0001, 10000.0], kernel_variance_prior=None, lengthscale_prior=None, lengthscales=None, kernel_variance=1.0, param_names=None, train_clf_on_init=True)[source]

Generic Classifier-GP class combining a GP with a classifier. The GP is trained on the data points that are within the GP threshold of the maximum value of the GP.

Parameters:
  • train_x (array-like, shape (n_samples, n_dim)) – Initial training points.

  • train_y (array-like, shape (n_samples,)) – Initial training values.

  • clf_type (str, optional) – Type of classifier (‘svm’, ‘nn’, ‘ellipsoid’, etc.). Default is ‘svm’.

  • clf_params (dict, optional) – Parameters specific to the chosen classifier. Default is None.

  • clf_use_size (int, optional) – Minimum number of points to start using the classifier. Default is 300.

  • clf_update_step (int, optional) – Update classifier every clf_update_step points after clf_use_size is reached. Default is 5.

  • probability_threshold (float, optional) – Threshold for classifier probability/score to consider a point feasible (important for nn, ellipsoid). Default is 0.5.

  • minus_inf (float, optional) – Value used for infeasible predictions. Default is -1e5.

  • clf_threshold (float, optional) – Threshold for initial classifier training labels (if used). If None, gp_threshold might be used or a default calculated.

  • gp_threshold (float, optional) – Threshold for adding points to the GP training set. Default is 5000.

  • noise – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • kernel – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • optimizer – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • kernel_variance_bounds – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • lengthscale_bounds – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • lengthscale_priors – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • lengthscales – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

  • kernel_variance – GP parameters (see DSLP_GP/SAAS_GP). Note: bounds are now in actual space, not log10.

train_classifier()[source]

Public method to train/retrain the classifier.

predict_mean_single(x)[source]

Single point prediction of mean

predict_var_single(x)[source]
predict_mean_batched(x)[source]
predict_var_batched(x)[source]
predict_single(x)[source]

Predicts the mean and variance of the GP at x but does not unstandardize it. To use with EI and the like.

fantasy_var(new_x, mc_points, k_train_mc)[source]

Computes the fantasy variance, see gp.py for more details. Classifier logic could potentially be added here if needed.

update(new_x, new_y)[source]

Updates the classifier and GP training sets. Retrains classifier/GP based on thresholds and steps.

kernel(x1, x2, lengthscales, kernel_variance, noise, include_noise=True)[source]

Returns the kernel function used by the GP.

get_random_point(rng=None, nstd=None)[source]

Returns a random point in the unit cube.

state_dict()[source]

Returns a dictionary containing the complete state of the GPwithClassifier. This can be used for saving, loading, or copying the GPwithClassifier.

Returns:

state – Dictionary containing all necessary information to reconstruct the GPwithClassifier

Return type:

dict

classmethod from_state_dict(state)[source]

Creates a GPwithClassifier instance from a state dictionary.

Parameters:

state (dict) – State dictionary returned by state_dict()

Returns:

gp_clf – The reconstructed GPwithClassifier object

Return type:

GPwithClassifier

save(filename='gp')[source]

Save the GPwithClassifier state to a file using state_dict.

Parameters:

filename (str) – The filename to save to (with or without .npz extension). Default is ‘gp’.

classmethod load(filename, **kwargs)[source]

Loads a GPwithClassifier from a file

Parameters:
  • filename (str) – The name of the file to load the GPwithClassifier from (with or without .npz extension)

  • **kwargs – Additional keyword arguments to pass to the GPwithClassifier constructor

Returns:

gp_clf – The loaded GPwithClassifier object

Return type:

GPwithClassifier

copy()[source]

Creates a deep copy of the GPwithClassifier using state_dict.

Returns:

gp_clf_copy – A deep copy of the current GPwithClassifier

Return type:

GPwithClassifier

property clf_data_size

Size of the classifier’s training inputs.

property npoints
BOBE.get_dimension_based_defaults(ndim)[source]

Compute reasonable default values for run() parameters based on problem dimension.

This method provides dimension-scaled defaults for parameters that should adapt to the complexity of the problem. Users can override these by providing explicit values to the run() method.

Returns:

Dictionary of default parameter values keyed by parameter name.

Return type:

dict

class BOBE.Likelihood(loglikelihood, param_list, param_labels=None, param_bounds=None, name=None, minus_inf=-10000000000.0)[source]

Bases: object

Base class for log-likelihoods with common evaluation logic.

Parameters:
  • loglikelihood (Callable) – Log-likelihood function that takes parameter array and returns float.

  • param_list (Optional[List[str]]) – List of parameter names.

  • param_labels (Optional[List[str]]) – LaTeX labels for parameters. Default is None (uses param_list).

  • param_bounds (Union[List, ndarray, None]) – Parameter bounds, shape (2, ndim). Default is None (unit cube).

  • name (Optional[str]) – Name for this likelihood. Default is “loglikelihood”.

  • minus_inf (float) – Value to return for failed evaluations. Default is -1e5.

__init__(loglikelihood, param_list, param_labels=None, param_bounds=None, name=None, minus_inf=-10000000000.0)[source]
__call__(X)[source]

Evaluate the likelihood function at a single point.

This method is designed to be called by pool.run_map_objective, which handles distributing tasks across MPI processes in parallel mode or iterating in serial mode.

Parameters:

X (Union[ndarray, List[float]]) – Single input parameter vector.

Returns:

val – Log-likelihood value at the input point.

Return type:

float

class BOBE.EI(optimizer='scipy', optimizer_options={})[source]

Bases: AcquisitionFunction

Expected Improvement acquisition function.

EI measures the expected improvement over the current best observed value. It balances exploitation (high mean) and exploration (high uncertainty).

The EI criterion is defined as:

EI(x) = E[max(f(x) - f_best - zeta, 0)]

where f_best is the best observed value and zeta is an exploration bonus.

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'EI'
__init__(optimizer='scipy', optimizer_options={})[source]
fun(x, gp, best_y, zeta)[source]

Compute Expected Improvement at point x.

Parameters:
  • x (jnp.ndarray) – Point at which to evaluate EI, shape (ndim,).

  • gp (GP) – Gaussian process model.

  • best_y (float) – Best observed function value.

  • zeta (float) – Exploration bonus parameter.

Returns:

Negative expected improvement (for minimization).

Return type:

float

get_next_point(gp, acq_kwargs, maxiter=250, n_restarts=20, verbose=True, early_stop_patience=25, rng=None)[source]

Optimize the acquisition function to obtain the next point to sample.

Parameters:
  • gp (GP) – Gaussian process model.

  • acq_kwargs (dict, optional) – Additional arguments for the acquisition function. Default is {}.

  • maxiter (int) – Maximum number of optimization iterations. Default is 500.

  • n_restarts (int) – Number of random restarts for optimization. Default is 8.

  • verbose (bool) – Whether to print optimization progress. Default is True.

  • early_stop_patience (int) – Patience for early stopping. Default is 25.

  • rng (np.random.Generator, optional) – Random number generator. Default is None.

Returns:

(best_point, best_value) where best_point is shape (ndim,) and best_value is the acquisition function value.

Return type:

tuple

class BOBE.LogEI(optimizer='scipy', optimizer_options={})[source]

Bases: EI

Log Expected Improvement acquisition function.

LogEI computes the logarithm of the Expected Improvement, providing better numerical stability compared to EI, especially when EI values are very small. Uses advanced numerical techniques for accurate computation in extreme cases.

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

References

[1] Ament, S., et al. (2023). “Unexpected Improvements to Expected Improvement

for Bayesian Optimization.” arXiv:2310.20708.

name: str = 'LogEI'
__init__(optimizer='scipy', optimizer_options={})[source]
fun(x, gp, best_y, zeta)[source]

Log Expected Improvement in pure JAX. Returns positive log-EI, so you can maximize directly.

class BOBE.WIPV(optimizer='scipy', optimizer_options={})[source]

Bases: WeightedIntegratedPosteriorBase

Weighted Integrated Posterior Variance acquisition function.

WIPV focuses on reducing uncertainty in regions weighted by posterior probability. It integrates the posterior variance over MC samples drawn from the GP posterior, making it particularly effective for Bayesian evidence estimation.

The criterion is defined as:

WIPV(x) = E_{x’ ~ p(x’ | D)}[Var[f(x) | D]]

where the expectation is over MC samples x’ from the posterior.

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'WIPV'
fun(x, gp, mc_points=None, k_train_mc=None)[source]
class BOBE.WIPStd(optimizer='scipy', optimizer_options={})[source]

Bases: WeightedIntegratedPosteriorBase

Weighted Integrated Posterior Standard Deviation acquisition function.

WIPStd is similar to WIPV but uses standard deviation instead of variance, which can provide different exploration characteristics. It integrates the posterior standard deviation over MC samples from the GP posterior.

The criterion is defined as:

WIPStd(x) = E_{x’ ~ p(x’ | D)}[Std[f(x) | D]]

Parameters:
  • optimizer (str) – Optimizer to use (‘scipy’ or ‘optax’). Default is ‘scipy’.

  • optimizer_options (Optional[Dict[str, Any]]) – Additional options for the optimizer. Default is {}.

name: str = 'WIPStd'
fun(x, gp, mc_points=None, k_train_mc=None)[source]
class BOBE.BOBEResults(param_names, param_labels, param_bounds, output_file='results', save_dir='./', settings=None, likelihood_name='unknown', resume_from_existing=False)[source]

Bases: object

Comprehensive results management for BOBE runs.

This class handles storing, organizing, and outputting results in formats compatible with standard nested sampling analysis tools.

__init__(param_names, param_labels, param_bounds, output_file='results', save_dir='./', settings=None, likelihood_name='unknown', resume_from_existing=False)[source]

Initialize the results manager.

Parameters:
  • output_file (str) – Base name for output files

  • param_names (List[str]) – List of parameter names

  • param_labels (List[str]) – List of parameter LaTeX labels

  • param_bounds (ndarray) – Parameter bounds array [n_params, 2]

  • settings (Optional[Dict[str, Any]]) – Dictionary of BOBE settings

  • likelihood_name (str) – Name of the likelihood function

  • resume_from_existing (bool) – If True, try to load existing results and continue from there

update_acquisition(iteration, acquisition_value, acquisition_function)[source]

Track acquisition function values throughout iterations.

Parameters:
  • iteration (int) – Current iteration number

  • acquisition_value (float) – Value of the acquisition function at the selected point

  • acquisition_function (str) – String name of the acquisition function used

update_gp_hyperparams(iteration, lengthscales, kernel_variance)[source]

Track GP hyperparameters evolution.

Parameters:
  • iteration (int) – Current iteration number

  • lengthscales (list) – List of lengthscale values (can be JAX arrays)

  • kernel_variance (float) – Kernel variance value

update_best_loglike(iteration, best_loglike)[source]

Track best loglikelihood evolution.

Parameters:
  • iteration (int) – Current iteration number

  • best_loglike (float) – Current best loglikelihood value

update_convergence(iteration, logz_dict, converged, threshold)[source]

Update convergence information from a nested sampling check.

Parameters:
  • iteration (int) – Current iteration number

  • logz_dict (Dict[str, float]) – Dictionary with logz information

  • converged (bool) – Whether convergence was achieved

  • threshold (float) – Convergence threshold used

update_kl_divergences(iteration, successive_kl=None)[source]

Update KL divergence tracking for convergence analysis.

Parameters:
  • iteration (int) – Current iteration number

  • successive_kl (Optional[Dict[str, float]]) – Optional KL divergence between successive iterations

get_last_iteration()[source]

Get the last iteration number from the results history.

Return type:

int

Returns:

Last iteration number, or 0 if no iterations have been recorded

is_resuming()[source]

Check if this is a resumed run (has existing data).

Return type:

bool

Returns:

True if this appears to be a resumed run

start_timing(phase_name)[source]

Start timing a specific phase.

end_timing(phase_name)[source]

End timing a specific phase and accumulate the time.

get_timing_summary()[source]

Get a summary of timing information.

Return type:

Dict[str, Any]

save_timing_data()[source]

Save timing data to JSON file.

get_gp_data()[source]

Get GP hyperparameter evolution data for plotting.

Return type:

Dict[str, list]

Returns:

Dictionary with ‘iterations’, ‘lengthscales’, and ‘kernel_variances’ keys

get_acquisition_data()[source]

Get acquisition function evolution data for plotting.

Return type:

Dict[str, list]

Returns:

Dictionary with ‘iterations’, ‘values’, and ‘functions’ keys

get_best_loglike_data()[source]

Get best loglikelihood evolution data for plotting.

Return type:

Dict[str, list]

Returns:

Dictionary with ‘iterations’ and ‘best_loglike’ keys

finalize(samples_dict={}, logz_dict=None, converged=False, termination_reason='Max iterations reached', gp_info=None, best_point=None, best_loglike=None, best_iteration=None)[source]

Finalize the results with final samples and metadata.

Parameters:
  • samples_dict (Dict[str, ndarray]) – Dictionary with ‘x’, ‘weights’, ‘logl’ keys for final samples

  • logz_dict (Optional[Dict[str, float]]) – Final evidence information

  • converged (bool) – Whether the run converged

  • termination_reason (str) – Reason for termination

  • gp_info (Optional[Dict[str, Any]]) – Dictionary containing GP and classifier information

  • best_point (Optional[ndarray]) – Best point found (physical parameter space)

  • best_loglike (Optional[float]) – Best log-likelihood value

  • best_iteration (Optional[int]) – Iteration where best point was found

get_results_dict()[source]

Get simplified results dictionary with only essential data.

Return type:

Dict[str, Any]

Returns:

Dictionary containing samples, weights, evidence evolution, and convergence info

save_all_formats()[source]

Save results in multiple formats for compatibility.

save_main_results()[source]

Save main comprehensive results file.

save_chain_files(samples_dict=None, filename=None)[source]

Save chain files in GetDist format using MCSamples.saveAsText method.

save_minimum_files()[source]

Save best point in GetDist minimum format.

Creates two files: - .minimum.txt: Simple table with best point - .minimum: Formatted text with parameter details

save_summary_stats()[source]

Save summary statistics in JSON format.

save_intermediate(gp, filename=None)[source]

Save intermediate results for crash recovery and resuming.

get_getdist_samples(samples_dict=None)[source]

Convert results to GetDist MCSamples object.

Return type:

Optional[MCSamples]

Returns:

GetDist MCSamples object if GetDist is available, None otherwise

classmethod load_results(output_file)[source]

Load results from saved files.

Parameters:

output_file (str) – Base name of the output files

Return type:

BOBEResults

Returns:

BOBEResults object with loaded data

class BOBE.BOBESummaryPlotter(results, figsize_scale=1.0)[source]

Bases: object

Comprehensive plotting class for BOBE run analysis and diagnostics.

__init__(results, figsize_scale=1.0)[source]

Initialize the plotter with BOBE results.

Parameters:
  • results (Union[BOBEResults, str]) – BOBEResults object or path to results file

  • figsize_scale (float) – Scale factor for figure sizes (default: 1.0)

plot_evidence_evolution(logz_data=None, ax=None, show_convergence=True)[source]

Plot the evolution of log evidence (logZ) with error bounds.

Parameters:
  • logz_data (Optional[Dict]) – Dictionary containing logZ evolution data (uses results if None)

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

  • show_convergence (bool) – Whether to mark convergence points

Return type:

Axes

Returns:

The matplotlib axes object

plot_gp_lengthscales(gp_data=None, ax=None)[source]

Plot evolution of GP lengthscales only.

Parameters:
  • gp_data (Optional[Dict]) – Dictionary containing GP hyperparameter evolution data

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_gp_kernel_variance(gp_data=None, ax=None)[source]

Plot evolution of GP kernel variance only.

Parameters:
  • gp_data (Optional[Dict]) – Dictionary containing GP hyperparameter evolution data

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_gp_hyperparameters(gp_data=None, ax=None)[source]

Plot evolution of GP hyperparameters (backward compatibility - now plots lengthscales only).

Parameters:
  • gp_data (Optional[Dict]) – Dictionary containing GP hyperparameter evolution data

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_best_loglike_evolution(best_loglike_data=None, ax=None, scatter_improvements=False)[source]

Plot evolution of the best log-likelihood found so far.

Parameters:
  • best_loglike_data (Optional[Dict]) – Dictionary with ‘iterations’ and ‘best_loglike’ keys

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_acquisition_evolution(acquisition_data=None, ax=None)[source]

Plot the evolution of acquisition function values throughout iterations.

Parameters:
  • acquisition_data (Optional[Dict]) – Dictionary with acquisition data (gets from results if None)

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_timing_breakdown(timing_data=None, ax=None)[source]

Plot timing breakdown of different phases of the algorithm.

Parameters:
  • timing_data (Optional[Dict]) – Dictionary with timing information for different phases

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_convergence_diagnostics(convergence_data=None, ax=None)[source]

Plot convergence diagnostics including thresholds and delta evolution.

Parameters:
  • convergence_data (Optional[Dict]) – Dictionary containing convergence history data (uses results if None)

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_kl_divergences(kl_data=None, ax=None, annotate=False)[source]

Plot successive KL divergences between NS iterations (reverse, forward, symmetric).

Parameters:
  • kl_data (Optional[Dict]) – Dictionary containing KL divergence data (uses results if None)

  • ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

plot_parameter_evolution(param_evolution_data=None, max_params=4)[source]

Plot evolution of parameter values during optimization.

Parameters:
  • param_evolution_data (Optional[Dict]) – Dictionary with parameter evolution data

  • max_params (int) – Maximum number of parameters to plot

Return type:

Figure

Returns:

The matplotlib figure object

create_summary_dashboard(logz_data=None, convergence_data=None, kl_data=None, gp_data=None, best_loglike_data=None, acquisition_data=None, timing_data=None, save_path=None, title=None)[source]

Create a comprehensive summary dashboard with all diagnostic plots.

Parameters:
  • logz_data (Optional[Dict]) – Log evidence evolution data

  • convergence_data (Optional[Dict]) – Convergence diagnostics data

  • kl_data (Optional[Dict]) – KL divergence data

  • gp_data (Optional[Dict]) – GP hyperparameter evolution data

  • best_loglike_data (Optional[Dict]) – Best log-likelihood evolution data

  • acquisition_data (Optional[Dict]) – Acquisition function evolution data

  • timing_data (Optional[Dict]) – Timing breakdown data

  • save_path (Optional[str]) – Path to save the figure (optional)

Return type:

Figure

Returns:

The matplotlib figure object

plot_summary_stats(ax=None)[source]

Plot key summary statistics as text.

Parameters:

ax (Optional[Axes]) – Matplotlib axes to plot on (creates new if None)

Return type:

Axes

Returns:

The matplotlib axes object

save_all_plots(output_dir=None, **data_kwargs)[source]

Save all individual plots and the summary dashboard.

Parameters:
  • output_dir (Optional[str]) – Directory to save plots (uses output_file base if None)

  • **data_kwargs – Data dictionaries for different plot types

BOBE.get_logger(name)[source]

Gets a logger. The root logger should be configured first via setup_logging.

BOBE.setup_logging(verbosity='INFO', log_file=None)[source]

Configure logging for serial or MPI runs.

Parameters:
  • verbosity – String level - ‘DEBUG’, ‘INFO’, ‘WARNING’, ‘ERROR’, ‘CRITICAL’, ‘QUIET’

  • log_file – Optional file to log to. In MPI runs, will be post-fixed with rank.

BOBE.scale_to_unit(x, param_bounds)[source]

Project from original domain to unit hypercube, X is N x d shaped, param_bounds are 2 x d

BOBE.scale_from_unit(x, param_bounds)[source]

Project from unit hypercube to original domain, X is N x d shaped, param_bounds are 2 x d

class BOBE.CobayaLikelihood(input_file_dict, confidence_for_unbounded=0.9999995, minus_inf=-10000000000.0, name='CobayaLikelihood')[source]

Bases: Likelihood

Likelihood wrapper for Cobaya models.

Parameters:
  • input_file_dict (Union[str, Dict[str, Any]]) – Cobaya input YAML file path or input dictionary.

  • confidence_for_unbounded (float) – Confidence level for unbounded priors. Default is 0.9999995.

  • minus_inf (float) – Value to return for failed/minus infinity evaluations. Default is -1e10.

  • name (str) – Name for this likelihood. Default is “CobayaLikelihood”.

__init__(input_file_dict, confidence_for_unbounded=0.9999995, minus_inf=-10000000000.0, name='CobayaLikelihood')[source]
__call__(X)[source]

Evaluate Cobaya likelihood with logprior volume correction. This is added to match Cobaya behaviour.

Return type:

float