BOBE 0.1.0

A Bayesian Optimization package for cosmology using JAX

BOBE

BOBE (Bayesian Optimisation for Bayesian Evidence) is a high-performance package for performing Bayesian model comparison with expensive likelihood functions, developed for applications to cosmology. It uses Bayesian Optimization to train a Gaussian process surrogate for the expensive likelihood function and runs Nested sampling/MCMC on the surrogate instead of the underlying likelihood. Training the surrogate typically requires around ~100x fewer true likelihood evaluations compared to running Nested sampling/MCMC on the true likelihood, leading to significant speed-ups for slow likelihoods (t>1s). BOBE uses acquisition functions that minimise the integrated uncertainty of the surrogate, prioritising regions that matter the most for the evidence. The algorithm is explained in more detail in our paper and code documentation is available at https://BOBE.readthedocs.io.

Use BOBE if:

• Your likelihood function is expensive to evaluate (t~1s or more)
• You need Bayesian evidence estimates and/or posterior samples efficiently

BOBE works best for problems with up to 15-20 parameters, although this can vary based on the specific problem and likelihood structure. It has been tested to work well upto 30 dimensions for simple multivariate Gaussian likelihoods and upto 16 dimensions for cosmological likelihoods (LCDM + curvature) with the Planck Camspec likelihood. BOBE may not be necessary if your likelihood already evaluates in milliseconds, as the overhead of training the GP surrogate and running Bayesian optimisation may not be worth the speedup. For such cases, there is generally no need to go beyond standard MCMC/nested sampling tools.

BOBE is under active development and welcome bug reports, patches, feature requests, and other comments via the GitHub issue tracker.

Key Features

• Efficient Bayesian Evidence Estimation: Train GP surrogates for expensive likelihoods and compute evidence and parameter posteriors
• JAX-Powered Performance: Leverages JAX for automatic differentiation, jit compilation (and in the future GPU/TPU acceleration...)
• MPI Parallelization: Distribute likelihood evaluations and GP fitting across multiple processes
• Flexible Likelihood Interface: Built-in support for cosmological likelihoods through Cobaya or your own custom functions

Requirements

• Python >=3.10 and <3.14
• JAX
• NumPyro
• scipy
• scikit-learn

Optional:

• cobaya (for cosmological likelihoods)
• mpi4py (for MPI parallelization)

See pyproject.toml and the documentation for full list of dependencies.

Installation

From PyPI

pip install BOBE

From source

git clone https://github.com/Ameek94/BOBE.git
cd BOBE
python -m pip install .

For an editable (dev) install do

python -m pip install -e .

from the package directory.

Optional Dependencies

BOBE has several optional dependencies for extended functionality. When installing from source, you can install them as follows:

• Cosmology Suite: Install with pip install -e '.[cosmo]'

  • This will install Cobaya, needed for several cosmological likelihoods. You will still need to download and install the data for some likelihoods - see the Cobaya documentation for more details.
  • It will also install mpi4py (you will need to have an MPI implementation such as openmpi or mpich). Enables parallel likelihood evaluation and GP fitting across multiple processes using mpi4py.

• Cobaya Only: Install with pip install -e '.[cobaya]'

• MPI Only: Install with pip install -e '.[mpi]'

• All Optional Dependencies: Install with pip install -e '.[all]'

Note: The -e flag installs in editable mode. For a regular install, use pip install '.[extra]' instead (replace extra with name of dependency).

Quick Start

from BOBE import BOBE

# Define your likelihood function
def my_likelihood(X):
    x, y = X[0], X[1]
    logpdf = -0.25 * (5 * (0.2 - x))**2 - (20 * (y/4 - x**4))**2 # Example: a likelihood with a curved degeneracy
    return logpdf 


# Initialize BOBE with setup parameters
sampler = BOBE(
    loglikelihood=my_likelihood,
    param_list=['x1', 'x2'], # list of parameter names
    param_bounds=np.array([[-1, 1], [-1, 2]]).T, # lower and upper bounds for parameters (2, ndim) shaped
    n_sobol_init=2, # number of initial Sobol samples to start the run from
    save_dir='./results',
)

# Run optimization with convergence settings.
# BOBE also has other run settings which you can specify, otherwise dimension-based defaults will be used. See the docs for more details.
results = sampler.run(
    logz_threshold=0.1, # target log-evidence uncertainty for convergence
)


# Access the evidence and posterior samples
print(f"Log Evidence: {results['logz']['mean']}")
samples = results['samples'] # dictionary containing keys 'x', 'logl', 'weights'

Cosmology Example with Cobaya

For cosmological likelihoods you will need to have Cobaya installed. Then, simply pass the Cobaya YAML file path in your script to the BOBE sampler.

from BOBE import BOBE

# Initialize BOBE with Cobaya YAML file, prepare it the same way you would for running MCMC/Nested Sampling through Cobaya
   sampler = BOBE(
       loglikelihood='path/to/cobaya_input.yaml',
       n_sobol_init=4,
       n_cobaya_init=4,  # We can also specify reference dists in the Cobaya yaml file to generate additional initial points
       likelihood_name='quickstart_cobaya_example', # name for output files
       save_dir='./results',
       use_clf=True # recommended to enable classifiers for cosmological examples where likelihood can sometimes return -inf values
   )
   
   # Run with optimization settings
   results = sampler.run(
       logz_threshold=0.1,
   )

# rest of the run remains the same as above

Full documentation is available at https://BOBE.readthedocs.io. The examples/ folder also contains several scripts on how to run the code with different likelihoods, including cosmological likelihoods interfaced through the Cobaya package or your own custom likelihoods. These scripts have been used to produce the results in our paper. Run them using

python your_chosen_example.py

MPI Parallelization

For expensive likelihoods (evaluation time > 1 second), you can also use MPI to parallelize likelihood evaluations at points proposed by the batched acquisition function. Make sure you have mpi4py installed.

mpirun -n 4 python your_chosen_example.py

where -n 4 specifies the number of MPI processes. In MPI mode, the code distributes the computation of the likelihood function at several candidate points across different MPI processes, significantly reducing wall-clock time for expensive likelihoods. It also distributes GP fitting by running multiple restarts across the different MPI processes. Note than on SLURM clusters you might need to substitute mpirun with srun.

Citation

If you use BOBE in your research, please cite:

@article{Cohen:2026iij,
    author = "Cohen, Nathan and Hamann, Jan and Malhotra, Ameek",
    title = "{Bayesian optimisation for Bayesian evidence (BOBE) -- a fast and efficient likelihood emulator for model selection}",
    eprint = "2601.11150",
    archivePrefix = "arXiv",
    primaryClass = "astro-ph.CO",
    month = "1",
    year = "2026"
}

Troubleshooting

In case you run into installation issues related to incompatible versions of packages, you can also set up an environment with the exact package versions BOBE was developed and tested with

# Create environment from the minimal essential packages
conda env create -f environment.yml

# Activate the environment
conda activate BOBE

# Install BOBE
pip install .

Support

• Documentation: https://BOBE.readthedocs.io
• Issues: https://github.com/Ameek94/BOBE/issues
• Repository: https://github.com/Ameek94/BOBE

License

BOBE is released under the MIT License.