hammer-and-sample

A simplistic MCMC sampler implementing the affine-invariant ensemble sampling of emcee with serial execution and optionally with parallel execution based on Rayon.

The implementation is relatively efficient, for example computing 1000 iterations of 100 walkers using the hierarchical model from hierarchical.rs takes approximately 1 min using emcee and multiprocessing versus 50 ms using this crate and Rayon, running on 8 hardware threads in both cases.

lib.rs:

Simplistic MCMC ensemble sampler based on emcee, the MCMC hammer

use hammer_and_sample::{sample, MinChainLen, Model, Serial};
use rand::{Rng, SeedableRng};
use rand_pcg::Pcg64;

fn estimate_bias(coin_flips: &[bool]) -> f64 {
    struct CoinFlips<'a>(&'a [bool]);

    impl Model for CoinFlips<'_> {
        type Params = [f64; 1];

        // likelihood of Bernoulli distribution and uninformative prior
        fn log_prob(&self, &[p]: &Self::Params) -> f64 {
            if p < 0. || p > 1. {
                return f64::NEG_INFINITY;
            }

            let ln_p = p.ln();
            let ln_1_p = (1. - p).ln();

            self.0
                .iter()
                .map(|coin_flip| if *coin_flip { ln_p } else { ln_1_p })
                .sum()
        }
    }

    let model = CoinFlips(coin_flips);

    let walkers = (0..10).map(|seed| {
        let mut rng = Pcg64::seed_from_u64(seed);

        let p = rng.gen_range(0.0..=1.0);

        ([p], rng)
    });

    let (chain, _accepted) = sample(&model, walkers, MinChainLen(10 * 1000), Serial);

    // 100 iterations of 10 walkers as burn-in
    let chain = &chain[10 * 100..];

    chain.iter().map(|&[p]| p).sum::<f64>() / chain.len() as f64
}