Crate nuts_rs 

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 PyTensor. The python wrapper for this sampler is nutpie.

Usage

use nuts_rs::{CpuLogpFunc, CpuMath, LogpError, DiagNutsSettings, Chain, Progress,
Settings, HasDims};
use thiserror::Error;
use rand::rng;
use std::collections::HashMap;

// Define a function that computes the unnormalized posterior density
// and its gradient.
#[derive(Debug)]
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 HasDims for PosteriorDensity {
    fn dim_sizes(&self) -> HashMap<String, u64> {
        vec![("unconstrained_parameter".to_string(), self.dim() as u64)].into_iter().collect()
    }
}

impl CpuLogpFunc for PosteriorDensity {
    type LogpError = PosteriorLogpError;
    type ExpandedVector = Vec<f64>;

    // Only used for transforming adaptation.
    type FlowParameters = ();

    // 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::LogpError> {
        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)
    }

    fn expand_vector<R: rand::Rng + ?Sized>(&mut self, rng: &mut R, position: &[f64]) -> Result<Vec<f64>, nuts_rs::CpuMathError> {
        Ok(position.to_vec())
    }
}

// We get the default sampler arguments
let mut settings = DiagNutsSettings::default();

// and modify as we like
settings.num_tune = 1000;
settings.maxdepth = 3;  // small value just for testing...

// We instanciate our posterior density function
let logp_func = PosteriorDensity {};
let math = CpuMath::new(logp_func);

let chain = 0;
let mut rng = rng();
let mut sampler = settings.new_chain(0, math, &mut rng);

// 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
for _ in 0..2000 {
    let (draw, info) = sampler.draw().expect("Unrecoverable error during sampling");
    trace.push(draw);
}

Users can also implement the Model trait for more control and parallel sampling.

See the examples directory in the repository for more examples.

Implementation details

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

Re-exports

pub use rand;

Structs

AdamOptions

Settings for Adam step size adaptation

ChainProgress

CpuMath

CsvConfig

Configuration for CSV-based MCMC storage.

CsvTraceStorage

Main CSV storage managing multiple chains

DiagAdaptExpSettings

Settings for mass matrix adaptation

DivergenceInfo

Details about a divergence that might have occured during sampling

EuclideanAdaptOptions

FlowSettings

HashMapConfig

LowRankSettings

MclmcChain

Single-chain MCLMC sampler.

MclmcInfo

Diagnostic information returned for a single MCLMC draw.

MclmcSettings

Settings for the unadjusted Microcanonical Langevin Monte Carlo (MCLMC) sampler.

MclmcStats

Per-draw sampler statistics for MclmcChain.

NutsSettings

Settings for the NUTS sampler

Progress

ProgressCallback

Sampler

StepSizeAdaptOptions

Options for step size adaptation

StepSizeSettings

Enums

CpuMathError

DateTimeUnit

HashMapValue

Container for different types of sample values in HashMaps

ItemType

KineticEnergyKind

Selects the kinetic-energy form (and thus the integrator) used by [TransformedHamiltonian].

MclmcTrajectoryKind

Selects which leapfrog integrator and partial momentum refresh are used by MclmcChain.

NutsError

SamplerWaitResult

StepSizeAdaptMethod

Method used for step size adaptation

Value

Traits

Chain

Draw samples from the posterior distribution using Hamiltonian MCMC.

CpuLogpFunc

HasDims

LogpError

Errors that happen when we evaluate the logp and gradient function

Math

Model

Trait for MCMC models with associated math backend and initialization.

SamplerStats

Settings

All sampler configurations implement this trait

Storable

Trait for types whose fields can be progressively written to a trace backend.

Functions

sample_sequentially

Type Aliases

DiagGradNutsSettingsDeprecated

Backwards-compatible alias for DiagNutsSettings.

DiagMclmcSettings

MCLMC settings with a diagonal mass matrix adaptation.

DiagNutsSettings

FlowMclmcSettings

MCLMC settings with a learned flow transformation.

FlowNutsSettings

LowRankMclmcSettings

MCLMC settings with a low-rank mass matrix adaptation.

LowRankNutsSettings

TransformedMclmcSettingsDeprecated

Backwards-compatible alias for FlowMclmcSettings.

TransformedNutsSettingsDeprecated

Backwards-compatible alias for FlowNutsSettings.

TransformedSettingsDeprecated

Backwards-compatible alias for FlowSettings.

Derive Macros

Storable