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

//! `Sma` — simple moving average over the last `length` values of one f64
//! input. The walking skeleton's first worked node: it proves the `aura-core`
//! `Node` contract is authorable from a downstream crate and evaluable with no
//! engine present (the test drives it by hand, as the sim loop later will).

use aura_core::{
    Ctx, FieldSpec, Firing, InputSpec, LeafFactory, Node, NodeSchema, ParamSpec, Scalar, ScalarKind,
};

/// Simple moving average over the last `length` values of one f64 input.
pub struct Sma {
    length: usize,
    out: [Scalar; 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, out: [Scalar::F64(0.0)] }
    }

    /// The param-generic recipe for a blueprint leaf: 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 factory() -> LeafFactory {
        LeafFactory::new(
            "SMA",
            vec![ParamSpec { name: "length".into(), kind: ScalarKind::I64 }],
            |p| Box::new(Sma::new(p[0].as_i64().expect("length slot is I64") as usize)),
        )
    }
}

impl Node for Sma {
    fn schema(&self) -> NodeSchema {
        NodeSchema {
            inputs: vec![InputSpec {
                kind: ScalarKind::F64,
                lookback: self.length,
                firing: Firing::Any,
            }],
            output: vec![FieldSpec { name: "value", kind: ScalarKind::F64 }],
            params: vec![ParamSpec { name: "length".into(), kind: ScalarKind::I64 }],
        }
    }

    fn eval(&mut self, ctx: Ctx<'_>) -> Option<&[Scalar]> {
        let w = ctx.f64_in(0);
        if w.len() < self.length {
            return None; // not yet warmed up
        }
        let mut sum = 0.0;
        for k in 0..self.length {
            sum += w[k]; // index 0 = newest (financial indexing)
        }
        self.out[0] = Scalar::F64(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, Timestamp};

    #[test]
    fn sma_warms_up_then_tracks_the_window_mean() {
        let mut sma = Sma::new(3);
        let schema = sma.schema();

        // size the input column from the schema, as the engine will at wiring
        let mut inputs = vec![AnyColumn::with_capacity(
            schema.inputs[0].kind,
            schema.inputs[0].lookback,
        )];

        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([Scalar::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([Scalar::F64(7.0)].as_slice()));

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

    #[test]
    fn labels_carry_identifying_params() {
        use crate::{Add, Exposure, 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!(Exposure::new(0.5).label(), "Exposure(0.5)");
        assert_eq!(SimBroker::new(0.0001).label(), "SimBroker(0.0001)");
        // bare-kind nodes (identity is not a mis-wiring axis here, per spec)
        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 factory_params_match_built_node_schema() {
        let f = Sma::factory();
        let built = f.build(&[Scalar::I64(3)]);
        assert_eq!(f.params(), built.schema().params.as_slice());
    }

    #[test]
    fn nodes_declare_expected_params() {
        use crate::{Add, Exposure, LinComb, Recorder, SimBroker, Sub};
        use aura_core::{Firing, ParamSpec, ScalarKind};
        // single scalar knobs
        assert_eq!(
            Sma::new(3).schema().params,
            vec![ParamSpec { name: "length".into(), kind: ScalarKind::I64 }],
        );
        assert_eq!(
            Exposure::new(0.5).schema().params,
            vec![ParamSpec { name: "scale".into(), kind: ScalarKind::F64 }],
        );
        // vector knob expands flat to N indexed F64 entries
        let lc = LinComb::new(vec![1.0, -1.0]).schema().params;
        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::new().schema().params.is_empty());
        assert!(Add::new().schema().params.is_empty());
        assert!(SimBroker::new(0.0001).schema().params.is_empty());
        let (tx, _rx) = std::sync::mpsc::channel();
        assert!(Recorder::new(&[ScalarKind::F64], Firing::Any, tx).schema().params.is_empty());
    }
}