diff --git a/crates/ml-features/src/alpha_pipeline.rs b/crates/ml-features/src/alpha_pipeline.rs index 2f032cffd..ebfc4404c 100644 --- a/crates/ml-features/src/alpha_pipeline.rs +++ b/crates/ml-features/src/alpha_pipeline.rs @@ -43,6 +43,7 @@ use std::collections::VecDeque; use data::providers::databento::mbp10::Mbp10Snapshot; +use rayon::prelude::*; use crate::OHLCVBar; use crate::Mbp10Trade; @@ -81,6 +82,16 @@ pub const ALPHA_FEATURE_DIM: usize = 134; /// Max trailing horizon for the trend scanner. Zero forward reads. const TREND_SCAN_MAX_HORIZON: usize = 40; +/// Bars of pre-roll per parallel chunk. Each chunk runs WARMUP_BARS leading +/// bars to warm its fresh aggregators (Hawkes excitation history, Bouchaud +/// trade-imbalance EMA, frac-diff FIR, LOB PCA covariance, microprice EMA, +/// spread_decomp running stats) before emitting feature rows. 2000 bars at +/// ~3-10 trades/bar gives 6k-20k trade events — sufficient to refill Hawkes's +/// 500-event history twice over and saturate the longer-memory EMAs. The +/// warmup overhead per chunk is bounded by `WARMUP_BARS / chunk_size` which +/// is well under 1% at production scale. +const WARMUP_BARS: usize = 2_000; + /// Extract alpha features for `n_output` bars starting at `bars[bar_start_offset]`. /// /// `bar_start_offset` is the WARMUP value used by the production pipeline @@ -90,6 +101,22 @@ const TREND_SCAN_MAX_HORIZON: usize = 40; /// /// Returns a `Vec>` of length `n_output`, each inner `Vec` having /// length `ALPHA_FEATURE_DIM`. +/// +/// # Parallelism +/// +/// The output range is split into `rayon::current_num_threads()` contiguous +/// chunks (one per CPU). Each chunk runs `process_chunk` with its own fresh +/// aggregator state, pre-rolling [`WARMUP_BARS`] leading bars without emission +/// so that the first emitted bar has well-saturated EMAs/FIRs/excitations. +/// Chunk 0 starts at `out_idx = 0` with cold state — identical to the original +/// sequential behaviour. Other chunks differ from a fully-serial run by +/// floating-point ε once warmup completes (state initialisers + EMA memory +/// loss are bounded by the pre-roll length). +/// +/// Trade-feed semantics are preserved: each trade is fed to a single chunk's +/// aggregators exactly once, partitioned by bar timestamp. `trade_cursor` is +/// seeded per chunk via `partition_point` on the previous warmup bar's +/// timestamp. pub fn extract_alpha_features( bars: &[OHLCVBar], snapshots: &[Mbp10Snapshot], @@ -97,6 +124,65 @@ pub fn extract_alpha_features( bar_start_offset: usize, n_output: usize, ) -> Vec> { + if n_output == 0 { + return Vec::new(); + } + + // ── Pre-index snapshots by timestamp (shared, read-only across chunks) ── + // Each bar gets the snapshot whose timestamp is the latest ≤ bar.timestamp. + let bar_snap_idx: Vec = bars + .iter() + .map(|bar| { + let bar_ts_ns = bar.timestamp.timestamp_nanos_opt().unwrap_or(0); + let idx = snapshots + .partition_point(|s| (s.timestamp as i64) <= bar_ts_ns); + idx.saturating_sub(1) + }) + .collect(); + + // ── Chunk the emit range ── + let n_chunks = rayon::current_num_threads().max(1).min(n_output); + let chunk_size = n_output.div_ceil(n_chunks); + let chunks: Vec<(usize, usize)> = (0..n_chunks) + .map(|k| { + let emit_start = k * chunk_size; + let emit_end = ((k + 1) * chunk_size).min(n_output); + (emit_start, emit_end) + }) + .filter(|(s, e)| s < e) + .collect(); + + // Parallel per-chunk computation; flatten into the final contiguous Vec. + chunks + .par_iter() + .flat_map_iter(|&(emit_start, emit_end)| { + process_chunk( + bars, + snapshots, + trades, + &bar_snap_idx, + bar_start_offset, + emit_start, + emit_end, + ) + .into_iter() + }) + .collect() +} + +/// Process one chunk [emit_start..emit_end) with fresh aggregator state, +/// pre-rolling [`WARMUP_BARS`] leading bars to saturate stateful blocks. +fn process_chunk( + bars: &[OHLCVBar], + snapshots: &[Mbp10Snapshot], + trades: &[Mbp10Trade], + bar_snap_idx: &[usize], + bar_start_offset: usize, + emit_start: usize, + emit_end: usize, +) -> Vec> { + let warmup_start = emit_start.saturating_sub(WARMUP_BARS); + // ── Stateful aggregators (persistent across bars) ── // Block A-E factory aggregators let mut bar_window: VecDeque = VecDeque::with_capacity(BAR_WINDOW); @@ -124,22 +210,27 @@ pub fn extract_alpha_features( // OFICalculator instance for stateful VPIN bucket tracking. let mut ofi_calc = OFICalculator::new(); - // ── Pre-index snapshots by timestamp for fast lookup ── - // Each bar gets the snapshot whose timestamp is the latest ≤ bar.timestamp. - let bar_snap_idx: Vec = bars - .iter() - .map(|bar| { - let bar_ts_ns = bar.timestamp.timestamp_nanos_opt().unwrap_or(0); - let idx = snapshots - .partition_point(|s| (s.timestamp as i64) <= bar_ts_ns); - idx.saturating_sub(1) + // ── Seed trade_cursor at the first trade belonging to this chunk ── + // For chunk 0 (warmup_start == 0) the cursor starts at 0 so the trade-feed + // inner loop replays *all* pre-window trades into the aggregators on the + // first iteration — preserving the original sequential semantics. + // + // For other chunks we partition past the previous warmup bar's timestamp + // so each trade is fed to exactly one chunk's aggregators. + let mut trade_cursor: usize = if warmup_start == 0 { + 0 + } else { + let prev_bar_ts = bars[warmup_start - 1 + bar_start_offset] + .timestamp + .timestamp_nanos_opt() + .unwrap_or(0); + trades.partition_point(|t| { + t.timestamp.timestamp_nanos_opt().unwrap_or(0) <= prev_bar_ts }) - .collect(); + }; + let mut features: Vec> = Vec::with_capacity(emit_end - emit_start); - let mut trade_cursor: usize = 0; - let mut features: Vec> = Vec::with_capacity(n_output); - - for out_idx in 0..n_output { + for out_idx in warmup_start..emit_end { let bar_idx = out_idx + bar_start_offset; let bar = &bars[bar_idx]; let bar_ts_ns = bar.timestamp.timestamp_nanos_opt().unwrap_or(0); @@ -395,7 +486,12 @@ pub fn extract_alpha_features( ALPHA_FEATURE_DIM ); - features.push(row); + // Suppress emission for the leading warmup bars; their job is to + // saturate Hawkes excitation / Bouchaud EMA / frac-diff FIR / LOB PCA + // / microprice / spread_decomp state before the first emit_start bar. + if out_idx >= emit_start { + features.push(row); + } } features