Aura/crates/aura-std/src/sma.rs

//! `Sma` — simple moving average over the last `length` values of one f64 input.
//!
//! Computed as an **O(1) incremental window sum**, not a per-tick re-sum: the node
//! owns a `length`-slot ring of the window and keeps a running `sum`, adding the
//! new sample and subtracting the evicted one each cycle (the ta-lib running-sum
//! shape). The running sum carries a **Kahan/Neumaier compensation** term so it
//! does not drift over millions of add/remove ops — the fix pandas' rolling mean
//! adopted after real float-drift bug reports; ta-lib omits it and drifts.
//!
//! `Ema` needs no equivalent (`ema.rs`): its recurrence is *contractive*
//! (`ema += alpha*(x - ema)` rescales old state by `1-alpha < 1` each tick), so a
//! past rounding error decays geometrically instead of accumulating. The SMA
//! running sum is *accumulative* — error has nowhere to go — which is exactly why
//! it, and not the EMA, needs compensation.
//!
//! Because the window lives in node state, `eval` reads only the newest sample and
//! `lookbacks()` is `1` (the input column is depth-1), exactly like `Ema`. O(1)
//! time, O(length) state, allocation-free on the hot path (the ring and output
//! buffer are sized once at construction).

use aura_core::{
    Cell, Ctx, FieldSpec, Firing, Node, NodeSchema, ParamSpec, PortSpec, PrimitiveBuilder,
    ScalarKind,
};

/// Simple moving average over the last `length` values of one f64 input,
/// maintained as an O(1) Kahan-compensated running sum over a node-owned ring.
pub struct Sma {
    length: usize,
    // The window, node-owned (so the input column is depth-1): `ring[pos]` is the
    // oldest value, the one the next push evicts. Sized once at construction (C7).
    ring: Box<[f64]>,
    pos: usize,
    // Samples seen so far — the warm-up gate (silent until `length`, like `Ema`).
    count: usize,
    // Running sum of the values currently in `ring`, plus its Kahan compensation
    // term (the low-order bits each add/remove dropped, folded into the next op).
    sum: f64,
    comp: f64,
    out: [Cell; 1],
}

impl Sma {
    /// Build an SMA of window `length` (must be >= 1).
    pub fn new(length: usize) -> Self {
        assert!(length >= 1, "SMA length must be >= 1");
        Self {
            length,
            ring: vec![0.0; length].into_boxed_slice(),
            pos: 0,
            count: 0,
            sum: 0.0,
            comp: 0.0,
            out: [Cell::from_f64(0.0)],
        }
    }

    /// The param-generic recipe for a blueprint primitive: declares `length` and builds
    /// through `Sma::new` (the single sizing/validation gate; the slice is
    /// kind-checked before `build` runs, so the typed read is total).
    pub fn builder() -> PrimitiveBuilder {
        PrimitiveBuilder::new(
            "SMA",
            NodeSchema {
                inputs: vec![PortSpec { kind: ScalarKind::F64, firing: Firing::Any, name: "series".into() }],
                output: vec![FieldSpec { name: "value".into(), kind: ScalarKind::F64 }],
                params: vec![ParamSpec { name: "length".into(), kind: ScalarKind::I64 }],
            },
            |p| Box::new(Sma::new(p[0].i64() as usize)),
        )
    }
}

/// Kahan/Neumaier compensated accumulation: fold `v` into `*sum`, carrying the
/// low-order bits lost on this step in `*comp` so a long sequence of adds (and
/// removes, which are adds of a negative) does not drift. The whole reason the SMA
/// running sum stays accurate over millions of ticks; `Ema`'s contractive
/// recurrence needs no such term (see the module docs).
fn kahan(sum: &mut f64, comp: &mut f64, v: f64) {
    let y = v - *comp;
    let t = *sum + y;
    *comp = (t - *sum) - y;
    *sum = t;
}

impl Node for Sma {
    // The window lives in node state, so only the newest sample is read each cycle
    // — `length` sizes the ring, not the input column (recursive, like `Ema`).
    fn lookbacks(&self) -> Vec<usize> {
        vec![1]
    }

