foxhunt/crates/ml-features/tests/ofi_parallel_bit_equiv_test.rs

//! Bit-equivalence test for `compute_ofi_per_bar_parallel` vs.
//! `compute_ofi_per_bar_sequential` (SP20, 2026-05-10).
//!
//! Asserts that the time-bucket-sharded parallel implementation produces
//! the SAME per-bar OFI features as the single-threaded reference, within
//! `1e-9` absolute tolerance per cell. In practice the result is
//! bit-identical because each shard's warmup phase replays the same
//! prefix of trades+snapshots into a freshly-initialised local
//! `OFICalculator`, so the rolling-window state at the warmup→core
//! boundary is identical to the sequential walk.
//!
//! Test strategy:
//!   1. Generate synthetic OHLCV bars + MBP-10 snapshots + trades with
//!      enough bars (>5000 + warmup) to exercise the rolling-window
//!      eviction path.
//!   2. Run `compute_ofi_per_bar_sequential` and capture the result.
//!   3. Run `compute_ofi_per_bar_parallel` with K = 4 and K = 8 shards.
//!   4. Assert max-abs-diff <= 1e-9 across all cells (32 dims × n bars).
//!
//! Why synthetic data: the test must be hermetic; we can't depend on
//! `test_data/futures-baseline-mbp10` being present in CI workspaces.
//! The synthetic generator constructs realistic-shaped MBP-10 books and
//! a Poisson-arrival trade tape with tick-rule sign so VPIN/Kyle/
//! trade-imb are all driven from real `feed_trade` calls.
//!
//! If this test fails: the parallel implementation has diverged from
//! sequential semantics. Do NOT relax the tolerance; instead, increase
//! the warmup size in `DEFAULT_OFI_PARALLEL_WARMUP_BARS` or audit the
//! shard boundaries in `compute_ofi_per_bar_parallel`. See the module
//! note at the top of `crates/ml-features/src/ofi_calculator.rs`.

use chrono::{DateTime, TimeZone, Utc};
use data::providers::databento::mbp10::{BidAskPair, Mbp10Snapshot};
use ml_features::{
    compute_ofi_per_bar_parallel, compute_ofi_per_bar_sequential, DbnTrade, OHLCVBar,
    PerBarOfiInputs, DEFAULT_OFI_PARALLEL_WARMUP_BARS, PER_BAR_OFI_DIM,
};

/// Generate synthetic OHLCV bars at 1-second cadence starting from a fixed epoch.
fn make_bars(n: usize, base_price: f64) -> Vec<OHLCVBar> {
    let mut bars = Vec::with_capacity(n);
    // Deterministic walk: ±1 tick per bar via index parity.
    let mut close = base_price;
    let start: DateTime<Utc> = Utc.with_ymd_and_hms(2025, 1, 1, 9, 30, 0).unwrap();
    for i in 0..n {
        let dir: f64 = if (i % 7) < 3 { 1.0 } else { -1.0 };
        let prev_close = close;
        close += dir * 0.25; // ES-like 0.25 tick
        let open = prev_close;
        let high = open.max(close) + 0.25;
        let low = open.min(close) - 0.25;
        bars.push(OHLCVBar {
            timestamp: start + chrono::Duration::seconds(i as i64),
            open,
            high,
            low,
            close,
            volume: 50.0 + (i % 17) as f64,
        });
    }
    bars
}

/// Generate MBP-10 snapshots aligned to bar timestamps. ~3 snapshots per bar
/// so the per-bar window has multiple snapshots (exercises micro_state
/// snapshot loop AND the OFIStats normalisation path with >0 history).
fn make_snapshots(bars: &[OHLCVBar], symbol: &str) -> Vec<Mbp10Snapshot> {
    let mut snaps = Vec::with_capacity(bars.len() * 3);
    for (i, bar) in bars.iter().enumerate() {
        let bar_ts_ns = bar.timestamp.timestamp_nanos_opt().unwrap_or(0) as u64;
        let prev_ts_ns = if i > 0 {
            bars[i - 1].timestamp.timestamp_nanos_opt().unwrap_or(0) as u64
        } else {
            bar_ts_ns.saturating_sub(1_000_000_000)
        };
        // Three offsets within the bar window (must be in (prev_ts, bar_ts]).
        let span = bar_ts_ns.saturating_sub(prev_ts_ns).max(3);
        let offsets = [span / 4, span / 2, span];
        for (j, off) in offsets.iter().enumerate() {
            let ts = prev_ts_ns + off;
            let mid = bar.close;
            let spread = 0.25;
            let mut levels = Vec::with_capacity(10);
            for level in 0..10 {
                let bid_px = ((mid - spread / 2.0 - 0.25 * level as f64) * 1e12) as i64;
                let ask_px = ((mid + spread / 2.0 + 0.25 * level as f64) * 1e12) as i64;
                // Asymmetric sizes so depth_imbalance, slopes, microprice are non-zero.
                let bid_sz = 100_u32 + (i as u32 % 23) + level as u32 * 5;
                let ask_sz = 100_u32 + ((i + 7) as u32 % 17) + level as u32 * 4;
                let bid_ct = 5_u32 + level as u32;
                let ask_ct = 6_u32 + level as u32;
                levels.push(BidAskPair {
                    bid_px,
                    bid_sz,
                    bid_ct,
                    ask_px,
                    ask_sz,
                    ask_ct,
                });
            }
            snaps.push(Mbp10Snapshot::new(
                symbol.to_string(),
                ts,
                levels,
                (i * 3 + j) as u32,
                1,
            ));
        }
    }
    snaps.sort_by_key(|s| s.timestamp);
    snaps
}

