Mini MCMC

A small (and growing) Rust library for Markov Chain Monte Carlo (MCMC) methods.

Installation

Once published on crates.io, add the following to your Cargo.toml:

[dependencies]
mini-mcmc = "0.2.0"

Then you can use mini_mcmc in your Rust code.

Example: Sampling From a 2D Gaussian

use mini_mcmc::core::ChainRunner;
use mini_mcmc::metropolis_hastings::MetropolisHastings;
use mini_mcmc::distributions::{Gaussian2D, IsotropicGaussian};

fn main() {
        let target = Gaussian2D {
            mean: [0.0, 0.0].into(),
            cov: [[1.0, 0.0], [0.0, 1.0]].into(),
        };
        let proposal = IsotropicGaussian::new(1.0);
        let initial_state = [0.0, 0.0];

        // Create a MH sampler with 4 parallel chains
        let mut mh = MetropolisHastings::new(target, proposal, &initial_state, 4);

        // Run the sampler for 1,000 steps, discarding the first 100 as burn-in
        let samples = mh.run(1000, 100);

        // We should have 900 * 4 = 3600 samples
        assert_eq!(samples.len(), 4);
        assert_eq!(samples[0].nrows(), 900); // samples[0] is a nalgebra::DMatrix
}

You can also find this example at examples/minimal_mh.rs.

Example: Sampling From a Custom Distribution

Below we define a custom Poisson distribution for nonnegative integer states (({0,1,2,\dots})) and a simple random-walk proposal. We then run Metropolis–Hastings to sample from this distribution, collecting frequencies of (k) after some burn-in:

use mini_mcmc::core::ChainRunner;
use mini_mcmc::distributions::{Proposal, Target};
use mini_mcmc::metropolis_hastings::MetropolisHastings;
use rand::Rng; // for thread_rng

/// A Poisson(\lambda) distribution, seen as a discrete target over k=0,1,2,...
#[derive(Clone)]
struct PoissonTarget {
    lambda: f64,
}

impl Target<usize, f64> for PoissonTarget {
    /// unnorm_log_prob(k) = log( p(k) ), ignoring normalizing constants if you wish.
    /// For Poisson(k|lambda) = exp(-lambda) * (lambda^k / k!)
    /// so log p(k) = -lambda + k*ln(lambda) - ln(k!)
    /// which is enough to do MH acceptance.
    fn unnorm_log_prob(&self, theta: &[usize]) -> f64 {
        let k = theta[0];
        let kf = k as f64;
        // If you like, you can omit -ln(k!) if you only need "unnormalized"—but including
        // it can improve acceptance ratio numerically. Here we keep the full log pmf.
        -self.lambda + kf * self.lambda.ln() - ln_factorial(k as u64)
    }
}

/// A simple random-walk proposal in the nonnegative integers:
/// - If current_state=0, propose 0 -> 1 always
/// - Otherwise propose x->x+1 or x->x-1 with p=0.5 each
#[derive(Clone)]
struct NonnegativeProposal;

impl Proposal<usize, f64> for NonnegativeProposal {
    fn sample(&mut self, current: &[usize]) -> Vec<usize> {
        let x = current[0];
        if x == 0 {
            // can't go negative; always move to 1
            vec![1]
        } else {
            // 50% chance to do x+1, 50% x-1
            let flip = rand::thread_rng().gen_bool(0.5);
            let next = if flip { x + 1 } else { x - 1 };
            vec![next]
        }
    }

    /// log_prob(x->y):
    ///  - if x=0 and y=1, p=1 => log p=0
    ///  - if x>0, then y in {x+1, x-1} => p=0.5 => log(0.5)
    ///  - otherwise => -∞ (impossible transition)
    fn log_prob(&self, from: &[usize], to: &[usize]) -> f64 {
        let x = from[0];
        let y = to[0];
        if x == 0 {
            if y == 1 {
                0.0 // ln(1.0)
            } else {
                f64::NEG_INFINITY
            }
        } else {
            // x>0
            if y == x + 1 || y + 1 == x {
                // y in {x+1, x-1} => prob=0.5 => ln(0.5)
                (0.5_f64).ln()
            } else {
                f64::NEG_INFINITY
            }
        }
    }

    fn set_seed(self, _seed: u64) -> Self {
        // no custom seeding logic here
        self
    }
}

// A small helper for computing ln(k!)
fn ln_factorial(k: u64) -> f64 {
    if k < 2 {
        0.0
    } else {
        let mut acc = 0.0;
        for i in 1..=k {
            acc += (i as f64).ln();
        }
        acc
    }
}

