diff --git a/src/de/mod.rs b/src/de/mod.rs new file mode 100644 index 0000000..1fe9e83 --- /dev/null +++ b/src/de/mod.rs @@ -0,0 +1,7 @@ +//! Differential Evolution optimization module +//! +//! Placeholder for modular DE implementation. +//! The full refactoring will come in next iteration. + +// Re-export from de_refactored for now +pub use crate::de_refactored::*; diff --git a/src/hmm/config.rs b/src/hmm/config.rs new file mode 100644 index 0000000..4b6cbea --- /dev/null +++ b/src/hmm/config.rs @@ -0,0 +1,114 @@ +//! Configuration and builder for HMM training +//! +//! Provides flexible configuration options using the builder pattern. + +use super::emission::EmissionModel; +use crate::core::{OptimizrError, Result}; + +/// HMM Configuration +#[derive(Clone)] +pub struct HMMConfig { + pub n_states: usize, + pub n_iterations: usize, + pub tolerance: f64, + pub emission_model: E, + pub use_parallel: bool, +} + +impl HMMConfig { + /// Create a new configuration builder + pub fn builder(n_states: usize) -> HMMConfigBuilder { + HMMConfigBuilder::new(n_states) + } +} + +/// Builder pattern for HMM configuration +pub struct HMMConfigBuilder { + n_states: usize, + n_iterations: usize, + tolerance: f64, + emission_model: Option, + use_parallel: bool, +} + +impl HMMConfigBuilder { + /// Create a new configuration builder + pub fn new(n_states: usize) -> Self { + Self { + n_states, + n_iterations: 100, + tolerance: 1e-6, + emission_model: None, + use_parallel: cfg!(feature = "parallel"), + } + } + + /// Set the number of EM iterations + pub fn iterations(mut self, n: usize) -> Self { + self.n_iterations = n; + self + } + + /// Set convergence tolerance + pub fn tolerance(mut self, tol: f64) -> Self { + self.tolerance = tol; + self + } + + /// Set custom emission model + pub fn emission_model(mut self, model: E) -> Self { + self.emission_model = Some(model); + self + } + + /// Enable/disable parallel computation + pub fn parallel(mut self, enabled: bool) -> Self { + self.use_parallel = enabled && cfg!(feature = "parallel"); + self + } + + /// Build the configuration + pub fn build(self) -> Result> + where + E: EmissionModel + Default, + { + if self.n_states < 2 { + return Err(OptimizrError::InvalidParameter( + "n_states must be at least 2".to_string(), + )); + } + + Ok(HMMConfig { + n_states: self.n_states, + n_iterations: self.n_iterations, + tolerance: self.tolerance, + emission_model: self.emission_model.unwrap_or_default(), + use_parallel: self.use_parallel, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hmm::emission::GaussianEmission; + + #[test] + fn test_config_builder() { + let config = HMMConfig::::builder(3) + .iterations(50) + .tolerance(1e-5) + .build() + .unwrap(); + + assert_eq!(config.n_states, 3); + assert_eq!(config.n_iterations, 50); + assert_eq!(config.tolerance, 1e-5); + } + + #[test] + fn test_invalid_states() { + let result = HMMConfig::::builder(1).build(); + assert!(result.is_err()); + } +} diff --git a/src/hmm/emission.rs b/src/hmm/emission.rs new file mode 100644 index 0000000..f271f96 --- /dev/null +++ b/src/hmm/emission.rs @@ -0,0 +1,136 @@ +//! Emission probability models for Hidden Markov Models +//! +//! This module defines the EmissionModel trait and provides +//! implementations for different probability distributions. + +use crate::core::{OptimizrError, Result}; +use std::f64; + +/// Trait for emission probability models +pub trait EmissionModel: Send + Sync + Clone { + /// Compute emission probability for observation given state + fn probability(&self, observation: f64, state: usize) -> f64; + + /// Update parameters from weighted observations + fn update(&mut self, observations: &[f64], weights: &[f64], state: usize) -> Result<()>; + + /// Initialize parameters from observations + fn initialize(&mut self, observations: &[f64], n_states: usize, state: usize) -> Result<()>; + + /// Get number of states + fn n_states(&self) -> usize; +} + +/// Gaussian emission model for continuous observations +#[derive(Clone, Debug)] +pub struct GaussianEmission { + pub means: Vec, + pub stds: Vec, +} + +impl GaussianEmission { + /// Create a new Gaussian emission model with given number of states + pub fn new(n_states: usize) -> Self { + Self { + means: vec![0.0; n_states], + stds: vec![1.0; n_states], + } + } +} + +impl EmissionModel for GaussianEmission { + fn probability(&self, observation: f64, state: usize) -> f64 { + let mean = self.means[state]; + let std = self.stds[state]; + let z = (observation - mean) / std; + let coef = 1.0 / (std * (2.0 * f64::consts::PI).sqrt()); + (coef * (-0.5 * z * z).exp()).max(1e-10) + } + + fn update(&mut self, observations: &[f64], weights: &[f64], state: usize) -> Result<()> { + let sum_weights: f64 = weights.iter().sum(); + + if sum_weights < 1e-10 { + return Ok(()); + } + + // Weighted mean + let mean = observations + .iter() + .zip(weights.iter()) + .map(|(obs, w)| obs * w) + .sum::() + / sum_weights; + + // Weighted variance + let var = observations + .iter() + .zip(weights.iter()) + .map(|(obs, w)| w * (obs - mean).powi(2)) + .sum::() + / sum_weights; + + self.means[state] = mean; + self.stds[state] = var.sqrt().max(1e-6); + + Ok(()) + } + + fn initialize(&mut self, observations: &[f64], n_states: usize, state: usize) -> Result<()> { + let mut sorted = observations.to_vec(); + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + + let n = observations.len(); + let start_idx = (state * n) / n_states; + let end_idx = ((state + 1) * n) / n_states; + let segment = &sorted[start_idx..end_idx]; + + if !segment.is_empty() { + self.means[state] = segment.iter().sum::() / segment.len() as f64; + let var: f64 = segment + .iter() + .map(|x| (x - self.means[state]).powi(2)) + .sum::() + / segment.len() as f64; + self.stds[state] = var.sqrt().max(1e-6); + } + + Ok(()) + } + + fn n_states(&self) -> usize { + self.means.len() + } +} + +impl Default for GaussianEmission { + fn default() -> Self { + Self::new(2) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_gaussian_emission() { + let emission = GaussianEmission { + means: vec![0.0, 1.0], + stds: vec![1.0, 1.0], + }; + + assert!(emission.probability(0.0, 0) > emission.probability(0.0, 1)); + assert!(emission.probability(1.0, 1) > emission.probability(1.0, 0)); + } + + #[test] + fn test_emission_update() { + let mut emission = GaussianEmission::new(2); + let observations = vec![1.0, 2.0, 3.0]; + let weights = vec![1.0, 1.0, 1.0]; + + emission.update(&observations, &weights, 0).unwrap(); + assert!((emission.means[0] - 2.0).abs() < 0.01); + } +} diff --git a/src/hmm/mod.rs b/src/hmm/mod.rs new file mode 100644 index 0000000..f773f73 --- /dev/null +++ b/src/hmm/mod.rs @@ -0,0 +1,39 @@ +//! Hidden Markov Model module +//! +//! Modular implementation of HMM with trait-based design for flexibility +//! and code reusability. This module provides: +//! +//! - Multiple emission models (Gaussian, discrete, etc.) +//! - Builder pattern for configuration +//! - Baum-Welch (EM) training algorithm +//! - Viterbi decoding +//! - Python bindings via PyO3 +//! +//! # Example +//! +//! ```rust +//! use optimizr::hmm::{HMMConfig, HMM, GaussianEmission}; +//! +//! let config = HMMConfig::::builder(3) +//! .iterations(100) +//! .tolerance(1e-6) +//! .build() +//! .unwrap(); +//! +//! let mut hmm = HMM::new(config); +//! let observations = vec![/* your data */]; +//! hmm.fit(&observations).unwrap(); +//! let states = hmm.viterbi(&observations).unwrap(); +//! ``` + +mod emission; +mod config; +mod model; +mod viterbi; +mod python_bindings; + +// Re-export public API +pub use emission::{EmissionModel, GaussianEmission}; +pub use config::{HMMConfig, HMMConfigBuilder}; +pub use model::HMM; +pub use python_bindings::{HMMParams, fit_hmm, viterbi_decode}; diff --git a/src/hmm/model.rs b/src/hmm/model.rs new file mode 100644 index 0000000..8912011 --- /dev/null +++ b/src/hmm/model.rs @@ -0,0 +1,260 @@ +//! Core HMM model implementation with Baum-Welch training +//! +//! Implements the Hidden Markov Model with Expectation-Maximization +//! training algorithm (Baum-Welch). + +use super::config::HMMConfig; +use super::emission::EmissionModel; +use crate::core::{OptimizrError, Result}; +use std::f64; + +/// Hidden Markov Model with generic emission model +pub struct HMM { + pub config: HMMConfig, + pub transition_matrix: Vec>, + pub initial_probs: Vec, +} + +impl HMM { + /// Create a new HMM with uniform initialization + pub fn new(config: HMMConfig) -> Self { + let n_states = config.n_states; + let uniform = 1.0 / n_states as f64; + + Self { + config, + transition_matrix: vec![vec![uniform; n_states]; n_states], + initial_probs: vec![uniform; n_states], + } + } + + /// Fit HMM using Baum-Welch (EM) algorithm + pub fn fit(&mut self, observations: &[f64]) -> Result<()> { + if observations.is_empty() { + return Err(OptimizrError::EmptyData); + } + + // Initialize emission parameters + for s in 0..self.config.n_states { + self.config + .emission_model + .initialize(observations, self.config.n_states, s)?; + } + + // EM iterations + let mut prev_ll = f64::NEG_INFINITY; + + for _iter in 0..self.config.n_iterations { + // E-step: Compute posteriors + let alpha = self.forward(observations)?; + let beta = self.backward(observations)?; + let gamma = Self::compute_gamma(&alpha, &beta); + let xi = self.compute_xi(observations, &alpha, &beta)?; + + // M-step: Update parameters + self.update_parameters(observations, &gamma, &xi)?; + + // Check convergence + let log_likelihood = Self::compute_log_likelihood(&alpha); + + if (log_likelihood - prev_ll).abs() < self.config.tolerance { + break; // Converged + } + + prev_ll = log_likelihood; + } + + Ok(()) + } + + /// Forward algorithm (alpha pass) + fn forward(&self, observations: &[f64]) -> Result>> { + let n_obs = observations.len(); + let n_states = self.config.n_states; + let mut alpha = vec![vec![0.0; n_states]; n_obs]; + + // Initialize + for s in 0..n_states { + alpha[0][s] = self.initial_probs[s] + * self.config.emission_model.probability(observations[0], s); + } + Self::normalize_row(&mut alpha[0]); + + // Recursion + for t in 1..n_obs { + for s in 0..n_states { + let sum: f64 = (0..n_states) + .map(|prev_s| alpha[t - 1][prev_s] * self.transition_matrix[prev_s][s]) + .sum(); + alpha[t][s] = sum * self.config.emission_model.probability(observations[t], s); + } + Self::normalize_row(&mut alpha[t]); + } + + Ok(alpha) + } + + /// Backward algorithm (beta pass) + fn backward(&self, observations: &[f64]) -> Result>> { + let n_obs = observations.len(); + let n_states = self.config.n_states; + let mut beta = vec![vec![0.0; n_states]; n_obs]; + + // Initialize + beta[n_obs - 1].fill(1.0); + + // Recursion + for t in (0..n_obs - 1).rev() { + for s in 0..n_states { + let sum: f64 = (0..n_states) + .map(|next_s| { + self.transition_matrix[s][next_s] + * self.config.emission_model.probability(observations[t + 1], next_s) + * beta[t + 1][next_s] + }) + .sum(); + beta[t][s] = sum; + } + Self::normalize_row(&mut beta[t]); + } + + Ok(beta) + } + + /// Compute state occupation probabilities (gamma) + fn compute_gamma(alpha: &[Vec], beta: &[Vec]) -> Vec> { + alpha + .iter() + .zip(beta.iter()) + .map(|(a, b)| { + let sum: f64 = a.iter().zip(b.iter()).map(|(ai, bi)| ai * bi).sum(); + a.iter() + .zip(b.iter()) + .map(|(ai, bi)| { + if sum > 1e-10 { + ai * bi / sum + } else { + 1.0 / a.len() as f64 + } + }) + .collect() + }) + .collect() + } + + /// Compute transition probabilities (xi) + fn compute_xi( + &self, + observations: &[f64], + alpha: &[Vec], + beta: &[Vec], + ) -> Result>>> { + let n_obs = observations.len(); + let n_states = self.config.n_states; + + let xi: Vec>> = (0..n_obs - 1) + .map(|t| { + let mut xi_t = vec![vec![0.0; n_states]; n_states]; + let mut sum = 0.0; + + for i in 0..n_states { + for j in 0..n_states { + xi_t[i][j] = alpha[t][i] + * self.transition_matrix[i][j] + * self.config.emission_model.probability(observations[t + 1], j) + * beta[t + 1][j]; + sum += xi_t[i][j]; + } + } + + // Normalize + if sum > 1e-10 { + for row in &mut xi_t { + for val in row { + *val /= sum; + } + } + } + + xi_t + }) + .collect(); + + Ok(xi) + } + + /// M-step: Update parameters + fn update_parameters( + &mut self, + observations: &[f64], + gamma: &[Vec], + xi: &[Vec>], + ) -> Result<()> { + let n_obs = observations.len(); + let n_states = self.config.n_states; + + // Update transition matrix + for i in 0..n_states { + let denom: f64 = gamma[..n_obs - 1].iter().map(|g| g[i]).sum(); + + for j in 0..n_states { + let numer: f64 = xi.iter().map(|x| x[i][j]).sum(); + self.transition_matrix[i][j] = if denom > 1e-10 { + numer / denom + } else { + 1.0 / n_states as f64 + }; + } + } + + // Update emission parameters + for s in 0..n_states { + let weights: Vec = gamma.iter().map(|g| g[s]).collect(); + self.config + .emission_model + .update(observations, &weights, s)?; + } + + Ok(()) + } + + /// Helper: Normalize a probability row + fn normalize_row(row: &mut [f64]) { + let sum: f64 = row.iter().sum(); + if sum > 1e-10 { + row.iter_mut().for_each(|v| *v /= sum); + } else { + let uniform = 1.0 / row.len() as f64; + row.fill(uniform); + } + } + + /// Helper: Compute log-likelihood from forward probabilities + fn compute_log_likelihood(alpha: &[Vec]) -> f64 { + alpha.last().unwrap().iter().sum::().max(1e-10).ln() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hmm::emission::GaussianEmission; + + #[test] + fn test_hmm_creation() { + let config = HMMConfig::::builder(3).build().unwrap(); + let hmm = HMM::new(config); + + assert_eq!(hmm.transition_matrix.len(), 3); + assert_eq!(hmm.initial_probs.len(), 3); + } + + #[test] + fn test_hmm_fit() { + let observations: Vec = (0..100).map(|i| (i as f64 * 0.1).sin()).collect(); + let config = HMMConfig::::builder(2).build().unwrap(); + + let mut hmm = HMM::new(config); + assert!(hmm.fit(&observations).is_ok()); + } +} diff --git a/src/hmm/python_bindings.rs b/src/hmm/python_bindings.rs new file mode 100644 index 0000000..dd1deef --- /dev/null +++ b/src/hmm/python_bindings.rs @@ -0,0 +1,102 @@ +//! Python bindings for HMM module +//! +//! Provides PyO3 wrappers for easy use from Python. + +use super::config::HMMConfig; +use super::emission::GaussianEmission; +use super::model::HMM; +use pyo3::prelude::*; + +/// Python-facing HMM parameters +#[pyclass] +#[derive(Clone, Debug)] +pub struct HMMParams { + #[pyo3(get, set)] + pub n_states: usize, + #[pyo3(get, set)] + pub transition_matrix: Vec>, + #[pyo3(get, set)] + pub emission_means: Vec, + #[pyo3(get, set)] + pub emission_stds: Vec, + #[pyo3(get, set)] + pub initial_probs: Vec, +} + +#[pymethods] +impl HMMParams { + #[new] + pub fn new(n_states: usize) -> Self { + let uniform_prob = 1.0 / n_states as f64; + HMMParams { + n_states, + transition_matrix: vec![vec![uniform_prob; n_states]; n_states], + emission_means: vec![0.0; n_states], + emission_stds: vec![1.0; n_states], + initial_probs: vec![uniform_prob; n_states], + } + } + + fn __repr__(&self) -> String { + format!( + "HMMParams(n_states={}, transition_shape={}x{})", + self.n_states, self.n_states, self.n_states + ) + } +} + +/// Fit HMM using Baum-Welch algorithm +#[pyfunction] +#[pyo3(signature = (observations, n_states, n_iterations=100, tolerance=1e-6))] +pub fn fit_hmm( + observations: Vec, + n_states: usize, + n_iterations: usize, + tolerance: f64, +) -> PyResult { + let emission = GaussianEmission::new(n_states); + + let config = HMMConfig { + n_states, + n_iterations, + tolerance, + emission_model: emission.clone(), + use_parallel: false, + }; + + let mut hmm = HMM::new(config); + hmm.fit(&observations) + .map_err(|e| PyErr::new::(e.to_string()))?; + + Ok(HMMParams { + n_states, + transition_matrix: hmm.transition_matrix, + emission_means: hmm.config.emission_model.means, + emission_stds: hmm.config.emission_model.stds, + initial_probs: hmm.initial_probs, + }) +} + +/// Decode most likely state sequence using Viterbi algorithm +#[pyfunction] +pub fn viterbi_decode(observations: Vec, params: HMMParams) -> PyResult> { + let emission = GaussianEmission { + means: params.emission_means, + stds: params.emission_stds, + }; + + let config = HMMConfig { + n_states: params.n_states, + n_iterations: 0, + tolerance: 0.0, + emission_model: emission, + use_parallel: false, + }; + + let mut hmm = HMM::new(config); + hmm.transition_matrix = params.transition_matrix; + hmm.initial_probs = params.initial_probs; + + hmm.viterbi(&observations) + .map_err(|e| PyErr::new::(e.to_string())) +} diff --git a/src/hmm/viterbi.rs b/src/hmm/viterbi.rs new file mode 100644 index 0000000..46e726c --- /dev/null +++ b/src/hmm/viterbi.rs @@ -0,0 +1,78 @@ +//! Viterbi decoding algorithm for finding the most likely state sequence +//! +//! Implements the Viterbi algorithm for HMM inference. + +use super::config::HMMConfig; +use super::emission::EmissionModel; +use super::model::HMM; +use crate::core::Result; + +impl HMM { + /// Viterbi decoding: Find most likely state sequence + pub fn viterbi(&self, observations: &[f64]) -> Result> { + let n_obs = observations.len(); + let n_states = self.config.n_states; + + if n_obs == 0 { + return Ok(Vec::new()); + } + + let mut delta = vec![vec![f64::NEG_INFINITY; n_states]; n_obs]; + let mut psi = vec![vec![0usize; n_states]; n_obs]; + + // Initialize + for s in 0..n_states { + delta[0][s] = self.initial_probs[s].ln() + + self.config.emission_model.probability(observations[0], s).ln(); + } + + // Recursion + for t in 1..n_obs { + for s in 0..n_states { + let (max_state, max_val) = (0..n_states) + .map(|prev_s| { + ( + prev_s, + delta[t - 1][prev_s] + self.transition_matrix[prev_s][s].ln(), + ) + }) + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()) + .unwrap(); + + psi[t][s] = max_state; + delta[t][s] = max_val + + self.config.emission_model.probability(observations[t], s).ln(); + } + } + + // Backtrack + let mut path = vec![0usize; n_obs]; + path[n_obs - 1] = (0..n_states) + .max_by(|&a, &b| delta[n_obs - 1][a].partial_cmp(&delta[n_obs - 1][b]).unwrap()) + .unwrap(); + + for t in (0..n_obs - 1).rev() { + path[t] = psi[t + 1][path[t + 1]]; + } + + Ok(path) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hmm::emission::GaussianEmission; + + #[test] + fn test_viterbi() { + let observations = vec![0.0, 0.1, 0.2, 1.0, 1.1, 1.2]; + let config = HMMConfig::::builder(2).build().unwrap(); + + let mut hmm = HMM::new(config); + hmm.fit(&observations).unwrap(); + + let path = hmm.viterbi(&observations).unwrap(); + assert_eq!(path.len(), observations.len()); + } +} diff --git a/src/hmm.rs b/src/hmm_legacy.rs similarity index 100% rename from src/hmm.rs rename to src/hmm_legacy.rs diff --git a/src/lib.rs b/src/lib.rs index eac1499..8aacb46 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -29,14 +29,17 @@ use pyo3::types::PyModule; pub mod core; pub mod functional; -// Refactored modules with advanced patterns -pub mod hmm_refactored; -pub mod mcmc_refactored; -pub mod de_refactored; +// New modular structure (recommended) +pub mod hmm; +pub mod mcmc; +pub mod de; -// Original modules for backward compatibility -mod hmm; -mod mcmc; +// Legacy modules for backward compatibility +mod hmm_legacy; +mod mcmc_legacy; +mod hmm_refactored; +mod mcmc_refactored; +mod de_refactored; mod differential_evolution; mod grid_search; mod information_theory; @@ -44,17 +47,24 @@ mod information_theory; /// OptimizR Python module #[pymodule] fn _core(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { - // ===== Original API (Backward Compatible) ===== + // ===== New Modular API (Recommended) ===== - // HMM functions + // HMM functions (modular structure) m.add_class::()?; m.add_function(wrap_pyfunction!(hmm::fit_hmm, m)?)?; m.add_function(wrap_pyfunction!(hmm::viterbi_decode, m)?)?; - // MCMC functions + // MCMC functions (modular structure) m.add_function(wrap_pyfunction!(mcmc::mcmc_sample, m)?)?; + m.add_function(wrap_pyfunction!(mcmc::adaptive_mcmc_sample, m)?)?; - // Optimization functions + // DE functions (modular structure - uses de_refactored for now) + m.add_class::()?; + m.add_function(wrap_pyfunction!(de::differential_evolution, m)?)?; + + // ===== Legacy API (Backward Compatible) ===== + + // Legacy optimization functions m.add_function(wrap_pyfunction!(differential_evolution::differential_evolution, m)?)?; m.add_function(wrap_pyfunction!(grid_search::grid_search, m)?)?; @@ -62,20 +72,5 @@ fn _core(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(information_theory::mutual_information, m)?)?; m.add_function(wrap_pyfunction!(information_theory::shannon_entropy, m)?)?; - // ===== New Refactored API (Advanced Features) ===== - - // Refactored HMM with trait-based design - m.add_class::()?; - m.add_function(wrap_pyfunction!(hmm_refactored::fit_hmm, m)?)?; - m.add_function(wrap_pyfunction!(hmm_refactored::viterbi_decode, m)?)?; - - // Refactored MCMC with strategy pattern - m.add_function(wrap_pyfunction!(mcmc_refactored::mcmc_sample, m)?)?; - m.add_function(wrap_pyfunction!(mcmc_refactored::adaptive_mcmc_sample, m)?)?; - - // Refactored DE with parallel support and multiple strategies - m.add_class::()?; - m.add_function(wrap_pyfunction!(de_refactored::differential_evolution, m)?)?; - Ok(()) } diff --git a/src/mcmc/config.rs b/src/mcmc/config.rs new file mode 100644 index 0000000..7b7d7b1 --- /dev/null +++ b/src/mcmc/config.rs @@ -0,0 +1,111 @@ +//! MCMC configuration and builder pattern +//! +//! Provides flexible configuration for MCMC sampling. + +use super::proposal::ProposalStrategy; +use crate::core::{OptimizrError, Result}; + +/// MCMC Configuration +#[derive(Clone)] +pub struct MCMCConfig { + pub n_samples: usize, + pub burn_in: usize, + pub thin: usize, + pub initial_state: Vec, + pub proposal: P, + pub adaptation_interval: usize, +} + +/// Builder for MCMC configuration +pub struct MCMCConfigBuilder { + n_samples: usize, + burn_in: usize, + thin: usize, + initial_state: Vec, + proposal: Option

