foxhunt/crates/ml/examples/hyperopt_mamba2_demo.rs

//! MAMBA-2 Hyperparameter Optimization Demo
//!
//! This example demonstrates how to use the argmin-based hyperparameter
//! optimization framework with MAMBA-2. It runs a small-scale optimization
//! to show the complete workflow.
//!
//! ## Usage
//!
//! ```bash
//! # Local training with small dataset (default /tmp)
//! cargo run -p ml --example hyperopt_mamba2_demo --release --features cuda -- \
//!   --parquet-file test_data/ES_FUT_small.parquet \
//!   --trials 4 --epochs 3
//!
//! # Runpod training with custom base dir
//! ./hyperopt_mamba2_demo \
//!   --parquet-file /runpod-volume/datasets/parquet/futures/ES_FUT_180d.parquet \
//!   --base-dir /runpod-volume \
//!   --trials 30 --epochs 50
//!
//! # Resume from specific run
//! ./hyperopt_mamba2_demo \
//!   --base-dir /runpod-volume \
//!   --run-id 20251028_223000_hyperopt \
//!   --parquet-file test_data/ES_FUT_180d.parquet \
//!   --trials 30 --epochs 50
//! ```
//!
//! ## Output
//!
//! The example will:
//! 1. Initialize MAMBA-2 trainer with specified Parquet file
//! 2. Run argmin optimization with Nelder-Mead simplex
//! 3. Display trial results including loss and parameter values
//! 4. Report best hyperparameters found
//! 5. Show expected improvement vs default parameters

use anyhow::Result;
use clap::Parser;
use ml::hyperopt::adapters::mamba2::Mamba2Trainer;
use ml::hyperopt::paths::{generate_run_id, TrainingPaths};
use ml::hyperopt::ArgminOptimizer;
use tracing::{info, Level};
use tracing_subscriber;

#[derive(Parser, Debug)]
#[command(name = "MAMBA-2 Hyperparameter Optimization Demo")]
#[command(about = "Demonstrates argmin-based hyperparameter optimization for MAMBA-2")]
struct Args {
    /// Path to Parquet file with OHLCV data
    #[arg(long)]
    parquet_file: String,

    /// Number of optimization trials (default: 10)
    #[arg(long, default_value = "10")]
    trials: usize,

    /// Number of training epochs per trial (default: 20)
    #[arg(long, default_value = "20")]
    epochs: usize,

    /// Number of initial random samples (default: 3)
    #[arg(long, default_value = "3")]
    n_initial: usize,

    /// Random seed for reproducibility (default: 42)
    #[arg(long, default_value = "42")]
    seed: u64,

    /// Minimum batch size (default: 4)
    #[arg(long, default_value = "4")]
    batch_size_min: usize,

    /// Maximum batch size for GPU memory constraints (default: 96 for RTX A4000 16GB)
    /// Examples: RTX 3050 Ti 4GB = 32, RTX A4000 16GB = 96, RTX 4090 24GB = 256
    #[arg(long, default_value = "96")]
    batch_size_max: usize,

    /// Base directory for training outputs (e.g., /runpod-volume)
    #[arg(long, default_value = "/tmp/ml_training")]
    base_dir: String,

    /// Run ID (auto-generated if not provided)
    #[arg(long)]
    run_id: Option<String>,

    /// Run type (hyperopt, production, test)
    #[arg(long, default_value = "hyperopt")]
    run_type: String,

    /// Early stopping patience (epochs without improvement before stopping)
    #[arg(long, default_value = "5")]
    early_stopping_patience: usize,

    /// Early stopping minimum epochs (minimum epochs before early stopping can trigger)
    #[arg(long, default_value = "5")]
    early_stopping_min_epochs: usize,
}

fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .with_target(false)
        .init();

    // Parse arguments
    let args = Args::parse();

    // Generate run ID if not provided
    let run_id = args
        .run_id
        .unwrap_or_else(|| generate_run_id(&args.run_type));

    // Create training paths
    let training_paths = TrainingPaths::new(&args.base_dir, "mamba2", &run_id);

    info!("========================================");
    info!("MAMBA-2 Hyperparameter Optimization Demo");
    info!("========================================");
    info!("Configuration:");
    info!("  Parquet file: {}", args.parquet_file);
    info!("  Trials: {}", args.trials);
    info!("  Epochs per trial: {}", args.epochs);
    info!("  Initial samples: {}", args.n_initial);
    info!("  Random seed: {}", args.seed);
    info!(
        "  Batch size bounds: [{}, {}]",
        args.batch_size_min, args.batch_size_max
    );
    info!("");
    info!("Training Paths:");
    info!("  Run ID: {}", run_id);
    info!("  Base directory: {}", args.base_dir);
    info!("  Run directory: {:?}", training_paths.run_dir());
    info!("  Checkpoints: {:?}", training_paths.checkpoints_dir());
    info!("  Logs: {:?}", training_paths.logs_dir());
    info!("  Hyperopt: {:?}", training_paths.hyperopt_dir());
    info!("");

    // Create trainer with training paths
    info!("Creating MAMBA-2 trainer...");
    let trainer = Mamba2Trainer::new(&args.parquet_file, args.epochs)?
        .with_batch_size_bounds(args.batch_size_min as f64, args.batch_size_max as f64)
        .with_early_stopping(args.early_stopping_patience, args.early_stopping_min_epochs)
        .with_training_paths(training_paths);

    // Create optimizer
    info!("Initializing argmin optimizer...");
    let optimizer = ArgminOptimizer::builder()
        .max_trials(args.trials)
        .n_initial(args.n_initial)
        .seed(args.seed)
        .build();

    // Run optimization
    info!("");
    info!("Starting optimization (this may take a while)...");
    info!(
        "Expected runtime: ~{} minutes",
        estimate_runtime(args.trials, args.epochs)
    );
    info!("");

    let result = optimizer.optimize(trainer)?;

    // Display results
    info!("");
    info!("========================================");
    info!("Optimization Complete!");
    info!("========================================");
    info!("");
    info!("Best Hyperparameters:");
    info!("  Learning rate: {:.6}", result.best_params.learning_rate);
    info!("  Batch size: {}", result.best_params.batch_size);
    info!("  Dropout: {:.3}", result.best_params.dropout);
    info!("  Weight decay: {:.6}", result.best_params.weight_decay);
    info!("");
    info!("Performance:");
    info!("  Best validation loss: {:.6}", result.best_objective);
    info!("  Total trials: {}", result.all_trials.len());

    // Find convergence trial (where best was found)
    let convergence_trial = result
        .all_trials
        .iter()
        .position(|t| (t.objective - result.best_objective).abs() < 1e-10)
        .unwrap_or(0);
    info!("  Convergence: {} trials to best", convergence_trial + 1);
    info!("");

    // Show top 5 trials
    if result.all_trials.len() >= 5 {
        info!("Top 5 Trials:");
        let mut sorted_trials = result.all_trials.clone();
        sorted_trials.sort_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap());

        for (i, trial) in sorted_trials.iter().take(5).enumerate() {
            info!(
                "  {}. Loss: {:.6} (LR: {:.6}, BS: {}, Dropout: {:.3})",
                i + 1,
                trial.objective,
                trial.params.learning_rate,
                trial.params.batch_size,
                trial.params.dropout
            );
        }
    }

    info!("");
    info!("========================================");
    info!("Next Steps:");
    info!("========================================");
    info!("1. Use best parameters for production training");
    info!("2. Run longer optimization (50+ trials) for better results");
    info!("3. Validate on holdout dataset");
    info!("4. Deploy optimized model to trading system");

    Ok(())
}

/// Estimate runtime based on trials and epochs
fn estimate_runtime(trials: usize, epochs: usize) -> usize {
    // Rough estimate: 2 min per 50 epochs on RTX 3050 Ti
    let minutes_per_trial = (epochs as f64 / 50.0) * 2.0;
    let total_minutes = (trials as f64 * minutes_per_trial).ceil() as usize;
    total_minutes
}