/// Generate ~5 trades per bar with deterministic side/volume so VPIN, Kyle,
/// and trade_imbalance are all driven by `feed_trade` (not the tick-rule
/// fallback).
fn make_trades(bars: &[OHLCVBar]) -> Vec<DbnTrade> {
    let mut trades = Vec::with_capacity(bars.len() * 5);
    for (i, bar) in bars.iter().enumerate() {
        let bar_ts_ns = bar.timestamp.timestamp_nanos_opt().unwrap_or(0) as u64;
        let prev_ts_ns = if i > 0 {
            bars[i - 1].timestamp.timestamp_nanos_opt().unwrap_or(0) as u64
        } else {
            bar_ts_ns.saturating_sub(1_000_000_000)
        };
        let span = bar_ts_ns.saturating_sub(prev_ts_ns).max(5);
        for k in 0..5 {
            let ts = prev_ts_ns + (span * (k + 1) as u64) / 6;
            let is_buy = ((i + k) % 3) != 0;
            let volume = 10_u64 + ((i + k * 11) as u64 % 20);
            let price = bar.close + if is_buy { 0.125 } else { -0.125 };
            trades.push(DbnTrade {
                timestamp: ts,
                price,
                volume,
                is_buy,
                instrument_id: 1,
            });
        }
    }
    trades.sort_by_key(|t| t.timestamp);
    trades
}

fn max_abs_diff(
    a: &[[f64; PER_BAR_OFI_DIM]],
    b: &[[f64; PER_BAR_OFI_DIM]],
) -> (f64, (usize, usize)) {
    assert_eq!(a.len(), b.len(), "row count mismatch");
    let mut max_diff = 0.0_f64;
    let mut where_at = (0_usize, 0_usize);
    for (i, (ra, rb)) in a.iter().zip(b.iter()).enumerate() {
        for k in 0..PER_BAR_OFI_DIM {
            let d = (ra[k] - rb[k]).abs();
            if d > max_diff {
                max_diff = d;
                where_at = (i, k);
            }
        }
    }
    (max_diff, where_at)
}

#[test]
fn parallel_matches_sequential_within_1e9_k4() {
    const WARMUP: usize = 50;
    // 5500 core bars + WARMUP ⇒ shard 0 sees no warmup overlap (it's the
    // start of the dataset), shards 1..3 each see 5000 bars of warmup.
    // This exercises the rolling-window eviction path.
    let n = 5500;
    let bars = make_bars(WARMUP + n, 5000.0);
    let snaps = make_snapshots(&bars, "ES.FUT");
    let trades = make_trades(&bars);

    let inputs = PerBarOfiInputs {
        bars: &bars,
        warmup: WARMUP,
        n,
        snapshots: &snaps,
        trades: Some(&trades),
    };

    let seq = compute_ofi_per_bar_sequential(&inputs);
    let par = compute_ofi_per_bar_parallel(&inputs, 4, DEFAULT_OFI_PARALLEL_WARMUP_BARS);

    assert_eq!(seq.len(), n);
    assert_eq!(par.len(), n);

    let (max_diff, (i, k)) = max_abs_diff(&seq, &par);
    assert!(
        max_diff <= 1e-9,
        "K=4: max abs diff {max_diff:e} exceeds 1e-9 at row {i}, dim {k} (seq={}, par={})",
        seq[i][k],
        par[i][k],
    );

    // Sanity: at least some non-zero output (otherwise the test is vacuous).
    let non_zero = seq.iter().filter(|r| r.iter().any(|&v| v != 0.0)).count();
    assert!(
        non_zero > n / 2,
        "expected most rows non-zero, got {non_zero}/{n}"
    );
}

