chronos_ts/
statespace.rs

1use crate::errors::{ChronosError, Result};
2use crate::linalg;
3use ndarray::{Array1, Array2, Axis};
4use serde::{Deserialize, Serialize};
5
6/// State-space specification:
7/// Y_t = Z * alpha_t + eps_t,        eps_t ~ N(0, H)
8/// alpha_{t+1} = T * alpha_t + R * eta_t,  eta_t ~ N(0, Q)
9#[derive(Serialize, Deserialize)]
10pub struct StateSpaceModel {
11    #[serde(with = "crate::utils::serde_array2")]
12    pub transition_matrix: Array2<f64>, // T (m x m)
13    #[serde(with = "crate::utils::serde_array2")]
14    pub selection_matrix: Array2<f64>, // R (m x r)
15    #[serde(with = "crate::utils::serde_array1")]
16    pub design_matrix: Array1<f64>, // Z (m)
17    #[serde(with = "crate::utils::serde_array2")]
18    pub state_cov: Array2<f64>, // Q (r x r)
19    pub obs_cov: f64, // H (scalar)
20}
21
22#[derive(Serialize, Deserialize)]
23pub struct FilterStepResult {
24    #[serde(with = "crate::utils::serde_array1")]
25    pub a_prior: Array1<f64>, // a_{t|t-1}
26    #[serde(with = "crate::utils::serde_array2")]
27    pub p_prior: Array2<f64>, // P_{t|t-1}
28    #[serde(with = "crate::utils::serde_array1")]
29    pub a_post: Array1<f64>, // a_{t|t}
30    #[serde(with = "crate::utils::serde_array2")]
31    pub p_post: Array2<f64>, // P_{t|t}
32    pub v: f64, // Innovation v_t
33    pub f: f64, // Innovation variance F_t
34    #[serde(with = "crate::utils::serde_array1")]
35    pub k: Array1<f64>, // Kalman gain K_t
36}
37
38pub struct KalmanFilterResult {
39    pub steps: Vec<FilterStepResult>,
40    pub log_likelihood: f64,
41}
42
43pub struct SmootherResult {
44    pub smoothed_states: Vec<Array1<f64>>, // a_{t|N}
45    pub smoothed_covs: Vec<Array2<f64>>,   // P_{t|N}
46}
47
48pub struct KalmanFilter<'a> {
49    pub model: &'a StateSpaceModel,
50}
51
52impl StateSpaceModel {
53    pub fn local_linear_trend(sigma_obs: f64, sigma_level: f64, sigma_trend: f64) -> Self {
54        use ndarray::{array, Array2};
55
56        // 1D observation vector (Z): [1.0, 0.0]
57        let design_matrix = array![1.0, 0.0];
58
59        // 2D transition matrix (T): [[1.0, 1.0], [0.0, 1.0]]
60        let transition_matrix = array![[1.0, 1.0], [0.0, 1.0]];
61
62        // Selection matrix (R): Identity matrix mapping state error terms
63        let selection_matrix = Array2::eye(2);
64
65        // State covariance matrix (Q)
66        let state_cov = array![[sigma_level.powi(2), 0.0], [0.0, sigma_trend.powi(2)]];
67
68        // Observation covariance variance scalar (sigma^2_obs)
69        let obs_cov = sigma_obs.powi(2);
70
71        Self {
72            transition_matrix,
73            selection_matrix,
74            design_matrix,
75            state_cov,
76            obs_cov,
77        }
78    }
79}
80
81impl<'a> KalmanFilter<'a> {
82    pub fn new(model: &'a StateSpaceModel) -> Self {
83        Self { model }
84    }
85
86    /// Performs exact Forward Kalman Filtering with optional Diffuse State Initialization
87    pub fn filter(
88        &self,
89        observations: &Array1<f64>,
90        a0: Option<Array1<f64>>,
91        p0: Option<Array2<f64>>,
92    ) -> Result<KalmanFilterResult> {
93        let n = observations.len();
94        let m = self.model.design_matrix.len();
95
96        // 1. Initialize states: If None, apply Diffuse Initializer (large covariance scalar kappa * I)
97        let mut a = a0.unwrap_or_else(|| Array1::zeros(m));
98        let mut p = p0.unwrap_or_else(|| {
99            let kappa = 1e6; // Diffuse initialization variance
100            Array2::eye(m) * kappa
101        });
102
103        let mut steps = Vec::with_capacity(n);
104        let mut log_like = 0.0;
105
106        let t = &self.model.transition_matrix;
107        let r = &self.model.selection_matrix;
108        let z = &self.model.design_matrix;
109        let q = &self.model.state_cov;
110        let h = self.model.obs_cov;
111
112        let rqr = r.dot(q).dot(&r.t());
113
114        for t_idx in 0..n {
115            let y_t = observations[t_idx];
116
117            let a_prior = a.clone();
118            let p_prior = p.clone();
119
120            // Measurement Prediction & Innovation
121            let y_hat = z.dot(&a_prior);
122            let v_t = y_t - y_hat;
123
124            // Innovation Variance F_t = Z * P * Z^T + H
125            let f_t = z.dot(&p_prior.dot(z)) + h;
126
127            if f_t <= 0.0 {
128                return Err(ChronosError::ConvergenceFailure(
129                    "Non-positive innovation variance encountered during filter pass".to_string(),
130                ));
131            }
132
133            // Kalman Gain K_t = P * Z^T / F_t
134            let k_t = p_prior.dot(z) / f_t;
135
136            // Posterior Updating
137            let a_post = &a_prior + &(&k_t * v_t);
138
139            // P_{t|t} = P_{t|t-1} - K_t * Z * P_{t|t-1}
140            let z_row = z.clone().insert_axis(Axis(0));
141            let k_col = k_t.clone().insert_axis(Axis(1));
142            let p_post = &p_prior - &k_col.dot(&z_row.dot(&p_prior));
143
144            // Replace LN_2_PI with std::f64::consts::TAU.ln()
145            let ln_2_pi = std::f64::consts::TAU.ln(); // ln(2 * PI)
146            log_like -= 0.5 * (ln_2_pi + f_t.ln() + (v_t * v_t) / f_t);
147
148            steps.push(FilterStepResult {
149                a_prior: a_prior.clone(),
150                p_prior: p_prior.clone(),
151                a_post: a_post.clone(),
152                p_post: p_post.clone(),
153                v: v_t,
154                f: f_t,
155                k: k_t,
156            });
157
158            // Time Update (Predict next state)
159            a = t.dot(&a_post);
160            p = t.dot(&p_post).dot(&t.t()) + &rqr;
161        }
162
163        Ok(KalmanFilterResult {
164            steps,
165            log_likelihood: log_like,
166        })
167    }
168
169    /// Rauch-Tung-Striebel (RTS) Backward Smoother Pass
170    /// Computes conditional state distributions given ALL observations: a_{t|N} and P_{t|N}
171    pub fn smooth(&self, filter_result: &KalmanFilterResult) -> Result<SmootherResult> {
172        let n = filter_result.steps.len();
173        if n == 0 {
174            return Err(ChronosError::InsufficientData {
175                required: 1,
176                found: 0,
177            });
178        }
179
180        let m = self.model.design_matrix.len();
181        let t = &self.model.transition_matrix;
182
183        let mut smoothed_states = vec![Array1::zeros(m); n];
184        let mut smoothed_covs = vec![Array2::zeros((m, m)); n];
185
186        // Terminal state at t = N is identical to the filtered state
187        smoothed_states[n - 1] = filter_result.steps[n - 1].a_post.clone();
188        smoothed_covs[n - 1] = filter_result.steps[n - 1].p_post.clone();
189
190        // Backward recursion from t = N-2 down to 0
191        for t_idx in (0..n - 1).rev() {
192            let step_curr = &filter_result.steps[t_idx];
193            let step_next = &filter_result.steps[t_idx + 1];
194
195            // Predicted covariance for next step: P_{t+1|t} = T * P_{t|t} * T^T + R*Q*R^T
196            let p_next_pred = &step_next.p_prior;
197
198            // Smoother Gain C_t = P_{t|t} * T^T * [P_{t+1|t}]^-1
199            let p_next_inv = linalg::inv(p_next_pred).map_err(ChronosError::LinalgError)?;
200            let c_t = step_curr.p_post.dot(&t.t()).dot(&p_next_inv);
201
202            // Smoothed State: a_{t|N} = a_{t|t} + C_t * (a_{t+1|N} - a_{t+1|t})
203            let a_diff = &smoothed_states[t_idx + 1] - &step_next.a_prior;
204            smoothed_states[t_idx] = &step_curr.a_post + &c_t.dot(&a_diff);
205
206            // Smoothed Covariance: P_{t|N} = P_{t|t} + C_t * (P_{t+1|N} - P_{t+1|t}) * C_t^T
207            let p_diff = &smoothed_covs[t_idx + 1] - p_next_pred;
208            smoothed_covs[t_idx] = &step_curr.p_post + &c_t.dot(&p_diff).dot(&c_t.t());
209        }
210
211        Ok(SmootherResult {
212            smoothed_states,
213            smoothed_covs,
214        })
215    }
216}
chronos_ts/statespace.rs

chronos_ts/
statespace.rs
