mod nuts

module nuts

NUTS Phase 3a: No-U-Turn Sampler trajectory builder.

Mirrors Stan’s recursive doubling tree from stan/mcmc/hmc/nuts/base_nuts.hpp:122-265. Each NUTS step draws fresh momentum, builds a random-doubling binary tree of leapfrog leaves, stops on a U-turn or divergence, and multinomial-samples a candidate from the visited leaves with the standard HMC-Metropolis probability.

Gaussian and q-Gaussian momentum both use the Momentum trait’s U-turn predicate, with divergence guarded by the Hamiltonian threshold.

Functions

fn nuts_step<G, M, Obj, R: Rng>(x: &Array1<f64>, u: f64, eps: f64, temp: f64, grad_fn: &G, momentum: &M, obj_fn: &Obj, max_depth: u32, rng: &mut R) -> NutsTransition
where
    G: Gradient<f64>,
    M: Momentum + ?Sized,
    Obj: Fn(&Array1<f64>) -> f64

One NUTS step. Builds the binary tree by doubling up to max_depth, then accepts the multinomial-sampled candidate.

Structs and Unions

struct NutsSaSampler<O, G, C, M>
where
    O: Objective<f64> + Send + Sync,
    G: Gradient<f64>,
    C: Cooling<f64>,
    M: Momentum

NUTS-driven SA sampler. Drops into run_rs and MultiChainSampler like HmcSaSampler since it impls Sampler<f64>.

obj: O

The objective.

gradient: G

The gradient (analytic or finite-difference).

cool: C

The cooling schedule.

momentum: M

The momentum kernel.

epsilon: f64

Base leapfrog step size; rescaled by sqrt(temp/temp_ref) per step.

temp_ref: f64

Reference temperature for the cooling rescaling.

max_depth: u32

Maximum doubling depth (max_n_leapfrog = 2^max_depth).

Implementations

impl<O, G, C, M> NutsSaSampler<O, G, C, M>
where
    O: Objective<f64> + Send + Sync,
    G: Gradient<f64>,
    C: Cooling<f64>,
    M: Momentum

Functions

fn new(obj: O, gradient: G, cool: C, momentum: M, epsilon: f64, temp_ref: f64, max_depth: u32) -> Self

Constructs a NUTS-SA sampler. temp_ref should typically equal cool.temperature(0).

Traits implemented

impl<O, G, C, M> Sampler<f64> for NutsSaSampler<O, G, C, M>
where
    O: Objective<f64> + Send + Sync,
    G: Gradient<f64>,
    C: Cooling<f64>,
    M: Momentum

struct NutsTransition

Diagnostic for a single NUTS step.

x: Array1<f64>

New position (the multinomial-sampled candidate).

n_leapfrog: usize

Number of leapfrog leaves visited (= 2^depth).

tree_depth: u32

Tree depth reached (0..max_depth).

diverged: bool

true if any leaf diverged (|delta_H| > max_delta_h).

accepted: bool

true if the candidate was accepted (always true in vanilla NUTS: candidate selection is itself the Metropolis step).