anofox-forecast 0.13.2

//! Ornstein-Uhlenbeck mean-reversion leaf.
//!
//! Discrete-time OU process:
//!
//! ```text
//!   y_{t+1} = y_t + θ · (μ − y_t) + σ · ε
//! ```
//!
//! Where `θ ∈ (0, 1)` is the reversion rate, `μ` is the long-run mean,
//! and `σ` is the innovation std. The h-step forecast is
//!
//! ```text
//!   ŷ_{t+h} = μ + (1 − θ)^h · (y_t − μ)
//! ```
//!
//! Mathematically equivalent to a mean-reverting AR(1) with `φ = 1 − θ`.
//! The distinct leaf exists because it estimates `θ` via method-of-moments
//! on centred first-differences (`Δy_t = θ(μ − y_{t-1}) + ε`) rather than
//! via streaming OLS on lagged levels — a specialisation that behaves
//! better on bounded / mean-reverting panels than the level-form AR(1),
//! particularly at longer horizons where the reversion asymptote matters.
//!
//! Predictive std uses `σ · √(Σ_{i=0..h} (1−θ)^{2i})` — the exact
//! stationary AR(1) variance sum, adapted to the OU parameterisation.

use crate::models::laplace::dist::Gaussian;
use crate::models::laplace::leaf::Leaf;

pub struct OuLeaf {
    alpha_mean: f64,
    /// Reversion rate `θ`; solved from running MoM sufficient stats.
    theta: f64,
    /// Long-run mean μ (EMA).
    mean: Option<f64>,
    last: Option<f64>,
    /// Σ (y_{t-1} − μ)² over observed history — denominator for θ MoM.
    s_xx: f64,
    /// Σ (y_{t-1} − μ)(y_t − y_{t-1}) — numerator (relates `θ` to reversion).
    s_xdy: f64,
    n: usize,
    ss: f64,
    mean_resid: f64,
}

impl OuLeaf {
    pub fn new(alpha_mean: f64) -> Self {
        Self {
            alpha_mean: alpha_mean.clamp(1e-3, 1.0 - 1e-3),
            theta: 0.0,
            mean: None,
            last: None,
            s_xx: 0.0,
            s_xdy: 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 OuLeaf {
    fn name(&self) -> &'static str {
        "ou"
    }

    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 = (1.0 - self.theta).clamp(-0.999, 0.999);
        let phi2 = phi * phi;
        (1..=horizon)
            .map(|h| {
                let mean = mu + phi.powi(h as i32) * (last - mu);
                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()
    }

    fn observe(&mut self, y: f64) {
        let mu_before = self.mean.unwrap_or(y);
        let last = self.last.unwrap_or(mu_before);
        let phi = (1.0 - self.theta).clamp(-0.999, 0.999);

        let predicted = mu_before + 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);

        // MoM stats: Δy_t = θ(μ − y_{t-1}) + ε, so
        // θ ≈ Σ (μ − y_{t-1})(Δy_t) / Σ (μ − y_{t-1})²  = −Σ x·Δy / Σ x²
        // We accumulate `s_xdy = Σ (y_{t-1} − μ)(y_t − y_{t-1})` so that
        // θ = − s_xdy / s_xx.
        let x = last - mu_before;
        let dy = y - last;
        self.s_xx += x * x;
        self.s_xdy += x * dy;
        if self.s_xx > 1e-12 {
            let theta = -self.s_xdy / self.s_xx;
            self.theta = theta.clamp(1e-3, 1.0 - 1e-3);
        }

        self.mean = Some(match self.mean {
            Some(m) => self.alpha_mean * y + (1.0 - self.alpha_mean) * m,
            None => y,
        });
        self.last = Some(y);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn feed_ou(leaf: &mut OuLeaf, theta: f64, mu: f64, n: usize) {
        let mut y = mu;
        for i in 0..n {
            let noise = ((i as f64 * 12.9898).sin() * 43758.5453).fract() - 0.5;
            y += theta * (mu - y) + noise;
            leaf.observe(y);
        }
    }

    #[test]
    fn theta_is_positive_on_mean_reverting_process() {
        // The MoM estimator centred on an EMA-tracked μ is biased low
        // when μ hasn't converged. We only assert direction (θ > 0 =
        // reversion detected) — the mixture softmax handles the
        // magnitude across candidates.
        let mut leaf = OuLeaf::new(0.05);
        feed_ou(&mut leaf, 0.3, 5.0, 1000);
        assert!(
            leaf.theta > 0.0 && leaf.theta < 1.0,
            "θ = {} not in (0, 1)",
            leaf.theta
        );
    }

    #[test]
    fn far_horizon_forecast_reverts_toward_mu() {
        let mut leaf = OuLeaf::new(0.05);
        feed_ou(&mut leaf, 0.4, 10.0, 500);
        // Now feed a jump — the leaf's `last` will be the jumped value.
        leaf.observe(50.0);
        let preds = leaf.predict(20);
        // h=1 should still be near 50 (small step back); h=20 should be
        // materially closer to μ than to 50.
        let mu_estimate = leaf.mean.unwrap();
        assert!(
            (preds[19].mean - mu_estimate).abs() < (50.0 - mu_estimate).abs() * 0.5,
            "h=20 forecast {} should have reverted toward μ≈{}",
            preds[19].mean,
            mu_estimate
        );
    }

    #[test]
    fn cold_start_produces_finite_predictions() {
        let mut leaf = OuLeaf::new(0.1);
        leaf.observe(3.0);
        leaf.observe(4.0);
        let preds = leaf.predict(5);
        for (h, p) in preds.iter().enumerate() {
            assert!(p.mean.is_finite(), "h={}: mean not finite", h + 1);
            assert!(p.std.is_finite() && p.std > 0.0, "h={}: std invalid", h + 1);
        }
    }
}