Skip to main content

quantwave_core/regimes/
gmm.rs

1//! Gaussian Mixture Models (Two Sigma 2021)
2//!
3//! Source: Two Sigma (2021). "A Machine Learning Approach to Regime Modeling."
4//! Foundational EM Algorithm: Dempster, A. P., Laird, N. M., & Rubin, D. B. (1977). 
5//! "Maximum Likelihood from Incomplete Data via the EM Algorithm." 
6//! Journal of the Royal Statistical Society: Series B (Methodological), 39(1), 1-22.
7//!
8//! Multi-variate clustering for latent market states using the Expectation-Maximization (EM) algorithm.
9//! This implementation uses diagonal covariance matrices for efficiency.
10
11use crate::traits::Next;
12use crate::regimes::MarketRegime;
13use serde::{Deserialize, Serialize};
14
15/// A Gaussian Mixture Model for multi-factor regime detection.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct GMM {
18    k: usize,
19    dims: usize,
20    /// Means for each component [k][dim]
21    means: Vec<Vec<f64>>,
22    /// Variances for each component [k][dim] (diagonal covariance)
23    vars: Vec<Vec<f64>>,
24    /// Mixing coefficients
25    weights: Vec<f64>,
26}
27
28impl GMM {
29    /// Creates a new GMM with pre-defined parameters.
30    pub fn new(means: Vec<Vec<f64>>, vars: Vec<Vec<f64>>, weights: Vec<f64>) -> Self {
31        let k = means.len();
32        let dims = means[0].len();
33        Self { k, dims, means, vars, weights }
34    }
35
36    /// Calculate multivariate Gaussian PDF (diagonal covariance)
37    fn pdf(&self, x: &[f64], k_idx: usize) -> f64 {
38        let mut prob = 1.0;
39        for d in 0..self.dims {
40            let mu = self.means[k_idx][d];
41            let var = self.vars[k_idx][d].max(1e-9);
42            let denom = (2.0 * std::f64::consts::PI * var).sqrt();
43            let exponent = -((x[d] - mu).powi(2)) / (2.0 * var);
44            prob *= exponent.exp() / denom;
45        }
46        prob
47    }
48
49    /// Batch fit using EM algorithm (simplified placeholder)
50    /// In a real implementation, this would iterate until convergence.
51    pub fn fit(&mut self, _data: &[Vec<f64>], _max_iter: usize) {
52        // TODO: Implement EM algorithm
53        // 1. E-step: Responsibilities
54        // 2. M-step: Update means, vars, weights
55    }
56}
57
58impl Next<&[f64]> for GMM {
59    type Output = MarketRegime;
60
61    fn next(&mut self, x: &[f64]) -> Self::Output {
62        let mut max_prob = -1.0;
63        let mut best_k = 0;
64
65        for k in 0..self.k {
66            let p = self.weights[k] * self.pdf(x, k);
67            if p > max_prob {
68                max_prob = p;
69                best_k = k;
70            }
71        }
72
73        match best_k {
74            0 => MarketRegime::Steady,
75            k if k == self.k - 1 => MarketRegime::Crisis,
76            _ => MarketRegime::Cluster(best_k as u8),
77        }
78    }
79}