7fe064c6e0
Replace .unwrap()/.expect() with safe alternatives across 51 files: - 41 `let _ = writeln!()` → `_ = writeln!()` (wildcard assignment) - 53 unwrap() in features/ → unwrap_or/match/early-return - 20 expect() in inference/metrics → module-level #[allow] for static init - 16 unwrap/expect in hyperopt/ → ?, map_err, unwrap_or - 20 unwrap in dqn/trainers/ → ?, map_err, unwrap_or - 28 unwrap in misc files → context-appropriate safe patterns Zero clippy errors remain across the entire workspace. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1298 lines
47 KiB
Rust
1298 lines
47 KiB
Rust
//! MAMBA-2 Hyperparameter Optimization Adapter
|
|
//!
|
|
//! This module provides a production-ready adapter for optimizing MAMBA-2
|
|
//! hyperparameters using the generic optimization framework. It implements:
|
|
//!
|
|
//! - Parameter space with log-scale handling for learning rates
|
|
//! - Training wrapper that integrates with existing MAMBA-2 pipeline
|
|
//! - Metrics extraction for validation loss optimization
|
|
//!
|
|
//! ## Usage Example
|
|
//!
|
|
//! ```rust,no_run
|
|
//! use ml::hyperopt::EgoboxOptimizer;
|
|
//! use ml::hyperopt::adapters::mamba2::{Mamba2Trainer, Mamba2Params};
|
|
//!
|
|
//! # async fn example() -> anyhow::Result<()> {
|
|
//! // Create trainer
|
|
//! let trainer = Mamba2Trainer::new(
|
|
//! "test_data/ES_FUT_180d.parquet",
|
|
//! 50, // epochs per trial
|
|
//! )?;
|
|
//!
|
|
//! // Run optimization
|
|
//! let optimizer = EgoboxOptimizer::with_trials(30, 5);
|
|
//! let result = optimizer.optimize(trainer)?;
|
|
//!
|
|
//! println!("Best learning rate: {}", result.best_params.learning_rate);
|
|
//! println!("Best batch size: {}", result.best_params.batch_size);
|
|
//! println!("Best validation loss: {:.6}", result.best_objective);
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
|
|
use anyhow::{Context, Result};
|
|
use candle_core::{Device, Tensor};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::fs::OpenOptions;
|
|
use std::io::Write as IoWrite;
|
|
use std::path::PathBuf;
|
|
use tracing::{info, warn};
|
|
|
|
use arrow::array::{Array, Float64Array, PrimitiveArray, UInt64Array};
|
|
use arrow::datatypes::TimestampNanosecondType;
|
|
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
|
use std::fs::File;
|
|
|
|
use crate::features::{extract_ml_features, FeatureConfig, OHLCVBar};
|
|
use crate::hyperopt::paths::TrainingPaths;
|
|
use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
|
|
use crate::mamba::{Mamba2Config, Mamba2SSM, OptimizerType};
|
|
use crate::MLError;
|
|
|
|
/// MAMBA-2 hyperparameter space
|
|
///
|
|
/// Defines the hyperparameters to optimize for MAMBA-2 training:
|
|
/// - Learning rate (log-scale: 1e-5 to 1e-2)
|
|
/// - Batch size (linear scale: 4 to 256, optimized for GPU utilization)
|
|
/// - Dropout rate (linear scale: 0.0 to 0.5)
|
|
/// - Weight decay (log-scale: 1e-6 to 1e-2)
|
|
///
|
|
/// ## Parameter Scaling
|
|
///
|
|
/// - **Log-scale**: Learning rate, weight decay (span multiple orders of magnitude)
|
|
/// - **Linear scale**: Batch size, dropout (span single order of magnitude)
|
|
///
|
|
/// This scaling ensures efficient exploration by egobox's Gaussian Process.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct Mamba2Params {
|
|
/// Learning rate for Adam optimizer (log-scale)
|
|
pub learning_rate: f64,
|
|
/// Batch size for training (linear scale, integer)
|
|
pub batch_size: usize,
|
|
/// Dropout rate for regularization (linear scale)
|
|
pub dropout: f64,
|
|
/// Weight decay for L2 regularization (log-scale)
|
|
pub weight_decay: f64,
|
|
/// Gradient clipping threshold (log-scale)
|
|
pub grad_clip: f64,
|
|
/// Warmup steps for learning rate schedule (linear scale)
|
|
pub warmup_steps: usize,
|
|
/// P0: Adam beta1 momentum parameter (linear scale)
|
|
pub adam_beta1: f64,
|
|
/// P1: Adam beta2 parameter (linear scale)
|
|
pub adam_beta2: f64,
|
|
/// P1: Adam epsilon (log-scale)
|
|
pub adam_epsilon: f64,
|
|
/// P2: Lookback window (sequence length) (linear scale, integer)
|
|
pub lookback_window: usize,
|
|
/// P2: Sequence stride for overlapping windows (linear scale, integer)
|
|
pub sequence_stride: usize,
|
|
/// P2: Normalization epsilon for layer norm (log-scale)
|
|
pub norm_eps: f64,
|
|
}
|
|
|
|
impl Default for Mamba2Params {
|
|
fn default() -> Self {
|
|
Self {
|
|
learning_rate: 1e-4,
|
|
batch_size: 32,
|
|
dropout: 0.1,
|
|
weight_decay: 1e-4,
|
|
grad_clip: 1.0,
|
|
warmup_steps: 100,
|
|
adam_beta1: 0.9,
|
|
adam_beta2: 0.999,
|
|
adam_epsilon: 1e-8,
|
|
lookback_window: 60,
|
|
sequence_stride: 1,
|
|
norm_eps: 1e-5,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ParameterSpace for Mamba2Params {
|
|
fn continuous_bounds() -> Vec<(f64, f64)> {
|
|
vec![
|
|
(1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log scale)
|
|
(4.0, 256.0), // batch_size (linear) - wide bounds, clamped by trainer config
|
|
(0.0, 0.5), // dropout (linear)
|
|
(1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log scale)
|
|
(0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log scale)
|
|
(100.0, 2000.0), // warmup_steps (linear)
|
|
(0.85, 0.95), // adam_beta1 (linear)
|
|
(0.98, 0.999), // adam_beta2 (linear)
|
|
(1e-9_f64.ln(), 1e-7_f64.ln()), // adam_epsilon (log scale)
|
|
(30.0, 120.0), // lookback_window (linear)
|
|
(1.0, 5.0), // sequence_stride (linear)
|
|
(1e-6_f64.ln(), 1e-4_f64.ln()), // norm_eps (log scale)
|
|
]
|
|
}
|
|
|
|
fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
|
|
if x.len() != 12 {
|
|
return Err(MLError::ConfigError {
|
|
reason: format!("Expected 12 parameters, got {}", x.len()),
|
|
});
|
|
}
|
|
|
|
Ok(Self {
|
|
learning_rate: x[0].exp(),
|
|
batch_size: x[1].round().max(1.0) as usize,
|
|
dropout: x[2].clamp(0.0, 0.5),
|
|
weight_decay: x[3].exp(),
|
|
grad_clip: x[4].exp(),
|
|
warmup_steps: x[5].round().max(1.0) as usize,
|
|
adam_beta1: x[6].clamp(0.85, 0.95),
|
|
adam_beta2: x[7].clamp(0.98, 0.999),
|
|
adam_epsilon: x[8].exp(),
|
|
lookback_window: x[9].round().max(30.0).min(120.0) as usize,
|
|
sequence_stride: x[10].round().max(1.0).min(5.0) as usize,
|
|
norm_eps: x[11].exp(),
|
|
})
|
|
}
|
|
|
|
fn to_continuous(&self) -> Vec<f64> {
|
|
vec![
|
|
self.learning_rate.ln(),
|
|
self.batch_size as f64,
|
|
self.dropout,
|
|
self.weight_decay.ln(),
|
|
self.grad_clip.ln(),
|
|
self.warmup_steps as f64,
|
|
self.adam_beta1,
|
|
self.adam_beta2,
|
|
self.adam_epsilon.ln(),
|
|
self.lookback_window as f64,
|
|
self.sequence_stride as f64,
|
|
self.norm_eps.ln(),
|
|
]
|
|
}
|
|
|
|
fn param_names() -> Vec<&'static str> {
|
|
vec![
|
|
"learning_rate",
|
|
"batch_size",
|
|
"dropout",
|
|
"weight_decay",
|
|
"grad_clip",
|
|
"warmup_steps",
|
|
"adam_beta1",
|
|
"adam_beta2",
|
|
"adam_epsilon",
|
|
"lookback_window",
|
|
"sequence_stride",
|
|
"norm_eps",
|
|
]
|
|
}
|
|
}
|
|
|
|
/// MAMBA-2 training metrics
|
|
///
|
|
/// Contains all relevant metrics from a MAMBA-2 training run.
|
|
/// The primary optimization target is validation loss.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Mamba2Metrics {
|
|
/// Final validation loss (optimization target)
|
|
pub val_loss: f64,
|
|
/// Final training loss
|
|
pub train_loss: f64,
|
|
/// Validation perplexity (exp(val_loss))
|
|
pub val_perplexity: f64,
|
|
/// Directional accuracy (% of correct direction predictions)
|
|
pub directional_accuracy: f64,
|
|
/// Mean Absolute Error
|
|
pub mae: f64,
|
|
/// Root Mean Squared Error
|
|
pub rmse: f64,
|
|
/// R² (Coefficient of Determination)
|
|
pub r_squared: f64,
|
|
/// Number of epochs completed
|
|
pub epochs_completed: usize,
|
|
}
|
|
|
|
/// MAMBA-2 trainer for hyperparameter optimization
|
|
///
|
|
/// This struct wraps the MAMBA-2 training pipeline and implements
|
|
/// `HyperparameterOptimizable` for use with `EgoboxOptimizer`.
|
|
///
|
|
/// ## Configuration
|
|
///
|
|
/// - **Parquet file**: Market data source (OHLCV bars)
|
|
/// - **Epochs**: Number of training epochs per trial
|
|
/// - **Device**: CUDA GPU (falls back to CPU if unavailable)
|
|
/// - **Features**: 51 features (WAVE 10: Reduced from 54→51 after Proxy OFI removal)
|
|
///
|
|
/// ## Fixed Architecture
|
|
///
|
|
/// The following parameters are fixed for consistency:
|
|
/// - `d_model`: 54 (WAVE 10: 51 market features + 3 portfolio features)
|
|
/// - `d_state`: 16
|
|
/// - `num_layers`: 6
|
|
/// - `sequence_length`: 60
|
|
///
|
|
/// ## Optimized Hyperparameters
|
|
///
|
|
/// The following are optimized by `Mamba2Params`:
|
|
/// - Learning rate
|
|
/// - Batch size
|
|
/// - Dropout
|
|
/// - Weight decay
|
|
#[derive(Debug)]
|
|
pub struct Mamba2Trainer {
|
|
parquet_file: PathBuf,
|
|
epochs: usize,
|
|
device: Device,
|
|
feature_config: FeatureConfig,
|
|
d_model: usize,
|
|
train_split: f64,
|
|
/// Target normalization parameters (set after data loading)
|
|
target_min: Option<f64>,
|
|
target_max: Option<f64>,
|
|
/// Minimum batch size (for GPU memory constraints)
|
|
batch_size_min: f64,
|
|
/// Maximum batch size (for GPU memory constraints)
|
|
batch_size_max: f64,
|
|
/// Enable async data loading (prefetch while GPU trains)
|
|
async_loading: bool,
|
|
/// Number of batches to prefetch (2-3 recommended)
|
|
prefetch_count: usize,
|
|
/// Training paths configuration (replaces hardcoded checkpoint_dir)
|
|
training_paths: TrainingPaths,
|
|
/// Early stopping patience (epochs without improvement before stopping)
|
|
early_stopping_patience: usize,
|
|
/// Early stopping minimum epochs (minimum epochs before early stopping can trigger)
|
|
early_stopping_min_epochs: usize,
|
|
/// Trial counter for hyperopt (tracks trial number across multiple trials)
|
|
trial_counter: usize,
|
|
}
|
|
|
|
impl Mamba2Trainer {
|
|
/// Create a new MAMBA-2 trainer
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `parquet_file` - Path to Parquet file with market data
|
|
/// * `epochs` - Number of training epochs per trial
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Configured trainer ready for optimization
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns error if:
|
|
/// - Parquet file doesn't exist
|
|
/// - CUDA device initialization fails (falls back to CPU)
|
|
pub fn new<P: Into<PathBuf>>(parquet_file: P, epochs: usize) -> Result<Self> {
|
|
let parquet_file = parquet_file.into();
|
|
|
|
if !parquet_file.exists() {
|
|
return Err(MLError::ConfigError {
|
|
reason: format!("Parquet file not found: {}", parquet_file.display()),
|
|
}
|
|
.into());
|
|
}
|
|
|
|
// Initialize device (CUDA preferred, CPU fallback)
|
|
let device = Device::new_cuda(0).unwrap_or_else(|e| {
|
|
warn!("CUDA unavailable ({}), falling back to CPU", e);
|
|
Device::Cpu
|
|
});
|
|
|
|
// Use WAVE 10 feature configuration (51 market + 3 portfolio = 54)
|
|
let feature_config = FeatureConfig::wave_d();
|
|
let d_model = feature_config.feature_count();
|
|
|
|
info!("MAMBA-2 Trainer initialized:");
|
|
info!(" Device: {:?}", device);
|
|
info!(" Features: {} (WAVE 10: 51 market + 3 portfolio)", d_model);
|
|
info!(" Epochs per trial: {}", epochs);
|
|
|
|
// Use temporary default paths - should be replaced with with_training_paths()
|
|
let training_paths = TrainingPaths::new("/tmp/ml_training", "mamba2", "default");
|
|
|
|
Ok(Self {
|
|
parquet_file,
|
|
epochs,
|
|
device,
|
|
feature_config,
|
|
d_model,
|
|
train_split: 0.8,
|
|
target_min: None,
|
|
target_max: None,
|
|
batch_size_min: 4.0, // Default minimum
|
|
batch_size_max: 96.0, // Default maximum (safe for RTX A4000 16GB)
|
|
async_loading: true, // Enable async loading by default
|
|
prefetch_count: 3, // Prefetch 3 batches (good balance)
|
|
training_paths,
|
|
early_stopping_patience: 5, // Default: 5 epochs (hyperopt optimized)
|
|
early_stopping_min_epochs: 5, // Default: 5 epochs (hyperopt optimized)
|
|
trial_counter: 0, // Start at trial 0
|
|
})
|
|
}
|
|
|
|
/// Configure early stopping parameters (overrides defaults)
|
|
pub fn with_early_stopping(mut self, patience: usize, min_epochs: usize) -> Self {
|
|
self.early_stopping_patience = patience;
|
|
self.early_stopping_min_epochs = min_epochs;
|
|
self
|
|
}
|
|
|
|
/// Set train/validation split ratio
|
|
pub fn with_train_split(mut self, split: f64) -> Self {
|
|
assert!(split > 0.0 && split < 1.0, "Split must be in (0, 1)");
|
|
self.train_split = split;
|
|
self
|
|
}
|
|
|
|
/// Set batch size bounds for GPU memory constraints
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `min` - Minimum batch size (must be >= 1)
|
|
/// * `max` - Maximum batch size (must be > min)
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```no_run
|
|
/// # use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
|
|
/// // Configure for RTX 4090 (24GB VRAM)
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_batch_size_bounds(4.0, 256.0);
|
|
///
|
|
/// // Configure for RTX 3050 Ti (4GB VRAM)
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_batch_size_bounds(4.0, 32.0);
|
|
/// # Ok::<(), anyhow::Error>(())
|
|
/// ```
|
|
pub fn with_batch_size_bounds(mut self, min: f64, max: f64) -> Self {
|
|
assert!(min >= 1.0, "Minimum batch size must be >= 1");
|
|
assert!(max > min, "Maximum batch size must be > minimum");
|
|
info!("Configuring batch_size bounds: [{}, {}]", min, max);
|
|
self.batch_size_min = min;
|
|
self.batch_size_max = max;
|
|
self
|
|
}
|
|
|
|
/// Enable or disable async data loading (prefetching)
|
|
///
|
|
/// When enabled, batches are prepared on CPU and transferred to GPU in a
|
|
/// background thread while the GPU trains on the current batch. This improves
|
|
/// GPU utilization from ~78% to ~90-95% and reduces training time by 20-30%.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `enabled` - Enable async loading
|
|
/// * `prefetch_count` - Number of batches to prefetch (2-3 recommended)
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```no_run
|
|
/// # use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
|
|
/// // Enable async loading with 3 batch prefetch
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_async_loading(true, 3);
|
|
///
|
|
/// // Disable async loading (sync mode)
|
|
/// let trainer = Mamba2Trainer::new("data.parquet", 50)?
|
|
/// .with_async_loading(false, 0);
|
|
/// # Ok::<(), anyhow::Error>(())
|
|
/// ```
|
|
pub fn with_async_loading(mut self, enabled: bool, prefetch_count: usize) -> Self {
|
|
if enabled {
|
|
assert!(
|
|
prefetch_count >= 2,
|
|
"Prefetch count must be >= 2 when async loading enabled"
|
|
);
|
|
assert!(
|
|
prefetch_count <= 10,
|
|
"Prefetch count must be <= 10 to avoid excessive memory"
|
|
);
|
|
}
|
|
info!(
|
|
"Configuring async data loading: enabled={}, prefetch={}",
|
|
enabled, prefetch_count
|
|
);
|
|
self.async_loading = enabled;
|
|
self.prefetch_count = prefetch_count;
|
|
self
|
|
}
|
|
|
|
/// Set checkpoint directory for saving best models
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `dir` - Path to checkpoint directory (will be created if it doesn't exist)
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Self for method chaining
|
|
///
|
|
/// # Deprecated
|
|
///
|
|
/// Use `with_training_paths()` instead for full path management
|
|
pub fn with_checkpoint_dir<P: Into<PathBuf>>(mut self, dir: P) -> Self {
|
|
// Update training_paths to use provided checkpoint directory
|
|
let checkpoint_dir = dir.into();
|
|
self.training_paths = TrainingPaths::new(
|
|
checkpoint_dir.parent().unwrap_or(&checkpoint_dir),
|
|
"mamba2",
|
|
"legacy",
|
|
);
|
|
info!(
|
|
"Checkpoint directory set to: {:?}",
|
|
self.training_paths.checkpoints_dir()
|
|
);
|
|
self
|
|
}
|
|
|
|
/// Set training paths configuration (recommended over with_checkpoint_dir)
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `paths` - Training paths configuration
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Self for method chaining
|
|
pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self {
|
|
info!("Training paths configured:");
|
|
info!(" Run directory: {:?}", paths.run_dir());
|
|
info!(" Checkpoints: {:?}", paths.checkpoints_dir());
|
|
info!(" Logs: {:?}", paths.logs_dir());
|
|
info!(" Hyperopt: {:?}", paths.hyperopt_dir());
|
|
self.training_paths = paths;
|
|
self
|
|
}
|
|
|
|
/// Denormalize a prediction from [0,1] to original price scale
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `normalized` - Normalized prediction in [0,1] range
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// `Some(price)` in original scale (e.g., $5000-6000 for ES futures),
|
|
/// or `None` if normalization params have not been set (i.e., training has not run yet).
|
|
pub fn denormalize_prediction(&self, normalized: f64) -> Option<f64> {
|
|
let min = self.target_min?;
|
|
let max = self.target_max?;
|
|
|
|
Some(normalized * (max - min) + min)
|
|
}
|
|
|
|
/// Load and prepare training data from Parquet
|
|
///
|
|
/// Reads OHLCV bars, extracts features, creates sequences.
|
|
fn load_and_prepare_data(
|
|
&self,
|
|
seq_len: usize,
|
|
_stride: usize, // P2 parameter - for future use with overlapping sequences
|
|
) -> Result<(Vec<(Tensor, Tensor)>, Vec<(Tensor, Tensor)>, f64, f64)> {
|
|
// Open Parquet file
|
|
let file = File::open(&self.parquet_file).with_context(|| {
|
|
format!(
|
|
"Failed to open Parquet file: {}",
|
|
self.parquet_file.display()
|
|
)
|
|
})?;
|
|
|
|
let builder = ParquetRecordBatchReaderBuilder::try_new(file)
|
|
.context("Failed to create Parquet reader")?;
|
|
|
|
let reader = builder.build().context("Failed to build Parquet reader")?;
|
|
|
|
// Read all OHLCV bars
|
|
let mut all_ohlcv_bars = Vec::new();
|
|
|
|
for batch_result in reader {
|
|
let batch = batch_result.context("Failed to read record batch")?;
|
|
|
|
let timestamps = batch
|
|
.column(9)
|
|
.as_any()
|
|
.downcast_ref::<PrimitiveArray<TimestampNanosecondType>>()
|
|
.context("Failed to downcast timestamp column")?;
|
|
|
|
let opens = batch
|
|
.column(3)
|
|
.as_any()
|
|
.downcast_ref::<Float64Array>()
|
|
.context("Failed to downcast open column")?;
|
|
|
|
let highs = batch
|
|
.column(4)
|
|
.as_any()
|
|
.downcast_ref::<Float64Array>()
|
|
.context("Failed to downcast high column")?;
|
|
|
|
let lows = batch
|
|
.column(5)
|
|
.as_any()
|
|
.downcast_ref::<Float64Array>()
|
|
.context("Failed to downcast low column")?;
|
|
|
|
let closes = batch
|
|
.column(6)
|
|
.as_any()
|
|
.downcast_ref::<Float64Array>()
|
|
.context("Failed to downcast close column")?;
|
|
|
|
let volumes = batch
|
|
.column(7)
|
|
.as_any()
|
|
.downcast_ref::<UInt64Array>()
|
|
.context("Failed to downcast volume column")?;
|
|
|
|
for i in 0..batch.num_rows() {
|
|
let timestamp_ns = timestamps.value(i);
|
|
let timestamp = chrono::DateTime::from_timestamp(
|
|
(timestamp_ns / 1_000_000_000) as i64,
|
|
(timestamp_ns.rem_euclid(1_000_000_000)) as u32,
|
|
)
|
|
.unwrap_or_else(|| chrono::Utc::now());
|
|
|
|
let bar = OHLCVBar {
|
|
timestamp,
|
|
open: opens.value(i),
|
|
high: highs.value(i),
|
|
low: lows.value(i),
|
|
close: closes.value(i),
|
|
volume: volumes.value(i) as f64,
|
|
};
|
|
|
|
all_ohlcv_bars.push(bar);
|
|
}
|
|
}
|
|
|
|
// P1 FIX: Validate minimum rows before processing
|
|
if all_ohlcv_bars.len() < seq_len + 1 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Insufficient data: {} bars (need at least {} for seq_len={})",
|
|
all_ohlcv_bars.len(),
|
|
seq_len + 1,
|
|
seq_len
|
|
))
|
|
.into());
|
|
}
|
|
|
|
// Extract features
|
|
let features =
|
|
extract_ml_features(&all_ohlcv_bars).context("Failed to extract features")?;
|
|
|
|
if features.is_empty() {
|
|
return Err(
|
|
MLError::ModelError("No features extracted from Parquet data".to_owned()).into(),
|
|
);
|
|
}
|
|
|
|
// P0 FIX: Collect all target prices for normalization
|
|
let mut all_target_prices = Vec::new();
|
|
for window_idx in 0..features.len().saturating_sub(seq_len) {
|
|
let target_price = all_ohlcv_bars[window_idx + seq_len].close;
|
|
all_target_prices.push(target_price);
|
|
}
|
|
|
|
// Compute normalization parameters
|
|
let target_min = all_target_prices
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::INFINITY, f64::min);
|
|
let target_max = all_target_prices
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::NEG_INFINITY, f64::max);
|
|
|
|
if (target_max - target_min).abs() < 1e-6 {
|
|
return Err(MLError::ModelError(
|
|
"Target prices have zero variance - cannot normalize".to_owned(),
|
|
)
|
|
.into());
|
|
}
|
|
|
|
info!(
|
|
"Target normalization: min={:.2}, max={:.2}, range={:.2}",
|
|
target_min,
|
|
target_max,
|
|
target_max - target_min
|
|
);
|
|
|
|
// FIX: Apply percentile clipping BEFORE normalization to prevent outliers
|
|
// (e.g., OBV features with extreme values) from crushing other features
|
|
let all_feature_values: Vec<f64> =
|
|
features.iter().flat_map(|f| f.iter().copied()).collect();
|
|
|
|
// Compute 1st and 99th percentiles
|
|
let mut sorted_features = all_feature_values.clone();
|
|
sorted_features.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
|
|
|
let p1_idx = (sorted_features.len() as f64 * 0.01).round() as usize;
|
|
let p99_idx = (sorted_features.len() as f64 * 0.99).round() as usize;
|
|
let p1 = sorted_features[p1_idx.min(sorted_features.len() - 1)];
|
|
let p99 = sorted_features[p99_idx.min(sorted_features.len() - 1)];
|
|
|
|
info!("Feature percentile clipping: p1={:.2}, p99={:.2}", p1, p99);
|
|
|
|
// Clip outliers to [p1, p99] range
|
|
let clipped_feature_values: Vec<f64> = all_feature_values
|
|
.iter()
|
|
.map(|&x| x.clamp(p1, p99))
|
|
.collect();
|
|
|
|
// Now compute normalization parameters from clipped data
|
|
let feature_min = clipped_feature_values
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::INFINITY, f64::min);
|
|
let feature_max = clipped_feature_values
|
|
.iter()
|
|
.copied()
|
|
.fold(f64::NEG_INFINITY, f64::max);
|
|
|
|
if (feature_max - feature_min).abs() < 1e-6 {
|
|
return Err(MLError::ModelError(
|
|
"Features have zero variance - cannot normalize".to_owned(),
|
|
)
|
|
.into());
|
|
}
|
|
|
|
info!(
|
|
"Feature normalization (after clipping): min={:.2}, max={:.2}, range={:.2}",
|
|
feature_min,
|
|
feature_max,
|
|
feature_max - feature_min
|
|
);
|
|
|
|
// Create sequences with normalized features and targets
|
|
let mut feature_sequences = Vec::new();
|
|
|
|
for (window_idx, &target_price) in all_target_prices.iter().enumerate() {
|
|
// FIX: Apply percentile clipping + normalization to [0, 1] range
|
|
let sequence: Vec<f64> = features[window_idx..window_idx + seq_len]
|
|
.iter()
|
|
.flat_map(|f| f.iter().copied())
|
|
.map(|val| {
|
|
// Clip to percentile range, then normalize
|
|
let clipped = val.clamp(p1, p99);
|
|
(clipped - feature_min) / (feature_max - feature_min)
|
|
})
|
|
.collect();
|
|
|
|
// Normalize target to [0,1]
|
|
let normalized_target = (target_price - target_min) / (target_max - target_min);
|
|
|
|
let input_tensor = Tensor::new(sequence.as_slice(), &Device::Cpu)?.reshape((
|
|
1,
|
|
seq_len,
|
|
self.d_model,
|
|
))?;
|
|
let target_tensor =
|
|
Tensor::new(&[normalized_target], &Device::Cpu)?.reshape((1, 1, 1))?;
|
|
|
|
feature_sequences.push((input_tensor, target_tensor));
|
|
}
|
|
|
|
// Split train/validation
|
|
let split_idx = (feature_sequences.len() as f64 * self.train_split) as usize;
|
|
let train_data = feature_sequences[..split_idx].to_vec();
|
|
let val_data = feature_sequences[split_idx..].to_vec();
|
|
|
|
// P1 FIX: Validate validation set size
|
|
if val_data.len() < 10 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Validation set too small: {} samples (need at least 10)",
|
|
val_data.len()
|
|
))
|
|
.into());
|
|
}
|
|
|
|
Ok((train_data, val_data, target_min, target_max))
|
|
}
|
|
|
|
/// Train model with async data loading (optimized for GPU utilization)
|
|
///
|
|
/// This method uses AsyncDataLoader to prefetch batches while GPU trains,
|
|
/// improving GPU utilization from ~78% to ~90-95% and reducing training
|
|
/// time by 20-30%.
|
|
///
|
|
/// # Architecture
|
|
///
|
|
/// ```text
|
|
/// CPU Thread (Background): GPU Thread (Main):
|
|
/// ┌─────────────────┐ ┌─────────────────┐
|
|
/// │ Load batch N+1 │ ────────> │ Train batch N │
|
|
/// │ Concat tensors │ │ Forward pass │
|
|
/// │ Transfer to GPU │ │ Backward pass │
|
|
/// └─────────────────┘ │ Optimizer step │
|
|
/// │ └─────────────────┘
|
|
/// ▼ │
|
|
/// ┌─────────────────┐ │
|
|
/// │ Load batch N+2 │ │
|
|
/// │ ... │ <───────────────────┘
|
|
/// └─────────────────┘
|
|
/// ```
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `model` - MAMBA-2 model to train
|
|
/// * `train_data` - Training data (already prepared tensors)
|
|
/// * `val_data` - Validation data
|
|
/// * `epochs` - Number of training epochs
|
|
/// * `batch_size` - Batch size for training
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// Training history with metrics per epoch
|
|
async fn train_with_async_loading(
|
|
&self,
|
|
model: &mut Mamba2SSM,
|
|
train_data: &[(Tensor, Tensor)],
|
|
val_data: &[(Tensor, Tensor)],
|
|
epochs: usize,
|
|
batch_size: usize,
|
|
checkpoint_dir: Option<&std::path::Path>,
|
|
) -> Result<Vec<crate::mamba::TrainingEpoch>, MLError> {
|
|
info!(
|
|
"Async data loading enabled (prefetch={}, batch_size={})",
|
|
self.prefetch_count, batch_size
|
|
);
|
|
|
|
// Call the new train_async() method with AsyncDataLoader
|
|
model
|
|
.train_async(
|
|
train_data,
|
|
val_data,
|
|
epochs,
|
|
batch_size,
|
|
self.prefetch_count,
|
|
checkpoint_dir,
|
|
)
|
|
.await
|
|
}
|
|
}
|
|
|
|
/// Write a log entry to the training log file
|
|
fn write_training_log_mamba2(
|
|
logs_dir: &std::path::Path,
|
|
message: &str,
|
|
) -> Result<(), std::io::Error> {
|
|
let log_file = logs_dir.join("training.log");
|
|
let mut file = OpenOptions::new()
|
|
.create(true)
|
|
.append(true)
|
|
.open(log_file)?;
|
|
|
|
let timestamp = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S");
|
|
writeln!(file, "[{}] {}", timestamp, message)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Write trial results to JSON file
|
|
fn write_trial_result_mamba2(
|
|
hyperopt_dir: &std::path::Path,
|
|
trial_result: &crate::hyperopt::traits::TrialResult<Mamba2Params>,
|
|
) -> Result<(), std::io::Error> {
|
|
let trials_file = hyperopt_dir.join("trials.json");
|
|
|
|
// Read existing trials (if any)
|
|
let mut all_trials = if trials_file.exists() {
|
|
let content = std::fs::read_to_string(&trials_file)?;
|
|
serde_json::from_str::<Vec<serde_json::Value>>(&content).unwrap_or_default()
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
// Append new trial
|
|
let trial_json = serde_json::to_value(trial_result)
|
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
|
|
all_trials.push(trial_json);
|
|
|
|
// Write back to file (pretty printed)
|
|
let content = serde_json::to_string_pretty(&all_trials)
|
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
|
|
std::fs::write(&trials_file, content)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
impl HyperparameterOptimizable for Mamba2Trainer {
|
|
type Params = Mamba2Params;
|
|
type Metrics = Mamba2Metrics;
|
|
|
|
#[allow(clippy::unwrap_in_result)]
|
|
fn train_with_params(&mut self, mut params: Self::Params) -> Result<Self::Metrics, MLError> {
|
|
// START: Add trial timing
|
|
let trial_start = std::time::Instant::now();
|
|
|
|
// Get current trial number and increment for next trial
|
|
let current_trial = self.trial_counter;
|
|
self.trial_counter += 1;
|
|
|
|
// Clamp batch_size to configured bounds (for GPU memory constraints)
|
|
let original_batch_size = params.batch_size;
|
|
let clamped_batch_size = (params.batch_size as f64)
|
|
.clamp(self.batch_size_min, self.batch_size_max)
|
|
.round() as usize;
|
|
|
|
if clamped_batch_size != original_batch_size {
|
|
warn!(
|
|
"Batch size clamped: {} → {} (bounds: [{}, {}])",
|
|
original_batch_size, clamped_batch_size, self.batch_size_min, self.batch_size_max
|
|
);
|
|
params.batch_size = clamped_batch_size;
|
|
}
|
|
|
|
info!("Training MAMBA-2 with 12 hyperparameters:");
|
|
info!(" Learning rate: {:.6}", params.learning_rate);
|
|
info!(
|
|
" Batch size: {} (bounds: [{}, {}])",
|
|
params.batch_size, self.batch_size_min, self.batch_size_max
|
|
);
|
|
info!(" Dropout: {:.3}", params.dropout);
|
|
info!(" Weight decay: {:.6}", params.weight_decay);
|
|
info!(" P0 - Grad clip: {:.3}", params.grad_clip);
|
|
info!(" P0 - Warmup steps: {}", params.warmup_steps);
|
|
info!(" P0 - Adam beta1: {:.4}", params.adam_beta1);
|
|
info!(" P1 - Adam beta2: {:.4}", params.adam_beta2);
|
|
info!(" P1 - Adam epsilon: {:.2e}", params.adam_epsilon);
|
|
info!(" P2 - Lookback window: {}", params.lookback_window);
|
|
info!(" P2 - Sequence stride: {}", params.sequence_stride);
|
|
info!(" P2 - Norm epsilon: {:.2e}", params.norm_eps);
|
|
|
|
// Log trial start (ensure directory exists first)
|
|
std::fs::create_dir_all(self.training_paths.logs_dir()).ok();
|
|
write_training_log_mamba2(
|
|
&self.training_paths.logs_dir(),
|
|
&format!("=== Starting MAMBA-2 Trial ===\nParams: {:#?}", params),
|
|
)
|
|
.ok();
|
|
|
|
// Load and prepare data FIRST (required to calculate total_decay_steps)
|
|
let (train_data, val_data, target_min, target_max) = self
|
|
.load_and_prepare_data(params.lookback_window, params.sequence_stride)
|
|
.map_err(|e| MLError::ModelError(format!("Data loading failed: {}", e)))?;
|
|
|
|
// Calculate total_decay_steps dynamically based on actual data size
|
|
let steps_per_epoch = (train_data.len() + params.batch_size - 1) / params.batch_size;
|
|
let total_decay_steps = self.epochs * steps_per_epoch;
|
|
|
|
info!(
|
|
"Calculated total_decay_steps: {} (epochs={}, steps_per_epoch={})",
|
|
total_decay_steps, self.epochs, steps_per_epoch
|
|
);
|
|
|
|
// Create MAMBA-2 config with trial hyperparameters (ALL 12 params + calculated total_decay_steps)
|
|
let mamba_config = Mamba2Config {
|
|
d_model: self.d_model,
|
|
d_state: 16,
|
|
d_head: self.d_model / 8,
|
|
num_heads: 8,
|
|
expand: 2,
|
|
num_layers: 6,
|
|
dropout: params.dropout,
|
|
use_ssd: true,
|
|
use_selective_state: true,
|
|
hardware_aware: true,
|
|
target_latency_us: 5,
|
|
max_seq_len: 120,
|
|
learning_rate: params.learning_rate,
|
|
weight_decay: params.weight_decay,
|
|
grad_clip: params.grad_clip, // P0
|
|
warmup_steps: params.warmup_steps, // P0
|
|
adam_beta1: params.adam_beta1, // P0
|
|
adam_beta2: params.adam_beta2, // P1
|
|
adam_epsilon: params.adam_epsilon, // P1
|
|
total_decay_steps, // P1 - CALCULATED DYNAMICALLY
|
|
batch_size: params.batch_size,
|
|
seq_len: params.lookback_window, // P2
|
|
shuffle_batches: false,
|
|
optimizer_type: OptimizerType::Adam,
|
|
sgd_momentum: 0.9,
|
|
sequence_stride: params.sequence_stride, // P2
|
|
norm_eps: params.norm_eps, // P2
|
|
early_stopping_enabled: true, // Enable early stopping for hyperopt
|
|
early_stopping_patience: self.early_stopping_patience,
|
|
early_stopping_min_delta: 1e-4, // Minimum improvement threshold
|
|
early_stopping_min_epochs: self.early_stopping_min_epochs,
|
|
};
|
|
|
|
// Store normalization params for inference
|
|
self.target_min = Some(target_min);
|
|
self.target_max = Some(target_max);
|
|
|
|
if train_data.is_empty() || val_data.is_empty() {
|
|
warn!("Empty training or validation data");
|
|
return Ok(Mamba2Metrics {
|
|
val_loss: 1000.0, // Penalty
|
|
train_loss: 1000.0,
|
|
val_perplexity: f64::INFINITY,
|
|
directional_accuracy: 0.0,
|
|
mae: 1000.0,
|
|
rmse: 1000.0,
|
|
r_squared: 0.0,
|
|
epochs_completed: 0,
|
|
});
|
|
}
|
|
|
|
// Create all training directories
|
|
self.training_paths.create_all().map_err(|e| {
|
|
MLError::ModelError(format!("Failed to create training directories: {}", e))
|
|
})?;
|
|
|
|
info!("Training directories created:");
|
|
info!(" Checkpoints: {:?}", self.training_paths.checkpoints_dir());
|
|
|
|
// Create and train model
|
|
let mut model = Mamba2SSM::new(mamba_config.clone(), &self.device)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to create model: {}", e)))?;
|
|
|
|
// P1 FIX: Wrap training in catch_unwind to handle CUDA OOM panics gracefully
|
|
// Run training (async or sync based on configuration)
|
|
let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
|
if self.async_loading {
|
|
info!(
|
|
"Using async data loading (prefetch={})",
|
|
self.prefetch_count
|
|
);
|
|
tokio::runtime::Runtime::new()
|
|
.map_err(|e| MLError::ModelError(format!("Failed to create async runtime: {}", e)))?
|
|
.block_on(self.train_with_async_loading(
|
|
&mut model,
|
|
&train_data,
|
|
&val_data,
|
|
self.epochs,
|
|
params.batch_size,
|
|
Some(&self.training_paths.checkpoints_dir()),
|
|
))
|
|
} else {
|
|
info!("Using synchronous data loading");
|
|
tokio::runtime::Runtime::new()
|
|
.map_err(|e| MLError::ModelError(format!("Failed to create async runtime: {}", e)))?
|
|
.block_on(model.train(
|
|
&train_data,
|
|
&val_data,
|
|
self.epochs,
|
|
Some(&self.training_paths.checkpoints_dir()),
|
|
))
|
|
}
|
|
}));
|
|
|
|
let training_history = match training_result {
|
|
Ok(Ok(history)) => history,
|
|
Ok(Err(e)) => {
|
|
warn!("Training failed (possibly OOM): {}", e);
|
|
// Return penalty metrics instead of propagating error
|
|
return Ok(Mamba2Metrics {
|
|
val_loss: 1000.0,
|
|
train_loss: 1000.0,
|
|
val_perplexity: f64::INFINITY,
|
|
directional_accuracy: 0.0,
|
|
mae: 1000.0,
|
|
rmse: 1000.0,
|
|
r_squared: 0.0,
|
|
epochs_completed: 0,
|
|
});
|
|
},
|
|
Err(_) => {
|
|
warn!("Training panicked (likely CUDA OOM or GPU error)");
|
|
// Return penalty metrics for optimizer
|
|
return Ok(Mamba2Metrics {
|
|
val_loss: 1000.0,
|
|
train_loss: 1000.0,
|
|
val_perplexity: f64::INFINITY,
|
|
directional_accuracy: 0.0,
|
|
mae: 1000.0,
|
|
rmse: 1000.0,
|
|
r_squared: 0.0,
|
|
epochs_completed: 0,
|
|
});
|
|
},
|
|
};
|
|
|
|
// Extract final metrics
|
|
let final_epoch = training_history
|
|
.last()
|
|
.ok_or_else(|| MLError::TrainingError("No training history".to_owned()))?;
|
|
|
|
// Map fields from TrainingEpoch (which only has loss/accuracy)
|
|
// to Mamba2Metrics (which expects detailed metrics)
|
|
let metrics = Mamba2Metrics {
|
|
val_loss: final_epoch.loss, // Use loss as val_loss
|
|
train_loss: final_epoch.loss, // Same value for train_loss (best available)
|
|
val_perplexity: final_epoch.loss.exp(),
|
|
directional_accuracy: final_epoch.accuracy, // Use accuracy as directional_accuracy
|
|
mae: final_epoch.loss, // Approximation
|
|
rmse: final_epoch.loss.sqrt(), // Approximation
|
|
r_squared: 1.0 - final_epoch.loss.min(1.0), // Approximation
|
|
epochs_completed: training_history.len(),
|
|
};
|
|
|
|
info!("Training completed:");
|
|
info!(" Training loss: {:.6}", metrics.train_loss);
|
|
info!(" Validation loss: {:.6}", metrics.val_loss);
|
|
info!(" Perplexity: {:.4}", metrics.val_perplexity);
|
|
info!(
|
|
" Directional accuracy: {:.2}%",
|
|
metrics.directional_accuracy * 100.0
|
|
);
|
|
info!(" MAE: {:.4}", metrics.mae);
|
|
info!(" RMSE: {:.4}", metrics.rmse);
|
|
info!(" R²: {:.4}", metrics.r_squared);
|
|
|
|
// END: Add trial completion logging
|
|
let duration_secs = trial_start.elapsed().as_secs_f64();
|
|
write_training_log_mamba2(
|
|
&self.training_paths.logs_dir(),
|
|
&format!(
|
|
"Training completed in {:.2}s: val_loss={:.6}, train_loss={:.6}, accuracy={:.2}%",
|
|
duration_secs,
|
|
metrics.val_loss,
|
|
metrics.train_loss,
|
|
metrics.directional_accuracy * 100.0
|
|
),
|
|
)
|
|
.ok();
|
|
|
|
// CRITICAL FIX: Explicit memory cleanup to prevent OOM between trials
|
|
// Drop model, data loaders, and sync CUDA to free GPU/RAM
|
|
info!("Cleaning up resources...");
|
|
drop(model);
|
|
drop(train_data);
|
|
drop(val_data);
|
|
|
|
// Sync CUDA to ensure GPU memory is freed
|
|
if self.device.is_cuda() {
|
|
use candle_core::Device;
|
|
if let Device::Cuda(_) = &self.device {
|
|
// Force CUDA synchronization to release GPU memory
|
|
std::thread::sleep(std::time::Duration::from_millis(100));
|
|
}
|
|
}
|
|
info!("Resource cleanup complete");
|
|
|
|
// Write trial result to JSON (ensure directory exists first)
|
|
let trial_result = crate::hyperopt::traits::TrialResult {
|
|
trial_num: current_trial,
|
|
params,
|
|
objective: Self::extract_objective(&metrics),
|
|
duration_secs,
|
|
};
|
|
|
|
std::fs::create_dir_all(self.training_paths.hyperopt_dir()).ok();
|
|
write_trial_result_mamba2(&self.training_paths.hyperopt_dir(), &trial_result).ok();
|
|
|
|
Ok(metrics)
|
|
}
|
|
|
|
fn extract_objective(metrics: &Self::Metrics) -> f64 {
|
|
metrics.val_loss
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_mamba2_params_roundtrip() {
|
|
let params = Mamba2Params {
|
|
learning_rate: 0.001,
|
|
batch_size: 64,
|
|
dropout: 0.2,
|
|
weight_decay: 0.0001,
|
|
grad_clip: 2.5,
|
|
warmup_steps: 500,
|
|
adam_beta1: 0.9,
|
|
adam_beta2: 0.999,
|
|
adam_epsilon: 1e-8,
|
|
lookback_window: 90,
|
|
sequence_stride: 3,
|
|
norm_eps: 5e-5,
|
|
};
|
|
|
|
let continuous = params.to_continuous();
|
|
assert_eq!(continuous.len(), 12, "Expected 12 continuous parameters");
|
|
|
|
let recovered = Mamba2Params::from_continuous(&continuous).unwrap();
|
|
|
|
assert!((recovered.learning_rate - params.learning_rate).abs() < 1e-10);
|
|
assert_eq!(recovered.batch_size, params.batch_size);
|
|
assert!((recovered.dropout - params.dropout).abs() < 1e-10);
|
|
assert!((recovered.weight_decay - params.weight_decay).abs() < 1e-10);
|
|
assert!((recovered.grad_clip - params.grad_clip).abs() < 1e-6);
|
|
assert_eq!(recovered.warmup_steps, params.warmup_steps);
|
|
assert!((recovered.adam_beta1 - params.adam_beta1).abs() < 1e-10);
|
|
assert!((recovered.adam_beta2 - params.adam_beta2).abs() < 1e-10);
|
|
assert!((recovered.adam_epsilon - params.adam_epsilon).abs() < 1e-12);
|
|
assert_eq!(recovered.lookback_window, params.lookback_window);
|
|
assert_eq!(recovered.sequence_stride, params.sequence_stride);
|
|
assert!((recovered.norm_eps - params.norm_eps).abs() < 1e-12);
|
|
}
|
|
|
|
#[test]
|
|
fn test_mamba2_params_bounds() {
|
|
let bounds = Mamba2Params::continuous_bounds();
|
|
assert_eq!(bounds.len(), 12, "Expected 12 parameter bounds");
|
|
|
|
// Check log-scale bounds are reasonable
|
|
assert!(bounds[0].0 < bounds[0].1); // learning_rate
|
|
assert!(bounds[3].0 < bounds[3].1); // weight_decay
|
|
assert!(bounds[4].0 < bounds[4].1); // grad_clip
|
|
assert!(bounds[8].0 < bounds[8].1); // adam_epsilon
|
|
assert!(bounds[11].0 < bounds[11].1); // norm_eps
|
|
|
|
// Check linear bounds
|
|
assert_eq!(bounds[1], (4.0, 256.0)); // batch_size (wide bounds for optimizer exploration)
|
|
assert_eq!(bounds[2], (0.0, 0.5)); // dropout
|
|
assert_eq!(bounds[5], (100.0, 2000.0)); // warmup_steps
|
|
assert_eq!(bounds[6], (0.85, 0.95)); // adam_beta1
|
|
assert_eq!(bounds[7], (0.98, 0.999)); // adam_beta2
|
|
assert_eq!(bounds[9], (30.0, 120.0)); // lookback_window
|
|
assert_eq!(bounds[10], (1.0, 5.0)); // sequence_stride
|
|
}
|
|
|
|
#[test]
|
|
fn test_param_names() {
|
|
let names = Mamba2Params::param_names();
|
|
assert_eq!(names.len(), 12, "Expected 12 parameter names");
|
|
assert_eq!(names[0], "learning_rate");
|
|
assert_eq!(names[1], "batch_size");
|
|
assert_eq!(names[2], "dropout");
|
|
assert_eq!(names[3], "weight_decay");
|
|
assert_eq!(names[4], "grad_clip");
|
|
assert_eq!(names[5], "warmup_steps");
|
|
assert_eq!(names[6], "adam_beta1");
|
|
assert_eq!(names[7], "adam_beta2");
|
|
assert_eq!(names[8], "adam_epsilon");
|
|
assert_eq!(names[9], "lookback_window");
|
|
assert_eq!(names[10], "sequence_stride");
|
|
assert_eq!(names[11], "norm_eps");
|
|
}
|
|
|
|
#[test]
|
|
fn test_target_normalization() {
|
|
// Test normalization/denormalization with realistic ES price ranges
|
|
let target_min = 5000.0;
|
|
let target_max = 6000.0;
|
|
|
|
// Test min value
|
|
let normalized_min: f64 = (target_min - target_min) / (target_max - target_min);
|
|
assert!(
|
|
(normalized_min - 0.0).abs() < 1e-10,
|
|
"Min should normalize to 0"
|
|
);
|
|
|
|
// Test max value
|
|
let normalized_max: f64 = (target_max - target_min) / (target_max - target_min);
|
|
assert!(
|
|
(normalized_max - 1.0).abs() < 1e-10,
|
|
"Max should normalize to 1"
|
|
);
|
|
|
|
// Test mid value
|
|
let mid_price = 5500.0;
|
|
let normalized_mid: f64 = (mid_price - target_min) / (target_max - target_min);
|
|
assert!(
|
|
(normalized_mid - 0.5).abs() < 1e-10,
|
|
"Midpoint should normalize to 0.5"
|
|
);
|
|
|
|
// Test denormalization recovers original
|
|
let denormalized: f64 = normalized_mid * (target_max - target_min) + target_min;
|
|
assert!(
|
|
(denormalized - mid_price).abs() < 1e-6,
|
|
"Denormalization should recover original price"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_denormalize_prediction() {
|
|
// Create a trainer with normalization params set
|
|
let trainer = Mamba2Trainer {
|
|
parquet_file: PathBuf::from("dummy.parquet"),
|
|
epochs: 1,
|
|
device: Device::Cpu,
|
|
feature_config: FeatureConfig::wave_d(),
|
|
d_model: 225,
|
|
train_split: 0.8,
|
|
target_min: Some(5000.0),
|
|
target_max: Some(6000.0),
|
|
batch_size_min: 4.0,
|
|
batch_size_max: 96.0,
|
|
async_loading: false,
|
|
prefetch_count: 0,
|
|
training_paths: TrainingPaths::new("/tmp/test", "mamba2", "test"),
|
|
early_stopping_patience: 10,
|
|
early_stopping_min_epochs: 5,
|
|
trial_counter: 0,
|
|
};
|
|
|
|
// Test denormalization
|
|
assert!((trainer.denormalize_prediction(0.0).unwrap_or(0.0) - 5000.0).abs() < 1e-6);
|
|
assert!((trainer.denormalize_prediction(1.0).unwrap_or(0.0) - 6000.0).abs() < 1e-6);
|
|
assert!((trainer.denormalize_prediction(0.5).unwrap_or(0.0) - 5500.0).abs() < 1e-6);
|
|
assert!((trainer.denormalize_prediction(0.25).unwrap_or(0.0) - 5250.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_denormalize_before_training_returns_none() {
|
|
// Create trainer without normalization params
|
|
let trainer = Mamba2Trainer {
|
|
parquet_file: PathBuf::from("dummy.parquet"),
|
|
epochs: 1,
|
|
device: Device::Cpu,
|
|
feature_config: FeatureConfig::wave_d(),
|
|
batch_size_min: 4.0,
|
|
batch_size_max: 96.0,
|
|
async_loading: false,
|
|
prefetch_count: 0,
|
|
d_model: 225,
|
|
train_split: 0.8,
|
|
target_min: None,
|
|
target_max: None,
|
|
training_paths: TrainingPaths::new("/tmp/test", "mamba2", "test"),
|
|
early_stopping_patience: 10,
|
|
early_stopping_min_epochs: 5,
|
|
trial_counter: 0,
|
|
};
|
|
|
|
// Should return None when normalization params not set
|
|
assert!(trainer.denormalize_prediction(0.5).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_normalized_targets_in_range() {
|
|
// Verify that normalized targets are always in [0,1]
|
|
let target_min = 5000.0;
|
|
let target_max = 6000.0;
|
|
|
|
let test_prices = vec![5000.0, 5100.0, 5500.0, 5900.0, 6000.0];
|
|
|
|
for price in test_prices {
|
|
let normalized: f64 = (price - target_min) / (target_max - target_min);
|
|
assert!(
|
|
normalized >= 0.0,
|
|
"Normalized target {} should be >= 0",
|
|
normalized
|
|
);
|
|
assert!(
|
|
normalized <= 1.0,
|
|
"Normalized target {} should be <= 1",
|
|
normalized
|
|
);
|
|
|
|
// Verify round-trip
|
|
let denormalized: f64 = normalized * (target_max - target_min) + target_min;
|
|
assert!(
|
|
(denormalized - price).abs() < 1e-6,
|
|
"Round-trip failed: {} -> {} -> {}",
|
|
price,
|
|
normalized,
|
|
denormalized
|
|
);
|
|
}
|
|
}
|
|
}
|