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

//! `Ema` — exponential moving average of one f64 input, smoothing constant
//! `alpha = 2 / (length + 1)` (the conventional length↔alpha identity, so an
//! `Ema` and an `Sma` of the same `length` are roughly comparable).
//!
//! Two design points worth stating:
//!
//! - **Seeding (ta-lib convention).** Like `Sma`, the EMA stays silent (`None`)
//!   until it has seen `length` samples, then seeds itself with the *SMA* of those
//!   first `length` values and runs the recurrence from there. This keeps warm-up
//!   semantics uniform across the moving-average nodes, matches the EMA traders
//!   know (ta-lib's MACD seeds its EMAs this way), and avoids claiming an average
//!   from a single observation. (A first-value seed — emit from sample 1 — is the
//!   other common convention; rejected here for the consistency/honesty reasons
//!   above.)
//! - **Recursive, so `lookback = 1`.** The running average lives in internal
//!   state, so the node reads only the newest sample each cycle — `length` sizes
//!   `alpha`, not the input window. `eval` is O(1) time, O(1) state, and allocates
//!   nothing on the hot path (the output buffer is reused).

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

/// Exponential moving average of one f64 input, smoothing `alpha = 2/(length+1)`,
/// seeded with the SMA of the first `length` samples.
pub struct Ema {
    alpha: f64,
    length: usize,
    // Warm-up accumulator: until `length` samples have arrived the node is silent
    // and sums them; on the `length`-th sample it locks in the SMA seed (`seeded`)
    // and `ema` carries the running value from then on.
    seeded: bool,
    warmup_sum: f64,
    count: usize,
    ema: f64,
    out: [Cell; 1],
}

impl Ema {
    /// Build an EMA of window `length` (must be >= 1); `alpha = 2/(length+1)`
    /// (`length == 1` gives `alpha == 1`, the identity).
    pub fn new(length: usize) -> Self {
        assert!(length >= 1, "EMA length must be >= 1");
        Self {
            alpha: 2.0 / (length as f64 + 1.0),
            length,
            seeded: false,
            warmup_sum: 0.0,
            count: 0,
            ema: 0.0,
            out: [Cell::from_f64(0.0)],
        }
    }

    /// The param-generic recipe for a blueprint primitive: declares `length` and builds
    /// through `Ema::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(
            "EMA",
            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 }],
                doc: "exponential moving average over the input series",
            },
            |p| Box::new(Ema::new(p[0].i64() as usize)),
        )
    }
}

impl Node for Ema {
    // recursive: the running average lives in internal state, so only the newest
    // sample is read — `length` sizes alpha, not the window.
    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)
        if !self.seeded {
            // warm up over `length` samples, then seed with their SMA (ta-lib
            // convention) — silent until ready, exactly like `Sma`.
            self.warmup_sum += x;
            self.count += 1;
            if self.count < self.length {
                return None;
            }
            self.ema = self.warmup_sum / self.length as f64;
            self.seeded = true;
        } else {
            // O(1) recurrence: one sub, one mul, one add.
            self.ema += self.alpha * (x - self.ema);
        }
        self.out[0] = Cell::from_f64(self.ema);
        Some(&self.out)
    }

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

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

    #[test]
    fn ema_warms_up_like_sma_then_smooths() {
        // length 3 -> alpha = 0.5: silent for the first 2 samples, seeds with the
        // SMA of the first 3, then runs the recurrence. The jump to 10 separates the
        // EMA from a plain SMA (EMA = 7.0 here, SMA(4,6,10) would be 6.667); all
        // magnitudes are exact in f64.
        let mut ema = Ema::new(3);
        let mut inputs = vec![AnyColumn::with_capacity(ScalarKind::F64, 1)];

        let feed = [2.0_f64, 4.0, 6.0, 10.0];
        // None, None, SMA(2,4,6) = 4.0, then 4.0 + 0.5*(10 - 4.0) = 7.0
        let expect = [None, None, Some(4.0), Some(7.0)];

        for (v, want) in feed.iter().zip(expect) {
            inputs[0].push(Scalar::f64(*v)).unwrap();
            let got = ema.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 ema_length_one_is_identity() {
        // alpha = 2/2 = 1.0: the seed locks in on the first sample (length 1) and
        // the recurrence reduces to ema = x.
        let mut ema = Ema::new(1);
        let mut inputs = vec![AnyColumn::with_capacity(ScalarKind::F64, 1)];

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

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

    #[test]
    fn label_carries_the_window() {
        assert_eq!(Ema::new(12).label(), "EMA(12)");
        assert_eq!(Ema::new(26).label(), "EMA(26)");
    }

    #[test]
    fn builder_declares_the_length_knob() {
        // the param-generic recipe carries the declared I64 `length` knob
        let b = Ema::builder();
        assert_eq!(b.params(), b.schema().params.as_slice());
        assert_eq!(b.schema().params[0].name, "length");
        assert_eq!(b.schema().params[0].kind, ScalarKind::I64);
    }

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