quantwave_core/regimes/
hsmm.rs1use crate::regimes::MarketRegime;
9use crate::traits::Next;
10use serde::{Deserialize, Serialize};
11use statrs::distribution::{Discrete, Poisson};
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
15pub enum DurationDistribution {
16 Poisson { lambda: f64 },
18 Fixed { duration: usize },
20}
21
22impl DurationDistribution {
23 pub fn p(&self, d: usize) -> f64 {
24 match self {
25 Self::Poisson { lambda } => Poisson::new(*lambda)
26 .map(|dist| dist.pmf(d as u64))
27 .unwrap_or(0.0),
28 Self::Fixed { duration } => {
29 if d == *duration {
30 1.0
31 } else {
32 0.0
33 }
34 }
35 }
36 }
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct HSMM {
42 pub n_states: usize,
43 pub a: Vec<Vec<f64>>,
45 pub means: Vec<f64>,
46 pub stds: Vec<f64>,
47 pub durations: Vec<DurationDistribution>,
48 current_duration: usize,
50 last_state: usize,
51 initialized: bool,
52}
53
54impl HSMM {
55 pub fn new(
56 a: Vec<Vec<f64>>,
57 means: Vec<f64>,
58 stds: Vec<f64>,
59 durations: Vec<DurationDistribution>,
60 ) -> Self {
61 Self {
62 n_states: a.len(),
63 a,
64 means,
65 stds,
66 durations,
67 current_duration: 0,
68 last_state: 0,
69 initialized: false,
70 }
71 }
72
73 fn gaussian_pdf(x: f64, mu: f64, sigma: f64) -> f64 {
74 let variance = sigma * sigma;
75 let denom = (2.0 * std::f64::consts::PI * variance).sqrt();
76 let exponent = -((x - mu).powi(2)) / (2.0 * variance);
77 exponent.exp() / denom
78 }
79}
80
81impl Next<f64> for HSMM {
82 type Output = MarketRegime;
83
84 fn next(&mut self, x: f64) -> Self::Output {
85 if !self.initialized {
86 self.initialized = true;
87 return MarketRegime::Steady;
89 }
90
91 self.current_duration += 1;
92
93 let prob_stay = self.durations[self.last_state].p(self.current_duration);
95
96 let mut max_prob;
97 let mut best_state = self.last_state;
98
99 let emission_stay =
101 Self::gaussian_pdf(x, self.means[self.last_state], self.stds[self.last_state]);
102 max_prob = prob_stay * emission_stay;
103
104 for j in 0..self.n_states {
106 if j == self.last_state {
107 continue;
108 }
109
110 let transition_prob = self.a[self.last_state][j];
111 let emission_j = Self::gaussian_pdf(x, self.means[j], self.stds[j]);
112 let prob_j = (1.0 - prob_stay) * transition_prob * self.durations[j].p(1) * emission_j;
114
115 if prob_j > max_prob {
116 max_prob = prob_j;
117 best_state = j;
118 }
119 }
120
121 if best_state != self.last_state {
122 self.last_state = best_state;
123 self.current_duration = 1;
124 }
125
126 match best_state {
127 0 => MarketRegime::Steady,
128 1 => MarketRegime::Crisis,
129 _ => MarketRegime::Cluster(best_state as u8),
130 }
131 }
132}