Skip to main content

quantwave_core/regimes/
hmm.rs

1//! Hidden Markov Models (Hamilton 1989)
2//!
3//! Source: Hamilton, J. D. (1989).
4//! "A New Approach to the Economic Analysis of Nonstationary Time Series and the Business Cycle."
5//! Econometrica, 57(2), 357-384.
6//!
7//! Implementation of a regime-switching Hidden Markov Model with Gaussian emissions.
8//! Includes the Viterbi algorithm for online state decoding and placeholder for Baum-Welch training.
9
10use crate::regimes::MarketRegime;
11use crate::traits::Next;
12use serde::{Deserialize, Serialize};
13
14/// A Hidden Markov Model for regime detection.
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct HMM {
17    n_states: usize,
18    /// Transition probability matrix [from_state][to_state]
19    a: Vec<Vec<f64>>,
20    /// Emission means for each state
21    means: Vec<f64>,
22    /// Emission standard deviations for each state
23    stds: Vec<f64>,
24    /// Initial state probabilities
25    pi: Vec<f64>,
26    /// Last Viterbi log-probabilities for each state
27    last_delta: Vec<f64>,
28    initialized: bool,
29}
30
31impl HMM {
32    /// Creates a new HMM with pre-defined parameters.
33    pub fn new(a: Vec<Vec<f64>>, means: Vec<f64>, stds: Vec<f64>, pi: Vec<f64>) -> Self {
34        let n_states = a.len();
35        Self {
36            n_states,
37            a,
38            means,
39            stds,
40            pi,
41            last_delta: vec![0.0; n_states],
42            initialized: false,
43        }
44    }
45
46    /// Default 2-state HMM (Bull/Bear)
47    pub fn bull_bear() -> Self {
48        Self::new(
49            vec![
50                vec![0.95, 0.05], // Bull -> Bull, Bull -> Bear
51                vec![0.10, 0.90], // Bear -> Bull, Bear -> Bear
52            ],
53            vec![0.001, -0.002], // Daily returns: 0.1% for Bull, -0.2% for Bear
54            vec![0.01, 0.02],    // Volatility: 1% for Bull, 2% for Bear
55            vec![0.5, 0.5],
56        )
57    }
58
59    /// Calculate Gaussian PDF: P(x | mu, sigma)
60    fn gaussian_pdf(x: f64, mu: f64, sigma: f64) -> f64 {
61        let variance = sigma * sigma;
62        let denom = (2.0 * std::f64::consts::PI * variance).sqrt();
63        let exponent = -((x - mu).powi(2)) / (2.0 * variance);
64        exponent.exp() / denom
65    }
66
67    /// Normalized state probabilities from Viterbi log-deltas (for soft regime features).
68    pub fn state_probabilities(&self) -> Vec<f64> {
69        if !self.initialized {
70            return self.pi.clone();
71        }
72        let max_log = self
73            .last_delta
74            .iter()
75            .cloned()
76            .fold(f64::NEG_INFINITY, f64::max);
77        let mut probs: Vec<f64> = self
78            .last_delta
79            .iter()
80            .map(|&d| (d - max_log).exp())
81            .collect();
82        let sum: f64 = probs.iter().sum();
83        if sum > 0.0 {
84            for p in &mut probs {
85                *p /= sum;
86            }
87        }
88        probs
89    }
90}
91
92impl Next<f64> for HMM {
93    type Output = MarketRegime;
94
95    fn next(&mut self, x: f64) -> Self::Output {
96        let mut next_delta = vec![0.0; self.n_states];
97        let mut best_state = 0;
98        let mut max_prob = -f64::INFINITY;
99
100        if !self.initialized {
101            for i in 0..self.n_states {
102                let emission = Self::gaussian_pdf(x, self.means[i], self.stds[i]);
103                next_delta[i] = (self.pi[i] * emission).ln();
104                if next_delta[i] > max_prob {
105                    max_prob = next_delta[i];
106                    best_state = i;
107                }
108            }
109            self.initialized = true;
110        } else {
111            for j in 0..self.n_states {
112                let mut max_prev = -f64::INFINITY;
113                for i in 0..self.n_states {
114                    // delta_j(t) = max_i [ delta_i(t-1) + ln(A_ij) ] + ln(P(x|j))
115                    let prob = self.last_delta[i] + self.a[i][j].ln();
116                    if prob > max_prev {
117                        max_prev = prob;
118                    }
119                }
120                let emission = Self::gaussian_pdf(x, self.means[j], self.stds[j]);
121                next_delta[j] = max_prev + emission.ln();
122
123                if next_delta[j] > max_prob {
124                    max_prob = next_delta[j];
125                    best_state = j;
126                }
127            }
128        }
129
130        self.last_delta = next_delta;
131
132        match best_state {
133            0 => MarketRegime::Bull,
134            1 => MarketRegime::Bear,
135            _ => MarketRegime::Cluster(best_state as u8),
136        }
137    }
138}