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 work-in-progress python wrapper for this sampler is nuts-py.

Usage

```rust use nutsrs::{CpuLogpFunc, LogpError, newsampler, 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 samplerargs.stepsizeadapt.target = 0.8; samplerargs.numtune = 1000; samplerargs.maxdepth = 3; // small value just for testing... samplerargs.massmatrixadapt.storemass_matrix = true;

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

let chain = 0; let seed = 42; let mut sampler = newsampler(logpfunc, sampler_args, chain, seed);

// Set to some initial position and start drawing samples. sampler.setposition(&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 _infovec = info.tovec(); // We can collect the stats in a Vec // Or get more detailed information about divergences if let Some(divinfo) = info.divergenceinfo() { println!("Divergence at position {:?}", divinfo.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: In a first window we sample using the identity as mass matrix and adapt the step size using the normal dual averaging algorithm. After discard_window draws we start computing a diagonal mass matrix using an exponentially decaying estimate for sqrt(sample_var / grad_var). After 2 * discard_window draws we switch to the entimated mass mass_matrix and keep adapting it live until stop_tune_at.