Skip to main content

quantwave_core/regimes/
hsmm.rs

1//! Hidden Semi-Markov Models (HSMM)
2//!
3//! Source: Yu (2010) "Hidden Semi-Markov Models"
4//!
5//! HSMMs extend HMMs by allowing explicit modeling of the duration of each state,
6//! avoiding the geometric duration distribution inherent in standard HMMs.
7
8use crate::regimes::MarketRegime;
9use crate::traits::Next;
10use serde::{Deserialize, Serialize};
11use statrs::distribution::{Discrete, Poisson};
12
13/// A distribution for the duration of a hidden state.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub enum DurationDistribution {
16    /// Poisson distribution with mean lambda
17    Poisson { lambda: f64 },
18    /// Fixed duration
19    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/// A Hidden Semi-Markov Model.
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct HSMM {
42    pub n_states: usize,
43    /// Transition matrix (diagonal must be zero)
44    pub a: Vec<Vec<f64>>,
45    pub means: Vec<f64>,
46    pub stds: Vec<f64>,
47    pub durations: Vec<DurationDistribution>,
48    /// Time spent in current state
49    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            // Start in first state by default
88            return MarketRegime::Steady;
89        }
90
91        self.current_duration += 1;
92
93        // Calculate probability of staying vs switching
94        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        // 1. Evaluate staying in current state
100        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        // 2. Evaluate switching to other states
105        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            // Probability of starting new state j (duration 1)
113            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}