1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

#![cfg_attr(feature = "simd_support", feature(stdsimd))]
#![cfg_attr(feature = "simd_support", feature(portable_simd))]
#![cfg_attr(feature = "simd_support", feature(slice_as_chunks))]
//! Sample from posterior distributions using the No U-turn Sampler (NUTS).
//! For details see the original [NUTS paper](https://arxiv.org/abs/1111.4246)
//! and the more recent [introduction](https://arxiv.org/abs/1701.02434).
//!
//! 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](https://github.com/pymc-devs/nuts-py).
//!
//! ## Usage
//!
//! ```
//! use nuts_rs::{CpuLogpFunc, LogpError, new_sampler, SamplerArgs, Chain, SampleStats};
//! use thiserror::Error;
//! use rand::thread_rng;
//!
//! // 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.num_tune = 1000;
//! sampler_args.maxdepth = 3;  // small value just for testing...
//!
//! // We instanciate our posterior density function
//! let logp_func = PosteriorDensity {};
//!
//! let chain = 0;
//! let mut rng = thread_rng();
//! let mut sampler = new_sampler(logp_func, sampler_args, chain, &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);
//!     // 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);
//! }
//! ```
//!
//! 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](https://mc-stan.org) and
//! [PyMC](https://docs.pymc.io/en/v3/), 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`.

pub(crate) mod adapt_strategy;
pub(crate) mod cpu_potential;
pub(crate) mod cpu_sampler;
pub(crate) mod cpu_state;
pub(crate) mod mass_matrix;
pub mod math;
pub(crate) mod nuts;
pub(crate) mod stepsize;

pub use adapt_strategy::DualAverageSettings;
pub use cpu_potential::CpuLogpFunc;
pub use cpu_sampler::test_logps;
pub use cpu_sampler::{
    new_sampler, sample_parallel, sample_sequentially, CpuLogpFuncMaker, InitPointFunc,
    JitterInitFunc, ParallelChainResult, ParallelSamplingError, SamplerArgs,
};
#[cfg(feature = "arrow")]
pub use nuts::{ArrowBuilder, ArrowRow};
pub use nuts::{Chain, DivergenceInfo, LogpError, NutsError, SampleStats};