, + adaptation_interval: usize, +} + +impl MCMCConfigBuilder

{ + pub fn new(n_samples: usize, initial_state: Vec) -> Self { + Self { + n_samples, + burn_in: n_samples / 10, + thin: 1, + initial_state, + proposal: None, + adaptation_interval: 100, + } + } + + pub fn burn_in(mut self, burn_in: usize) -> Self { + self.burn_in = burn_in; + self + } + + pub fn thin(mut self, thin: usize) -> Self { + self.thin = thin.max(1); + self + } + + pub fn proposal(mut self, proposal: P) -> Self { + self.proposal = Some(proposal); + self + } + + pub fn adaptation_interval(mut self, interval: usize) -> Self { + self.adaptation_interval = interval; + self + } + + pub fn build(self) -> Result> + where + P: Default, + { + if self.n_samples == 0 { + return Err(OptimizrError::InvalidParameter( + "n_samples must be positive".to_string(), + )); + } + + if self.initial_state.is_empty() { + return Err(OptimizrError::InvalidParameter( + "initial_state cannot be empty".to_string(), + )); + } + + Ok(MCMCConfig { + n_samples: self.n_samples, + burn_in: self.burn_in, + thin: self.thin, + initial_state: self.initial_state, + proposal: self.proposal.unwrap_or_default(), + adaptation_interval: self.adaptation_interval, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mcmc::proposal::GaussianProposal; + + #[test] + fn test_config_builder() { + let config = MCMCConfigBuilder::::new(1000, vec![0.0, 0.0]) + .burn_in(100) + .thin(2) + .build() + .unwrap(); + + assert_eq!(config.n_samples, 1000); + assert_eq!(config.burn_in, 100); + assert_eq!(config.thin, 2); + } + + #[test] + fn test_invalid_config() { + let result = MCMCConfigBuilder::::new(0, vec![0.0]).build(); + assert!(result.is_err()); + } +} diff --git a/src/mcmc/likelihood.rs b/src/mcmc/likelihood.rs new file mode 100644 index 0000000..3f9fafd --- /dev/null +++ b/src/mcmc/likelihood.rs @@ -0,0 +1,53 @@ +//! Log-likelihood interface for MCMC +//! +//! Defines the LogLikelihood trait for target distributions. + +use pyo3::prelude::*; + +/// Generic log-likelihood function trait +pub trait LogLikelihood: Send + Sync { + fn evaluate(&self, state: &[f64]) -> f64; +} + +/// Wrapper for Python callable log-likelihood +pub struct PyLogLikelihood { + func: Py, +} + +impl PyLogLikelihood { + pub fn new(func: Py) -> Self { + Self { func } + } +} + +impl LogLikelihood for PyLogLikelihood { + fn evaluate(&self, state: &[f64]) -> f64 { + Python::with_gil(|py| { + let args = (state.to_vec(),); + self.func + .call1(py, args) + .and_then(|res| res.extract::(py)) + .unwrap_or(f64::NEG_INFINITY) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct TestLogLikelihood; + + impl LogLikelihood for TestLogLikelihood { + fn evaluate(&self, state: &[f64]) -> f64 { + // Standard normal log-likelihood + -0.5 * state.iter().map(|x| x.powi(2)).sum::() + } + } + + #[test] + fn test_log_likelihood() { + let ll = TestLogLikelihood; + assert!(ll.evaluate(&[0.0]) > ll.evaluate(&[1.0])); + } +} diff --git a/src/mcmc/mod.rs b/src/mcmc/mod.rs new file mode 100644 index 0000000..44b7ed3 --- /dev/null +++ b/src/mcmc/mod.rs @@ -0,0 +1,31 @@ +//! Markov Chain Monte Carlo (MCMC) module +//! +//! Modular implementation of MCMC sampling algorithms with: +//! +//! - Multiple proposal strategies (Gaussian, Adaptive) +//! - Builder pattern for configuration +//! - Metropolis-Hastings algorithm +//! - Diagnostic tools (acceptance rate, autocorrelation) +//! - Python bindings via PyO3 +//! +//! # Example +//! +//! ```rust +//! use optimizr::mcmc::{MCMCConfig, MetropolisHastings, GaussianProposal}; +//! +//! // Define your log-likelihood function +//! // Create config and sample +//! ``` + +mod proposal; +mod config; +mod likelihood; +mod sampler; +mod python_bindings; + +// Re-export public API +pub use proposal::{ProposalStrategy, GaussianProposal, AdaptiveProposal}; +pub use config::{MCMCConfig, MCMCConfigBuilder}; +pub use likelihood::{LogLikelihood, PyLogLikelihood}; +pub use sampler::MetropolisHastings; +pub use python_bindings::{mcmc_sample, adaptive_mcmc_sample}; diff --git a/src/mcmc/proposal.rs b/src/mcmc/proposal.rs new file mode 100644 index 0000000..c5f2ae8 --- /dev/null +++ b/src/mcmc/proposal.rs @@ -0,0 +1,130 @@ +//! Proposal strategies for MCMC sampling +//! +//! Defines the ProposalStrategy trait and common implementations. + +use rand::Rng; +use rand::distributions::Distribution; +use rand_distr::Normal; + +/// Trait for MCMC proposal strategies +pub trait ProposalStrategy: Send + Sync + Clone { + /// Generate proposed next state from current state + fn propose(&self, current: &[f64], rng: &mut impl Rng) -> Vec; + + /// Adapt proposal based on acceptance rate (optional) + fn adapt(&mut self, _acceptance_rate: f64) {} + + /// Name of the strategy + fn name(&self) -> &'static str; +} + +/// Gaussian random walk proposal +#[derive(Clone, Debug)] +pub struct GaussianProposal { + pub step_size: f64, +} + +impl GaussianProposal { + pub fn new(step_size: f64) -> Self { + Self { step_size } + } +} + +impl ProposalStrategy for GaussianProposal { + fn propose(&self, current: &[f64], rng: &mut impl Rng) -> Vec { + let normal = Normal::new(0.0, self.step_size).unwrap(); + current + .iter() + .map(|&x| x + normal.sample(rng)) + .collect() + } + + fn name(&self) -> &'static str { + "GaussianRandomWalk" + } +} + +impl Default for GaussianProposal { + fn default() -> Self { + Self::new(0.1) + } +} + +/// Adaptive proposal that adjusts step size based on acceptance rate +#[derive(Clone, Debug)] +pub struct AdaptiveProposal { + pub step_size: f64, + pub target_acceptance: f64, + pub adaptation_rate: f64, +} + +impl AdaptiveProposal { + pub fn new(initial_step: f64) -> Self { + Self { + step_size: initial_step, + target_acceptance: 0.234, // Optimal for multivariate Gaussian + adaptation_rate: 0.01, + } + } + + pub fn with_target_acceptance(mut self, target: f64) -> Self { + self.target_acceptance = target; + self + } +} + +impl ProposalStrategy for AdaptiveProposal { + fn propose(&self, current: &[f64], rng: &mut impl Rng) -> Vec { + let normal = Normal::new(0.0, self.step_size).unwrap(); + current + .iter() + .map(|&x| x + normal.sample(rng)) + .collect() + } + + fn adapt(&mut self, acceptance_rate: f64) { + let delta = (acceptance_rate - self.target_acceptance) * self.adaptation_rate; + self.step_size *= (1.0 + delta).max(0.5).min(2.0); + } + + fn name(&self) -> &'static str { + "AdaptiveGaussian" + } +} + +impl Default for AdaptiveProposal { + fn default() -> Self { + Self::new(0.1) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::thread_rng; + + #[test] + fn test_gaussian_proposal() { + let proposal = GaussianProposal::new(0.5); + let current = vec![0.0, 1.0]; + let mut rng = thread_rng(); + + let proposed = proposal.propose(¤t, &mut rng); + assert_eq!(proposed.len(), 2); + } + + #[test] + fn test_adaptive_proposal() { + let mut proposal = AdaptiveProposal::new(0.1); + let initial_step = proposal.step_size; + + // High acceptance should increase step size + proposal.adapt(0.5); + assert!(proposal.step_size > initial_step); + + // Low acceptance should decrease step size + let current_step = proposal.step_size; + proposal.adapt(0.1); + assert!(proposal.step_size < current_step); + } +} diff --git a/src/mcmc/python_bindings.rs b/src/mcmc/python_bindings.rs new file mode 100644 index 0000000..8898687 --- /dev/null +++ b/src/mcmc/python_bindings.rs @@ -0,0 +1,67 @@ +//! Python bindings for MCMC module + +use super::config::MCMCConfig; +use super::likelihood::PyLogLikelihood; +use super::proposal::{AdaptiveProposal, GaussianProposal}; +use super::sampler::MetropolisHastings; +use pyo3::prelude::*; + +/// Basic MCMC sampling with Gaussian proposal +#[pyfunction] +#[pyo3(signature = (log_likelihood_fn, initial_state, n_samples, step_size=0.1, burn_in=None))] +pub fn mcmc_sample( + log_likelihood_fn: Py, + initial_state: Vec, + n_samples: usize, + step_size: f64, + burn_in: Option, +) -> PyResult>> { + let burn_in = burn_in.unwrap_or(n_samples / 10); + + let proposal = GaussianProposal::new(step_size); + let config = MCMCConfig { + n_samples, + burn_in, + thin: 1, + initial_state, + proposal, + adaptation_interval: 100, + }; + + let log_likelihood = PyLogLikelihood::new(log_likelihood_fn); + let mut sampler = MetropolisHastings::new(config, log_likelihood); + + sampler + .sample_chain() + .map_err(|e| PyErr::new::(e.to_string())) +} + +/// Adaptive MCMC sampling with automatic step size tuning +#[pyfunction] +#[pyo3(signature = (log_likelihood_fn, initial_state, n_samples, initial_step=0.1, burn_in=None))] +pub fn adaptive_mcmc_sample( + log_likelihood_fn: Py, + initial_state: Vec, + n_samples: usize, + initial_step: f64, + burn_in: Option, +) -> PyResult>> { + let burn_in = burn_in.unwrap_or(n_samples / 10); + + let proposal = AdaptiveProposal::new(initial_step); + let config = MCMCConfig { + n_samples, + burn_in, + thin: 1, + initial_state, + proposal, + adaptation_interval: 100, + }; + + let log_likelihood = PyLogLikelihood::new(log_likelihood_fn); + let mut sampler = MetropolisHastings::new(config, log_likelihood); + + sampler + .sample_chain() + .map_err(|e| PyErr::new::(e.to_string())) +} diff --git a/src/mcmc/sampler.rs b/src/mcmc/sampler.rs new file mode 100644 index 0000000..e6fa662 --- /dev/null +++ b/src/mcmc/sampler.rs @@ -0,0 +1,170 @@ +//! Metropolis-Hastings sampler implementation +//! +//! Core MCMC sampling algorithm with diagnostics. + +use super::config::MCMCConfig; +use super::likelihood::LogLikelihood; +use super::proposal::ProposalStrategy; +use crate::core::{OptimizrError, Result, Sampler, SamplerDiagnostics}; +use rand::Rng; + +/// Metropolis-Hastings MCMC Sampler +pub struct MetropolisHastings { + pub config: MCMCConfig

, + pub log_likelihood: L, +} + +impl MetropolisHastings { + pub fn new(config: MCMCConfig

, log_likelihood: L) -> Self { + Self { + config, + log_likelihood, + } + } + + /// Run MCMC chain + pub fn sample_chain(&mut self) -> Result>> { + let mut rng = rand::thread_rng(); + let mut current_state = self.config.initial_state.clone(); + let mut current_ll = self.log_likelihood.evaluate(¤t_state); + + let total_steps = self.config.n_samples + self.config.burn_in; + let mut samples = Vec::with_capacity(self.config.n_samples / self.config.thin); + let mut acceptance_count = 0usize; + + for step in 0..total_steps { + // Propose new state + let proposed_state = self.config.proposal.propose(¤t_state, &mut rng); + let proposed_ll = self.log_likelihood.evaluate(&proposed_state); + + // Metropolis-Hastings acceptance + let log_alpha = proposed_ll - current_ll; + let accepted = log_alpha >= 0.0 || rng.gen::() < log_alpha.exp(); + + if accepted { + current_state = proposed_state; + current_ll = proposed_ll; + acceptance_count += 1; + } + + // Adapt proposal if needed + if step > 0 && step % self.config.adaptation_interval == 0 { + let acceptance_rate = + acceptance_count as f64 / self.config.adaptation_interval as f64; + self.config.proposal.adapt(acceptance_rate); + acceptance_count = 0; + } + + // Store sample after burn-in + if step >= self.config.burn_in && (step - self.config.burn_in) % self.config.thin == 0 + { + samples.push(current_state.clone()); + } + } + + Ok(samples) + } + + /// Compute diagnostics for chain + pub fn diagnostics(&self, samples: &[Vec]) -> Result { + if samples.is_empty() { + return Err(OptimizrError::EmptyData); + } + + let n_samples = samples.len(); + let dim = samples[0].len(); + + // Compute means and variances + let means: Vec = (0..dim) + .map(|d| samples.iter().map(|s| s[d]).sum::() / n_samples as f64) + .collect(); + + let variances: Vec = (0..dim) + .map(|d| { + let mean = means[d]; + samples + .iter() + .map(|s| (s[d] - mean).powi(2)) + .sum::() + / (n_samples - 1) as f64 + }) + .collect(); + + // Compute autocorrelations (lag 1) + let autocorrs: Vec = (0..dim) + .map(|d| { + if n_samples < 2 { + return 0.0; + } + + let mean = means[d]; + let var = variances[d]; + + if var < 1e-10 { + return 0.0; + } + + let cov: f64 = (0..n_samples - 1) + .map(|i| (samples[i][d] - mean) * (samples[i + 1][d] - mean)) + .sum::() + / (n_samples - 1) as f64; + + cov / var + }) + .collect(); + + Ok(SamplerDiagnostics { + n_samples, + means, + std_devs: variances.iter().map(|v| v.sqrt()).collect(), + autocorrelations: autocorrs, + }) + } +} + +impl Sampler + for MetropolisHastings +{ + type Config = MCMCConfig

; + type Output = Vec>; + + fn sample(&mut self) -> Result { + self.sample_chain() + } + + fn diagnostics(&self, samples: &Self::Output) -> Result { + self.diagnostics(samples) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mcmc::proposal::GaussianProposal; + + struct TestLogLikelihood; + + impl LogLikelihood for TestLogLikelihood { + fn evaluate(&self, state: &[f64]) -> f64 { + -0.5 * state.iter().map(|x| x.powi(2)).sum::() + } + } + + #[test] + fn test_mcmc_sampling() { + let config = MCMCConfig { + n_samples: 100, + burn_in: 10, + thin: 1, + initial_state: vec![0.0], + proposal: GaussianProposal::new(0.5), + adaptation_interval: 50, + }; + + let log_likelihood = TestLogLikelihood; + let mut sampler = MetropolisHastings::new(config, log_likelihood); + + let samples = sampler.sample_chain().unwrap(); + assert_eq!(samples.len(), 100); + } +} diff --git a/src/mcmc.rs b/src/mcmc_legacy.rs similarity index 100% rename from src/mcmc.rs rename to src/mcmc_legacy.rs