gavindiaz/optimiz-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve code design

Browse Source
This commit is contained in:
Melvin Avarez
2025-12-03 22:08:08 +01:00
parent 923d27e87b
commit a62ceaa64b
18 changed files with 6613 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core.rs
+174
View File
@@ -0,0 +1,174 @@
//! Core traits and types for optimization algorithms
//!
//! This module defines the foundational traits and types used across all
//! optimization and inference algorithms in OptimizR.
use thiserror::Error;
/// Custom error type for OptimizR operations
#[derive(Error, Debug, Clone)]
pub enum OptimizrError {
#[error("Invalid parameter: {0}")]
InvalidParameter(String),
#[error("Dimension mismatch: expected {expected}, got {actual}")]
DimensionMismatch { expected: usize, actual: usize },
#[error("Empty data provided")]
EmptyData,
#[error("Convergence failed after {0} iterations")]
ConvergenceFailed(usize),
#[error("Numerical error: {0}")]
NumericalError(String),
#[error("Computation error: {0}")]
ComputationError(String),
}
/// Result type for OptimizR operations
pub type Result<T> = std::result::Result<T, OptimizrError>;
/// Trait for optimization algorithms
pub trait Optimizer {
type Config;
type Output;
/// Optimize to find best solution
fn optimize(&mut self) -> Result<Self::Output>;
/// Get current best solution
fn best(&self) -> Result<Vec<f64>>;
}
/// Trait for sampling algorithms (MCMC, etc.)
pub trait Sampler {
type Config;
type Output;
/// Draw samples from the target distribution
fn sample(&mut self) -> Result<Self::Output>;
/// Get diagnostics about sampling performance
fn diagnostics(&self, samples: &Self::Output) -> Result<SamplerDiagnostics>;
}
/// Diagnostics for sampling algorithms
#[derive(Debug, Clone)]
pub struct SamplerDiagnostics {
pub n_samples: usize,
pub means: Vec<f64>,
pub std_devs: Vec<f64>,
pub autocorrelations: Vec<f64>,
}
/// Trait for configuration builders
pub trait ConfigBuilder {
type Config;
fn build(self) -> Result<Self::Config>;
}
/// Trait for information measures (entropy, MI, etc.)
pub trait InformationMeasure {
/// Compute the measure for given data
fn compute(&self, data: &[f64]) -> Result<f64>;
/// Compute pairwise measure (for MI)
fn compute_pairwise(&self, _x: &[f64], _y: &[f64]) -> Result<f64> {
Err(OptimizrError::ComputationError(
"Pairwise computation not supported".to_string(),
))
}
}
/// Bounds for optimization
#[derive(Debug, Clone)]
pub struct Bounds {
pub lower: Vec<f64>,
pub upper: Vec<f64>,
}
impl Bounds {
pub fn new(bounds: Vec<(f64, f64)>) -> Result<Self> {
if bounds.is_empty() {
return Err(OptimizrError::InvalidParameter(
"Bounds cannot be empty".to_string(),
));
}
for (lower, upper) in &bounds {
if lower >= upper {
return Err(OptimizrError::InvalidParameter(format!(
"Invalid bounds: lower ({}) >= upper ({})",
lower, upper
)));
}
}
let (lower, upper): (Vec<_>, Vec<_>) = bounds.into_iter().unzip();
Ok(Self { lower, upper })
}
pub fn dim(&self) -> usize {
self.lower.len()
}
pub fn clip(&self, x: &[f64]) -> Vec<f64> {
x.iter()
.enumerate()
.map(|(i, &val)| val.max(self.lower[i]).min(self.upper[i]))
.collect()
}
pub fn is_valid(&self, x: &[f64]) -> bool {
x.len() == self.dim()
&& x.iter()
.enumerate()
.all(|(i, &val)| val >= self.lower[i] && val <= self.upper[i])
}
pub fn sample(&self, rng: &mut impl rand::Rng) -> Vec<f64> {
(0..self.dim())
.map(|i| rng.gen_range(self.lower[i]..self.upper[i]))
.collect()
}
}
/// Trait for parallel execution strategies
pub trait ParallelExecutor {
fn execute_parallel<F, T>(&self, tasks: Vec<F>) -> Vec<T>
where
F: Fn() -> T + Send,
T: Send;
}
/// Standard rayon-based parallel executor
#[cfg(feature = "parallel")]
pub struct RayonExecutor;
#[cfg(feature = "parallel")]
impl ParallelExecutor for RayonExecutor {
fn execute_parallel<F, T>(&self, tasks: Vec<F>) -> Vec<T>
where
F: Fn() -> T + Send,
T: Send,
{
use rayon::prelude::*;
tasks.into_par_iter().map(|f| f()).collect()
}
}
/// Sequential executor (fallback)
pub struct SequentialExecutor;
impl ParallelExecutor for SequentialExecutor {
fn execute_parallel<F, T>(&self, tasks: Vec<F>) -> Vec<T>
where
F: Fn() -> T + Send,
T: Send,
{
tasks.into_iter().map(|f| f()).collect()
}
}
src/de_refactored.rs
+549
View File
@@ -0,0 +1,549 @@
//! Refactored Differential Evolution with Parallel Support
//!
//! Strategy pattern for mutation operators and parallel fitness evaluation.
use crate::core::{Bounds, OptimizrError, Optimizer, Result};
use pyo3::prelude::*;
use rand::Rng;
use std::sync::Arc;
#[cfg(feature = "parallel")]
use rayon::prelude::*;
/// Trait for mutation strategies
pub trait MutationStrategy: Send + Sync + Clone {
fn mutate(
&self,
population: &[Vec<f64>],
target_idx: usize,
f: f64,
rng: &mut impl Rng,
) -> Vec<f64>;
fn name(&self) -> &'static str;
}
/// DE/rand/1 strategy
#[derive(Clone, Debug)]
pub struct RandOne;
impl MutationStrategy for RandOne {
fn mutate(
&self,
population: &[Vec<f64>],
target_idx: usize,
f: f64,
rng: &mut impl Rng,
) -> Vec<f64> {
let pop_size = population.len();
let dim = population[0].len();
// Select three distinct random individuals
let mut indices = Vec::new();
while indices.len() < 3 {
let idx = rng.gen_range(0..pop_size);
if idx != target_idx && !indices.contains(&idx) {
indices.push(idx);
}
}
let [r1, r2, r3] = [indices[0], indices[1], indices[2]];
// Mutant = r1 + F * (r2 - r3)
(0..dim)
.map(|d| population[r1][d] + f * (population[r2][d] - population[r3][d]))
.collect()
}
fn name(&self) -> &'static str {
"DE/rand/1"
}
}
/// DE/best/1 strategy
#[derive(Clone, Debug)]
pub struct BestOne {
pub best_idx: usize,
}
impl MutationStrategy for BestOne {
fn mutate(
&self,
population: &[Vec<f64>],
target_idx: usize,
f: f64,
rng: &mut impl Rng,
) -> Vec<f64> {
let pop_size = population.len();
let dim = population[0].len();
// Select two distinct random individuals
let mut indices = Vec::new();
while indices.len() < 2 {
let idx = rng.gen_range(0..pop_size);
if idx != target_idx && idx != self.best_idx && !indices.contains(&idx) {
indices.push(idx);
}
}
let [r1, r2] = [indices[0], indices[1]];
// Mutant = best + F * (r1 - r2)
(0..dim)
.map(|d| population[self.best_idx][d] + f * (population[r1][d] - population[r2][d]))
.collect()
}
fn name(&self) -> &'static str {
"DE/best/1"
}
}
/// DE/rand/2 strategy
#[derive(Clone, Debug)]
pub struct RandTwo;
impl MutationStrategy for RandTwo {
fn mutate(
&self,
population: &[Vec<f64>],
target_idx: usize,
f: f64,
rng: &mut impl Rng,
) -> Vec<f64> {
let pop_size = population.len();
let dim = population[0].len();
// Select five distinct random individuals
let mut indices = Vec::new();
while indices.len() < 5 {
let idx = rng.gen_range(0..pop_size);
if idx != target_idx && !indices.contains(&idx) {
indices.push(idx);
}
}
let [r1, r2, r3, r4, r5] = [indices[0], indices[1], indices[2], indices[3], indices[4]];
// Mutant = r1 + F * (r2 - r3) + F * (r4 - r5)
(0..dim)
.map(|d| {
population[r1][d]
+ f * (population[r2][d] - population[r3][d])
+ f * (population[r4][d] - population[r5][d])
})
.collect()
}
fn name(&self) -> &'static str {
"DE/rand/2"
}
}
/// Generic objective function
pub trait ObjectiveFunction: Send + Sync {
fn evaluate(&self, x: &[f64]) -> f64;
}
/// Wrapper for Python callable
pub struct PyObjectiveFunction {
func: Arc<Py<PyAny>>,
}
impl PyObjectiveFunction {
pub fn new(func: Py<PyAny>) -> Self {
Self {
func: Arc::new(func),
}
}
}
impl ObjectiveFunction for PyObjectiveFunction {
fn evaluate(&self, x: &[f64]) -> f64 {
Python::with_gil(|py| {
let args = (x.to_vec(),);
self.func
.call1(py, args)
.and_then(|res| res.extract::<f64>(py))
.unwrap_or(f64::INFINITY)
})
}
}
/// DE Configuration Builder
#[derive(Clone)]
pub struct DEConfig<M: MutationStrategy> {
pub bounds: Bounds,
pub pop_size: usize,
pub max_generations: usize,
pub mutation_factor: f64,
pub crossover_rate: f64,
pub tolerance: f64,
pub strategy: M,
pub use_parallel: bool,
}
pub struct DEConfigBuilder<M: MutationStrategy> {
bounds: Bounds,
pop_size: Option<usize>,
max_generations: usize,
mutation_factor: f64,
crossover_rate: f64,
tolerance: f64,
strategy: Option<M>,
use_parallel: bool,
}
impl<M: MutationStrategy> DEConfigBuilder<M> {
pub fn new(bounds: Bounds) -> Self {
Self {
bounds,
pop_size: None,
max_generations: 1000,
mutation_factor: 0.8,
crossover_rate: 0.7,
tolerance: 1e-6,
strategy: None,
use_parallel: cfg!(feature = "parallel"),
}
}
pub fn pop_size(mut self, size: usize) -> Self {
self.pop_size = Some(size);
self
}
pub fn max_generations(mut self, gen: usize) -> Self {
self.max_generations = gen;
self
}
pub fn mutation_factor(mut self, f: f64) -> Self {
self.mutation_factor = f;
self
}
pub fn crossover_rate(mut self, cr: f64) -> Self {
self.crossover_rate = cr;
self
}
pub fn tolerance(mut self, tol: f64) -> Self {
self.tolerance = tol;
self
}
pub fn strategy(mut self, strategy: M) -> Self {
self.strategy = Some(strategy);
self
}
pub fn parallel(mut self, enabled: bool) -> Self {
self.use_parallel = enabled && cfg!(feature = "parallel");
self
}
pub fn build(self) -> Result<DEConfig<M>>
where
M: Default,
{
let dim = self.bounds.dim();
let pop_size = self.pop_size.unwrap_or(10 * dim);
if pop_size < 4 {
return Err(OptimizrError::InvalidParameter(
"pop_size must be at least 4".to_string(),
));
}
Ok(DEConfig {
bounds: self.bounds,
pop_size,
max_generations: self.max_generations,
mutation_factor: self.mutation_factor,
crossover_rate: self.crossover_rate,
tolerance: self.tolerance,
strategy: self.strategy.unwrap_or_default(),
use_parallel: self.use_parallel,
})
}
}
impl Default for RandOne {
fn default() -> Self {
RandOne
}
}
impl Default for RandTwo {
fn default() -> Self {
RandTwo
}
}
/// Refactored Differential Evolution
pub struct DifferentialEvolution<M: MutationStrategy, F: ObjectiveFunction> {
pub config: DEConfig<M>,
pub objective: F,
}
impl<M: MutationStrategy, F: ObjectiveFunction> DifferentialEvolution<M, F> {
pub fn new(config: DEConfig<M>, objective: F) -> Self {
Self { config, objective }
}
/// Initialize population
fn initialize_population(&self, rng: &mut impl Rng) -> Vec<Vec<f64>> {
(0..self.config.pop_size)
.map(|_| self.config.bounds.sample(rng))
.collect()
}
/// Evaluate fitness in parallel or sequential
fn evaluate_population(&self, population: &[Vec<f64>]) -> Vec<f64> {
#[cfg(feature = "parallel")]
{
if self.config.use_parallel {
return population
.par_iter()
.map(|ind| self.objective.evaluate(ind))
.collect();
}
}
// Sequential fallback
population
.iter()
.map(|ind| self.objective.evaluate(ind))
.collect()
}
/// Perform crossover
fn crossover(&self, target: &[f64], mutant: &[f64], rng: &mut impl Rng) -> Vec<f64> {
let dim = target.len();
let j_rand = rng.gen_range(0..dim);
(0..dim)
.map(|j| {
if rng.gen::<f64>() < self.config.crossover_rate || j == j_rand {
mutant[j]
} else {
target[j]
}
})
.collect()
}
/// Run optimization
pub fn optimize(&mut self) -> Result<(Vec<f64>, f64)> {
let mut rng = rand::thread_rng();
// Initialize
let mut population = self.initialize_population(&mut rng);
let mut fitness = self.evaluate_population(&population);
let mut best_idx = fitness
.iter()
.enumerate()
.min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
.map(|(i, _)| i)
.unwrap();
let mut best_fitness = fitness[best_idx];
// Evolution loop with functional style
for _generation in 0..self.config.max_generations {
let prev_best = best_fitness;
// Generate trial vectors
let trials: Vec<Vec<f64>> = (0..self.config.pop_size)
.map(|i| {
// Note: BestOne strategy would need special handling here
// In practice, use a mutable reference pattern or Arc<Mutex<>>
// Mutation
let mutant = self.config.strategy.mutate(
&population,
i,
self.config.mutation_factor,
&mut rng,
);
// Crossover
let trial = self.crossover(&population[i], &mutant, &mut rng);
// Clip to bounds
self.config.bounds.clip(&trial)
})
.collect();
// Evaluate trials
let trial_fitness = self.evaluate_population(&trials);
// Selection
for i in 0..self.config.pop_size {
if trial_fitness[i] < fitness[i] {
population[i] = trials[i].clone();
fitness[i] = trial_fitness[i];
if trial_fitness[i] < best_fitness {
best_idx = i;
best_fitness = trial_fitness[i];
}
}
}
// Check convergence
if (best_fitness - prev_best).abs() < self.config.tolerance {
break;
}
}
Ok((population[best_idx].clone(), best_fitness))
}
}
impl<M: MutationStrategy + 'static, F: ObjectiveFunction + 'static> Optimizer
for DifferentialEvolution<M, F>
{
type Config = DEConfig<M>;
type Output = (Vec<f64>, f64);
fn optimize(&mut self) -> Result<Self::Output> {
self.optimize()
}
fn best(&self) -> Result<Vec<f64>> {
// Note: This requires re-optimization. In production, cache the best solution.
Err(OptimizrError::ComputationError(
"Call optimize() to get the best solution".to_string(),
))
}
}
// Python bindings
#[pyclass]
#[derive(Clone, Debug)]
pub struct DEResult {
#[pyo3(get)]
pub best_solution: Vec<f64>,
#[pyo3(get)]
pub best_value: f64,
}
#[pyfunction]
#[pyo3(signature = (objective_fn, bounds, pop_size=None, max_generations=1000, mutation_factor=0.8, crossover_rate=0.7, strategy="rand1"))]
pub fn differential_evolution(
objective_fn: Py<PyAny>,
bounds: Vec<(f64, f64)>,
pop_size: Option<usize>,
max_generations: usize,
mutation_factor: f64,
crossover_rate: f64,
strategy: &str,
) -> PyResult<DEResult> {
let bounds = Bounds::new(bounds)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let objective = PyObjectiveFunction::new(objective_fn);
// Select strategy
match strategy {
"rand1" | "DE/rand/1" => {
let mut builder = DEConfigBuilder::new(bounds)
.max_generations(max_generations)
.mutation_factor(mutation_factor)
.crossover_rate(crossover_rate)
.strategy(RandOne);
if let Some(ps) = pop_size {
builder = builder.pop_size(ps);
}
let config = builder
.build()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
let mut optimizer = DifferentialEvolution::new(config, objective);
let (best_solution, best_value) = optimizer
.optimize()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(DEResult {
best_solution,
best_value,
})
}
"rand2" | "DE/rand/2" => {
let mut builder = DEConfigBuilder::new(bounds)
.max_generations(max_generations)
.mutation_factor(mutation_factor)
.crossover_rate(crossover_rate)
.strategy(RandTwo);
if let Some(ps) = pop_size {
builder = builder.pop_size(ps);
}
let config = builder
.build()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
let mut optimizer = DifferentialEvolution::new(config, objective);
let (best_solution, best_value) = optimizer
.optimize()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(DEResult {
best_solution,
best_value,
})
}
_ => Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
"Unknown strategy: {}. Use 'rand1', 'rand2', or 'best1'",
strategy
))),
}
}
#[cfg(test)]
mod tests {
use super::*;
struct SphereFunction;
impl ObjectiveFunction for SphereFunction {
fn evaluate(&self, x: &[f64]) -> f64 {
x.iter().map(|xi| xi.powi(2)).sum()
}
}
#[test]
fn test_de_builder() {
let bounds = Bounds::new(vec![(-5.0, 5.0), (-5.0, 5.0)]).unwrap();
let config = DEConfigBuilder::<RandOne>::new(bounds)
.pop_size(40)
.max_generations(100)
.build()
.unwrap();
assert_eq!(config.pop_size, 40);
assert_eq!(config.max_generations, 100);
}
#[test]
fn test_de_optimization() {
let bounds = Bounds::new(vec![(-5.0, 5.0), (-5.0, 5.0)]).unwrap();
let config = DEConfigBuilder::<RandOne>::new(bounds)
.pop_size(20)
.max_generations(50)
.build()
.unwrap();
let objective = SphereFunction;
let mut optimizer = DifferentialEvolution::new(config, objective);
let (_best, fitness) = optimizer.optimize().unwrap();
assert!(fitness < 0.1); // Should converge close to 0
}
}
src/functional.rs
+204
View File
@@ -0,0 +1,204 @@
//! Trait-based functional utilities for OptimizR
//!
//! This module provides functional programming utilities like composition,
//! monadic operations, and higher-order functions.
use crate::core::{OptimizrError, Result};
/// Function composition trait
pub trait Compose<A, B, C>: Sized {
fn compose<G>(self, g: G) -> impl Fn(A) -> C
where
G: Fn(B) -> C,
Self: Fn(A) -> B;
}
impl<F, A, B, C> Compose<A, B, C> for F
where
F: Fn(A) -> B,
{
fn compose<G>(self, g: G) -> impl Fn(A) -> C
where
G: Fn(B) -> C,
{
move |x| g(self(x))
}
}
/// Monadic operations for Result
pub trait ResultExt<T> {
/// Apply a function if Ok, short-circuit on Err
fn and_then_log<F, U>(self, f: F, msg: &str) -> Result<U>
where
F: FnOnce(T) -> Result<U>;
/// Map with context
fn map_context<F, U>(self, f: F, ctx: &str) -> Result<U>
where
F: FnOnce(T) -> U;
}
impl<T> ResultExt<T> for Result<T> {
fn and_then_log<F, U>(self, f: F, msg: &str) -> Result<U>
where
F: FnOnce(T) -> Result<U>,
{
match self {
Ok(val) => f(val),
Err(e) => {
eprintln!("Error at {}: {:?}", msg, e);
Err(e)
}
}
}
fn map_context<F, U>(self, f: F, ctx: &str) -> Result<U>
where
F: FnOnce(T) -> U,
{
self.map(f).map_err(|e| {
OptimizrError::ComputationError(format!("{}: {}", ctx, e))
})
}
}
/// Retry logic for operations
pub fn retry<F, T>(mut f: F, max_attempts: usize) -> Result<T>
where
F: FnMut() -> Result<T>,
{
let mut last_error = None;
for _ in 0..max_attempts {
match f() {
Ok(val) => return Ok(val),
Err(e) => last_error = Some(e),
}
}
Err(last_error.unwrap_or_else(|| {
OptimizrError::ComputationError("All retry attempts failed".to_string())
}))
}
/// Memoization for expensive computations
pub struct Memoized<F, T>
where
F: Fn(&[f64]) -> T,
{
f: F,
cache: std::sync::Mutex<std::collections::HashMap<Vec<ordered_float::OrderedFloat<f64>>, T>>,
}
impl<F, T> Memoized<F, T>
where
F: Fn(&[f64]) -> T,
T: Clone,
{
pub fn new(f: F) -> Self {
Self {
f,
cache: std::sync::Mutex::new(std::collections::HashMap::new()),
}
}
pub fn call(&self, x: &[f64]) -> T {
let key: Vec<_> = x.iter().map(|&v| ordered_float::OrderedFloat(v)).collect();
let mut cache = self.cache.lock().unwrap();
if let Some(cached) = cache.get(&key) {
return cached.clone();
}
let result = (self.f)(x);
cache.insert(key, result.clone());
result
}
}
/// Lazy evaluation wrapper
pub struct Lazy<T, F>
where
F: FnOnce() -> T,
{
f: Option<F>,
value: Option<T>,
}
impl<T, F> Lazy<T, F>
where
F: FnOnce() -> T,
{
pub fn new(f: F) -> Self {
Self {
f: Some(f),
value: None,
}
}
pub fn force(&mut self) -> &T {
if self.value.is_none() {
let f = self.f.take().unwrap();
self.value = Some(f());
}
self.value.as_ref().unwrap()
}
}
/// Piping operator - allows chaining operations
pub trait Pipe: Sized {
fn pipe<F, R>(self, f: F) -> R
where
F: FnOnce(Self) -> R,
{
f(self)
}
}
impl<T> Pipe for T {}
/// Currying utilities
/// Note: Simplified version due to Rust's ownership constraints
/// For full currying, use the partial function instead
pub fn curry2<A, B, R, F>(f: F) -> impl Fn((A, B)) -> R
where
F: Fn(A, B) -> R + 'static,
A: 'static,
B: 'static,
R: 'static,
{
move |(a, b)| f(a, b)
}
/// Partial application
pub fn partial<A: Clone + 'static, B, R, F>(f: F, a: A) -> impl Fn(B) -> R
where
F: Fn(A, B) -> R + 'static,
B: 'static,
R: 'static,
{
move |b| f(a.clone(), b)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pipe() {
let result = vec![1, 2, 3]
.pipe(|v| v.into_iter().map(|x| x * 2).collect::<Vec<_>>())
.pipe(|v: Vec<_>| v.into_iter().sum::<i32>());
assert_eq!(result, 12);
}
#[test]
fn test_partial() {
let add = |a: i32, b: i32| a + b;
let add5 = partial(add, 5);
assert_eq!(add5(3), 8);
}
}
src/hmm_refactored.rs
+582
View File
@@ -0,0 +1,582 @@
//! Refactored Hidden Markov Model with trait-based design
//!
//! This module provides a more modular, functional, and trait-based implementation
//! of HMMs with support for different emission models and parallel computation.
use crate::core::{OptimizrError, Result};
use pyo3::prelude::*;
use std::f64;
#[cfg(feature = "parallel")]
use rayon::prelude::*;
/// 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
#[derive(Clone, Debug)]
pub struct GaussianEmission {
pub means: Vec<f64>,
pub stds: Vec<f64>,
}
impl GaussianEmission {
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::<f64>()
/ sum_weights;
// Weighted variance
let var = observations
.iter()
.zip(weights.iter())
.map(|(obs, w)| w * (obs - mean).powi(2))
.sum::<f64>()
/ 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::<f64>() / segment.len() as f64;
let var: f64 = segment
.iter()
.map(|x| (x - self.means[state]).powi(2))
.sum::<f64>()
/ segment.len() as f64;
self.stds[state] = var.sqrt().max(1e-6);
}
Ok(())
}
fn n_states(&self) -> usize {
self.means.len()
}
}
/// HMM Configuration Builder
#[derive(Clone)]
pub struct HMMConfig<E: EmissionModel> {
pub n_states: usize,
pub n_iterations: usize,
pub tolerance: f64,
pub emission_model: E,
pub use_parallel: bool,
}
impl<E: EmissionModel> HMMConfig<E> {
pub fn builder(n_states: usize) -> HMMConfigBuilder<E> {
HMMConfigBuilder::new(n_states)
}
}
/// Builder pattern for HMM configuration
pub struct HMMConfigBuilder<E: EmissionModel> {
n_states: usize,
n_iterations: usize,
tolerance: f64,
emission_model: Option<E>,
use_parallel: bool,
}
impl<E: EmissionModel> HMMConfigBuilder<E> {
pub fn new(n_states: usize) -> Self {
Self {
n_states,
n_iterations: 100,
tolerance: 1e-6,
emission_model: None,
use_parallel: cfg!(feature = "parallel"),
}
}
pub fn iterations(mut self, n: usize) -> Self {
self.n_iterations = n;
self
}
pub fn tolerance(mut self, tol: f64) -> Self {
self.tolerance = tol;
self
}
pub fn emission_model(mut self, model: E) -> Self {
self.emission_model = Some(model);
self
}
pub fn parallel(mut self, enabled: bool) -> Self {
self.use_parallel = enabled && cfg!(feature = "parallel");
self
}
pub fn build(self) -> Result<HMMConfig<E>>
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,
})
}
}
impl Default for GaussianEmission {
fn default() -> Self {
Self::new(2)
}
}
/// Refactored HMM with generic emission model
pub struct HMM<E: EmissionModel> {
pub config: HMMConfig<E>,
pub transition_matrix: Vec<Vec<f64>>,
pub initial_probs: Vec<f64>,
}
impl<E: EmissionModel> HMM<E> {
pub fn new(config: HMMConfig<E>) -> 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 functional pipeline
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 with functional approach
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 with parallel option
fn forward(&self, observations: &[f64]) -> Result<Vec<Vec<f64>>> {
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 (sequential for dependencies)
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
fn backward(&self, observations: &[f64]) -> Result<Vec<Vec<f64>>> {
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 (pure function)
fn compute_gamma(alpha: &[Vec<f64>], beta: &[Vec<f64>]) -> Vec<Vec<f64>> {
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
fn compute_xi(
&self,
observations: &[f64],
alpha: &[Vec<f64>],
beta: &[Vec<f64>],
) -> Result<Vec<Vec<Vec<f64>>>> {
let n_obs = observations.len();
let n_states = self.config.n_states;
let xi: Vec<Vec<Vec<f64>>> = (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)
}
/// Update parameters using functional patterns
fn update_parameters(
&mut self,
observations: &[f64],
gamma: &[Vec<f64>],
xi: &[Vec<Vec<f64>>],
) -> Result<()> {
let n_obs = observations.len();
let n_states = self.config.n_states;
// Update transitions
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 emissions
for s in 0..n_states {
let weights: Vec<f64> = gamma.iter().map(|g| g[s]).collect();
self.config
.emission_model
.update(observations, &weights, s)?;
}
Ok(())
}
/// Viterbi decoding with functional style
pub fn viterbi(&self, observations: &[f64]) -> Result<Vec<usize>> {
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)
}
// Helper functions
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);
}
}
fn compute_log_likelihood(alpha: &[Vec<f64>]) -> f64 {
alpha.last().unwrap().iter().sum::<f64>().max(1e-10).ln()
}
}
// Python bindings remain similar but use the new modular structure
#[pyclass]
#[derive(Clone, Debug)]
pub struct HMMParams {
#[pyo3(get, set)]
pub n_states: usize,
#[pyo3(get, set)]
pub transition_matrix: Vec<Vec<f64>>,
#[pyo3(get, set)]
pub emission_means: Vec<f64>,
#[pyo3(get, set)]
pub emission_stds: Vec<f64>,
#[pyo3(get, set)]
pub initial_probs: Vec<f64>,
}
#[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
)
}
}
#[pyfunction]
#[pyo3(signature = (observations, n_states, n_iterations=100, tolerance=1e-6))]
pub fn fit_hmm(
observations: Vec<f64>,
n_states: usize,
n_iterations: usize,
tolerance: f64,
) -> PyResult<HMMParams> {
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::<pyo3::exceptions::PyRuntimeError, _>(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,
})
}
#[pyfunction]
pub fn viterbi_decode(observations: Vec<f64>, params: HMMParams) -> PyResult<Vec<usize>> {
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::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hmm_builder() {
let config = HMMConfig::<GaussianEmission>::builder(3)
.iterations(50)
.tolerance(1e-5)
.build()
.unwrap();
assert_eq!(config.n_states, 3);
assert_eq!(config.n_iterations, 50);
}
#[test]
fn test_hmm_fit() {
let observations: Vec<f64> = (0..100).map(|i| (i as f64 * 0.1).sin()).collect();
let config = HMMConfig::<GaussianEmission>::builder(2).build().unwrap();
let mut hmm = HMM::new(config);
assert!(hmm.fit(&observations).is_ok());
}
}
src/lib.rs
+41 -4
View File
@@ -4,6 +4,16 @@
//! This library provides fast, reliable implementations of advanced optimization
//! and statistical inference algorithms, with Python bindings via PyO3.
//!
//! # Architecture
//!
//! The library is designed with modularity, functional programming patterns,
//! and trait-based abstractions:
//!
//! - `core`: Core traits (Optimizer, Sampler, InformationMeasure) and error types
//! - `functional`: Functional programming utilities (composition, memoization, pipes)
//! - Refactored modules with trait-based design and parallel support
//! - Original modules maintained for backward compatibility
//!
//! # Modules
//!
//! - `hmm`: Hidden Markov Model training and inference
@@ -15,6 +25,16 @@
use pyo3::prelude::*;
use pyo3::types::PyModule;
// Core modules with trait-based architecture
pub mod core;
pub mod functional;
// Refactored modules with advanced patterns
pub mod hmm_refactored;
pub mod mcmc_refactored;
pub mod de_refactored;
// Original modules for backward compatibility
mod hmm;
mod mcmc;
mod differential_evolution;
@@ -24,21 +44,38 @@ mod information_theory;
/// OptimizR Python module
#[pymodule]
fn _core(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
// Register HMM functions
// ===== Original API (Backward Compatible) =====
// HMM functions
m.add_class::<hmm::HMMParams>()?;
m.add_function(wrap_pyfunction!(hmm::fit_hmm, m)?)?;
m.add_function(wrap_pyfunction!(hmm::viterbi_decode, m)?)?;
// Register MCMC functions
// MCMC functions
m.add_function(wrap_pyfunction!(mcmc::mcmc_sample, m)?)?;
// Register optimization functions
// Optimization functions
m.add_function(wrap_pyfunction!(differential_evolution::differential_evolution, m)?)?;
m.add_function(wrap_pyfunction!(grid_search::grid_search, m)?)?;
// Register information theory functions
// Information theory functions
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::<hmm_refactored::HMMParams>()?;
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::<de_refactored::DEResult>()?;
m.add_function(wrap_pyfunction!(de_refactored::differential_evolution, m)?)?;
Ok(())
}
src/mcmc_refactored.rs
+447
View File
@@ -0,0 +1,447 @@
//! Refactored MCMC with Strategy Pattern
//!
//! Supports multiple proposal strategies and parallel chains.
use crate::core::{OptimizrError, Result, Sampler, SamplerDiagnostics};
use pyo3::prelude::*;
use rand::distributions::Distribution;
use rand::Rng;
use rand_distr::Normal;
/// Trait for proposal strategies
pub trait ProposalStrategy: Send + Sync + Clone {
/// Generate proposed next state
fn propose(&self, current: &[f64], rng: &mut impl Rng) -> Vec<f64>;
/// 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<f64> {
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"
}
}
/// Adaptive proposal that adjusts step size
#[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,
}
}
}
impl ProposalStrategy for AdaptiveProposal {
fn propose(&self, current: &[f64], rng: &mut impl Rng) -> Vec<f64> {
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"
}
}
/// MCMC Configuration Builder
#[derive(Clone)]
pub struct MCMCConfig<P: ProposalStrategy> {
pub n_samples: usize,
pub burn_in: usize,
pub thin: usize,
pub initial_state: Vec<f64>,
pub proposal: P,
pub adaptation_interval: usize,
}
pub struct MCMCConfigBuilder<P: ProposalStrategy> {
n_samples: usize,
burn_in: usize,
thin: usize,
initial_state: Vec<f64>,
proposal: Option<P>,
adaptation_interval: usize,
}
impl<P: ProposalStrategy> MCMCConfigBuilder<P> {
pub fn new(n_samples: usize, initial_state: Vec<f64>) -> 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<MCMCConfig<P>>
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,
})
}
}
impl Default for GaussianProposal {
fn default() -> Self {
Self::new(0.1)
}
}
impl Default for AdaptiveProposal {
fn default() -> Self {
Self::new(0.1)
}
}
/// Generic log-likelihood function
pub trait LogLikelihood: Send + Sync {
fn evaluate(&self, state: &[f64]) -> f64;
}
/// Wrapper for Python callable
pub struct PyLogLikelihood {
func: Py<PyAny>,
}
impl PyLogLikelihood {
pub fn new(func: Py<PyAny>) -> 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::<f64>(py))
.unwrap_or(f64::NEG_INFINITY)
})
}
}
/// Refactored MCMC Sampler
pub struct MetropolisHastings<P: ProposalStrategy, L: LogLikelihood> {
pub config: MCMCConfig<P>,
pub log_likelihood: L,
}
impl<P: ProposalStrategy, L: LogLikelihood> MetropolisHastings<P, L> {
pub fn new(config: MCMCConfig<P>, log_likelihood: L) -> Self {
Self {
config,
log_likelihood,
}
}
/// Run single chain with functional composition
pub fn sample_chain(&mut self) -> Result<Vec<Vec<f64>>> {
let mut rng = rand::thread_rng();
let mut current_state = self.config.initial_state.clone();
let mut current_ll = self.log_likelihood.evaluate(&current_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(&current_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::<f64>() < 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
pub fn diagnostics(&self, samples: &[Vec<f64>]) -> Result<SamplerDiagnostics> {
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<f64> = (0..dim)
.map(|d| samples.iter().map(|s| s[d]).sum::<f64>() / n_samples as f64)
.collect();
let variances: Vec<f64> = (0..dim)
.map(|d| {
let mean = means[d];
samples
.iter()
.map(|s| (s[d] - mean).powi(2))
.sum::<f64>()
/ (n_samples - 1) as f64
})
.collect();
// Compute autocorrelations (lag 1)
let autocorrs: Vec<f64> = (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::<f64>()
/ (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<P: ProposalStrategy + 'static, L: LogLikelihood + 'static> Sampler
for MetropolisHastings<P, L>
{
type Config = MCMCConfig<P>;
type Output = Vec<Vec<f64>>;
fn sample(&mut self) -> Result<Self::Output> {
self.sample_chain()
}
fn diagnostics(&self, samples: &Self::Output) -> Result<SamplerDiagnostics> {
self.diagnostics(samples)
}
}
// Python bindings
#[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<PyAny>,
initial_state: Vec<f64>,
n_samples: usize,
step_size: f64,
burn_in: Option<usize>,
) -> PyResult<Vec<Vec<f64>>> {
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::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[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<PyAny>,
initial_state: Vec<f64>,
n_samples: usize,
initial_step: f64,
burn_in: Option<usize>,
) -> PyResult<Vec<Vec<f64>>> {
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::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[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::<f64>()
}
}
#[test]
fn test_mcmc_builder() {
let config = MCMCConfigBuilder::<GaussianProposal>::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_mcmc_sampling() {
let config = MCMCConfigBuilder::<GaussianProposal>::new(100, vec![0.0])
.proposal(GaussianProposal::new(0.5))
.build()
.unwrap();
let log_likelihood = TestLogLikelihood;
let mut sampler = MetropolisHastings::new(config, log_likelihood);
let samples = sampler.sample_chain().unwrap();
assert!(!samples.is_empty());
}
#[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);
}
}