foxhunt/crates/ml-alpha/tests/trunk_forward.rs

//! CfcTrunk forward — invariants on the sequentially-dispatched perception path.

use approx::assert_relative_eq;
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::cfc::{CfcConfig, CfcTrunk};
use ml_alpha::heads::{N_HORIZONS, PROJ_DIM};
use ml_core::device::MlDevice;

fn test_device() -> MlDevice {
    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
}

fn synthetic_input(dt_ns: u64) -> Mbp10RawInput {
    let mut bid_px = [0.0f32; 10];
    let mut bid_sz = [0.0f32; 10];
    let mut ask_px = [0.0f32; 10];
    let mut ask_sz = [0.0f32; 10];
    for i in 0..10 {
        bid_px[i] = 5500.00 - 0.25 * i as f32;
        ask_px[i] = 5500.25 + 0.25 * i as f32;
        bid_sz[i] = 12.0 + i as f32;
        ask_sz[i] = 10.0 + i as f32 * 1.5;
    }
    Mbp10RawInput {
        bid_px,
        bid_sz,
        ask_px,
        ask_sz,
        prev_mid: 5499.875,
        trade_signed_vol: 4.0,
        trade_count: 7,
        ts_ns: dt_ns,
        prev_ts_ns: 0,
    }
}

#[test]
fn forward_returns_probs_in_unit_interval() {
    let dev = test_device();
    let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xABCD).expect("init");
    let (probs, proj) = trunk.forward_snapshot(&synthetic_input(1_000_000)).expect("forward");
    assert_eq!(probs.len(), N_HORIZONS);
    assert_eq!(proj.len(), PROJ_DIM);
    for &p in &probs {
        assert!((0.0..=1.0).contains(&p), "head prob {p} out of [0,1]");
    }
}

#[test]
fn forward_updates_hidden_state() {
    let dev = test_device();
    let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xCAFE).expect("init");
    let h0 = trunk.snapshot_hidden().expect("h0");
    trunk
        .forward_snapshot(&synthetic_input(20_000_000))
        .expect("forward");
    let h1 = trunk.snapshot_hidden().expect("h1");
    let diff: f32 = h0.iter().zip(&h1).map(|(a, b)| (a - b).abs()).sum();
    assert!(
        diff > 1e-6,
        "hidden state must change after a forward step (got total |Δh| = {diff})"
    );
}

#[test]
fn forward_outputs_are_finite() {
    let dev = test_device();
    let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0x1234).expect("init");
    let (probs, proj) = trunk.forward_snapshot(&synthetic_input(50_000_000)).expect("forward");
    for &p in &probs {
        assert!(p.is_finite(), "prob {p} not finite");
    }
    for &v in &proj {
        assert!(v.is_finite(), "proj {v} not finite");
    }
}

#[test]
fn forward_layer_norm_proj_has_near_zero_mean() {
    let dev = test_device();
    let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0x5EED).expect("init");
    let (_, proj) = trunk.forward_snapshot(&synthetic_input(20_000_000)).expect("forward");
    let mean: f32 = proj.iter().sum::<f32>() / PROJ_DIM as f32;
    assert_relative_eq!(mean, 0.0, epsilon = 1e-3);
}

#[test]
fn forward_multi_step_advances_state_consistently() {
    // After many forward steps, the hidden state should still be finite and
    // within the bounded range expected by the CfC formula (tanh of pre).
    let dev = test_device();
    let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xBEEF).expect("init");
    let mut last_ts = 0u64;
    for k in 0..50 {
        last_ts += 20_000_000;
        let mut input = synthetic_input(last_ts);
        input.prev_ts_ns = last_ts - 20_000_000;
        // Slowly mutate trade_signed_vol so the input changes each step.
        input.trade_signed_vol = k as f32 * 0.1;
        trunk.forward_snapshot(&input).expect("forward");
    }
    let h = trunk.snapshot_hidden().expect("snapshot");
    for (i, v) in h.iter().enumerate() {
        assert!(v.is_finite(), "hidden[{i}] = {v} not finite after 50 steps");
    }
}