COMPASS: Comparison Of Models using Probabilistic Assessment in Simulation-based Settings

COMPASS is a Python package designed for Bayesian Model Comparison in simulation-based settings. By comparing the predictive power of various models, it aims to identify the most suitable model for a given dataset.
It is especially suited for fields like astrophysics and computational biology, where simulation is integral to the modeling process.

Features

• Perform Bayesian model comparison in simulation-based settings with ModelTransfuser
• Perform Simulation-Based Inference with ScoreBasedInferenceModel

Installation

Install the package using pip:

pip install bayes-compass

Usage

There are two examples provided in the tutorials folder, demonstrating the model comparison workflow and the parameter inference with compass.

Model Comparison Example

The ModelTransfuser class provides a framework for the model comparison workflow.
It uses the ScoreBasedInferenceModel class to perform the simulation-based inference, predicts the posterior distribution and samples from the Likelihood function with the inferred parameters.
With a Gaussian Kernel Density Estimator, it evaluates the observed data at the Likelihood function and computes the posterior model probabilities.

from compass import ModelTransfuser 

# Initialize the ModelTransfuser
mtf = ModelTransfuser()

# Add data from simulators
mtf.add_data(model_name="Model1", train_data=data_1, val_data=val_data_1)
mtf.add_data(model_name="Model2", train_data=data_2, val_data=val_data_2)

# Initialize ScoreBasedInferenceModels
mtf.init_models()

# Train the models
mtf.train_models()

# Compare Posterior Model Probabilities
observations = load_your_observations
mtf.compare(x=observations, err=observations_err)

stats = mtf.stats

# Plot results
mtf.plot_comparison()
mtf.plot_attention()

mtf.stats is a dictionary containing all inferred parameters and their uncertainties, the posterior model probabilities, and the log-likelihood values for each model.

Simulation-Based Inference Model

The ScoreBasedInferenceModel is a Diffusion Model with a score predicting Transformer network.
It is able to sample from the posterior and likelihood function by utalizing the attention mechanism of the Transformer architecture.
The model is also able to handle observational uncertainties and is designed to run on all available GPUs.

from compass import ScoreBasedInferenceModel as SBIm

sbimodel = SBIm(nodes_size, sigma, depth, hidden_size, num_heads, mlp_ratio)

sbimodel.train(train_data, val_data=val_data, path=path, device="cuda")

nodes_size - The number of parameters $\theta$ and the number of data points $x$.
sigma - The noise level of the initial noise distribution of the diffusion model.
depth - The number of layers in the Transformer
hidden_size - The size of the embedding in the Transformer
num_heads - The number of attention heads in the Transformer
mlp_ratio - The ratio of the hidden size to the number of Nodes in the MLPs in the Transformer

To sample from the posterior distribution, you can use the sample() function and provide the observations x and the observational uncertainties err (optional):

# Sample from the posterior distribution
posterior_samples = sbimodel.sample(x=observations, err=observations_err, timesteps=100)

To sample from the Likelihood function, you can again use the sample() function and provide the Maximum-A-Posteriori $\hat\theta$ values in theta and the standard deviation err (optional):

# Samlpe from the likelihood function
likelihood_samples = sbimodel.sample(theta=theta_hat, err=std_theta_hat, timesteps=100)

Contributing

Contributions are welcome! Feel free to open issues or submit pull requests to improve this package.