Files
foxhunt/crates/ml-features/src/mbp10_loader.rs
jgrusewski 828417b523 fix: test reads threshold from dqn-smoketest.toml — no hardcoded values
Integration test now loads imbalance_bar_threshold and ewma_alpha from
config/training/dqn-smoketest.toml. Single source of truth for all
config values. Production threshold lowered to 0.5 for maximum bar yield.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-31 12:37:25 +02:00

1081 lines
38 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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, 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, |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, |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,
}
/// 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).
///
/// 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]) -> 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,
});
}
}
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,
});
}
}
}
}
}
Ok(None) => break,
Err(e) => {
return Err(MLError::InsufficientData(format!(
"DBN decode error: {}",
e
)));
}
}
}
Ok(trades)
}
/// 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
}
/// 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 all trades from all files
let mut all_trades: Vec<Mbp10Trade> = Vec::new();
for file_path in &dbn_files {
match extract_trades_from_dbn_file(file_path) {
Ok(trades) => {
tracing::debug!(
" {} -> {} trades",
file_path.file_name().unwrap_or_default().to_string_lossy(),
trades.len()
);
all_trades.extend(trades);
}
Err(e) => {
tracing::warn!(
"Failed to extract trades from {}: {}",
file_path.display(),
e
);
}
}
}
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));
tracing::info!(
"Extracted {} trade ticks, feeding into ImbalanceBarSampler (threshold={}, alpha={})",
all_trades.len(),
threshold,
ewma_alpha,
);
// Initialize sampler with first trade
let first = &all_trades[0];
let mut sampler =
ImbalanceBarSampler::new_with_ewma(first.price, threshold, first.timestamp, ewma_alpha);
let mut bars: Vec<OHLCVBar> = Vec::new();
for trade in &all_trades {
if let Some(bar) = sampler.update(trade.price, trade.volume, trade.timestamp) {
bars.push(bar);
}
}
// Sort by timestamp (should already be ordered)
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 (TODO: optimize with binary search for large datasets)
/// - 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);
// 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);
}
#[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());
}
}