Aura/crates/aura-engine/src/sweep.rs

//! Grid param-sweep (C12.1): enumerate a cartesian grid over a blueprint's
//! param-space and run each point disjointly (C1). This module owns enumeration
//! (`GridSpace`), execution (`sweep`), and collection (`SweepFamily`); the
//! per-point run-to-metrics closure is the author's (harness-specific sink glue
//! the engine cannot generically own — C8/C18).

use aura_core::{ParamSpec, Scalar, ScalarKind};
use crate::RunReport;
use std::sync::atomic::{AtomicUsize, Ordering};

/// A validated cartesian grid over a blueprint's param-space: one discrete
/// value-list per param slot, in `param_space()` order. Enumerates a family of
/// points (C12.1 grid axis).
#[derive(Debug)]
pub struct GridSpace {
    axes: Vec<Vec<Scalar>>,
}

impl GridSpace {
    /// Validate `axes` against `space` (the blueprint's `param_space()`): one
    /// axis per slot (`Arity`), every value the slot's declared kind
    /// (`KindMismatch`), no empty axis (`EmptyAxis`, which would yield zero
    /// points). On success the grid enumerates `∏ |axis_i|` points.
    pub fn new(space: &[ParamSpec], axes: Vec<Vec<Scalar>>) -> Result<Self, SweepError> {
        if axes.len() != space.len() {
            return Err(SweepError::Arity { expected: space.len(), got: axes.len() });
        }
        for (slot, (axis, ps)) in axes.iter().zip(space).enumerate() {
            if axis.is_empty() {
                return Err(SweepError::EmptyAxis { slot });
            }
            for (value_index, v) in axis.iter().enumerate() {
                if v.kind() != ps.kind {
                    return Err(SweepError::KindMismatch {
                        slot,
                        value_index,
                        expected: ps.kind,
                        got: v.kind(),
                    });
                }
            }
        }
        Ok(Self { axes })
    }

    /// The number of grid points (`∏ |axis_i|`), always `>= 1` (a valid grid has
    /// no empty axis; a zero-param grid is the single empty point).
    pub fn len(&self) -> usize {
        self.axes.iter().map(|a| a.len()).product()
    }

    /// Always `false` — a valid `GridSpace` rejects empty axes, so it has at
    /// least one point. Present to satisfy clippy's `len_without_is_empty`.
    pub fn is_empty(&self) -> bool {
        false
    }

    /// The cartesian product, in odometer order: the **last** axis varies
    /// fastest. Deterministic — the same grid yields the same point sequence.
    pub fn points(&self) -> Vec<Vec<Scalar>> {
        let mut out = Vec::with_capacity(self.len());
        let mut idx = vec![0usize; self.axes.len()];
        loop {
            out.push(self.axes.iter().zip(&idx).map(|(a, &i)| a[i]).collect());
            // odometer increment from the last axis
            let mut k = self.axes.len();
            loop {
                if k == 0 {
                    return out;
                }
                k -= 1;
                idx[k] += 1;
                if idx[k] < self.axes[k].len() {
                    break;
                }
                idx[k] = 0;
            }
        }
    }
}

/// A structural fault constructing a `GridSpace` — the typed gate before any run.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SweepError {
    /// The number of axes does not equal the param-space slot count.
    Arity { expected: usize, got: usize },
    /// A grid value's kind does not match its slot's declared kind. `slot` is the
    /// flat param-space index; `value_index` is the position within that axis.
    KindMismatch { slot: usize, value_index: usize, expected: ScalarKind, got: ScalarKind },
    /// A slot was given no values (would collapse the product to zero points).
    EmptyAxis { slot: usize },
}

/// One enumerated point and the full `RunReport` its run produced.
/// Self-describing: `params` is the point's coordinate in `param_space()` order,
/// and `report` carries the run's `(manifest, metrics)` — the unit the run
/// registry indexes (C18).
#[derive(Clone, Debug, PartialEq)]
pub struct SweepPoint {
    pub params: Vec<Scalar>,
    pub report: RunReport,
}

