use crate::models::laplace::dist::Gaussian;
use crate::models::laplace::leaf::Leaf;
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Ar1Leaf {
alpha_mean: f64,
mean: Option<f64>,
last: Option<f64>,
s_xx: f64, s_xy: f64, phi: f64,
n: usize,
ss: f64,
mean_resid: f64,
}
impl Ar1Leaf {
pub fn new(alpha_mean: f64) -> Self {
Self {
alpha_mean: alpha_mean.clamp(1e-3, 1.0 - 1e-3),
mean: None,
last: None,
s_xx: 0.0,
s_xy: 0.0,
phi: 0.0,
n: 0,
ss: 0.0,
mean_resid: 0.0,
}
}
fn sigma(&self) -> f64 {
if self.n < 2 {
return 1.0;
}
(self.ss / (self.n as f64 - 1.0)).sqrt().max(1e-9)
}
}
impl Leaf for Ar1Leaf {
fn name(&self) -> &'static str {
"ar1"
}
fn predict(&self, horizon: usize) -> Vec<Gaussian> {
let mu = self.mean.unwrap_or(0.0);
let last = self.last.unwrap_or(mu);
let sigma = self.sigma();
let phi = self.phi.clamp(-0.999, 0.999);
(1..=horizon)
.map(|h| {
let phi_h = phi.powi(h as i32);
let mean = mu + phi_h * (last - mu);
let phi2 = phi * phi;
let var_scale = if (1.0 - phi2).abs() < 1e-12 {
h as f64
} else {
(1.0 - phi2.powi(h as i32)) / (1.0 - phi2)
};
Gaussian::new(mean, sigma * var_scale.sqrt())
})
.collect()
}
#[inline]
fn predict_one(&self) -> Gaussian {
let mu = self.mean.unwrap_or(0.0);
let last = self.last.unwrap_or(mu);
let phi = self.phi.clamp(-0.999, 0.999);
Gaussian::new(mu + phi * (last - mu), self.sigma())
}
fn observe(&mut self, y: f64) {
let mu_before = self.mean.unwrap_or(y);
let last = self.last.unwrap_or(mu_before);
let predicted = mu_before + self.phi * (last - mu_before);
let resid = y - predicted;
self.n += 1;
let delta = resid - self.mean_resid;
self.mean_resid += delta / self.n as f64;
self.ss += delta * (resid - self.mean_resid);
let x = last - mu_before;
let z = y - mu_before;
self.s_xx += x * x;
self.s_xy += x * z;
if self.s_xx > 1e-12 {
self.phi = (self.s_xy / self.s_xx).clamp(-0.999, 0.999);
}
self.mean = Some(match self.mean {
Some(m) => self.alpha_mean * y + (1.0 - self.alpha_mean) * m,
None => y,
});
self.last = Some(y);
}
}