quantwave_core/regimes/
hmm.rs1use crate::regimes::MarketRegime;
11use crate::traits::Next;
12use serde::{Deserialize, Serialize};
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct HMM {
17 n_states: usize,
18 a: Vec<Vec<f64>>,
20 means: Vec<f64>,
22 stds: Vec<f64>,
24 pi: Vec<f64>,
26 last_delta: Vec<f64>,
28 initialized: bool,
29}
30
31impl HMM {
32 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 pub fn bull_bear() -> Self {
48 Self::new(
49 vec![
50 vec![0.95, 0.05], vec![0.10, 0.90], ],
53 vec![0.001, -0.002], vec![0.01, 0.02], vec![0.5, 0.5],
56 )
57 }
58
59 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 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 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}