/// The ordered result family of a sweep — one `SweepPoint` per grid point, in
/// enumeration (odometer) order, independent of thread completion order.
#[derive(Clone, Debug, PartialEq)]
pub struct SweepFamily {
    pub points: Vec<SweepPoint>,
}

/// Run `run_one` over every grid point, disjointly in parallel (C1), and collect
/// the family in enumeration order. `run_one` builds + bootstraps + runs +
/// summarizes one point; it shares nothing mutable, so it is `Sync` and the runs
/// are lock-free. Parallelism is `available_parallelism()` workers — `std` only,
/// via `std::thread::scope`.
pub fn sweep<F>(space: &GridSpace, run_one: F) -> SweepFamily
where
    F: Fn(&[Scalar]) -> RunReport + Sync,
{
    let nthreads = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1);
    sweep_with_threads(space, nthreads, run_one)
}

/// The thread-count-explicit core of [`sweep`]. Module-private: the public
/// `sweep` derives the count, while the tests drive it at 1 and at N to pin
/// determinism under concurrency (C1). Workers pull point indices from a shared
/// atomic cursor (work-stealing load-balances uneven per-point cost); each tags
/// its result with the point index, and the family is assembled by sorting on
/// that index after the scope joins — so the order is the enumeration order, not
/// the completion order. Only the cursor is shared; the results side is lock-free.
fn sweep_with_threads<F>(space: &GridSpace, nthreads: usize, run_one: F) -> SweepFamily
where
    F: Fn(&[Scalar]) -> RunReport + Sync,
{
    let points = space.points();
    // `points.len() >= 1` always (GridSpace rejects empty axes), so the clamp
    // lower bound never exceeds the upper; this also coerces a 0 to 1 worker.
    let nthreads = nthreads.clamp(1, points.len());
    let cursor = AtomicUsize::new(0);

    let mut results: Vec<(usize, RunReport)> = std::thread::scope(|scope| {
        let handles: Vec<_> = (0..nthreads)
            .map(|_| {
                scope.spawn(|| {
                    let mut local: Vec<(usize, RunReport)> = Vec::new();
                    loop {
                        let i = cursor.fetch_add(1, Ordering::Relaxed);
                        if i >= points.len() {
                            break;
                        }
                        local.push((i, run_one(&points[i])));
                    }
                    local
                })
            })
            .collect();
        handles.into_iter().flat_map(|h| h.join().unwrap()).collect()
    });

    results.sort_by_key(|&(i, _)| i);
    SweepFamily {
        points: results
            .into_iter()
            .map(|(i, report)| SweepPoint { params: points[i].clone(), report })
            .collect(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        f64_field, summarize, BlueprintNode, Composite, Edge, OutField, Role, RunManifest, Target,
    };
    use aura_core::{Firing, ParamSpec, Scalar, ScalarKind, Timestamp};
    use aura_std::{Exposure, Recorder, SimBroker, Sma, Sub};
    use std::sync::mpsc;

    fn i64_space(n: usize) -> Vec<ParamSpec> {
        (0..n)
            .map(|i| ParamSpec { name: format!("p{i}"), kind: ScalarKind::I64 })
            .collect()
    }

    #[test]
    fn points_enumerate_in_odometer_order() {
        let space = i64_space(2);
        let grid = GridSpace::new(
            &space,
            vec![
                vec![Scalar::I64(2), Scalar::I64(3)],
                vec![Scalar::I64(4), Scalar::I64(5)],
            ],
        )
        .expect("valid grid");
        assert_eq!(
            grid.points(),
            vec![
                vec![Scalar::I64(2), Scalar::I64(4)],
                vec![Scalar::I64(2), Scalar::I64(5)],
                vec![Scalar::I64(3), Scalar::I64(4)],
                vec![Scalar::I64(3), Scalar::I64(5)],
            ],
        );
    }

    #[test]
    fn len_is_product_of_axis_lengths() {
        let space = vec![
            ParamSpec { name: "a".into(), kind: ScalarKind::I64 },
            ParamSpec { name: "b".into(), kind: ScalarKind::I64 },
            ParamSpec { name: "c".into(), kind: ScalarKind::F64 },
        ];
        let grid = GridSpace::new(
            &space,
            vec![
                vec![Scalar::I64(2), Scalar::I64(3)],
                vec![Scalar::I64(4), Scalar::I64(5)],
                vec![Scalar::F64(0.5)],
            ],
        )
        .expect("valid grid");
        assert_eq!(grid.len(), 4);
        assert!(!grid.is_empty());
    }

    #[test]
    fn arity_mismatch_is_an_error() {
        let space = i64_space(2);
        let err = GridSpace::new(&space, vec![vec![Scalar::I64(2)]]).unwrap_err();
        assert_eq!(err, SweepError::Arity { expected: 2, got: 1 });
    }

    #[test]
    fn wrong_kind_is_a_kind_mismatch() {
        let space = i64_space(1);
        let err = GridSpace::new(&space, vec![vec![Scalar::F64(1.0)]]).unwrap_err();
        assert_eq!(
            err,
            SweepError::KindMismatch {
                slot: 0,
                value_index: 0,
                expected: ScalarKind::I64,
                got: ScalarKind::F64,
            },
        );
    }

    #[test]
    fn empty_axis_is_an_error() {
        let space = i64_space(1);
        let err = GridSpace::new(&space, vec![vec![]]).unwrap_err();
        assert_eq!(err, SweepError::EmptyAxis { slot: 0 });
    }

    fn synthetic_prices() -> Vec<(Timestamp, Scalar)> {
        [
            (1_i64, 1.0000_f64),
            (2, 1.0010),
            (3, 1.0030),
            (4, 1.0060),
            (5, 1.0040),
            (6, 1.0010),
            (7, 0.9990),
        ]
        .iter()
        .map(|&(t, p)| (Timestamp(t), Scalar::F64(p)))
        .collect()
    }

    /// The SMA-cross signal as a reusable value-empty composite (duplicated from
    /// `blueprint.rs`'s test module — test fixtures stay module-local, the same
    /// way `report.rs` copies its harness helper).
    fn sma_cross() -> Composite {
        Composite::new(
            "sma_cross",
            vec![
                Sma::builder().named("fast").into(),
                Sma::builder().named("slow").into(),
                Sub::builder().into(),
            ],
            vec![
                Edge { from: 0, to: 2, slot: 0, from_field: 0 },
                Edge { from: 1, to: 2, slot: 1, from_field: 0 },
            ],
            vec![Role {
                name: "price".into(),
                targets: vec![Target { node: 0, slot: 0 }, Target { node: 1, slot: 0 }],
                source: None,
            }],
            vec![OutField { node: 2, field: 0, name: "out".into() }],
        )
    }

    /// The signal-quality harness as a value-empty composite blueprint with two
    /// recording sinks (duplicated from `blueprint.rs`'s test module).
    #[allow(clippy::type_complexity)]
    fn composite_sma_cross_harness() -> (
        Composite,
        mpsc::Receiver<(Timestamp, Vec<Scalar>)>,
        mpsc::Receiver<(Timestamp, Vec<Scalar>)>,
    ) {
        let (tx_eq, rx_eq) = mpsc::channel();
        let (tx_ex, rx_ex) = mpsc::channel();
        let bp = Composite::new(
            "root",
            vec![
                BlueprintNode::Composite(sma_cross()),
                Exposure::builder().into(),
                SimBroker::builder(0.0001).into(),
                Recorder::builder(vec![ScalarKind::F64], Firing::Any, tx_eq).into(),
                Recorder::builder(vec![ScalarKind::F64], Firing::Any, tx_ex).into(),
            ],
            vec![
                Edge { from: 0, to: 1, slot: 0, from_field: 0 }, // composite out -> Exposure
                Edge { from: 1, to: 2, slot: 0, from_field: 0 }, // exposure -> broker slot 0
                Edge { from: 2, to: 3, slot: 0, from_field: 0 }, // equity -> sink
                Edge { from: 1, to: 4, slot: 0, from_field: 0 }, // exposure -> sink
            ],
            vec![Role {
                name: "src".into(),
                targets: vec![Target { node: 0, slot: 0 }, Target { node: 2, slot: 1 }],
                source: Some(ScalarKind::F64),
            }],
            vec![], // output: the root ends in sinks
        );
        (bp, rx_eq, rx_ex)
    }

    /// Build + bootstrap + run + drain + summarize one grid point into a
    /// `RunReport`. A free `fn` (Copy + Sync) so it serves as the `sweep` closure
    /// AND a direct reference for the "sweep == N independent runs" comparison.
    /// The manifest is a minimal fixed fixture — the metrics are the run's, and
    /// determinism makes `run_point` reproduce a point's report exactly.
    fn run_point(point: &[Scalar]) -> RunReport {
        let (bp, rx_eq, rx_ex) = composite_sma_cross_harness();
        let mut h = bp
            .bootstrap_with_params(point.to_vec())
            .expect("grid points are kind-checked against param_space");
        h.run(vec![synthetic_prices()]);
        let equity = f64_field(&rx_eq.try_iter().collect::<Vec<_>>(), 0);
        let exposure = f64_field(&rx_ex.try_iter().collect::<Vec<_>>(), 0);
        RunReport {
            manifest: RunManifest {
                commit: "test".to_string(),
                params: Vec::new(),
                window: (Timestamp(0), Timestamp(0)),
                seed: 0,
                broker: "test".to_string(),
            },
            metrics: summarize(&equity, &exposure),
        }
    }

    fn sma_cross_grid() -> GridSpace {
        let space = composite_sma_cross_harness().0.param_space();
        GridSpace::new(
            &space,
            vec![
                vec![Scalar::I64(2), Scalar::I64(3)], // fast  ∈ {2, 3}
                vec![Scalar::I64(4), Scalar::I64(5)], // slow  ∈ {4, 5}
                vec![Scalar::F64(0.5)],               // scale ∈ {0.5}
            ],
        )
        .expect("grid matches the sample param-space")
    }

    #[test]
    fn sweep_equals_n_independent_runs() {
        let grid = sma_cross_grid();
        let family = sweep(&grid, run_point);
        assert_eq!(family.points.len(), 4);
        // params carried in enumeration (odometer) order, self-describing
        assert_eq!(
            family.points[0].params,
            vec![Scalar::I64(2), Scalar::I64(4), Scalar::F64(0.5)],
        );
        assert_eq!(
            family.points[3].params,
            vec![Scalar::I64(3), Scalar::I64(5), Scalar::F64(0.5)],
        );
        // each point's metrics equal a direct, independent run of the same point:
        // the sweep adds enumeration + execution, never a metrics change (C1).
        for pt in &family.points {
            assert_eq!(pt.report, run_point(&pt.params));
            assert!(pt.report.metrics.total_pips.is_finite());
        }
    }

    #[test]
    fn family_is_deterministic_across_thread_counts() {
        let grid = sma_cross_grid();
        let one = sweep_with_threads(&grid, 1, run_point);
        let many = sweep_with_threads(&grid, 8, run_point);
        // same family at 1 worker and at N (C1: order = enumeration, not completion)
        assert_eq!(one, many);
        // and identical to the public `sweep`, which derives the worker count
        assert_eq!(one, sweep(&grid, run_point));
    }

    #[test]
    fn distinct_points_produce_distinct_metrics() {
        let family = sweep(&sma_cross_grid(), run_point);
        // not a constant family: differing SMA lengths produce differing equity
        let first = family.points[0].report.metrics.total_pips;
        assert!(
            family.points.iter().any(|p| p.report.metrics.total_pips != first),
            "a 4-point SMA-length grid must not collapse to one metric",
        );
    }
}