nuts_rs/
lib.rs

1//! Sample from posterior distributions using the No U-turn Sampler (NUTS).
2//! For details see the original [NUTS paper](https://arxiv.org/abs/1111.4246)
3//! and the more recent [introduction](https://arxiv.org/abs/1701.02434).
4//!
5//! This crate was developed as a faster replacement of the sampler in PyMC,
6//! to be used with the new numba backend of PyTensor. The python wrapper
7//! for this sampler is [nutpie](https://github.com/pymc-devs/nutpie).
8//!
9//! ## Usage
10//!
11//! ```
12//! use nuts_rs::{CpuLogpFunc, CpuMath, LogpError, DiagGradNutsSettings, Chain, Progress,
13//! Settings, HasDims};
14//! use thiserror::Error;
15//! use rand::thread_rng;
16//! use std::collections::HashMap;
17//!
18//! // Define a function that computes the unnormalized posterior density
19//! // and its gradient.
20//! #[derive(Debug)]
21//! struct PosteriorDensity {}
22//!
23//! // The density might fail in a recoverable or non-recoverable manner...
24//! #[derive(Debug, Error)]
25//! enum PosteriorLogpError {}
26//! impl LogpError for PosteriorLogpError {
27//!     fn is_recoverable(&self) -> bool { false }
28//! }
29//!
30//! impl HasDims for PosteriorDensity {
31//!     fn dim_sizes(&self) -> HashMap<String, u64> {
32//!         vec![("unconstrained_parameter".to_string(), self.dim() as u64)].into_iter().collect()
33//!     }
34//! }
35//!
36//! impl CpuLogpFunc for PosteriorDensity {
37//!     type LogpError = PosteriorLogpError;
38//!     type ExpandedVector = Vec<f64>;
39//!
40//!     // Only used for transforming adaptation.
41//!     type FlowParameters = ();
42//!
43//!     // We define a 10 dimensional normal distribution
44//!     fn dim(&self) -> usize { 10 }
45//!
46//!     // The normal likelihood with mean 3 and its gradient.
47//!     fn logp(&mut self, position: &[f64], grad: &mut [f64]) -> Result<f64, Self::LogpError> {
48//!         let mu = 3f64;
49//!         let logp = position
50//!             .iter()
51//!             .copied()
52//!             .zip(grad.iter_mut())
53//!             .map(|(x, grad)| {
54//!                 let diff = x - mu;
55//!                 *grad = -diff;
56//!                 -diff * diff / 2f64
57//!             })
58//!             .sum();
59//!         return Ok(logp)
60//!     }
61//!
62//!     fn expand_vector<R: rand::Rng + ?Sized>(&mut self, rng: &mut R, position: &[f64]) -> Result<Vec<f64>, nuts_rs::CpuMathError> {
63//!         Ok(position.to_vec())
64//!     }
65//! }
66//!
67//! // We get the default sampler arguments
68//! let mut settings = DiagGradNutsSettings::default();
69//!
70//! // and modify as we like
71//! settings.num_tune = 1000;
72//! settings.maxdepth = 3;  // small value just for testing...
73//!
74//! // We instanciate our posterior density function
75//! let logp_func = PosteriorDensity {};
76//! let math = CpuMath::new(logp_func);
77//!
78//! let chain = 0;
79//! let mut rng = thread_rng();
80//! let mut sampler = settings.new_chain(0, math, &mut rng);
81//!
82//! // Set to some initial position and start drawing samples.
83//! sampler.set_position(&vec![0f64; 10]).expect("Unrecoverable error during init");
84//! let mut trace = vec![];  // Collection of all draws
85//! for _ in 0..2000 {
86//!     let (draw, info) = sampler.draw().expect("Unrecoverable error during sampling");
87//!     trace.push(draw);
88//! }
89//! ```
90//!
91//! Users can also implement the `Model` trait for more control and parallel sampling.
92//!
93//! See the examples directory in the repository for more examples.
94//!
95//! ## Implementation details
96//!
97//! This crate mostly follows the implementation of NUTS in [Stan](https://mc-stan.org) and
98//! [PyMC](https://docs.pymc.io/en/v3/), only tuning of mass matrix and step size differs
99//! somewhat.
100
101mod adapt_strategy;
102mod chain;
103mod cpu_math;
104mod euclidean_hamiltonian;
105mod hamiltonian;
106mod mass_matrix;
107mod math;
108mod math_base;
109mod model;
110mod nuts;
111mod sampler;
112mod sampler_stats;
113mod state;
114mod stepsize;
115mod storage;
116mod transform_adapt_strategy;
117mod transformed_hamiltonian;
118
119pub use nuts_derive::Storable;
120pub use nuts_storable::{HasDims, ItemType, Storable, Value};
121
122pub use adapt_strategy::EuclideanAdaptOptions;
123pub use chain::Chain;
124pub use cpu_math::{CpuLogpFunc, CpuMath, CpuMathError};
125pub use hamiltonian::DivergenceInfo;
126pub use math_base::{LogpError, Math};
127pub use model::Model;
128pub use nuts::NutsError;
129pub use sampler::{
130    ChainProgress, DiagGradNutsSettings, LowRankNutsSettings, NutsSettings, Progress,
131    ProgressCallback, Sampler, SamplerWaitResult, Settings, TransformedNutsSettings,
132    sample_sequentially,
133};
134pub use sampler_stats::SamplerStats;
135
136pub use mass_matrix::DiagAdaptExpSettings;
137pub use mass_matrix::LowRankSettings;
138pub use stepsize::{AdamOptions, StepSizeAdaptMethod, StepSizeAdaptOptions, StepSizeSettings};
139pub use transform_adapt_strategy::TransformedSettings;
140
141#[cfg(feature = "zarr")]
142pub use storage::{ZarrAsyncConfig, ZarrAsyncTraceStorage, ZarrConfig, ZarrTraceStorage};
143
144pub use storage::{CsvConfig, CsvTraceStorage};
145pub use storage::{HashMapConfig, HashMapValue};
146#[cfg(feature = "ndarray")]
147pub use storage::{NdarrayConfig, NdarrayTrace, NdarrayValue};
148
149#[cfg(feature = "arrow")]
150pub use storage::{ArrowConfig, ArrowTrace, ArrowTraceStorage};
nuts_rs/lib.rs

nuts_rs/
lib.rs
