#statistics #bayes

nuts-rs

Sample from unnormalized densities using Hamiltonian MCMC

10 releases (6 breaking)

0.8.0 Sep 20, 2023
0.6.0 Jul 21, 2023
0.4.1 Jan 2, 2023
0.4.0 Nov 12, 2022
0.2.1 Jul 20, 2022

#1 in #nuts

Download history 108/week @ 2023-11-03 17/week @ 2023-11-10 29/week @ 2023-11-17 45/week @ 2023-11-24 54/week @ 2023-12-01 41/week @ 2023-12-08 64/week @ 2023-12-15 71/week @ 2023-12-22 117/week @ 2023-12-29 119/week @ 2024-01-05 47/week @ 2024-01-12 52/week @ 2024-01-19 38/week @ 2024-01-26 65/week @ 2024-02-02 46/week @ 2024-02-09 553/week @ 2024-02-16

709 downloads per month

MIT license

265KB
3K SLoC

Workflow Status dependency status

Sample from posterior distributions using the No U-turn Sampler (NUTS). For details see the original NUTS paper and the more recent introduction.

This crate was developed as a faster replacement of the sampler in PyMC, to be used with the new numba backend of aesara. The python wrapper for this sampler is nutpie.

Usage

use nuts_rs::{CpuLogpFunc, LogpError, new_sampler, SamplerArgs, Chain, SampleStats};
use thiserror::Error;

// Define a function that computes the unnormalized posterior density
// and its gradient.
struct PosteriorDensity {}

// The density might fail in a recoverable or non-recoverable manner...
#[derive(Debug, Error)]
enum PosteriorLogpError {}
impl LogpError for PosteriorLogpError {
    fn is_recoverable(&self) -> bool { false }
}

impl CpuLogpFunc for PosteriorDensity {
    type Err = PosteriorLogpError;

    // We define a 10 dimensional normal distribution
    fn dim(&self) -> usize { 10 }

    // The normal likelihood with mean 3 and its gradient.
    fn logp(&mut self, position: &[f64], grad: &mut [f64]) -> Result<f64, Self::Err> {
        let mu = 3f64;
        let logp = position
            .iter()
            .copied()
            .zip(grad.iter_mut())
            .map(|(x, grad)| {
                let diff = x - mu;
                *grad = -diff;
                -diff * diff / 2f64
            })
            .sum();
        return Ok(logp)
    }
}

// We get the default sampler arguments
let mut sampler_args = SamplerArgs::default();

// and modify as we like
sampler_args.step_size_adapt.target = 0.8;
sampler_args.num_tune = 1000;
sampler_args.maxdepth = 3;  // small value just for testing...
sampler_args.mass_matrix_adapt.store_mass_matrix = true;

// We instanciate our posterior density function
let logp_func = PosteriorDensity {};

let chain = 0;
let seed = 42;
let mut sampler = new_sampler(logp_func, sampler_args, chain, seed);

// Set to some initial position and start drawing samples.
sampler.set_position(&vec![0f64; 10]).expect("Unrecoverable error during init");
let mut trace = vec![];  // Collection of all draws
let mut stats = vec![];  // Collection of statistics like the acceptance rate for each draw
for _ in 0..2000 {
    let (draw, info) = sampler.draw().expect("Unrecoverable error during sampling");
    trace.push(draw);
    let _info_vec = info.to_vec();  // We can collect the stats in a Vec
    // Or get more detailed information about divergences
    if let Some(div_info) = info.divergence_info() {
        println!("Divergence at position {:?}", div_info.start_location());
    }
    dbg!(&info);
    stats.push(info);
}

Sampling several chains in parallel so that samples are accessable as they are generated is implemented in sample_parallel.

Implementation details

This crate mostly follows the implementation of NUTS in Stan and PyMC, only tuning of mass matrix and step size differs somewhat.

Dependencies

~4–5.5MB
~116K SLoC