#[test]
fn parallel_matches_sequential_within_1e9_k8() {
    const WARMUP: usize = 50;
    let n = 5500;
    let bars = make_bars(WARMUP + n, 5000.0);
    let snaps = make_snapshots(&bars, "ES.FUT");
    let trades = make_trades(&bars);

    let inputs = PerBarOfiInputs {
        bars: &bars,
        warmup: WARMUP,
        n,
        snapshots: &snaps,
        trades: Some(&trades),
    };

    let seq = compute_ofi_per_bar_sequential(&inputs);
    let par = compute_ofi_per_bar_parallel(&inputs, 8, DEFAULT_OFI_PARALLEL_WARMUP_BARS);

    let (max_diff, (i, k)) = max_abs_diff(&seq, &par);
    assert!(
        max_diff <= 1e-9,
        "K=8: max abs diff {max_diff:e} exceeds 1e-9 at row {i}, dim {k} (seq={}, par={})",
        seq[i][k],
        par[i][k],
    );
}

#[test]
fn parallel_k1_is_sequential() {
    const WARMUP: usize = 50;
    let n = 1500;
    let bars = make_bars(WARMUP + n, 5000.0);
    let snaps = make_snapshots(&bars, "ES.FUT");
    let trades = make_trades(&bars);

    let inputs = PerBarOfiInputs {
        bars: &bars,
        warmup: WARMUP,
        n,
        snapshots: &snaps,
        trades: Some(&trades),
    };

    let seq = compute_ofi_per_bar_sequential(&inputs);
    let par = compute_ofi_per_bar_parallel(&inputs, 1, DEFAULT_OFI_PARALLEL_WARMUP_BARS);

    let (max_diff, _) = max_abs_diff(&seq, &par);
    // K=1 must be byte-identical (same code path).
    assert_eq!(max_diff, 0.0, "K=1 should be byte-identical to sequential");
}

#[test]
fn parallel_handles_no_trades() {
    const WARMUP: usize = 50;
    let n = 800;
    let bars = make_bars(WARMUP + n, 5000.0);
    let snaps = make_snapshots(&bars, "ES.FUT");

    let inputs = PerBarOfiInputs {
        bars: &bars,
        warmup: WARMUP,
        n,
        snapshots: &snaps,
        trades: None,
    };

    let seq = compute_ofi_per_bar_sequential(&inputs);
    let par = compute_ofi_per_bar_parallel(&inputs, 4, DEFAULT_OFI_PARALLEL_WARMUP_BARS);

    let (max_diff, (i, k)) = max_abs_diff(&seq, &par);
    assert!(
        max_diff <= 1e-9,
        "no-trades: max abs diff {max_diff:e} at row {i}, dim {k}"
    );
}

/// Speedup smoke — `#[ignore]`'d because it's bench-shaped, not
/// assertion-shaped. Run via `cargo test ... -- --ignored --nocapture`.
///
/// At production scale (~50k bars × ~3 snapshots/bar × ~5 trades/bar) the
/// 5000-bar warmup amortises across each shard's core range so K=8 gives
/// meaningful speedup. For small N (≤8k) the warmup overhead dominates
/// and parallel can be slightly slower than sequential — this is by
/// design (warmup must be ≥ rolling-window depth for bit-equivalence).
///
/// `precompute_features.rs` typically runs at N ≥ 20k bars where the
/// speedup is realised.
#[test]
#[ignore = "bench-shaped; run with --ignored --nocapture for human review"]
fn parallel_speedup_smoke() {
    const WARMUP: usize = 50;
    for &n in &[8000_usize, 30_000, 80_000] {
        let bars = make_bars(WARMUP + n, 5000.0);
        let snaps = make_snapshots(&bars, "ES.FUT");
        let trades = make_trades(&bars);

        let inputs = PerBarOfiInputs {
            bars: &bars,
            warmup: WARMUP,
            n,
            snapshots: &snaps,
            trades: Some(&trades),
        };

        let t0 = std::time::Instant::now();
        let seq = compute_ofi_per_bar_sequential(&inputs);
        let dt_seq = t0.elapsed();

        let t1 = std::time::Instant::now();
        let par =
            compute_ofi_per_bar_parallel(&inputs, 8, DEFAULT_OFI_PARALLEL_WARMUP_BARS);
        let dt_par = t1.elapsed();

        let (max_diff, _) = max_abs_diff(&seq, &par);
        assert!(max_diff <= 1e-9, "speedup test diverged: {max_diff:e}");

        println!(
            "OFI per-bar smoke: n={n} bars, seq={:.2}ms, par(K=8)={:.2}ms, speedup={:.2}x",
            dt_seq.as_secs_f64() * 1000.0,
            dt_par.as_secs_f64() * 1000.0,
            dt_seq.as_secs_f64() / dt_par.as_secs_f64().max(1e-9),
        );
    }
}