BOBE package
Submodules
acquisition
- class BOBE.acquisition.AcquisitionFunction(optimizer='scipy', optimizer_options={})[source]
Bases:
objectBase 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.
- acq_optimize
Optimization function (optimize_scipy or optimize_optax).
- Type:
callable
- 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:
- class BOBE.acquisition.EI(optimizer='scipy', optimizer_options={})[source]
Bases:
AcquisitionFunctionExpected 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:
- 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:
- class BOBE.acquisition.LogEI(optimizer='scipy', optimizer_options={})[source]
Bases:
EILog 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:
References
- [1] Ament, S., et al. (2023). “Unexpected Improvements to Expected Improvement
for Bayesian Optimization.” arXiv:2310.20708.
- class BOBE.acquisition.WeightedIntegratedPosteriorBase(optimizer='scipy', optimizer_options={})[source]
Bases:
AcquisitionFunctionBase 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:
- 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:
- class BOBE.acquisition.WIPV(optimizer='scipy', optimizer_options={})[source]
Bases:
WeightedIntegratedPosteriorBaseWeighted 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:
- class BOBE.acquisition.WIPStd(optimizer='scipy', optimizer_options={})[source]
Bases:
WeightedIntegratedPosteriorBaseWeighted 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:
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.
- 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:
- 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 fixednprocsand 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 (notablyTASK_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:
- 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.
- 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:
- 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:
- 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:
- 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.
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:
- Returns:
Decision function value (scalar). Sign of this value gives the predicted class.
- 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.
- 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
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.
- 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.
- 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:
- 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:
- 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:
- 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:
- property clf_data_size
Size of the classifier’s training inputs.
- property npoints
gp
- class BOBE.gp.DummyDistribution[source]
Bases:
objectA dummy distribution that always returns log_prob = 0.0
- 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.
- BOBE.gp.gp_mll(k, train_y, num_points)[source]
Computes the negative marginal log likelihood of the GP
- 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:
- Returns:
result – Dictionary containing the best ‘mll’ and corresponding ‘params’ (log space) found.
- Return type:
- update_hyperparams(hyperparams)[source]
Update the GP hyperparameters and recompute the Cholesky and alphas.
- 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.
- update(new_x, new_y)[source]
Updates the GP with new training points and refits the GP if refit is True.
- 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
- 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:
- 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:
- property npoints
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:
ABCAbstract base class for all kernels in BOBE.
- lengthscales
Lengthscale parameters for each dimension, shape (D,)
- Type:
jnp.ndarray
- __init__(lengthscales, kernel_variance, noise=1e-08)[source]
Initialize kernel with hyperparameters.
- 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
- class BOBE.kernels.RBFKernel(lengthscales, kernel_variance, noise=1e-08)[source]
Bases:
KernelRadial 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:
KernelMaté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:
objectBase class for log-likelihoods with common evaluation logic.
- Parameters:
loglikelihood (
Callable) – Log-likelihood function that takes parameter array and returns float.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]
- class BOBE.likelihood.CobayaLikelihood(input_file_dict, confidence_for_unbounded=0.9999995, minus_inf=-10000000000.0, name='CobayaLikelihood')[source]
Bases:
LikelihoodLikelihood 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”.
samplers
- 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 modelndim (
int) – Number of dimensionsdlogz (
float) – Log evidence goaldynamic (
bool) – Use dynamic nested sampling, see Dynesty documentation for more detailslogz_std (bool) – Compute the upper and lower bounds on logZ using the GP uncertainty
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 weightssample_method (
str) – Sampling method for dynestyrng (random number generator) – Random number generator
- Return type:
- 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:
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
- 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
- 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 minimizefun_args (
Optional[Tuple]) – Additional arguments to pass to the functionfun_kwargs (
Optional[dict]) – Additional keyword arguments to pass to the functionnum_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:
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 fixednprocsand 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 (notablyTASK_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:
- 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.
- 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:
- 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:
- 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:
- 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.
- 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:
- Returns:
result – Dictionary containing the best ‘mll’ and corresponding ‘params’ (log space) found.
- Return type:
- update_hyperparams(hyperparams)[source]
Update the GP hyperparameters and recompute the Cholesky and alphas.
- 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.
- update(new_x, new_y)[source]
Updates the GP with new training points and refits the GP if refit is True.
- 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
- 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:
- 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:
- property npoints
- 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.
- 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.
- 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:
- 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:
- 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:
- 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:
- 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:
- class BOBE.Likelihood(loglikelihood, param_list, param_labels=None, param_bounds=None, name=None, minus_inf=-10000000000.0)[source]
Bases:
objectBase class for log-likelihoods with common evaluation logic.
- Parameters:
loglikelihood (
Callable) – Log-likelihood function that takes parameter array and returns float.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]
- class BOBE.EI(optimizer='scipy', optimizer_options={})[source]
Bases:
AcquisitionFunctionExpected 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:
- 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:
- class BOBE.LogEI(optimizer='scipy', optimizer_options={})[source]
Bases:
EILog 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:
References
- [1] Ament, S., et al. (2023). “Unexpected Improvements to Expected Improvement
for Bayesian Optimization.” arXiv:2310.20708.
- class BOBE.WIPV(optimizer='scipy', optimizer_options={})[source]
Bases:
WeightedIntegratedPosteriorBaseWeighted 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:
- class BOBE.WIPStd(optimizer='scipy', optimizer_options={})[source]
Bases:
WeightedIntegratedPosteriorBaseWeighted 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:
- 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:
objectComprehensive 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 filesparam_bounds (
ndarray) – Parameter bounds array [n_params, 2]settings (
Optional[Dict[str,Any]]) – Dictionary of BOBE settingslikelihood_name (
str) – Name of the likelihood functionresume_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.
- update_gp_hyperparams(iteration, lengthscales, kernel_variance)[source]
Track GP hyperparameters evolution.
- update_convergence(iteration, logz_dict, converged, threshold)[source]
Update convergence information from a nested sampling check.
- update_kl_divergences(iteration, successive_kl=None)[source]
Update KL divergence tracking for convergence analysis.
- get_last_iteration()[source]
Get the last iteration number from the results history.
- Return type:
- 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:
- Returns:
True if this appears to be a resumed run
- 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 sampleslogz_dict (
Optional[Dict[str,float]]) – Final evidence informationconverged (
bool) – Whether the run convergedtermination_reason (
str) – Reason for terminationgp_info (
Optional[Dict[str,Any]]) – Dictionary containing GP and classifier informationbest_point (
Optional[ndarray]) – Best point found (physical parameter space)best_iteration (
Optional[int]) – Iteration where best point was found
- 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_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
- class BOBE.BOBESummaryPlotter(results, figsize_scale=1.0)[source]
Bases:
objectComprehensive 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 filefigsize_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.
- plot_gp_hyperparameters(gp_data=None, ax=None)[source]
Plot evolution of GP hyperparameters (backward compatibility - now plots lengthscales only).
- plot_best_loglike_evolution(best_loglike_data=None, ax=None, scatter_improvements=False)[source]
Plot evolution of the best log-likelihood found so far.
- plot_acquisition_evolution(acquisition_data=None, ax=None)[source]
Plot the evolution of acquisition function values throughout iterations.
- plot_timing_breakdown(timing_data=None, ax=None)[source]
Plot timing breakdown of different phases of the algorithm.
- plot_convergence_diagnostics(convergence_data=None, ax=None)[source]
Plot convergence diagnostics including thresholds and delta evolution.
- plot_kl_divergences(kl_data=None, ax=None, annotate=False)[source]
Plot successive KL divergences between NS iterations (reverse, forward, symmetric).
- plot_parameter_evolution(param_evolution_data=None, max_params=4)[source]
Plot evolution of parameter values during optimization.
- 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:
- Return type:
Figure- Returns:
The matplotlib figure object
- 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:
LikelihoodLikelihood 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”.