78a9e08358
Replaces Option<u32> instrument_id_filter with InstrumentFilter enum {All,
Id(u32), FrontMonth}. FrontMonth runs a two-pass detect over the DBN
stream: pass 1 counts instrument_ids and collects SymbolMapping records,
picks the dominant id, validates it resolves to an ES contract via regex
ES[FGHJKMNQUVXZ]\d{1,2}; pass 2 streams the filtered records.
Motivated by alpha-perception-k54wd: a single-id filter on parent-symbol
ES.FUT data caught Q1 2024 (kept=73M) but kept=0 for Q2-Q9 because ES
front-month rolls quarterly (ESH4 -> ESM4 -> ESU4 -> ESZ4 ...). FrontMonth
self-tunes across the rolls without needing a per-file id table.
Sidecar keys distinguish modes: mbp10 / mbp10_instr<id> / mbp10_front_month.
CLI flag renamed --instrument-id -> --instrument-mode {all,id=N,front-month}
with matching parameter rename in argo-alpha-perception.sh + template.
1442 lines
55 KiB
Rust
1442 lines
55 KiB
Rust
//! MBP-10 Data Loader Helper
|
||
//!
|
||
//! Simple integration layer between DBN MBP-10 parser and OFI feature extraction.
|
||
//! Provides synchronous loading wrapper and snapshot window selection for real-time
|
||
//! Order Flow Imbalance (OFI) calculation.
|
||
//!
|
||
//! # Purpose
|
||
//!
|
||
//! - Load MBP-10 snapshots from DBN files
|
||
//! - Provide snapshot windows for OFI calculation
|
||
//! - Abstract async DBN parser complexity
|
||
//!
|
||
//! # Example
|
||
//!
|
||
//! ```no_run
|
||
//! use ml::features::mbp10_loader::load_mbp10_snapshots_sync;
|
||
//! use ml::features::ofi_calculator::OFICalculator;
|
||
//! use std::path::Path;
|
||
//!
|
||
//! let snapshots = load_mbp10_snapshots_sync(Path::new("test_data/ES.FUT.mbp10.dbn"))?;
|
||
//! let mut calculator = OFICalculator::new();
|
||
//!
|
||
//! for snapshot in &snapshots {
|
||
//! let features = calculator.calculate(snapshot)?;
|
||
//! // Use features for model input
|
||
//! }
|
||
//! ```
|
||
|
||
use data::providers::databento::{
|
||
dbn_parser::{DbnParser, InstrumentFilter},
|
||
mbp10::{BidAskPair, Mbp10Snapshot},
|
||
};
|
||
use std::path::Path;
|
||
|
||
use chrono::{DateTime, Utc};
|
||
|
||
use crate::alternative_bars::ImbalanceBarSampler;
|
||
use crate::ofi_calculator::OFICalculator;
|
||
use crate::trades_loader::load_trades_sync;
|
||
use crate::OHLCVBar;
|
||
use crate::MLError;
|
||
|
||
/// Load MBP-10 snapshots from DBN file (synchronous wrapper)
|
||
///
|
||
/// This is a blocking wrapper around the async `DbnParser::parse_mbp10_file()`.
|
||
/// Use this for simple synchronous contexts like testing or batch processing.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `file_path` - Path to the DBN file containing MBP-10 data
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// * `Ok(Vec<Mbp10Snapshot>)` - All snapshots from the file
|
||
/// * `Err(MLError)` - If file cannot be read or parsed
|
||
///
|
||
/// # Performance
|
||
///
|
||
/// - Target: <10ms for typical 180-day files
|
||
/// - Memory: ~1KB per snapshot (10 levels × 48 bytes)
|
||
///
|
||
pub fn load_mbp10_snapshots_sync(file_path: &Path) -> Result<Vec<Mbp10Snapshot>, MLError> {
|
||
let parser = DbnParser::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
|
||
})?;
|
||
|
||
// Use tokio runtime for async operation
|
||
let runtime = tokio::runtime::Runtime::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create async runtime: {}", e))
|
||
})?;
|
||
|
||
runtime.block_on(async {
|
||
parser
|
||
.parse_mbp10_file(file_path)
|
||
.await
|
||
.map_err(|e| MLError::InsufficientData(format!("Failed to parse MBP-10 file: {}", e)))
|
||
})
|
||
}
|
||
|
||
/// Compute OFI features from an MBP-10 file in a single streaming pass.
|
||
///
|
||
/// This is ~10x faster than `load_mbp10_snapshots_sync` + separate OFI iteration
|
||
/// because it avoids materializing millions of intermediate `Mbp10Snapshot` clones.
|
||
/// The parser feeds snapshots by reference directly to the OFI calculator.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `file_path` - Path to the .dbn or .dbn.zst MBP-10 file
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// * `Ok(Vec<[f64; 8]>)` - OFI feature arrays (one per aggregated snapshot)
|
||
/// * `Err(MLError)` - If file cannot be read or parsed
|
||
///
|
||
pub fn compute_ofi_from_file(file_path: &Path) -> Result<Vec<[f64; 8]>, MLError> {
|
||
let parser = DbnParser::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
|
||
})?;
|
||
|
||
let mut calculator = OFICalculator::new();
|
||
let mut features = Vec::new();
|
||
|
||
let snapshot_count = parser
|
||
.parse_mbp10_streaming(file_path, 100, InstrumentFilter::All, |snapshot| {
|
||
if let Ok(f) = calculator.calculate(snapshot) {
|
||
features.push(f.to_array());
|
||
} // Skip failed calculations (e.g. first snapshot with no prev)
|
||
})
|
||
.map_err(|e| MLError::InsufficientData(format!("MBP-10 streaming parse failed: {}", e)))?;
|
||
|
||
tracing::debug!(
|
||
"Computed {} OFI features from {} snapshots in {:?}",
|
||
features.len(),
|
||
snapshot_count,
|
||
file_path.file_name().unwrap_or_default()
|
||
);
|
||
|
||
Ok(features)
|
||
}
|
||
|
||
/// Compute OFI features from MBP-10 file enriched with trade data.
|
||
///
|
||
/// Interleaves trade events (from a `.dbn.zst` trades file) into the MBP-10
|
||
/// streaming pipeline. Before each MBP-10 snapshot is processed, all trades
|
||
/// with timestamps up to that snapshot are fed into the OFI calculator,
|
||
/// populating VPIN, Kyle's Lambda, and trade imbalance with real trade data
|
||
/// instead of leaving them at zero.
|
||
///
|
||
/// Falls back to `compute_ofi_from_file()` if the trade file doesn't exist
|
||
/// or fails to load.
|
||
pub fn compute_ofi_with_trades(
|
||
mbp10_file: &Path,
|
||
trades_file: &Path,
|
||
) -> Result<Vec<[f64; 8]>, MLError> {
|
||
// Load all trades upfront (sorted by timestamp). Trade files are much
|
||
// smaller than MBP-10 (~120 MB vs ~5+ GB), so this is fine in memory.
|
||
let trades = match load_trades_sync(trades_file) {
|
||
Ok(t) => {
|
||
tracing::info!(
|
||
"Loaded {} trades from {} for OFI enrichment",
|
||
t.len(),
|
||
trades_file.file_name().unwrap_or_default().to_string_lossy()
|
||
);
|
||
t
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
"Failed to load trades from {}: {} \u{2014} falling back to MBP-10 only",
|
||
trades_file.display(),
|
||
e
|
||
);
|
||
return compute_ofi_from_file(mbp10_file);
|
||
}
|
||
};
|
||
|
||
let parser = DbnParser::new().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create DBN parser: {}", e))
|
||
})?;
|
||
|
||
let mut calculator = OFICalculator::new();
|
||
let mut features = Vec::new();
|
||
let mut trade_cursor: usize = 0;
|
||
|
||
let snapshot_count = parser
|
||
.parse_mbp10_streaming(mbp10_file, 100, InstrumentFilter::All, |snapshot| {
|
||
// Feed all trades with timestamp <= this snapshot's timestamp.
|
||
// This ensures VPIN/Kyle's Lambda/trade_imbalance are populated
|
||
// before the OFI calculator processes the snapshot.
|
||
while trade_cursor < trades.len()
|
||
&& trades[trade_cursor].timestamp <= snapshot.timestamp
|
||
{
|
||
let t = &trades[trade_cursor];
|
||
calculator.feed_trade(t.price, t.volume, t.is_buy);
|
||
trade_cursor += 1;
|
||
}
|
||
|
||
if let Ok(f) = calculator.calculate(snapshot) {
|
||
features.push(f.to_array());
|
||
} // Skip failed calculations (e.g. first snapshot)
|
||
})
|
||
.map_err(|e| MLError::InsufficientData(format!("MBP-10 streaming parse failed: {}", e)))?;
|
||
|
||
tracing::info!(
|
||
"Computed {} OFI features from {} snapshots + {} trades (cursor at {}/{})",
|
||
features.len(),
|
||
snapshot_count,
|
||
trades.len(),
|
||
trade_cursor,
|
||
trades.len()
|
||
);
|
||
|
||
Ok(features)
|
||
}
|
||
|
||
/// Find matching trade file for an MBP-10 file in the trades directory.
|
||
///
|
||
/// Matches by filename stem: `ES.FUT_2024-Q1.dbn.zst` in MBP-10 dir
|
||
/// → looks for `ES.FUT_2024-Q1.dbn.zst` in trades dir (maintaining
|
||
/// the same subdirectory structure).
|
||
fn find_matching_trade_file(mbp10_file: &Path, trades_dir: &Path) -> Option<std::path::PathBuf> {
|
||
let file_name = mbp10_file.file_name()?;
|
||
|
||
// Try exact match: trades_dir/SYMBOL/filename
|
||
// MBP-10 files are typically under mbp10_dir/ES.FUT/ES.FUT_2024-Q1.dbn.zst
|
||
// Trades would be under trades_dir/ES.FUT/ES.FUT_2024-Q1.dbn.zst
|
||
if let Some(parent) = mbp10_file.parent() {
|
||
if let Some(symbol_dir) = parent.file_name() {
|
||
let candidate = trades_dir.join(symbol_dir).join(file_name);
|
||
if candidate.exists() {
|
||
return Some(candidate);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Fallback: flat structure — trades_dir/filename
|
||
let candidate = trades_dir.join(file_name);
|
||
if candidate.exists() {
|
||
return Some(candidate);
|
||
}
|
||
|
||
None
|
||
}
|
||
|
||
/// Load OFI features from all MBP-10 files in a directory using parallel streaming.
|
||
///
|
||
/// Files are processed concurrently with rayon. Each file streams through the DBN parser
|
||
/// and computes OFI features inline — no intermediate `Vec<Mbp10Snapshot>` allocation.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `mbp10_dir` - Directory containing .dbn or .dbn.zst MBP-10 files
|
||
/// * `file_collector` - Function to collect files from the directory (supports both
|
||
/// flat and recursive collection patterns used by DQN/PPO adapters)
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// * `Some(Vec<[f64; 8]>)` - Concatenated OFI features from all files (file order preserved)
|
||
/// * `None` - If directory doesn't exist, has no files, or all files failed
|
||
///
|
||
#[allow(clippy::cognitive_complexity)]
|
||
pub fn load_ofi_features_parallel<F>(
|
||
mbp10_dir: &Path,
|
||
trades_dir: Option<&Path>,
|
||
file_collector: F,
|
||
) -> Option<Vec<[f64; 8]>>
|
||
where
|
||
F: FnOnce(&Path) -> Vec<std::path::PathBuf>,
|
||
{
|
||
use rayon::prelude::*;
|
||
|
||
if !mbp10_dir.exists() {
|
||
tracing::info!(
|
||
"No MBP10 directory at {}, OFI features will be zero-padded",
|
||
mbp10_dir.display()
|
||
);
|
||
return None;
|
||
}
|
||
|
||
let mbp10_files = file_collector(mbp10_dir);
|
||
if mbp10_files.is_empty() {
|
||
tracing::info!("No MBP10 .dbn files found in {}", mbp10_dir.display());
|
||
return None;
|
||
}
|
||
|
||
tracing::info!(
|
||
"Loading OFI features from {} MBP10 files (parallel streaming{})",
|
||
mbp10_files.len(),
|
||
if trades_dir.is_some() { " + trades" } else { "" }
|
||
);
|
||
|
||
let start = std::time::Instant::now();
|
||
|
||
// Process files in parallel — each file gets its own OFI calculator (independent state).
|
||
// If a trades directory is provided, each MBP-10 file looks for a matching trade file
|
||
// to enrich VPIN, Kyle's Lambda, and trade imbalance with real trade data.
|
||
let per_file_results: Vec<_> = mbp10_files
|
||
.par_iter()
|
||
.map(|file| {
|
||
let result = if let Some(t_dir) = trades_dir {
|
||
if let Some(trade_file) = find_matching_trade_file(file, t_dir) {
|
||
compute_ofi_with_trades(file, &trade_file)
|
||
} else {
|
||
tracing::debug!(
|
||
"No matching trade file for {}, using MBP-10 only",
|
||
file.file_name().unwrap_or_default().to_string_lossy()
|
||
);
|
||
compute_ofi_from_file(file)
|
||
}
|
||
} else {
|
||
compute_ofi_from_file(file)
|
||
};
|
||
|
||
match result {
|
||
Ok(features) => {
|
||
tracing::info!(
|
||
" {} \u{2192} {} OFI features",
|
||
file.file_name().unwrap_or_default().to_string_lossy(),
|
||
features.len()
|
||
);
|
||
features
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!("Failed to compute OFI from {}: {}", file.display(), e);
|
||
Vec::new()
|
||
}
|
||
}
|
||
})
|
||
.collect();
|
||
|
||
// Concatenate results in file order
|
||
let total_features: usize = per_file_results.iter().map(|v| v.len()).sum();
|
||
let mut all_ofi = Vec::with_capacity(total_features);
|
||
for features in per_file_results {
|
||
all_ofi.extend_from_slice(&features);
|
||
}
|
||
|
||
if all_ofi.is_empty() {
|
||
return None;
|
||
}
|
||
|
||
let elapsed = start.elapsed();
|
||
tracing::info!(
|
||
"Computed {} OFI feature vectors from {} files in {:.1}s ({:.0} features/sec)",
|
||
all_ofi.len(),
|
||
mbp10_files.len(),
|
||
elapsed.as_secs_f64(),
|
||
all_ofi.len() as f64 / elapsed.as_secs_f64()
|
||
);
|
||
|
||
Some(all_ofi)
|
||
}
|
||
|
||
/// Trade tick extracted from MBP-10 data.
|
||
///
|
||
/// MBP-10 records include order book updates AND trades. When an MBP-10 record
|
||
/// has `action == b'T'` it represents a trade execution. We extract the price,
|
||
/// size, timestamp, and buy/sell classification for feeding into the
|
||
/// `ImbalanceBarSampler`.
|
||
#[derive(Debug, Clone)]
|
||
pub struct Mbp10Trade {
|
||
/// Trade price (f64, converted from fixed-point)
|
||
pub price: f64,
|
||
/// Trade volume in contracts
|
||
pub volume: f64,
|
||
/// Event timestamp
|
||
pub timestamp: DateTime<Utc>,
|
||
/// True if buyer-initiated (side == 'B'), false otherwise
|
||
pub is_buy: bool,
|
||
/// Databento `instrument_id` for the contract this trade belongs to.
|
||
/// Used by `mbp10_to_imbalance_bars` to filter to the front-month
|
||
/// (highest-volume) contract per file, mirroring the per-file pattern
|
||
/// in `precompute_features.rs:354`. Populated from `mbp10.hd.instrument_id`
|
||
/// in `extract_trades_from_dbn_file`; supplied by the caller in
|
||
/// `extract_trades_from_snapshots` (snapshot records don't carry per-
|
||
/// record IDs, so the caller passes the stream's contract id explicitly).
|
||
pub instrument_id: u32,
|
||
}
|
||
|
||
/// Extract trade ticks from MBP-10 snapshots using consecutive snapshot diffs.
|
||
///
|
||
/// MBP-10 snapshots don't directly carry trade action/side info (those live on
|
||
/// the raw `dbn::Mbp10Msg` records). This function uses a heuristic:
|
||
/// when `trade_count` increases between consecutive snapshots, a trade occurred.
|
||
/// Direction is inferred by comparing mid-prices (tick rule).
|
||
///
|
||
/// `instrument_id` is required — `Mbp10Snapshot` does NOT carry per-record
|
||
/// instrument IDs, so the caller must supply the contract this snapshot
|
||
/// stream belongs to (typically captured at the data-acquisition site
|
||
/// alongside `symbol`). Stamped onto every emitted `Mbp10Trade` so downstream
|
||
/// front-month filtering (`filter_front_month_mbp10`) works uniformly across
|
||
/// both extraction paths.
|
||
///
|
||
/// For higher-fidelity trade extraction, use [`extract_trades_from_dbn_file`]
|
||
/// which reads raw MBP-10 records with their action/side fields.
|
||
pub fn extract_trades_from_snapshots(
|
||
snapshots: &[Mbp10Snapshot],
|
||
instrument_id: u32,
|
||
) -> Vec<Mbp10Trade> {
|
||
let mut trades = Vec::new();
|
||
let mut prev_trade_count: u32 = 0;
|
||
let mut prev_mid: f64 = 0.0;
|
||
|
||
for snap in snapshots {
|
||
let mid = snap.mid_price();
|
||
if mid <= 0.0 {
|
||
continue;
|
||
}
|
||
|
||
let new_trades = snap.trade_count.saturating_sub(prev_trade_count);
|
||
if new_trades > 0 && prev_mid > 0.0 {
|
||
// Tick rule: price up => buy, price down => sell, unchanged => previous direction
|
||
let is_buy = mid >= prev_mid;
|
||
// Best ask size as proxy for trade volume (actual volume not in snapshot)
|
||
let volume = if !snap.levels.is_empty() {
|
||
snap.levels[0].bid_sz.max(1) as f64
|
||
} else {
|
||
1.0
|
||
};
|
||
|
||
let ts_secs = (snap.timestamp / 1_000_000_000) as i64;
|
||
let ts_nanos = (snap.timestamp % 1_000_000_000) as u32;
|
||
if let Some(dt) = DateTime::<Utc>::from_timestamp(ts_secs, ts_nanos) {
|
||
trades.push(Mbp10Trade {
|
||
price: mid,
|
||
volume,
|
||
timestamp: dt,
|
||
is_buy,
|
||
instrument_id,
|
||
});
|
||
}
|
||
}
|
||
|
||
prev_trade_count = snap.trade_count;
|
||
prev_mid = mid;
|
||
}
|
||
|
||
trades
|
||
}
|
||
|
||
/// Extract trade ticks directly from a DBN MBP-10 file.
|
||
///
|
||
/// Reads raw `Mbp10Msg` records and filters for trade actions (`action == b'T'`).
|
||
/// This provides higher-fidelity trades than [`extract_trades_from_snapshots`]
|
||
/// because it uses the actual action and side fields from each record.
|
||
pub fn extract_trades_from_dbn_file(file_path: &Path) -> Result<Vec<Mbp10Trade>, MLError> {
|
||
use dbn::decode::{DbnDecoder, DecodeRecordRef};
|
||
use dbn::RecordRefEnum;
|
||
use std::fs::File;
|
||
use std::io::BufReader;
|
||
|
||
let file = File::open(file_path).map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to open MBP-10 file {:?}: {}", file_path, e))
|
||
})?;
|
||
|
||
let is_zstd = file_path.to_string_lossy().ends_with(".dbn.zst");
|
||
let reader: Box<dyn std::io::Read> = if is_zstd {
|
||
Box::new(
|
||
zstd::Decoder::new(BufReader::new(file)).map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create zstd decoder: {}", e))
|
||
})?,
|
||
)
|
||
} else {
|
||
Box::new(BufReader::new(file))
|
||
};
|
||
|
||
let mut decoder = DbnDecoder::new(reader).map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to create DBN decoder: {}", e))
|
||
})?;
|
||
|
||
let mut trades = Vec::new();
|
||
let mut last_mid: f64 = 0.0;
|
||
|
||
loop {
|
||
match decoder.decode_record_ref() {
|
||
Ok(Some(record)) => {
|
||
let record_enum = record.as_enum().map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to convert record: {}", e))
|
||
})?;
|
||
|
||
if let RecordRefEnum::Mbp10(mbp10) = record_enum {
|
||
let action = mbp10.action as u8;
|
||
|
||
// Compute mid-price from this record's top-of-book levels
|
||
// MBP-10 records carry 10 levels; index 0 is best bid/ask
|
||
let bid_px = BidAskPair::price_to_f64(mbp10.levels[0].bid_px);
|
||
let ask_px = BidAskPair::price_to_f64(mbp10.levels[0].ask_px);
|
||
let mid = if bid_px > 0.0 && ask_px > 0.0 {
|
||
(bid_px + ask_px) / 2.0
|
||
} else {
|
||
last_mid
|
||
};
|
||
if mid > 0.0 {
|
||
last_mid = mid;
|
||
}
|
||
|
||
// Trade action: 'T' (0x54) = trade execution
|
||
if action == b'T' {
|
||
let price = mbp10.price as f64 * 1e-9;
|
||
let volume = mbp10.size as f64;
|
||
let is_buy = mbp10.side == b'B' as i8;
|
||
|
||
let ts_secs = (mbp10.hd.ts_event / 1_000_000_000) as i64;
|
||
let ts_nanos = (mbp10.hd.ts_event % 1_000_000_000) as u32;
|
||
if let Some(dt) = DateTime::<Utc>::from_timestamp(ts_secs, ts_nanos) {
|
||
if price > 0.0 && volume > 0.0 {
|
||
trades.push(Mbp10Trade {
|
||
price,
|
||
volume,
|
||
timestamp: dt,
|
||
is_buy,
|
||
instrument_id: mbp10.hd.instrument_id,
|
||
});
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
Ok(None) => break,
|
||
Err(e) => {
|
||
return Err(MLError::InsufficientData(format!(
|
||
"DBN decode error: {}",
|
||
e
|
||
)));
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(trades)
|
||
}
|
||
|
||
/// Filter `Mbp10Trade`s to keep only the front-month (highest cumulative
|
||
/// volume) contract. Mirrors `trades_loader::filter_front_month` for
|
||
/// `DbnTrade` — same logic on `Mbp10Trade.instrument_id`. Apply per-file
|
||
/// because each quarterly file has a different front-month contract.
|
||
fn filter_front_month_mbp10(trades: &[Mbp10Trade]) -> Vec<Mbp10Trade> {
|
||
if trades.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
let mut volume_by_id: std::collections::HashMap<u32, u64> =
|
||
std::collections::HashMap::new();
|
||
for trade in trades {
|
||
*volume_by_id.entry(trade.instrument_id).or_insert(0) += trade.volume as u64;
|
||
}
|
||
let front_month_id = volume_by_id
|
||
.into_iter()
|
||
.max_by_key(|&(_, vol)| vol)
|
||
.map(|(id, _)| id)
|
||
.unwrap_or(0);
|
||
trades
|
||
.iter()
|
||
.filter(|t| t.instrument_id == front_month_id)
|
||
.cloned()
|
||
.collect()
|
||
}
|
||
|
||
/// Recursively collect all .dbn and .dbn.zst files from a directory.
|
||
fn collect_dbn_files(dir: &Path) -> Vec<std::path::PathBuf> {
|
||
let mut files = Vec::new();
|
||
if let Ok(entries) = std::fs::read_dir(dir) {
|
||
for entry in entries.flatten() {
|
||
let path = entry.path();
|
||
if path.is_dir() {
|
||
files.extend(collect_dbn_files(&path));
|
||
} else if path.extension().and_then(|s| s.to_str()) == Some("dbn")
|
||
|| path.to_string_lossy().ends_with(".dbn.zst")
|
||
{
|
||
files.push(path);
|
||
}
|
||
}
|
||
}
|
||
files
|
||
}
|
||
|
||
// ══════════════════════════════════════════════════════════════════════════════
|
||
// Imbalance Bar Cache
|
||
//
|
||
// Computing imbalance bars from MBP-10 data requires decompressing ~19GB of
|
||
// .dbn.zst files (~450s). This cache stores the resulting `Vec<OHLCVBar>` to
|
||
// `/tmp/.foxhunt_imbalance_cache/<hash>.bars.bin` using bincode.
|
||
//
|
||
// Cache key: DefaultHasher of (mbp10_dir, symbol, threshold bits, alpha bits).
|
||
// Invalidation: cache mtime vs. newest .dbn/.dbn.zst file in the source dir.
|
||
// ══════════════════════════════════════════════════════════════════════════════
|
||
|
||
/// Compute a cache path for an imbalance bar set.
|
||
fn imbalance_cache_path(mbp10_dir: &Path, symbol: &str, threshold: f64, alpha: f64) -> std::path::PathBuf {
|
||
use std::collections::hash_map::DefaultHasher;
|
||
use std::hash::{Hash, Hasher};
|
||
let mut hasher = DefaultHasher::new();
|
||
mbp10_dir.hash(&mut hasher);
|
||
symbol.hash(&mut hasher);
|
||
threshold.to_bits().hash(&mut hasher);
|
||
alpha.to_bits().hash(&mut hasher);
|
||
let hash = hasher.finish();
|
||
std::path::PathBuf::from(format!("/tmp/.foxhunt_imbalance_cache/{:016x}.bars.bin", hash))
|
||
}
|
||
|
||
/// Check whether the cache file is newer than every `.dbn`/`.dbn.zst` file under `source_dir`.
|
||
fn is_imbalance_cache_valid(cache: &Path, source_dir: &Path) -> bool {
|
||
let cache_mtime = match std::fs::metadata(cache) {
|
||
Ok(m) => match m.modified() {
|
||
Ok(t) => t,
|
||
Err(_) => return false,
|
||
},
|
||
Err(_) => return false,
|
||
};
|
||
|
||
fn walk_dbn_mtimes(dir: &Path, newest: &mut std::time::SystemTime) -> std::io::Result<()> {
|
||
for entry in std::fs::read_dir(dir)? {
|
||
let entry = entry?;
|
||
let path = entry.path();
|
||
if path.is_dir() {
|
||
walk_dbn_mtimes(&path, newest)?;
|
||
} else if path.extension().and_then(|s| s.to_str()) == Some("dbn")
|
||
|| path.to_string_lossy().ends_with(".dbn.zst")
|
||
{
|
||
if let Ok(meta) = std::fs::metadata(&path) {
|
||
if let Ok(mtime) = meta.modified() {
|
||
if mtime > *newest {
|
||
*newest = mtime;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
let mut newest_source = std::time::SystemTime::UNIX_EPOCH;
|
||
if walk_dbn_mtimes(source_dir, &mut newest_source).is_err() {
|
||
return false;
|
||
}
|
||
|
||
// Cache is valid if it was created after all source files were last modified
|
||
cache_mtime >= newest_source
|
||
}
|
||
|
||
/// Default minimum bars-per-rayon-task for parallel imbalance-bar OHLCV
|
||
/// reduction. Below this size the parallel overhead (rayon task spawning +
|
||
/// bar-vec allocation) dominates and sequential is faster. Empirically tuned
|
||
/// against `IMBALANCE_BAR_MIN_BARS_PER_TASK` env override.
|
||
pub const DEFAULT_IMBALANCE_BAR_MIN_BARS_PER_TASK: usize = 256;
|
||
|
||
/// Run the imbalance-bar sampler over a contiguous, pre-sorted trade slice
|
||
/// using a **two-pass** parallel decomposition.
|
||
///
|
||
/// Why two-pass (NOT time-bucket sharding)
|
||
/// =======================================
|
||
/// The OFI parallelisation in `compute_ofi_per_bar_parallel` works because
|
||
/// per-bar OFI features are derived from BOUNDED rolling windows
|
||
/// (VPIN ≤50, Kyle ≤100, etc.). After ≥window-size warmup updates, two
|
||
/// rolling buffers started from different initial states converge to
|
||
/// bit-identical contents.
|
||
///
|
||
/// `ImbalanceBarSampler` is FUNDAMENTALLY DIFFERENT: its `cumulative_imbalance`
|
||
/// is a path-integral that resets only when crossing ±`threshold`. There is
|
||
/// NO bounded-lookback window that determines the state. Two replays of the
|
||
/// same trade tape starting from different cum_imb offsets emit at DIFFERENT
|
||
/// trade indices, and the offset is NOT guaranteed to vanish at any future
|
||
/// trade — even after many emissions, a bounded phase difference can persist
|
||
/// indefinitely. Empirically this surfaces as a ±1-bar count drift at low
|
||
/// emission density (high `threshold`), as confirmed by the K=4 high-threshold
|
||
/// run of `imbalance_bars_parallel_high_threshold_few_emissions` against an
|
||
/// earlier time-bucket-sharded prototype.
|
||
///
|
||
/// Two-pass decomposition
|
||
/// ----------------------
|
||
/// **Pass 1 (sequential, lightweight)**: Walk the trade tape ONCE tracking
|
||
/// only `cum_imb`, `prev_price`, `last_direction`. Record the trade index at
|
||
/// the END of every bar emission (the trade that triggered the emission). No
|
||
/// OHLCV state, no Vec<OHLCVBar> allocation, no per-trade conditional bar
|
||
/// construction. This pass is `O(N)` simple arithmetic — for 50k–500k trades
|
||
/// it runs in a few ms.
|
||
///
|
||
/// **Pass 2 (parallel rayon)**: Each emission segment `[boundary[i-1]+1,
|
||
/// boundary[i]]` produces exactly one bar via independent OHLCV reduction
|
||
/// over its trade slice (open=first.price, close=last.price, high=max,
|
||
/// low=min, volume=sum, timestamp=first.timestamp). Segments are fully
|
||
/// independent — `par_iter()` over segment ranges is trivially correct.
|
||
///
|
||
/// Bit-equivalence
|
||
/// ---------------
|
||
/// Pass 1 mirrors the sequential `ImbalanceBarSampler::update` direction +
|
||
/// imbalance arithmetic exactly (same tie-break, same comparisons), so the
|
||
/// emission boundary set is identical to the sequential walk. Pass 2's per-
|
||
/// segment reduction matches the sampler's per-bar OHLCV update verbatim
|
||
/// (open = first non-zero-volume trade in segment, high/low/close from same
|
||
/// trade scan, volume = sum, timestamp = open's timestamp). See
|
||
/// `tests/imbalance_bars_parallel_bit_equiv_test.rs` for end-to-end
|
||
/// verification at K=4 / K=8 against `imbalance_bars_sequential`.
|
||
///
|
||
/// Zero-volume trade handling
|
||
/// --------------------------
|
||
/// `ImbalanceBarSampler::update` early-returns on `volume == 0` (line 484-486
|
||
/// of `alternative_bars.rs`). Pass 1 must replicate this exactly — zero-
|
||
/// volume trades are SKIPPED (no direction update, no imbalance update, no
|
||
/// `prev_price` update). Pass 2 also skips zero-volume trades when computing
|
||
/// OHLCV (in particular the segment's "first non-zero-volume trade" defines
|
||
/// `open` and `timestamp`).
|
||
///
|
||
/// `shard_count` is unused (kept in the signature for symmetry with
|
||
/// `compute_ofi_per_bar_parallel`); rayon's `par_iter` uses the global thread
|
||
/// pool. `min_bars_per_task` controls the parallel/sequential cutoff: below
|
||
/// this segment count, fall back to a sequential reduction. Pass `0` for
|
||
/// the default.
|
||
///
|
||
/// Returns bars in trade-index order (== chronological order for pre-sorted
|
||
/// input). Empty trade slice returns empty vec.
|
||
pub fn imbalance_bars_parallel(
|
||
trades: &[Mbp10Trade],
|
||
threshold: f64,
|
||
_shard_count: usize,
|
||
min_bars_per_task: usize,
|
||
) -> Vec<OHLCVBar> {
|
||
use rayon::prelude::*;
|
||
|
||
if trades.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
let min_bars_per_task = if min_bars_per_task == 0 {
|
||
DEFAULT_IMBALANCE_BAR_MIN_BARS_PER_TASK
|
||
} else {
|
||
min_bars_per_task
|
||
};
|
||
|
||
// ── Pass 1: find emission boundaries (sequential, lightweight) ──────
|
||
//
|
||
// Mirrors `ImbalanceBarSampler::update` direction + imbalance logic
|
||
// EXACTLY. Records the trade index of the emission-triggering trade as
|
||
// the end of each segment. Segments are inclusive of both endpoints.
|
||
//
|
||
// Init mirrors `ImbalanceBarSampler::new(trades[0].price, ...)`:
|
||
// prev_price = Some(trades[0].price), last_direction = 0,
|
||
// cum_imb = 0.
|
||
//
|
||
// Segment definitions
|
||
// -------------------
|
||
// - `segments[i] = (start_idx, end_idx)`: closed-closed range over
|
||
// non-empty trade slice.
|
||
// - First segment starts at the first non-zero-volume trade. Subsequent
|
||
// segments start at the trade IMMEDIATELY AFTER the previous segment's
|
||
// `end_idx`. The last segment is the open trailing range from
|
||
// `last_emission + 1` (or 0) to `trades.len() - 1` IFF there were any
|
||
// non-zero-volume trades AFTER the last emission AND that range
|
||
// produced a non-empty bar in sequential — but `ImbalanceBarSampler`
|
||
// does NOT flush trailing trades into a final bar (no `flush()` call
|
||
// in `mbp10_to_imbalance_bars`), so we DO NOT emit a bar for the
|
||
// trailing range. This matches sequential exactly.
|
||
let mut segments: Vec<(usize, usize)> = Vec::new();
|
||
let mut cum_imb = 0.0_f64;
|
||
let mut prev_price: Option<f64> = Some(trades[0].price);
|
||
let mut last_direction: i8 = 0;
|
||
let mut segment_start: Option<usize> = None;
|
||
for (i, trade) in trades.iter().enumerate() {
|
||
// Mirror `update()` line 484-486: zero-volume trades are no-ops.
|
||
if trade.volume == 0.0 {
|
||
continue;
|
||
}
|
||
// Track the first non-zero-volume trade as the start of segment 0.
|
||
if segment_start.is_none() {
|
||
segment_start = Some(i);
|
||
}
|
||
|
||
// Mirror `update()` line 495-506: classify direction.
|
||
let direction: i8 = if let Some(pp) = prev_price {
|
||
if trade.price > pp {
|
||
1
|
||
} else if trade.price < pp {
|
||
-1
|
||
} else {
|
||
last_direction
|
||
}
|
||
} else {
|
||
0
|
||
};
|
||
|
||
// Mirror `update()` line 508-510: imbalance update.
|
||
cum_imb += direction as f64 * trade.volume;
|
||
|
||
// Mirror `update()` line 519-520: state for next tick.
|
||
prev_price = Some(trade.price);
|
||
last_direction = direction;
|
||
|
||
// Mirror `update()` line 522-523: emission check.
|
||
if cum_imb.abs() >= threshold {
|
||
// Segment ends at this trade index (inclusive). Reset for next.
|
||
let s = segment_start.expect("segment_start set on first non-zero volume");
|
||
segments.push((s, i));
|
||
cum_imb = 0.0;
|
||
// `prev_price` and `last_direction` are KEPT (matches `reset()`
|
||
// in `alternative_bars.rs:558-566`).
|
||
segment_start = None;
|
||
}
|
||
}
|
||
// Trailing trades after the last emission are DISCARDED — sequential
|
||
// sampler does not flush a final partial bar in `mbp10_to_imbalance_bars`.
|
||
|
||
if segments.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
// ── Pass 2: build OHLCV bars from segments (parallel rayon) ────────
|
||
let build_bar = |&(s, e): &(usize, usize)| -> OHLCVBar {
|
||
// s..=e is guaranteed non-empty and to contain at least one
|
||
// non-zero-volume trade (segment_start was set on encountering one).
|
||
// open = first non-zero-volume trade's price; timestamp = same.
|
||
// close = LAST non-zero-volume trade's price (the emission trigger,
|
||
// which by Pass 1 invariant has volume > 0).
|
||
// high/low = max/min over all non-zero-volume trades in the segment.
|
||
// volume = sum of all (non-zero) volumes (zero-volume contributes 0
|
||
// so the filter is implicit).
|
||
let mut open: f64 = 0.0;
|
||
let mut open_ts = trades[s].timestamp;
|
||
let mut high = f64::NEG_INFINITY;
|
||
let mut low = f64::INFINITY;
|
||
let mut close = trades[e].price;
|
||
let mut volume_sum = 0.0_f64;
|
||
let mut found_open = false;
|
||
for trade in &trades[s..=e] {
|
||
if trade.volume == 0.0 {
|
||
continue;
|
||
}
|
||
if !found_open {
|
||
open = trade.price;
|
||
open_ts = trade.timestamp;
|
||
found_open = true;
|
||
}
|
||
high = high.max(trade.price);
|
||
low = low.min(trade.price);
|
||
close = trade.price;
|
||
volume_sum += trade.volume;
|
||
}
|
||
OHLCVBar {
|
||
timestamp: open_ts,
|
||
open,
|
||
high,
|
||
low,
|
||
close,
|
||
volume: volume_sum,
|
||
}
|
||
};
|
||
|
||
if segments.len() < min_bars_per_task {
|
||
// Sequential reduction — small enough that rayon overhead dominates.
|
||
segments.iter().map(build_bar).collect()
|
||
} else {
|
||
segments.par_iter().map(build_bar).collect()
|
||
}
|
||
}
|
||
|
||
/// Sequential reference walk — the golden against which `imbalance_bars_parallel`
|
||
/// must be bit-identical. Public for use as the K=1 fast-path and for the
|
||
/// bit-equivalence test in `tests/imbalance_bars_parallel_bit_equiv_test.rs`.
|
||
///
|
||
/// Wraps `ImbalanceBarSampler::new(trades[0].price, threshold,
|
||
/// trades[0].timestamp)` and the same per-trade `update()` loop as
|
||
/// `mbp10_to_imbalance_bars`'s legacy path.
|
||
pub fn imbalance_bars_sequential(trades: &[Mbp10Trade], threshold: f64) -> Vec<OHLCVBar> {
|
||
if trades.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
let first = &trades[0];
|
||
let mut sampler =
|
||
ImbalanceBarSampler::new(first.price, threshold, first.timestamp);
|
||
let mut bars: Vec<OHLCVBar> = Vec::new();
|
||
for trade in trades {
|
||
if let Some(bar) = sampler.update(trade.price, trade.volume, trade.timestamp) {
|
||
bars.push(bar);
|
||
}
|
||
}
|
||
bars
|
||
}
|
||
|
||
/// Load MBP-10 files and convert to imbalance bars via `ImbalanceBarSampler`.
|
||
///
|
||
/// Scans `mbp10_dir/symbol/` for `.dbn` and `.dbn.zst` files, loads each,
|
||
/// extracts trades (action == 'T'), feeds into adaptive `ImbalanceBarSampler`.
|
||
///
|
||
/// Results are cached to `/tmp/.foxhunt_imbalance_cache/<hash>.bars.bin`.
|
||
/// Subsequent calls with the same parameters return in <1s if no source files
|
||
/// have changed.
|
||
///
|
||
/// Returns sorted `OHLCVBar`s or error if directory doesn't exist.
|
||
/// No fallback -- fails loudly if the data isn't there.
|
||
pub fn mbp10_to_imbalance_bars(
|
||
mbp10_dir: &Path,
|
||
symbol: &str,
|
||
threshold: f64,
|
||
ewma_alpha: f64,
|
||
) -> Result<Vec<OHLCVBar>, MLError> {
|
||
let symbol_dir = mbp10_dir.join(symbol);
|
||
if !symbol_dir.exists() {
|
||
return Err(MLError::InsufficientData(format!(
|
||
"MBP-10 directory not found: {}. Set data_source='ohlcv' or provide MBP-10 data.",
|
||
symbol_dir.display()
|
||
)));
|
||
}
|
||
|
||
// ── Cache check ──────────────────────────────────────────────────────
|
||
let cache_file = imbalance_cache_path(mbp10_dir, symbol, threshold, ewma_alpha);
|
||
if is_imbalance_cache_valid(&cache_file, &symbol_dir) {
|
||
let load_start = std::time::Instant::now();
|
||
match load_imbalance_cache(&cache_file) {
|
||
Ok(bars) => {
|
||
let elapsed_ms = load_start.elapsed().as_secs_f64() * 1000.0;
|
||
tracing::info!(
|
||
"Imbalance bar cache HIT: loaded {} bars in {:.1}ms",
|
||
bars.len(),
|
||
elapsed_ms
|
||
);
|
||
return Ok(bars);
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!("Imbalance bar cache read failed (recomputing): {}", e);
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── Cache miss — compute from scratch ────────────────────────────────
|
||
let mut dbn_files = collect_dbn_files(&symbol_dir);
|
||
dbn_files.sort();
|
||
|
||
if dbn_files.is_empty() {
|
||
return Err(MLError::InsufficientData(format!(
|
||
"No .dbn/.dbn.zst files found in {}",
|
||
symbol_dir.display()
|
||
)));
|
||
}
|
||
|
||
tracing::info!(
|
||
"Imbalance bar cache MISS: computing from {} MBP-10 files...",
|
||
dbn_files.len()
|
||
);
|
||
|
||
// Collect trades from all files, filtering to the front-month (highest-
|
||
// volume) contract per file. Without per-file filtering the resulting
|
||
// bar stream interleaves contract months — when the dominant ES contract
|
||
// rolls over (ESZ24 → ESH25 → ESM25), price jumps appear at every
|
||
// rollover, failing the `precompute_features.rs:371-376` price-continuity
|
||
// gate. Mirror the per-file pattern in `precompute_features.rs:347-365`
|
||
// (uses `filter_front_month` on `DbnTrade` slices); same logic, applied
|
||
// to `Mbp10Trade.instrument_id`.
|
||
//
|
||
// SP20 parallelism: `extract_trades_from_dbn_file` is a pure
|
||
// file → Vec<Mbp10Trade> function and the 9 quarterly files are fully
|
||
// independent (each carries its own contract universe). Mirror the
|
||
// par_iter pattern in `precompute_features.rs:551-565` so a 28-core
|
||
// CPU node can decode the files concurrently. Front-month filter is
|
||
// applied per-file as before — semantics unchanged. Final ordering
|
||
// is enforced by the `all_trades.sort_by(|a, b| ...)` step below, so
|
||
// file-level reduce order is irrelevant for correctness.
|
||
let per_file_trades: Vec<Vec<Mbp10Trade>> = {
|
||
use rayon::prelude::*;
|
||
dbn_files
|
||
.par_iter()
|
||
.filter_map(|file_path| match extract_trades_from_dbn_file(file_path) {
|
||
Ok(trades) => {
|
||
let raw_count = trades.len();
|
||
let filtered = filter_front_month_mbp10(&trades);
|
||
tracing::info!(
|
||
" {} -> {} trades, {} front-month",
|
||
file_path.file_name().unwrap_or_default().to_string_lossy(),
|
||
raw_count,
|
||
filtered.len()
|
||
);
|
||
Some(filtered)
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
"Failed to extract trades from {}: {}",
|
||
file_path.display(),
|
||
e
|
||
);
|
||
None
|
||
}
|
||
})
|
||
.collect()
|
||
};
|
||
let mut all_trades: Vec<Mbp10Trade> =
|
||
Vec::with_capacity(per_file_trades.iter().map(Vec::len).sum());
|
||
for chunk in per_file_trades {
|
||
all_trades.extend(chunk);
|
||
}
|
||
|
||
if all_trades.is_empty() {
|
||
return Err(MLError::InsufficientData(format!(
|
||
"No trade ticks extracted from MBP-10 files in {}. \
|
||
Files may contain only book updates without trade events.",
|
||
symbol_dir.display()
|
||
)));
|
||
}
|
||
|
||
// Sort by timestamp (files should be chronological but ensure correctness)
|
||
all_trades.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
|
||
|
||
// wgdc8 experiment: EWMA adaptation bypassed. With α=0.1 the configured
|
||
// threshold is washed out within ~5 bars (recursion: T_new = 0.1·T_old +
|
||
// 0.9·observed; fixed point is the data's natural imbalance scale, not
|
||
// the configured value). To make `imbalance_bar_threshold` actually
|
||
// honored — and the resolution-hypothesis smoke meaningful — use the
|
||
// fixed-threshold constructor. ewma_alpha=0.1 is logged for cache-key
|
||
// continuity but does not drive bar formation on this branch.
|
||
tracing::warn!(
|
||
"wgdc8: EWMA adaptation BYPASSED. Feeding {} trade ticks into \
|
||
ImbalanceBarSampler with FIXED threshold={} (config alpha={} ignored)",
|
||
all_trades.len(),
|
||
threshold,
|
||
ewma_alpha,
|
||
);
|
||
|
||
// SP20 parallelism (Bottleneck B): two-pass parallel decomposition.
|
||
// Pass 1 walks the trade tape sequentially to find emission boundaries
|
||
// (cheap arithmetic, no allocation). Pass 2 builds OHLCV bars from each
|
||
// independent emission segment via rayon `par_iter` reduction. See
|
||
// `imbalance_bars_parallel` doc for the bit-equivalence argument and
|
||
// `tests/imbalance_bars_parallel_bit_equiv_test.rs` for K=4/K=8/high-
|
||
// threshold verification at exact 0.0 diff.
|
||
//
|
||
// Note: time-bucket sharding (the OFI pattern) does NOT bit-equate to
|
||
// sequential here because `ImbalanceBarSampler` has unbounded path-
|
||
// dependent state (cum_imb between emissions, no rolling window).
|
||
// See module note in `imbalance_bars_parallel`.
|
||
let mut bars = imbalance_bars_parallel(
|
||
&all_trades,
|
||
threshold,
|
||
rayon::current_num_threads(),
|
||
DEFAULT_IMBALANCE_BAR_MIN_BARS_PER_TASK,
|
||
);
|
||
|
||
// Sort by timestamp (should already be ordered — shards concatenate in
|
||
// trade-index order, and trades were sorted by timestamp above).
|
||
bars.sort_by(|a, b| a.timestamp.cmp(&b.timestamp));
|
||
|
||
tracing::info!(
|
||
"ImbalanceBarSampler produced {} bars from {} trade ticks (ratio 1:{:.0})",
|
||
bars.len(),
|
||
all_trades.len(),
|
||
if bars.is_empty() {
|
||
0.0
|
||
} else {
|
||
all_trades.len() as f64 / bars.len() as f64
|
||
},
|
||
);
|
||
|
||
// ── Save to cache ────────────────────────────────────────────────────
|
||
if let Err(e) = save_imbalance_cache(&cache_file, &bars) {
|
||
tracing::warn!("Failed to write imbalance bar cache (non-fatal): {}", e);
|
||
} else {
|
||
tracing::info!("Cached {} bars to {}", bars.len(), cache_file.display());
|
||
}
|
||
|
||
Ok(bars)
|
||
}
|
||
|
||
/// Deserialize cached imbalance bars from a bincode file.
|
||
fn load_imbalance_cache(path: &Path) -> Result<Vec<OHLCVBar>, MLError> {
|
||
use std::io::BufReader;
|
||
let file = std::fs::File::open(path).map_err(|e| {
|
||
MLError::InsufficientData(format!("Failed to open imbalance cache {:?}: {}", path, e))
|
||
})?;
|
||
let reader = BufReader::new(file);
|
||
bincode::deserialize_from(reader).map_err(|e| {
|
||
MLError::InsufficientData(format!(
|
||
"Failed to deserialize imbalance cache {:?}: {}",
|
||
path, e
|
||
))
|
||
})
|
||
}
|
||
|
||
/// Serialize imbalance bars to a bincode file, creating the cache directory if needed.
|
||
fn save_imbalance_cache(path: &Path, bars: &[OHLCVBar]) -> Result<(), MLError> {
|
||
use std::io::BufWriter;
|
||
if let Some(parent) = path.parent() {
|
||
std::fs::create_dir_all(parent).map_err(|e| {
|
||
MLError::InsufficientData(format!(
|
||
"Failed to create imbalance cache dir {:?}: {}",
|
||
parent, e
|
||
))
|
||
})?;
|
||
}
|
||
let file = std::fs::File::create(path).map_err(|e| {
|
||
MLError::InsufficientData(format!(
|
||
"Failed to create imbalance cache {:?}: {}",
|
||
path, e
|
||
))
|
||
})?;
|
||
let writer = BufWriter::new(file);
|
||
bincode::serialize_into(writer, bars).map_err(|e| {
|
||
MLError::InsufficientData(format!(
|
||
"Failed to serialize imbalance cache {:?}: {}",
|
||
path, e
|
||
))
|
||
})
|
||
}
|
||
|
||
/// Get snapshots for a specific timestamp window
|
||
///
|
||
/// Returns a slice of snapshots starting from the first snapshot at or before
|
||
/// the target timestamp. Used for calculating OFI features that require
|
||
/// forward-looking context (e.g., next N snapshots after a given timestamp).
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `snapshots` - All available snapshots (must be sorted by timestamp)
|
||
/// * `target_ts` - Target timestamp (nanoseconds since Unix epoch)
|
||
/// * `window_size` - Number of snapshots to include in window
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Slice of up to `window_size` snapshots starting from the first snapshot
|
||
/// with timestamp >= `target_ts`. Returns empty slice if no snapshots exist
|
||
/// at or after the target timestamp.
|
||
///
|
||
/// # Algorithm
|
||
///
|
||
/// - Simple linear search over the sorted snapshot slice.
|
||
/// - Returns snapshots starting from the first one with timestamp >= `target_ts`.
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```ignore
|
||
/// // Snapshots at times: [1000, 2000, 3000, 4000]
|
||
/// // target_ts = 2000, window_size = 2
|
||
/// // Returns: [2000, 3000]
|
||
/// ```
|
||
///
|
||
pub fn get_snapshots_for_timestamp(
|
||
snapshots: &[Mbp10Snapshot],
|
||
target_ts: u64,
|
||
window_size: usize,
|
||
) -> &[Mbp10Snapshot] {
|
||
if snapshots.is_empty() {
|
||
return &[];
|
||
}
|
||
|
||
// Binary search for the first snapshot with timestamp >= target_ts
|
||
// Snapshots are pre-sorted by timestamp (from DBN file ordering)
|
||
let idx = snapshots.partition_point(|s| s.timestamp < target_ts);
|
||
if idx >= snapshots.len() {
|
||
return &[]; // Target timestamp is after all snapshots
|
||
}
|
||
let end_idx = (idx + window_size).min(snapshots.len());
|
||
&snapshots[idx..end_idx]
|
||
}
|
||
|
||
/// Get the most recent N snapshots ending at the given index
|
||
///
|
||
/// Useful for calculating features that require a rolling window of recent data.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `snapshots` - All available snapshots
|
||
/// * `end_idx` - Index of the last snapshot to include (exclusive)
|
||
/// * `window_size` - Number of snapshots to include
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Slice of up to `window_size` snapshots ending at `end_idx`
|
||
///
|
||
pub fn get_recent_snapshots(
|
||
snapshots: &[Mbp10Snapshot],
|
||
end_idx: usize,
|
||
window_size: usize,
|
||
) -> &[Mbp10Snapshot] {
|
||
if snapshots.is_empty() || end_idx == 0 {
|
||
return &[];
|
||
}
|
||
|
||
let actual_end = end_idx.min(snapshots.len());
|
||
let start_idx = actual_end.saturating_sub(window_size);
|
||
|
||
&snapshots[start_idx..actual_end]
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use data::providers::databento::mbp10::BidAskPair;
|
||
|
||
fn create_test_snapshot(timestamp: u64, bid_px: i64) -> Mbp10Snapshot {
|
||
let levels = vec![BidAskPair {
|
||
bid_px,
|
||
bid_sz: 100,
|
||
bid_ct: 5,
|
||
ask_px: bid_px + 1000000000, // 1 tick higher
|
||
ask_sz: 120,
|
||
ask_ct: 6,
|
||
}];
|
||
|
||
Mbp10Snapshot::new("ES.FUT".to_owned(), timestamp, levels, 0, 0)
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_exact_match() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
create_test_snapshot(4000, 150030000000000),
|
||
];
|
||
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2000, 2);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_between() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
];
|
||
|
||
// Target timestamp between 2000 and 3000
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2500, 2);
|
||
assert_eq!(window.len(), 1); // Only snapshot at 3000 remains
|
||
assert_eq!(window[0].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_empty() {
|
||
let snapshots: Vec<Mbp10Snapshot> = vec![];
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2000, 2);
|
||
assert_eq!(window.len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_before_all() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(2000, 150000000000000),
|
||
create_test_snapshot(3000, 150010000000000),
|
||
];
|
||
|
||
// Target is before all snapshots, should return first 2 snapshots
|
||
let window = get_snapshots_for_timestamp(&snapshots, 1000, 2);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_for_timestamp_after_all() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
];
|
||
|
||
let window = get_snapshots_for_timestamp(&snapshots, 5000, 2);
|
||
assert_eq!(window.len(), 0); // No snapshots after target
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_snapshots_window_clipping() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
];
|
||
|
||
// Request 10 snapshots but only 2 available at/after target (2000, 3000)
|
||
let window = get_snapshots_for_timestamp(&snapshots, 2000, 10);
|
||
assert_eq!(window.len(), 2); // 2 snapshots at/after 2000: [2000, 3000]
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_full_window() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
create_test_snapshot(3000, 150020000000000),
|
||
create_test_snapshot(4000, 150030000000000),
|
||
];
|
||
|
||
let window = get_recent_snapshots(&snapshots, 3, 2);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 2000);
|
||
assert_eq!(window[1].timestamp, 3000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_partial_window() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
];
|
||
|
||
// Request 5 snapshots but only 2 available
|
||
let window = get_recent_snapshots(&snapshots, 2, 5);
|
||
assert_eq!(window.len(), 2);
|
||
assert_eq!(window[0].timestamp, 1000);
|
||
assert_eq!(window[1].timestamp, 2000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_empty() {
|
||
let snapshots: Vec<Mbp10Snapshot> = vec![];
|
||
let window = get_recent_snapshots(&snapshots, 0, 2);
|
||
assert_eq!(window.len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_zero_end_idx() {
|
||
let snapshots = vec![
|
||
create_test_snapshot(1000, 150000000000000),
|
||
create_test_snapshot(2000, 150010000000000),
|
||
];
|
||
|
||
let window = get_recent_snapshots(&snapshots, 0, 2);
|
||
assert_eq!(window.len(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_get_recent_snapshots_one_element() {
|
||
let snapshots = vec![create_test_snapshot(1000, 150000000000000)];
|
||
|
||
let window = get_recent_snapshots(&snapshots, 1, 1);
|
||
assert_eq!(window.len(), 1);
|
||
assert_eq!(window[0].timestamp, 1000);
|
||
}
|
||
|
||
#[test]
|
||
fn test_extract_trades_from_snapshots() {
|
||
// Create snapshots with increasing trade_count and rising prices
|
||
let snapshots = vec![
|
||
{
|
||
let mut s = create_test_snapshot(1_000_000_000, 150000000000000); // 150.0
|
||
s.trade_count = 0;
|
||
s
|
||
},
|
||
{
|
||
let mut s = create_test_snapshot(2_000_000_000, 150010000000000); // 150.01
|
||
s.trade_count = 1;
|
||
s
|
||
},
|
||
{
|
||
let mut s = create_test_snapshot(3_000_000_000, 150020000000000); // 150.02
|
||
s.trade_count = 3;
|
||
s
|
||
},
|
||
{
|
||
let mut s = create_test_snapshot(4_000_000_000, 149990000000000); // 149.99
|
||
s.trade_count = 4;
|
||
s
|
||
},
|
||
];
|
||
|
||
let trades = extract_trades_from_snapshots(&snapshots, 12345);
|
||
// 3 transitions where trade_count increases (snap 0->1, 1->2, 2->3)
|
||
assert_eq!(trades.len(), 3);
|
||
// First trade: price went up => is_buy=true
|
||
assert!(trades[0].is_buy);
|
||
// Third trade: price went down => is_buy=false
|
||
assert!(!trades[2].is_buy);
|
||
// All trades carry the caller-supplied instrument_id (verifies the
|
||
// explicit-param contract — no implicit 0 default).
|
||
assert!(trades.iter().all(|t| t.instrument_id == 12345));
|
||
}
|
||
|
||
#[test]
|
||
fn test_filter_front_month_mbp10_picks_highest_volume() {
|
||
let ts = DateTime::<Utc>::from_timestamp(1_700_000_000, 0).unwrap();
|
||
let mk = |id: u32, vol: f64| Mbp10Trade {
|
||
price: 100.0,
|
||
volume: vol,
|
||
timestamp: ts,
|
||
is_buy: true,
|
||
instrument_id: id,
|
||
};
|
||
// Three contracts with different total volumes; id=200 is the
|
||
// highest-volume contract (10 + 5 = 15).
|
||
let trades = vec![
|
||
mk(100, 3.0),
|
||
mk(200, 10.0),
|
||
mk(300, 1.0),
|
||
mk(200, 5.0),
|
||
mk(100, 2.0),
|
||
];
|
||
let filtered = filter_front_month_mbp10(&trades);
|
||
assert_eq!(filtered.len(), 2, "should keep only id=200 trades");
|
||
assert!(filtered.iter().all(|t| t.instrument_id == 200));
|
||
}
|
||
|
||
#[test]
|
||
fn test_filter_front_month_mbp10_empty() {
|
||
let filtered = filter_front_month_mbp10(&[]);
|
||
assert!(filtered.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn test_mbp10_to_imbalance_bars_missing_dir() {
|
||
let result = mbp10_to_imbalance_bars(
|
||
Path::new("/nonexistent/path"),
|
||
"ES.FUT",
|
||
100.0,
|
||
0.1,
|
||
);
|
||
assert!(result.is_err());
|
||
let err_msg = format!("{}", result.unwrap_err());
|
||
assert!(err_msg.contains("MBP-10 directory not found"), "Error: {err_msg}");
|
||
}
|
||
|
||
/// Integration test: loads real MBP-10 data and produces imbalance bars.
|
||
/// Reads threshold/alpha from config/training/dqn-smoketest.toml (single source of truth).
|
||
/// Requires test_data/futures-baseline-mbp10/ES.FUT/ with .dbn.zst files.
|
||
/// Run with: cargo test -p ml-features --lib -- test_mbp10_to_imbalance_bars_real --ignored --nocapture
|
||
#[test]
|
||
#[ignore]
|
||
fn test_mbp10_to_imbalance_bars_real() {
|
||
// Resolve workspace root from CARGO_MANIFEST_DIR
|
||
let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap();
|
||
let workspace = std::path::Path::new(&manifest)
|
||
.parent() // crates/
|
||
.and_then(|p| p.parent()) // repo root
|
||
.expect("Failed to find workspace root");
|
||
|
||
// Load threshold/alpha from TOML config (no hardcoded values)
|
||
let toml_path = workspace.join("config/training/dqn-smoketest.toml");
|
||
let toml_str = std::fs::read_to_string(&toml_path)
|
||
.unwrap_or_else(|_| panic!("Failed to read {}", toml_path.display()));
|
||
let toml: toml::Value = toml_str.parse()
|
||
.unwrap_or_else(|e| panic!("Failed to parse TOML: {e}"));
|
||
let training = toml.get("training").expect("[training] section missing");
|
||
let threshold = training.get("imbalance_bar_threshold")
|
||
.and_then(|v| v.as_float()).unwrap_or(1.0);
|
||
let alpha = training.get("imbalance_bar_ewma_alpha")
|
||
.and_then(|v| v.as_float()).unwrap_or(0.1);
|
||
println!("Config: threshold={threshold}, alpha={alpha} (from dqn-smoketest.toml)");
|
||
|
||
let mbp10_dir = workspace.join("test_data/futures-baseline-mbp10");
|
||
assert!(mbp10_dir.exists(), "MBP-10 test data not found at {}", mbp10_dir.display());
|
||
|
||
let bars = mbp10_to_imbalance_bars(&mbp10_dir, "ES.FUT", threshold, alpha)
|
||
.expect("mbp10_to_imbalance_bars must succeed with real data");
|
||
|
||
assert!(!bars.is_empty(), "Must produce imbalance bars from MBP-10 data");
|
||
println!("Produced {} imbalance bars from MBP-10 data", bars.len());
|
||
println!(" First bar: {:?}", bars.first().unwrap().timestamp);
|
||
println!(" Last bar: {:?}", bars.last().unwrap().timestamp);
|
||
|
||
// Verify time-ordering
|
||
for w in bars.windows(2) {
|
||
assert!(
|
||
w[0].timestamp <= w[1].timestamp,
|
||
"Bars must be time-ordered: {:?} > {:?}",
|
||
w[0].timestamp, w[1].timestamp
|
||
);
|
||
}
|
||
|
||
// Verify reasonable bar properties
|
||
for bar in &bars {
|
||
assert!(bar.high >= bar.low, "high must >= low");
|
||
assert!(bar.volume > 0.0, "volume must be positive");
|
||
assert!(bar.close > 0.0, "close must be positive");
|
||
}
|
||
println!("All {} bars valid: time-ordered, positive volume, valid OHLC", bars.len());
|
||
}
|
||
}
|