Mini MCMC

A compact Rust library for Markov Chain Monte Carlo (MCMC) methods with GPU support.

Installation

Add the following to your Cargo.toml:

[dependencies]
mini-mcmc = "0.4.2"

Then use mini_mcmc in your Rust code.

Example: Sampling From a 2D Gaussian

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

fn main() {
    let target = Gaussian2D {
        mean: arr1(&[0.0, 0.0]),
        cov: arr2(&[[1.0, 0.0], [0.0, 1.0]]),
    };
    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,100 steps, discarding the first 100 as burn-in
    let samples = mh.run(1000, 100).unwrap();

    // We should have 1000 * 4 = 3600 samples
    assert_eq!(samples.shape()[0], 4);
    assert_eq!(samples.shape()[1], 1000);
}

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 basic 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);

    // Collect 10,000 samples and use 1,000 for burn-in (not returned).
    let samples = mh
        .run(10_000, 1_000)
        .expect("Expected generating samples to succeed");
    let chain0 = samples.to_shape(10_000).unwrap();
    println!("Elements in chain: {}", chain0.len());

    // 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.iter() {
        let k = *row;
        if k <= cutoff {
            counts[k] += 1;
        }
    }

    let total = chain0.len();
    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$$
  The log form of it is $\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-14$ with probability $0.5$ each.
  • log_prob returns $\ln(0.5)$ for the possible moves, or $-\infty$ for impossible moves.

• Usage:
  We start the chain at $k=0$, run 11,000 iterations discarding 1,000 as burn-in, and tally the final sample frequencies for $k=0 \dots 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.

Below is an additional documentation section that you can add to your README. It first gives a minimal version of the rosenbrock3d_hmc.rs example for sampling using HMC. (Note that the full example also plots the sampled data interactively using Plotly.)

Example: Sampling from a 3D Rosenbrock Distribution Using HMC

The following minimal example demonstrates how to create and run an HMC sampler to sample from a 3D Rosenbrock distribution. In this example, we construct an HMC sampler, run it for a fixed number of iterations, and print the shape of the collected samples. The corresponding file can also be found at examples/minimal_hmc.rs. For a complete example—including interactive 3D plotting with Plotly, refer to examples/rosenbrock3d_hmc.rs.

use burn::tensor::Element;
use burn::{backend::Autodiff, prelude::Tensor};
use mini_mcmc::hmc::{GradientTarget, HMC};
use num_traits::Float;

/// The 3D Rosenbrock distribution.
///
/// For a point x = (x₁, x₂, x₃), the log probability is defined as the negative of
/// the sum of two Rosenbrock terms:
///
///   f(x) = 100*(x₂ - x₁²)² + (1 - x₁)² + 100*(x₃ - x₂²)² + (1 - x₂)²
///
/// This implementation generalizes to d dimensions, but here we use it for 3D.
struct RosenbrockND {}

impl<T, B> GradientTarget<T, B> for RosenbrockND
where
    T: Float + std::fmt::Debug + Element,
    B: burn::tensor::backend::AutodiffBackend,
{
    fn log_prob_batch(&self, positions: &Tensor<B, 2>) -> Tensor<B, 1> {
        // Assume positions has shape [n_chains, d] with d = 3.
        let k = positions.dims()[0] as i64;
        let n = positions.dims()[1] as i64;
        let low = positions.clone().slice([(0, k), (0, n - 1)]);
        let high = positions.clone().slice([(0, k), (1, n)]);
        let term_1 = (high - low.clone().powi_scalar(2))
            .powi_scalar(2)
            .mul_scalar(100);
        let term_2 = low.neg().add_scalar(1).powi_scalar(2);
        -(term_1 + term_2).sum_dim(1).squeeze(1)
    }
}

fn main() {
    // Use the CPU backend wrapped in Autodiff (e.g., NdArray).
    type BackendType = Autodiff<burn::backend::NdArray>;

    // Create the 3D Rosenbrock target.
    let target = RosenbrockND {};

    // Define initial positions for 6 chains (each a 3D point).
    let initial_positions = vec![vec![1.0_f32, 2.0_f32, 3.0_f32]; 6];

    // Create the HMC sampler with a step size of 0.01 and 50 leapfrog steps.
    let mut sampler = HMC::<f32, BackendType, RosenbrockND>::new(
        target,
        initial_positions,
        0.032,
        50,
    );

    // Run the sampler for 1100 iterations, discard 100
    let samples = sampler.run(1000, 100);

    // Print the shape of the collected samples.
    println!("Collected samples with shape: {:?}", samples.dims());
}

Overview

This library provides implementations of

• Hamiltonian Monte Carlo (HMC): an MCMC method that efficiently samples by simulating Hamiltonian dynamics using gradients of the target distribution.
• Metropolis-Hastings: an MCMC algorithm that samples from a distribution by proposing candidates and probabilistically accepting or rejecting them.
• Gibbs Sampling: an MCMC method that iteratively samples each variable from its conditional distribution given all other variables.

with

• Implementations of Common Distributions: handy Gaussian and isotropic Gaussian implementations, along with traits for defining custom log-prob functions.
• Parallelization: run multiple chains in parallel.
• Progress Bars: show progress of MCMC algorithms with convergence statistics and acceptance rates.
• Support for Discrete & Continuous Distributions: Metropolis-Hastings and Gibbs sampling support continuous and discrete distributions.
• Generic Datatypes: Support sampling vectors of various integer or floating point types.

Roadmap

• No-U-Turn Sampler (NUTS): An extension of HMC that removes the need to choose path lengths.
• Rank Normalized Rhat: Modern convergence diagnostic, see paper.
• Ensemble Slice Sampling (ESS): Efficient gradient-free sampler, see paper.
• Effective Size Estimation: Online estimation of effective sample size for early stopping.

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/: Examples on how to use this library.
• src/io/arrow.rs: Helper functions for saving samples as Apache Arrow files. Enable via arrow feature.
• src/io/parquet.rs: Helper functions for saving samples as Apache Parquet files. Enable via parquet feature.
• src/io/csv.rs: Helper functions for saving samples as Apache Parquet files. Enable via csv feature.

Usage (Local)

1. Build (Library + Demo):

   cargo build --release
   

2. Run the Demo:

   cargo run --release --example gauss_mh --features parquet
   

   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.
• wgpu: Enables sampling from gradient based samplers using burn's WGPU backend.
• By default, all features are disabled.

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.