    fn eval(&mut self, ctx: Ctx<'_>) -> Option<&[Cell]> {
        let w = ctx.f64_in(0);
        if w.is_empty() {
            return None; // no sample yet
        }
        let x = w[0]; // index 0 = newest (financial indexing)

        // Add the newest into the running sum; once the window is full, remove the
        // value this push evicts (the ring slot about to be overwritten). Both go
        // through Kahan so the running sum tracks the true window sum.
        kahan(&mut self.sum, &mut self.comp, x);
        if self.count >= self.length {
            let evicted = self.ring[self.pos];
            kahan(&mut self.sum, &mut self.comp, -evicted);
        }
        self.ring[self.pos] = x;
        self.pos = (self.pos + 1) % self.length;
        if self.count < self.length {
            self.count += 1;
        }
        if self.count < self.length {
            return None; // not yet warmed up
        }

        self.out[0] = Cell::from_f64(self.sum / self.length as f64);
        Some(&self.out)
    }

    fn label(&self) -> String {
        format!("SMA({})", self.length)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aura_core::{AnyColumn, Scalar, Timestamp};

    #[test]
    fn sma_warms_up_then_tracks_the_window_mean() {
        let sma_for_depth = Sma::new(3);

        // size the input column from the node's lookback, as bootstrap will at wiring
        let mut inputs =
            vec![AnyColumn::with_capacity(ScalarKind::F64, sma_for_depth.lookbacks()[0])];
        let mut sma = sma_for_depth;

        let feed = [1.0_f64, 2.0, 3.0, 4.0, 5.0];
        // means of [1,2,3], [2,3,4], [3,4,5] once warmed up
        let expect = [None, None, Some(2.0), Some(3.0), Some(4.0)];

        for (v, want) in feed.iter().zip(expect) {
            inputs[0].push(Scalar::f64(*v)).unwrap();
            let got = sma.eval(Ctx::new(&inputs, Timestamp(0)));
            match want {
                None => assert_eq!(got, None),
                Some(m) => assert_eq!(got, Some([Cell::from_f64(m)].as_slice())),
            }
        }
    }

    #[test]
    fn sma_length_one_is_identity() {
        let mut sma = Sma::new(1);
        let mut inputs = vec![AnyColumn::with_capacity(ScalarKind::F64, 1)];

        inputs[0].push(Scalar::f64(7.0)).unwrap();
        assert_eq!(sma.eval(Ctx::new(&inputs, Timestamp(0))), Some([Cell::from_f64(7.0)].as_slice()));

        inputs[0].push(Scalar::f64(9.0)).unwrap();
        assert_eq!(sma.eval(Ctx::new(&inputs, Timestamp(0))), Some([Cell::from_f64(9.0)].as_slice()));
    }

    #[test]
    fn lookback_is_one_window_lives_in_node_state() {
        // the incremental SMA owns its window ring, so it reads only the newest
        // sample each cycle — input column depth drops from `length` to 1 (like Ema).
        assert_eq!(Sma::new(20).lookbacks(), vec![1]);
        assert_eq!(Sma::new(1).lookbacks(), vec![1]);
    }

    #[test]
    fn incremental_matches_full_resum_within_tolerance() {
        // the running Kahan sum must track the true window mean across a long, noisy
        // f64 series. Drive the node and a reference full re-sum side by side; they
        // agree to a tight tolerance (Kahan keeps drift near machine epsilon — a
        // bare running sum would slowly diverge, which is the bug pandas fixed).
        let length = 50;
        let mut sma = Sma::new(length);
        let mut inputs = vec![AnyColumn::with_capacity(ScalarKind::F64, 1)];

        // a deterministic non-integer series mixing magnitudes, so cancellation bites
        let series: Vec<f64> =
            (0..5_000).map(|i| 1.1234 + (i as f64) * 1e-3 + ((i % 7) as f64) * 0.37).collect();

        let mut hist: Vec<f64> = Vec::new();
        for &x in &series {
            inputs[0].push(Scalar::f64(x)).unwrap();
            let got = sma.eval(Ctx::new(&inputs, Timestamp(0))).map(|r| r[0].f64());
            hist.push(x);
            if hist.len() >= length {
                let want: f64 = hist[hist.len() - length..].iter().sum::<f64>() / length as f64;
                let g = got.expect("warmed up after `length` samples");
                assert!((g - want).abs() < 1e-9, "incremental {g} vs re-sum {want}");
            } else {
                assert_eq!(got, None, "silent until warmed up");
            }
        }
    }

    #[test]
    fn labels_carry_identifying_params() {
        use crate::{Add, Bias, LinComb, Recorder, SimBroker, Sub};
        use aura_core::{Firing, ScalarKind};

        // the load-bearing payoff: two SMAs disambiguate by window
        assert_eq!(Sma::new(2).label(), "SMA(2)");
        assert_eq!(Sma::new(4).label(), "SMA(4)");
        // param-carrying single nodes
        assert_eq!(Bias::new(0.5).label(), "Bias(0.5)");
        assert_eq!(SimBroker::new(0.0001).label(), "SimBroker(0.0001)");
        // bare-kind nodes (identity is not a mis-wiring axis here)
        assert_eq!(Sub::new().label(), "Sub");
        assert_eq!(Add::new().label(), "Add");
        assert_eq!(LinComb::new(vec![1.0, -1.0]).label(), "LinComb");
        let (tx, _rx) = std::sync::mpsc::channel();
        assert_eq!(Recorder::new(&[ScalarKind::F64], Firing::Any, tx).label(), "Recorder");
    }

    #[test]
    fn nodes_declare_expected_params() {
        use crate::{Add, Bias, LinComb, Recorder, SimBroker, Sub};
        use aura_core::{Firing, ParamSpec, ScalarKind};
        // single scalar knobs (declared on the param-generic builder, pre-build)
        assert_eq!(
            Sma::builder().schema().params,
            vec![ParamSpec { name: "length".into(), kind: ScalarKind::I64 }],
        );
        assert_eq!(
            Bias::builder().schema().params,
            vec![ParamSpec { name: "scale".into(), kind: ScalarKind::F64 }],
        );
        // vector knob expands flat to N indexed F64 entries
        let lc = LinComb::builder(2).schema().params.clone();
        assert_eq!(lc.len(), 2);
        assert_eq!(lc[0].name, "weights[0]");
        assert_eq!(lc[1].name, "weights[1]");
        assert!(lc.iter().all(|p| p.kind == ScalarKind::F64));
        // param-less nodes declare empty
        assert!(Sub::builder().schema().params.is_empty());
        assert!(Add::builder().schema().params.is_empty());
        assert!(SimBroker::builder(0.0001).schema().params.is_empty());
        let (tx, _rx) = std::sync::mpsc::channel();
        assert!(
            Recorder::builder(vec![ScalarKind::F64], Firing::Any, tx)
                .schema()
                .params
                .is_empty()
        );
    }

    #[test]
    fn input_slot_is_named_series() {
        assert_eq!(Sma::builder().schema().inputs[0].name, "series");
    }

    #[test]
    fn bind_removes_slot_from_param_space() {
        // a bound param-bearing node reports an empty param surface — parity with the
        // SimBroker precedent (nodes_declare_expected_params, this file)
        let sma2 = Sma::builder().named("bias").bind("length", Scalar::i64(2));
        assert!(sma2.schema().params.is_empty());
        // contrast: the length-generic SMA keeps `length` open
        assert_eq!(Sma::builder().named("bias").params().len(), 1);
    }

    #[test]
    fn bound_node_builds_with_injected_value() {
        // built with an empty open slice, the bound builder yields SMA(2)
        let node = Sma::builder().bind("length", Scalar::i64(2)).build(&[]);
        assert_eq!(node.label(), "SMA(2)");
    }
}