Minimal demo of using BayesKit + PosteriorDB + BridgeStan, plus a Dockerfile for a container to run it in

demo_hmc.py

import numpy as np
import tempfile
import json
import posteriordb
from tqdm.auto import trange
from bayes_kit import HMCDiag
from bayes_kit.model_types import HessianModel
from bridgestan.model import StanModel


class BayesKitCompatibleStanModel(StanModel, HessianModel):
    """A bridgestan.StanModel that implements the bayes_kit.HessianModel interface.

    Note that the functions log_density, log_density_gradient, and log_density_hessian are already implemented in
    StanModel, so we only need to implement the dims function.
    """

    def dims(self):
        return self.param_unc_num()


# Pick a problem from posteriordb; note location of the DB is hardcoded in the Dockerfile
pdb = posteriordb.PosteriorDatabase("/posteriordb/posterior_database/")
print("Available PDB models are", "\n\t".join(sorted(pdb.posterior_names())))
posterior = pdb.posterior("pilots-pilots")


# Model creation fails using posterior.data.file_path() directly because bridgestan explicitly looks for the ".json"
# file extension. So data_to_json here is basically a copy of posterior.data.file_path() with the ".json" suffix added.
def data_to_json(posterior):
    data_file = tempfile.NamedTemporaryFile(suffix=".json", delete=False)
    with open(data_file.name, "w") as f:
        json.dump(posterior.data.values(), f)
    return data_file.name

# Instantiate a Stan model from the problem
print("Creating the model")
model = BayesKitCompatibleStanModel.from_stan_file(posterior.model.stan_code_file_path(), data_to_json(posterior))

# Run HMC
print("Running HMC")
hmc = HMCDiag(model=model, stepsize=1e-2, steps=100)
thetas, logps = [], []
for _ in trange(1000, desc="Samples", leave=False):
    theta, logp = hmc.sample()
    thetas.append(theta)
    logps.append(logp)
samples = np.array(thetas)
logps = np.array(logps)
print("done")

Dockerfile

# Dockerfile for container with Python, Stan, Bridgestan, BayesKit, and PosteriorDB
FROM python:3.10.11-bullseye

# Install git
RUN apt-get install -y git

# Install BayesKit
RUN pip install --upgrade pip \
    && pip install git+https://github.com/flatironinstitute/bayes-kit.git

# Install posteriordb and its python interface
RUN git clone https://github.com/stan-dev/posteriordb /posteriordb
RUN pip install posteriordb

# Download Bridgestan and `make` an example model to trigger stan tools compilation
# See https://roualdes.github.io/bridgestan/latest/getting-started.html#testing-the-installation
RUN git clone --recurse-submodules https://github.com/roualdes/bridgestan.git \
    && cd /bridgestan \
    && make test_models/multi/multi_model.so

# Install python interface to bridgestan
RUN cd /bridgestan && pip install -e python

# Set BRIDGESTAN env variable so it knows where to look for compilation tools etc
ENV BRIDGESTAN=/bridgestan

# Add all other packages from requirements.txt
COPY requirements.txt /requirements.txt
RUN pip install -r /requirements.txt