fn main() {
    // We'll do Poisson with lambda=4.0, for instance
    let target = PoissonTarget { lambda: 4.0 };

    // We'll have a random-walk in nonnegative integers
    let proposal = NonnegativeProposal;

    // Start the chain at k=0
    let initial_state = [0usize];

    // Create Metropolis–Hastings with 1 chain (or more, up to you)
    let mut mh = MetropolisHastings::new(target, proposal, &initial_state, 1);

    // Run 10_000 steps, discarding first 1_000
    let samples = mh.run(10_000, 1_000);
    let chain0 = &samples[0];
    println!("Chain shape = {} x {}", chain0.nrows(), chain0.ncols());

    // Tally frequencies of each k up to some cutoff
    let cutoff = 20; // enough to see the mass near lambda=4
    let mut counts = vec![0usize; cutoff + 1];
    for row in chain0.row_iter() {
        let k = row[0];
        if k <= cutoff {
            counts[k] += 1;
        }
    }

    let total = chain0.nrows();
    println!("Frequencies for k=0..{cutoff}, from chain after burn-in:");
    for (k, &cnt) in counts.iter().enumerate() {
        let freq = cnt as f64 / total as f64;
        println!("k={k:2}: freq ~ {freq:.3}");
    }

    // We might compare these frequencies to the theoretical Poisson(4.0) pmf
    // in a quick check.
    println!("Done sampling Poisson(4).");
}

You can also find this example at examples/poisson_mh.rs.

Explanation

• PoissonTarget implements Target<usize, f64> for a discrete Poisson((\lambda)) distribution: [ p(k) = e^{{-\lambda}, \frac{\lambda}k}{k!},\quad k=0,1,2,\ldots ] In log form, (\log p(k) = -\lambda + k \log(\lambda) - \log(k!)).

• NonnegativeProposal provides a random-walk in the set ({0,1,2,\dots}):

  • If (x=0), propose (1) with probability 1.
  • If (x>0), propose (x+1) or (x-1) with probability 0.5 each.
  • log_prob returns (\ln(0.5)) for the valid moves, or (-\infty) for invalid moves.

• Usage:
  We start the chain at (k=0), run 10,000 iterations discarding 1,000 as burn-in, and tally the final sample frequencies for (k=0..20). They should approximate the Poisson(4.0) distribution (peak around (k=4)).

With this example, you can see how to use mini_mcmc for unbounded discrete distributions via a custom random-walk proposal and a log‐PMF.

Overview

This library provides:

• Metropolis-Hastings: A generic implementation suitable for various target distributions and proposal mechanisms.
• Distributions: Handy Gaussian and isotropic Gaussian implementations, along with traits for defining custom log-prob functions.

Roadmap

• Parallel Chains: Run multiple Metropolis-Hastings Markov chains in parallel. ✅
• Discrete & Continuous Distributions: Get Metropolis-Hastings to work for continuous and discrete distributions. ✅
• Generic Datatypes: Support sampling vectors of arbitrary integer or floating point types. ✅
• Gibbs Sampling: A component-wise MCMC approach for higher-dimensional problems. ✅
• Hamiltonian Monte Carlo (HMC): A gradient-based method for efficient exploration.
• No-U-Turn Sampler (NUTS): An extension of HMC that removes the need to choose path lengths.
• Ensemble Slice Sampling (ESS): Efficient gradient-free sampler, see paper.

Structure

• src/lib.rs: The main library entry point—exports MCMC functionality.
• src/distributions.rs: Target distributions (e.g., multivariate Gaussians) and proposal distributions.
• src/metropolis_hastings.rs: The Metropolis-Hastings algorithm implementation.
• src/gibbs.rs: The Gibbs sampling algorithm implementation.
• examples/demo.rs: Example usage demonstrating 2D Gaussian sampling and plotting.

Usage (Local)

1. Build (Library + Demo):

   cargo build --release
   

2. Run the Demo:

   cargo run --release --bin gauss_mh
   

   Prints basic statistics of the MCMC chain (e.g., estimated mean). Saves a scatter plot of sampled points in scatter_plot.png and a Parquet file samples.parquet.

Optional Features

• csv: Enables CSV I/O for samples.
• arrow / parquet: Enables Apache Arrow / Parquet I/O.
• By default, all features are enabled. You can disable them if you want a smaller build.

License

Licensed under the Apache License, Version 2.0. See LICENSE for details.
This project includes code from the kolmogorov_smirnov project, licensed under Apache 2.0 as noted in NOTICE.