Adaptive Gaussian random-walk Metropolis with Mici

adaptive_rwm.py

import numpy as np
import mici

def get_rwm_sampler_and_adapters(
    neg_log_posterior_density, rng, target_accept_stat=0.234
):
    """Get a Mici sampler and adapters for adaptive Gaussian random-walk Metropolis proposals."""
    system = mici.systems.EuclideanMetricSystem(
        neg_log_dens=neg_log_posterior_density,
        grad_neg_log_dens=lambda q: q * 0,
    )
    integrator = mici.integrators.LeapfrogIntegrator(system)
    sampler = mici.samplers.StaticMetropolisHMC(system, integrator, rng, n_step=1)
    adapters = [
        mici.adapters.DualAveragingStepSizeAdapter(target_accept_stat),
        mici.adapters.OnlineCovarianceMetricAdapter()
    ]
    return sampler, adapters

if __name__ == "__main__":

    seed = 1234
    n_chain = 4
    n_dim = 2
    n_warm_up_iter = 500
    n_main_iter = 1000
    rng = np.random.default_rng(seed)

    def neg_log_posterior_density(q):
        """Independent standard normal distribution on all dimensions."""
        return q @ q / 2

    sampler, adapters = get_rwm_sampler_and_adapters(neg_log_posterior_density, rng)
    init_states = rng.standard_normal((n_chain, n_dim))
    final_states, traces, stats = sampler.sample_chains(
        n_warm_up_iter, n_main_iter, init_states, adapters=adapters, n_process=n_chain
    )