mini_mcmc/lib.rs
1//! # Mini MCMC
2//!
3//! A compact Rust library offering **Markov Chain Monte Carlo (MCMC)** methods,
4//! including **Hamiltonian Monte Carlo (HMC)**, **Metropolis–Hastings**, and
5//! **Gibbs Sampling** for both discrete and continuous targets.
6//!
7//! ## Example 1: Sampling a 2D Gaussian (Metropolis–Hastings)
8//!
9//! ```rust
10//! use mini_mcmc::core::ChainRunner;
11//! use mini_mcmc::distributions::{Gaussian2D, IsotropicGaussian};
12//! use mini_mcmc::metropolis_hastings::MetropolisHastings;
13//! use ndarray::{arr1, arr2};
14//!
15//! let target = Gaussian2D {
16//! mean: arr1(&[0.0, 0.0]),
17//! cov: arr2(&[[1.0, 0.0], [0.0, 1.0]]),
18//! };
19//! let proposal = IsotropicGaussian::new(1.0);
20//! let initial_state = [0.0, 0.0];
21//!
22//! let mut mh = MetropolisHastings::new(target, proposal, &initial_state, 4);
23//! let samples = mh.run(1000, 100).unwrap();
24//! println!("Metropolis–Hastings samples shape: {:?}", samples.shape());
25//! ```
26//!
27//! ## Example 2: Sampling a 3D Rosenbrock (HMC)
28//!
29//! ```rust
30//! use burn::tensor::Element;
31//! use burn::{backend::Autodiff, prelude::Tensor};
32//! use mini_mcmc::hmc::{GradientTarget, HMC};
33//! use num_traits::Float;
34//!
35//! /// The 3D Rosenbrock distribution.
36//! ///
37//! /// For a point x = (x₁, x₂, x₃), the log probability is defined as the negative of
38//! /// the sum of two Rosenbrock terms:
39//! ///
40//! /// f(x) = 100*(x₂ - x₁²)² + (1 - x₁)² + 100*(x₃ - x₂²)² + (1 - x₂)²
41//! ///
42//! /// This implementation generalizes to d dimensions, but here we use it for 3D.
43//! struct RosenbrockND {}
44//!
45//! impl<T, B> GradientTarget<T, B> for RosenbrockND
46//! where
47//! T: Float + std::fmt::Debug + Element,
48//! B: burn::tensor::backend::AutodiffBackend,
49//! {
50//! fn log_prob_batch(&self, positions: &Tensor<B, 2>) -> Tensor<B, 1> {
51//! // Assume positions has shape [n_chains, d] with d = 3.
52//! let k = positions.dims()[0] as i64;
53//! let n = positions.dims()[1] as i64;
54//! let low = positions.clone().slice([(0, k), (0, n - 1)]);
55//! let high = positions.clone().slice([(0, k), (1, n)]);
56//! let term_1 = (high - low.clone().powi_scalar(2))
57//! .powi_scalar(2)
58//! .mul_scalar(100);
59//! let term_2 = low.neg().add_scalar(1).powi_scalar(2);
60//! -(term_1 + term_2).sum_dim(1).squeeze(1)
61//! }
62//! }
63//!
64//! // Use the CPU backend wrapped in Autodiff (e.g., NdArray).
65//! type BackendType = Autodiff<burn::backend::NdArray>;
66//!
67//! // Create the 3D Rosenbrock target.
68//! let target = RosenbrockND {};
69//!
70//! // Define initial positions for 6 chains (each a 3D point).
71//! let initial_positions = vec![vec![1.0_f32, 2.0_f32, 3.0_f32]; 6];
72//!
73//! // Create the HMC sampler with a step size of 0.01 and 50 leapfrog steps.
74//! let mut sampler = HMC::<f32, BackendType, RosenbrockND>::new(
75//! target,
76//! initial_positions,
77//! 0.032,
78//! 50,
79//! );
80//!
81//! // Run the sampler for 1100 iterations, discard 100
82//! let samples = sampler.run(1000, 100);
83//!
84//! // Print the shape of the collected samples.
85//! println!("Collected samples with shape: {:?}", samples.dims());
86//! ```
87//!
88//! ## Example 3: Sampling a Poisson Distribution (Discrete)
89//!
90//! ```rust
91//! use mini_mcmc::core::ChainRunner;
92//! use mini_mcmc::distributions::{Proposal, Target};
93//! use mini_mcmc::metropolis_hastings::MetropolisHastings;
94//! use rand::Rng;
95//!
96//! #[derive(Clone)]
97//! struct PoissonTarget {
98//! lambda: f64,
99//! }
100//!
101//! impl Target<usize, f64> for PoissonTarget {
102//! fn unnorm_log_prob(&self, theta: &[usize]) -> f64 {
103//! let k = theta[0];
104//! -self.lambda + (k as f64) * self.lambda.ln() - ln_factorial(k as u64)
105//! }
106//! }
107//!
108//! #[derive(Clone)]
109//! struct NonnegativeProposal;
110//!
111//! impl Proposal<usize, f64> for NonnegativeProposal {
112//! fn sample(&mut self, current: &[usize]) -> Vec<usize> {
113//! let x = current[0];
114//! if x == 0 {
115//! vec![1]
116//! } else {
117//! let step_up = rand::thread_rng().gen_bool(0.5);
118//! vec![if step_up { x + 1 } else { x - 1 }]
119//! }
120//! }
121//!
122//! fn log_prob(&self, from: &[usize], to: &[usize]) -> f64 {
123//! let (x, y) = (from[0], to[0]);
124//! if x == 0 && y == 1 {
125//! 0.0
126//! } else if x > 0 && (y == x + 1 || y == x - 1) {
127//! (0.5_f64).ln()
128//! } else {
129//! f64::NEG_INFINITY
130//! }
131//! }
132//! fn set_seed(self, _seed: u64) -> Self {
133//! self
134//! }
135//! }
136//!
137//! fn ln_factorial(k: u64) -> f64 {
138//! (1..=k).map(|v| (v as f64).ln()).sum()
139//! }
140//!
141//! let target = PoissonTarget { lambda: 4.0 };
142//! let proposal = NonnegativeProposal;
143//! let initial_state = [0];
144//!
145//! let mut mh = MetropolisHastings::new(target, proposal, &initial_state, 1);
146//! let samples = mh.run(5000, 100).unwrap();
147//! println!("Poisson samples shape: {:?}", samples.shape());
148//! ```
149//!
150//! For more complete implementations (including Gibbs sampling and I/O helpers),
151//! see the `examples/` directory.
152//!
153//! ## Features
154//! - **Parallel Chains** for improved throughput
155//! - **Progress Indicators** (acceptance rates, iteration counts)
156//! - **Common Distributions** (e.g. Gaussian) plus easy traits for custom log‐prob
157//! - **Optional I/O** (CSV, Arrow, Parquet) and GPU sampling (WGPU)
158//!
159//! ## Roadmap
160//! - No-U-Turn Sampler (NUTS)
161//! - Rank-Normalized R-hat diagnostics
162//! - Ensemble Slice Sampling (ESS)
163//! - Effective Sample Size estimation
164
165pub mod core;
166mod dev_tools;
167pub mod distributions;
168pub mod gibbs;
169pub mod hmc;
170pub mod io;
171pub mod ks_test;
172pub mod metropolis_hastings;
173pub mod stats;