From b09ebfc3cb2ffc475c181506dd53cfce308a50fe Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 26 Feb 2026 23:47:52 +0100 Subject: [PATCH] feat(ml): add HyperparameterOptimizable trainers for all 8 supervised models Extend hyperopt infrastructure to support all 10 ML models (DQN, PPO + 8 supervised). Previously only TFT and Mamba2 had hyperopt trainers. - Add HyperparameterOptimizable impl for Liquid, TGGN, TLOB, KAN, xLSTM, Diffusion - Create shared_data.rs with common data prep utilities (build_flat_pairs, build_sequence_pairs, write_trial_result_json) - Extend hyperopt_baseline_supervised binary to dispatch all 8 models (individual, "both" for tft+mamba2, "all" for all 8) - Add CI jobs: 7 train-validate + 10 hyperopt jobs for all models - Fix DiffusionMetrics NaN default, XLSTMMetrics serde, safe indexing Co-Authored-By: Claude Opus 4.6 --- .gitlab-ci.yml | 513 ++++++++++++++++++ .../examples/hyperopt_baseline_supervised.rs | 374 +++++++++++-- crates/ml/src/hyperopt/adapters/diffusion.rs | 265 ++++++++- crates/ml/src/hyperopt/adapters/kan.rs | 252 ++++++++- crates/ml/src/hyperopt/adapters/liquid.rs | 280 +++++++++- crates/ml/src/hyperopt/adapters/mod.rs | 12 +- crates/ml/src/hyperopt/adapters/tggn.rs | 255 ++++++++- crates/ml/src/hyperopt/adapters/tlob.rs | 246 ++++++++- crates/ml/src/hyperopt/adapters/xlstm.rs | 272 +++++++++- crates/ml/src/hyperopt/mod.rs | 1 + crates/ml/src/hyperopt/shared_data.rs | 300 ++++++++++ 11 files changed, 2694 insertions(+), 76 deletions(-) create mode 100644 crates/ml/src/hyperopt/shared_data.rs diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 422fab66f..2d33cad8c 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -548,6 +548,519 @@ hyperopt-ppo: when: always expire_in: 90 days +# --- Train-validate: remaining supervised models --- + +train-validate-mamba2: + extends: .train-validate-base + needs: [build-training] + script: + - /usr/local/bin/train_baseline_supervised + --model mamba2 + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --max-steps-per-epoch 50 + - echo "=== Mamba2 training complete, running evaluation ===" + - /usr/local/bin/evaluate_supervised + --model mamba2 + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/mamba2_eval_report.json + --symbol ES.FUT + - echo "Mamba2 validation + evaluation passed" + - cat ${CI_PROJECT_DIR}/output/mamba2_eval_report.json + artifacts: + paths: + - output/mamba2_eval_report.json + when: always + expire_in: 30 days + +train-validate-liquid: + extends: .train-validate-base + needs: [build-training] + script: + - /usr/local/bin/train_baseline_supervised + --model liquid + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --max-steps-per-epoch 50 + - echo "=== Liquid training complete, running evaluation ===" + - /usr/local/bin/evaluate_supervised + --model liquid + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/liquid_eval_report.json + --symbol ES.FUT + - echo "Liquid validation + evaluation passed" + - cat ${CI_PROJECT_DIR}/output/liquid_eval_report.json + artifacts: + paths: + - output/liquid_eval_report.json + when: always + expire_in: 30 days + +train-validate-tggn: + extends: .train-validate-base + needs: [build-training] + script: + - /usr/local/bin/train_baseline_supervised + --model tggn + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --max-steps-per-epoch 50 + - echo "=== TGGN training complete, running evaluation ===" + - /usr/local/bin/evaluate_supervised + --model tggn + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/tggn_eval_report.json + --symbol ES.FUT + - echo "TGGN validation + evaluation passed" + - cat ${CI_PROJECT_DIR}/output/tggn_eval_report.json + artifacts: + paths: + - output/tggn_eval_report.json + when: always + expire_in: 30 days + +train-validate-tlob: + extends: .train-validate-base + needs: [build-training] + script: + - /usr/local/bin/train_baseline_supervised + --model tlob + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --max-steps-per-epoch 50 + - echo "=== TLOB training complete, running evaluation ===" + - /usr/local/bin/evaluate_supervised + --model tlob + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/tlob_eval_report.json + --symbol ES.FUT + - echo "TLOB validation + evaluation passed" + - cat ${CI_PROJECT_DIR}/output/tlob_eval_report.json + artifacts: + paths: + - output/tlob_eval_report.json + when: always + expire_in: 30 days + +train-validate-kan: + extends: .train-validate-base + needs: [build-training] + script: + - /usr/local/bin/train_baseline_supervised + --model kan + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --max-steps-per-epoch 50 + - echo "=== KAN training complete, running evaluation ===" + - /usr/local/bin/evaluate_supervised + --model kan + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/kan_eval_report.json + --symbol ES.FUT + - echo "KAN validation + evaluation passed" + - cat ${CI_PROJECT_DIR}/output/kan_eval_report.json + artifacts: + paths: + - output/kan_eval_report.json + when: always + expire_in: 30 days + +train-validate-xlstm: + extends: .train-validate-base + needs: [build-training] + script: + - /usr/local/bin/train_baseline_supervised + --model xlstm + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --max-steps-per-epoch 50 + - echo "=== xLSTM training complete, running evaluation ===" + - /usr/local/bin/evaluate_supervised + --model xlstm + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/xlstm_eval_report.json + --symbol ES.FUT + - echo "xLSTM validation + evaluation passed" + - cat ${CI_PROJECT_DIR}/output/xlstm_eval_report.json + artifacts: + paths: + - output/xlstm_eval_report.json + when: always + expire_in: 30 days + +train-validate-diffusion: + extends: .train-validate-base + needs: [build-training] + script: + - /usr/local/bin/train_baseline_supervised + --model diffusion + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --max-steps-per-epoch 50 + - echo "=== Diffusion training complete, running evaluation ===" + - /usr/local/bin/evaluate_supervised + --model diffusion + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/diffusion_eval_report.json + --symbol ES.FUT + - echo "Diffusion validation + evaluation passed" + - cat ${CI_PROJECT_DIR}/output/diffusion_eval_report.json + artifacts: + paths: + - output/diffusion_eval_report.json + when: always + expire_in: 30 days + +# --- Hyperopt: all models --- + +hyperopt-dqn: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_rl + --model dqn + --trials 20 + --n-initial 5 + --epochs 15 + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/dqn_hyperopt_results.json + - echo "=== DQN Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/dqn_hyperopt_results.json + - echo "=== Training DQN with best params ===" + - /usr/local/bin/train_baseline_rl + --model dqn + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/dqn_hyperopt_results.json + - echo "=== Evaluating DQN with best params ===" + - /usr/local/bin/evaluate_baseline + --model dqn + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/dqn_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/dqn_eval_report.json + artifacts: + paths: + - output/dqn_hyperopt_results.json + - output/dqn_eval_report.json + when: always + expire_in: 90 days + +hyperopt-tft: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model tft + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/tft_hyperopt_results.json + - echo "=== TFT Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/tft_hyperopt_results.json + - echo "=== Training TFT with best params ===" + - /usr/local/bin/train_baseline_supervised + --model tft + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/tft_hyperopt_results.json + - echo "=== Evaluating TFT with best params ===" + - /usr/local/bin/evaluate_supervised + --model tft + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/tft_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/tft_eval_report.json + artifacts: + paths: + - output/tft_hyperopt_results.json + - output/tft_eval_report.json + when: always + expire_in: 90 days + +hyperopt-mamba2: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model mamba2 + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/mamba2_hyperopt_results.json + - echo "=== Mamba2 Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/mamba2_hyperopt_results.json + - echo "=== Training Mamba2 with best params ===" + - /usr/local/bin/train_baseline_supervised + --model mamba2 + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/mamba2_hyperopt_results.json + - echo "=== Evaluating Mamba2 with best params ===" + - /usr/local/bin/evaluate_supervised + --model mamba2 + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/mamba2_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/mamba2_eval_report.json + artifacts: + paths: + - output/mamba2_hyperopt_results.json + - output/mamba2_eval_report.json + when: always + expire_in: 90 days + +hyperopt-liquid: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model liquid + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/liquid_hyperopt_results.json + - echo "=== Liquid Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/liquid_hyperopt_results.json + - echo "=== Training Liquid with best params ===" + - /usr/local/bin/train_baseline_supervised + --model liquid + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/liquid_hyperopt_results.json + - echo "=== Evaluating Liquid with best params ===" + - /usr/local/bin/evaluate_supervised + --model liquid + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/liquid_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/liquid_eval_report.json + artifacts: + paths: + - output/liquid_hyperopt_results.json + - output/liquid_eval_report.json + when: always + expire_in: 90 days + +hyperopt-tggn: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model tggn + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/tggn_hyperopt_results.json + - echo "=== TGGN Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/tggn_hyperopt_results.json + - echo "=== Training TGGN with best params ===" + - /usr/local/bin/train_baseline_supervised + --model tggn + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/tggn_hyperopt_results.json + - echo "=== Evaluating TGGN with best params ===" + - /usr/local/bin/evaluate_supervised + --model tggn + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/tggn_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/tggn_eval_report.json + artifacts: + paths: + - output/tggn_hyperopt_results.json + - output/tggn_eval_report.json + when: always + expire_in: 90 days + +hyperopt-tlob: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model tlob + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/tlob_hyperopt_results.json + - echo "=== TLOB Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/tlob_hyperopt_results.json + - echo "=== Training TLOB with best params ===" + - /usr/local/bin/train_baseline_supervised + --model tlob + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/tlob_hyperopt_results.json + - echo "=== Evaluating TLOB with best params ===" + - /usr/local/bin/evaluate_supervised + --model tlob + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/tlob_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/tlob_eval_report.json + artifacts: + paths: + - output/tlob_hyperopt_results.json + - output/tlob_eval_report.json + when: always + expire_in: 90 days + +hyperopt-kan: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model kan + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/kan_hyperopt_results.json + - echo "=== KAN Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/kan_hyperopt_results.json + - echo "=== Training KAN with best params ===" + - /usr/local/bin/train_baseline_supervised + --model kan + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/kan_hyperopt_results.json + - echo "=== Evaluating KAN with best params ===" + - /usr/local/bin/evaluate_supervised + --model kan + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/kan_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/kan_eval_report.json + artifacts: + paths: + - output/kan_hyperopt_results.json + - output/kan_eval_report.json + when: always + expire_in: 90 days + +hyperopt-xlstm: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model xlstm + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/xlstm_hyperopt_results.json + - echo "=== xLSTM Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/xlstm_hyperopt_results.json + - echo "=== Training xLSTM with best params ===" + - /usr/local/bin/train_baseline_supervised + --model xlstm + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/xlstm_hyperopt_results.json + - echo "=== Evaluating xLSTM with best params ===" + - /usr/local/bin/evaluate_supervised + --model xlstm + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/xlstm_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/xlstm_eval_report.json + artifacts: + paths: + - output/xlstm_hyperopt_results.json + - output/xlstm_eval_report.json + when: always + expire_in: 90 days + +hyperopt-diffusion: + extends: .train-validate-base + needs: [build-training] + timeout: 4h + script: + - /usr/local/bin/hyperopt_baseline_supervised + --model diffusion + --trials 20 + --n-initial 5 + --epochs 15 + --data-dir /mnt/training-data/futures-baseline + --base-dir ${CI_PROJECT_DIR}/output/hyperopt + --output ${CI_PROJECT_DIR}/output/diffusion_hyperopt_results.json + - echo "=== Diffusion Hyperopt complete ===" + - cat ${CI_PROJECT_DIR}/output/diffusion_hyperopt_results.json + - echo "=== Training Diffusion with best params ===" + - /usr/local/bin/train_baseline_supervised + --model diffusion + --symbol ES.FUT + --data-dir /mnt/training-data/futures-baseline + --output-dir ${CI_PROJECT_DIR}/output + --hyperopt-params ${CI_PROJECT_DIR}/output/diffusion_hyperopt_results.json + - echo "=== Evaluating Diffusion with best params ===" + - /usr/local/bin/evaluate_supervised + --model diffusion + --models-dir ${CI_PROJECT_DIR}/output + --data-dir /mnt/training-data/futures-baseline + --output ${CI_PROJECT_DIR}/output/diffusion_eval_report.json + --symbol ES.FUT + - cat ${CI_PROJECT_DIR}/output/diffusion_eval_report.json + artifacts: + paths: + - output/diffusion_hyperopt_results.json + - output/diffusion_eval_report.json + when: always + expire_in: 90 days + # -------------------------------------------------------------------------- # IaC: Terragrunt plan on MR (runs on gitlab pool) # -------------------------------------------------------------------------- diff --git a/crates/ml/examples/hyperopt_baseline_supervised.rs b/crates/ml/examples/hyperopt_baseline_supervised.rs index a8e2aef6f..301f0a410 100644 --- a/crates/ml/examples/hyperopt_baseline_supervised.rs +++ b/crates/ml/examples/hyperopt_baseline_supervised.rs @@ -1,8 +1,9 @@ //! Hyperopt Runner for Supervised Models on DBN Data //! //! Runs hyperparameter optimization using Particle Swarm Optimization (PSO) for -//! TFT, MAMBA-2, or both models. Uses the same --data-dir / --symbol interface -//! as the RL hyperopt and training binaries. +//! any of the 8 supervised models: TFT, Mamba2, Liquid, TGGN, TLOB, KAN, xLSTM, +//! and Diffusion. Uses the same --data-dir / --symbol interface as the RL hyperopt +//! and training binaries. //! //! ## Usage //! @@ -13,8 +14,14 @@ //! # Run Mamba2 hyperopt //! hyperopt_baseline_supervised --model mamba2 --data-dir /data/NQ.FUT --trials 20 --epochs 10 //! -//! # Run both models +//! # Run a specific model +//! hyperopt_baseline_supervised --model liquid --data-dir /data/ES.FUT +//! +//! # Run TFT + Mamba2 (backward-compatible alias) //! hyperopt_baseline_supervised --model both --data-dir /data/ES.FUT +//! +//! # Run all 8 supervised models +//! hyperopt_baseline_supervised --model all --data-dir /data/ES.FUT //! ``` //! //! ## Output @@ -30,17 +37,24 @@ use std::path::PathBuf; use std::time::Instant; use tracing::{error, info, warn, Level}; +use ml::hyperopt::adapters::diffusion::DiffusionTrainer; +use ml::hyperopt::adapters::kan::KANTrainer; +use ml::hyperopt::adapters::liquid::LiquidTrainer; use ml::hyperopt::adapters::mamba2::Mamba2Trainer; use ml::hyperopt::adapters::tft::TFTTrainer; +use ml::hyperopt::adapters::tggn::TGGNTrainer; +use ml::hyperopt::adapters::tlob::TLOBTrainer; +use ml::hyperopt::adapters::xlstm::XLSTMTrainer; use ml::hyperopt::paths::{generate_run_id, TrainingPaths}; use ml::hyperopt::ArgminOptimizer; /// Hyperparameter optimization runner for supervised models #[derive(Parser, Debug)] #[command(name = "hyperopt-baseline-supervised")] -#[command(about = "Run hyperparameter optimization for supervised models (TFT, Mamba2)")] +#[command(about = "Run hyperparameter optimization for supervised models (tft, mamba2, liquid, tggn, tlob, kan, xlstm, diffusion, both, all)")] struct Args { - /// Model to optimize: "tft", "mamba2", or "both" + /// Model to optimize: "tft", "mamba2", "liquid", "tggn", "tlob", "kan", + /// "xlstm", "diffusion", "both" (tft+mamba2), or "all" (all 8 models) #[arg(long, default_value = "both")] model: String, @@ -184,6 +198,292 @@ fn run_mamba2_hyperopt(args: &Args) -> Result { )) } +fn run_liquid_hyperopt(args: &Args) -> Result { + info!("========================================"); + info!(" Liquid Hyperparameter Optimization"); + info!("========================================"); + + let run_id = generate_run_id("hyperopt-liquid"); + let training_paths = TrainingPaths::new(&args.base_dir, "liquid", &run_id); + + info!("Run ID: {}", run_id); + info!("Data dir: {}", args.data_dir.display()); + info!("Epochs per trial: {}", args.epochs); + info!("Trials: {}", args.trials); + + let trainer = LiquidTrainer::new(&args.data_dir, args.epochs) + .context("Failed to create Liquid trainer")? + .with_early_stopping(args.early_stopping_patience) + .with_training_paths(training_paths); + + let optimizer = ArgminOptimizer::builder() + .max_trials(args.trials) + .n_initial(args.n_initial) + .seed(args.seed) + .build(); + + let start = Instant::now(); + let result = optimizer + .optimize(trainer) + .context("Liquid hyperopt optimization failed")?; + let elapsed = start.elapsed().as_secs_f64(); + + info!("Liquid hyperopt complete:"); + info!(" Best objective: {:.6}", result.best_objective); + info!(" Total trials: {}", result.all_trials.len()); + info!(" Elapsed: {:.1}s", elapsed); + + let best_params_json = serde_json::to_value(&result.best_params) + .ok() + .unwrap_or(Value::Null); + + Ok(build_model_result( + result.best_objective, + best_params_json, + result.all_trials.len(), + elapsed, + )) +} + +fn run_tggn_hyperopt(args: &Args) -> Result { + info!("========================================"); + info!(" TGGN Hyperparameter Optimization"); + info!("========================================"); + + let run_id = generate_run_id("hyperopt-tggn"); + let training_paths = TrainingPaths::new(&args.base_dir, "tggn", &run_id); + + info!("Run ID: {}", run_id); + info!("Data dir: {}", args.data_dir.display()); + info!("Epochs per trial: {}", args.epochs); + info!("Trials: {}", args.trials); + + let trainer = TGGNTrainer::new(&args.data_dir, args.epochs) + .context("Failed to create TGGN trainer")? + .with_early_stopping(args.early_stopping_patience) + .with_training_paths(training_paths); + + let optimizer = ArgminOptimizer::builder() + .max_trials(args.trials) + .n_initial(args.n_initial) + .seed(args.seed) + .build(); + + let start = Instant::now(); + let result = optimizer + .optimize(trainer) + .context("TGGN hyperopt optimization failed")?; + let elapsed = start.elapsed().as_secs_f64(); + + info!("TGGN hyperopt complete:"); + info!(" Best objective: {:.6}", result.best_objective); + info!(" Total trials: {}", result.all_trials.len()); + info!(" Elapsed: {:.1}s", elapsed); + + let best_params_json = serde_json::to_value(&result.best_params) + .ok() + .unwrap_or(Value::Null); + + Ok(build_model_result( + result.best_objective, + best_params_json, + result.all_trials.len(), + elapsed, + )) +} + +fn run_tlob_hyperopt(args: &Args) -> Result { + info!("========================================"); + info!(" TLOB Hyperparameter Optimization"); + info!("========================================"); + + let run_id = generate_run_id("hyperopt-tlob"); + let training_paths = TrainingPaths::new(&args.base_dir, "tlob", &run_id); + + info!("Run ID: {}", run_id); + info!("Data dir: {}", args.data_dir.display()); + info!("Epochs per trial: {}", args.epochs); + info!("Trials: {}", args.trials); + + let trainer = TLOBTrainer::new(&args.data_dir, args.epochs) + .context("Failed to create TLOB trainer")? + .with_early_stopping(args.early_stopping_patience) + .with_training_paths(training_paths); + + let optimizer = ArgminOptimizer::builder() + .max_trials(args.trials) + .n_initial(args.n_initial) + .seed(args.seed) + .build(); + + let start = Instant::now(); + let result = optimizer + .optimize(trainer) + .context("TLOB hyperopt optimization failed")?; + let elapsed = start.elapsed().as_secs_f64(); + + info!("TLOB hyperopt complete:"); + info!(" Best objective: {:.6}", result.best_objective); + info!(" Total trials: {}", result.all_trials.len()); + info!(" Elapsed: {:.1}s", elapsed); + + let best_params_json = serde_json::to_value(&result.best_params) + .ok() + .unwrap_or(Value::Null); + + Ok(build_model_result( + result.best_objective, + best_params_json, + result.all_trials.len(), + elapsed, + )) +} + +fn run_kan_hyperopt(args: &Args) -> Result { + info!("========================================"); + info!(" KAN Hyperparameter Optimization"); + info!("========================================"); + + let run_id = generate_run_id("hyperopt-kan"); + let training_paths = TrainingPaths::new(&args.base_dir, "kan", &run_id); + + info!("Run ID: {}", run_id); + info!("Data dir: {}", args.data_dir.display()); + info!("Epochs per trial: {}", args.epochs); + info!("Trials: {}", args.trials); + + let trainer = KANTrainer::new(&args.data_dir, args.epochs) + .context("Failed to create KAN trainer")? + .with_early_stopping(args.early_stopping_patience) + .with_training_paths(training_paths); + + let optimizer = ArgminOptimizer::builder() + .max_trials(args.trials) + .n_initial(args.n_initial) + .seed(args.seed) + .build(); + + let start = Instant::now(); + let result = optimizer + .optimize(trainer) + .context("KAN hyperopt optimization failed")?; + let elapsed = start.elapsed().as_secs_f64(); + + info!("KAN hyperopt complete:"); + info!(" Best objective: {:.6}", result.best_objective); + info!(" Total trials: {}", result.all_trials.len()); + info!(" Elapsed: {:.1}s", elapsed); + + let best_params_json = serde_json::to_value(&result.best_params) + .ok() + .unwrap_or(Value::Null); + + Ok(build_model_result( + result.best_objective, + best_params_json, + result.all_trials.len(), + elapsed, + )) +} + +fn run_xlstm_hyperopt(args: &Args) -> Result { + info!("========================================"); + info!(" xLSTM Hyperparameter Optimization"); + info!("========================================"); + + let run_id = generate_run_id("hyperopt-xlstm"); + let training_paths = TrainingPaths::new(&args.base_dir, "xlstm", &run_id); + + info!("Run ID: {}", run_id); + info!("Data dir: {}", args.data_dir.display()); + info!("Epochs per trial: {}", args.epochs); + info!("Trials: {}", args.trials); + + let trainer = XLSTMTrainer::new(&args.data_dir, args.epochs) + .context("Failed to create xLSTM trainer")? + .with_early_stopping(args.early_stopping_patience) + .with_training_paths(training_paths); + + let optimizer = ArgminOptimizer::builder() + .max_trials(args.trials) + .n_initial(args.n_initial) + .seed(args.seed) + .build(); + + let start = Instant::now(); + let result = optimizer + .optimize(trainer) + .context("xLSTM hyperopt optimization failed")?; + let elapsed = start.elapsed().as_secs_f64(); + + info!("xLSTM hyperopt complete:"); + info!(" Best objective: {:.6}", result.best_objective); + info!(" Total trials: {}", result.all_trials.len()); + info!(" Elapsed: {:.1}s", elapsed); + + let best_params_json = serde_json::to_value(&result.best_params) + .ok() + .unwrap_or(Value::Null); + + Ok(build_model_result( + result.best_objective, + best_params_json, + result.all_trials.len(), + elapsed, + )) +} + +fn run_diffusion_hyperopt(args: &Args) -> Result { + info!("========================================"); + info!(" Diffusion Hyperparameter Optimization"); + info!("========================================"); + + let run_id = generate_run_id("hyperopt-diffusion"); + let training_paths = TrainingPaths::new(&args.base_dir, "diffusion", &run_id); + + info!("Run ID: {}", run_id); + info!("Data dir: {}", args.data_dir.display()); + info!("Epochs per trial: {}", args.epochs); + info!("Trials: {}", args.trials); + + let trainer = DiffusionTrainer::new(&args.data_dir, args.epochs) + .context("Failed to create Diffusion trainer")? + .with_early_stopping(args.early_stopping_patience) + .with_training_paths(training_paths); + + let optimizer = ArgminOptimizer::builder() + .max_trials(args.trials) + .n_initial(args.n_initial) + .seed(args.seed) + .build(); + + let start = Instant::now(); + let result = optimizer + .optimize(trainer) + .context("Diffusion hyperopt optimization failed")?; + let elapsed = start.elapsed().as_secs_f64(); + + info!("Diffusion hyperopt complete:"); + info!(" Best objective: {:.6}", result.best_objective); + info!(" Total trials: {}", result.all_trials.len()); + info!(" Elapsed: {:.1}s", elapsed); + + let best_params_json = serde_json::to_value(&result.best_params) + .ok() + .unwrap_or(Value::Null); + + Ok(build_model_result( + result.best_objective, + best_params_json, + result.all_trials.len(), + elapsed, + )) +} + +const VALID_MODELS: &[&str] = &[ + "tft", "mamba2", "liquid", "tggn", "tlob", "kan", "xlstm", "diffusion", +]; + fn main() -> Result<()> { tracing_subscriber::fmt() .with_max_level(Level::INFO) @@ -210,15 +510,22 @@ fn main() -> Result<()> { ); } - let run_tft = args.model == "tft" || args.model == "both"; - let run_mamba2 = args.model == "mamba2" || args.model == "both"; - - if !run_tft && !run_mamba2 { + // Resolve model selection: "both" = tft+mamba2 (backward compat), "all" = all 8 models + let models: Vec<&str> = if args.model == "both" { + vec!["tft", "mamba2"] + } else if args.model == "all" { + VALID_MODELS.to_vec() + } else if VALID_MODELS.contains(&args.model.as_str()) { + vec![args.model.as_str()] + } else { anyhow::bail!( - "Invalid --model value '{}'. Must be 'tft', 'mamba2', or 'both'.", - args.model + "Invalid --model value '{}'. Must be one of: {}, 'both', or 'all'.", + args.model, + VALID_MODELS.join(", ") ); - } + }; + + info!("Models to optimize: {:?}", models); if args.trials <= args.n_initial { anyhow::bail!( @@ -235,32 +542,31 @@ fn main() -> Result<()> { let mut results = serde_json::Map::new(); - if run_tft { - match run_tft_hyperopt(&args) { - Ok(tft_result) => { - results.insert("tft".to_owned(), tft_result); + for model_name in &models { + let result = match *model_name { + "tft" => run_tft_hyperopt(&args), + "mamba2" => run_mamba2_hyperopt(&args), + "liquid" => run_liquid_hyperopt(&args), + "tggn" => run_tggn_hyperopt(&args), + "tlob" => run_tlob_hyperopt(&args), + "kan" => run_kan_hyperopt(&args), + "xlstm" => run_xlstm_hyperopt(&args), + "diffusion" => run_diffusion_hyperopt(&args), + other => { + error!("Unknown model: {}", other); + continue; + } + }; + + match result { + Ok(model_result) => { + results.insert(model_name.to_string(), model_result); } Err(e) => { - error!("TFT hyperopt failed: {:#}", e); + error!("{} hyperopt failed: {:#}", model_name, e); warn!("Continuing with remaining models..."); results.insert( - "tft".to_owned(), - serde_json::json!({ "error": format!("{:#}", e) }), - ); - } - } - } - - if run_mamba2 { - match run_mamba2_hyperopt(&args) { - Ok(mamba2_result) => { - results.insert("mamba2".to_owned(), mamba2_result); - } - Err(e) => { - error!("Mamba2 hyperopt failed: {:#}", e); - warn!("Continuing..."); - results.insert( - "mamba2".to_owned(), + model_name.to_string(), serde_json::json!({ "error": format!("{:#}", e) }), ); } diff --git a/crates/ml/src/hyperopt/adapters/diffusion.rs b/crates/ml/src/hyperopt/adapters/diffusion.rs index c8314f3b6..ed4a55e5c 100644 --- a/crates/ml/src/hyperopt/adapters/diffusion.rs +++ b/crates/ml/src/hyperopt/adapters/diffusion.rs @@ -1,13 +1,23 @@ //! Hyperopt adapter for the Diffusion model. //! -//! Defines `DiffusionParams` (ParameterSpace) for hyperparameter optimization -//! and `DiffusionMetrics` for tracking training results. +//! Defines `DiffusionParams` (ParameterSpace), `DiffusionMetrics`, +//! and `DiffusionTrainer` (HyperparameterOptimizable) for hyperparameter +//! optimization of the Diffusion model via the unified framework. +use candle_core::Device; +use std::path::PathBuf; +use tracing::{info, warn}; + +use crate::diffusion::config::{DiffusionConfig, NoiseSchedule}; +use crate::diffusion::trainable::DiffusionTrainableAdapter; +use crate::features::extract_ml_features; +use crate::hyperopt::paths::TrainingPaths; +use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace}; +use crate::training::unified_trainer::UnifiedTrainable; use crate::MLError; -use crate::hyperopt::traits::ParameterSpace; /// Hyperparameters for Diffusion model hyperopt tuning. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct DiffusionParams { /// Learning rate (log scale). pub learning_rate: f64, @@ -103,7 +113,7 @@ impl ParameterSpace for DiffusionParams { } /// Metrics returned from Diffusion hyperopt training. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct DiffusionMetrics { /// Validation noise prediction loss. pub val_loss: f64, @@ -116,13 +126,251 @@ pub struct DiffusionMetrics { impl Default for DiffusionMetrics { fn default() -> Self { Self { - val_loss: f64::NAN, - train_loss: f64::NAN, + val_loss: 0.0, + train_loss: 0.0, epochs_completed: 0, } } } +/// Diffusion model trainer for hyperparameter optimization. +/// +/// Wraps the `DiffusionTrainableAdapter` (which implements `UnifiedTrainable`) +/// and drives it through the `HyperparameterOptimizable` interface. +/// +/// The Diffusion model generates price paths via denoising diffusion. +/// Training uses sequence input: `(1, seq_len, feature_dim)`. +#[derive(Debug)] +pub struct DiffusionTrainer { + data_dir: PathBuf, + epochs: usize, + device: Device, + training_paths: TrainingPaths, + early_stopping_patience: usize, + trial_counter: usize, +} + +impl DiffusionTrainer { + /// Create a new Diffusion trainer. + /// + /// # Arguments + /// * `data_dir` - Directory containing .dbn/.dbn.zst files + /// * `epochs` - Number of training epochs per trial + pub fn new>(data_dir: P, epochs: usize) -> Result { + let data_dir = data_dir.into(); + if !data_dir.exists() { + return Err(MLError::ConfigError { + reason: format!("Data directory not found: {}", data_dir.display()), + }); + } + + let device = Device::new_cuda(0).unwrap_or_else(|e| { + warn!("CUDA unavailable ({}), falling back to CPU", e); + Device::Cpu + }); + + info!( + "Diffusion Trainer initialized: Device={:?}, Data={}, Epochs={}", + device, data_dir.display(), epochs + ); + + Ok(Self { + data_dir, + epochs, + device, + training_paths: TrainingPaths::new("/tmp/ml_training", "diffusion", "default"), + early_stopping_patience: 10, + trial_counter: 0, + }) + } + + /// Set training paths configuration. + pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self { + self.training_paths = paths; + self + } + + /// Configure early stopping patience. + pub fn with_early_stopping(mut self, patience: usize) -> Self { + self.early_stopping_patience = patience; + self + } +} + +impl HyperparameterOptimizable for DiffusionTrainer { + type Params = DiffusionParams; + type Metrics = DiffusionMetrics; + + fn train_with_params(&mut self, params: Self::Params) -> Result { + let trial_start = std::time::Instant::now(); + let current_trial = self.trial_counter; + self.trial_counter += 1; + + info!( + "Training Diffusion: lr={:.6}, timesteps={}, sampling={}, hidden={}, layers={}, batch={}", + params.learning_rate, params.num_timesteps, params.sampling_steps, + params.hidden_dim, params.num_layers, params.batch_size + ); + + self.training_paths.create_all().map_err(|e| { + MLError::ModelError(format!("Failed to create training directories: {}", e)) + })?; + + let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir) + .map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?; + + let features = extract_ml_features(&bars) + .map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?; + + if features.len() < 2 { + return Err(MLError::ModelError("Insufficient features extracted".to_owned())); + } + + let feature_dim = 51; + + // Diffusion uses flat input: data_dim = seq_len * feature_dim + // For hyperopt, we use feature_dim=1 and seq_len=feature_dim (51) + // so data_dim = 51, matching our feature vectors. + let pairs = crate::hyperopt::shared_data::build_flat_pairs( + &features, &bars, feature_dim, &self.device, + )?; + + let split = (pairs.len() as f64 * 0.8) as usize; + let train_data = &pairs[..split]; + let val_data = &pairs[split..]; + + if val_data.len() < 2 { + return Err(MLError::ModelError("Validation set too small".to_owned())); + } + + let config = DiffusionConfig { + num_timesteps: params.num_timesteps, + sampling_steps: params.sampling_steps, + seq_len: feature_dim, // Treat feature vector as a sequence + feature_dim: 1, // Univariate per position + hidden_dim: params.hidden_dim, + num_layers: params.num_layers, + time_embed_dim: params.time_embed_dim, + schedule: NoiseSchedule::Cosine, + learning_rate: params.learning_rate, + weight_decay: params.weight_decay, + grad_clip: params.grad_clip, + }; + + let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result { + let mut model = DiffusionTrainableAdapter::new(config, self.device.clone())?; + + let mut best_val_loss = f64::MAX; + let mut patience_counter = 0_usize; + let mut last_train_loss = 0.0_f64; + + for epoch in 0..self.epochs { + let mut epoch_loss = 0.0_f64; + let mut batch_count = 0_usize; + + for (input, target) in train_data { + let pred = model.forward(input)?; + let loss = model.compute_loss(&pred, target)?; + let loss_val = loss + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Loss scalar: {}", e)))? + as f64; + + if loss_val.is_nan() || loss_val.is_infinite() { + return Ok(DiffusionMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + epochs_completed: epoch, + }); + } + + model.backward(&loss)?; + model.optimizer_step()?; + epoch_loss += loss_val; + batch_count += 1; + } + + last_train_loss = if batch_count > 0 { + epoch_loss / batch_count as f64 + } else { + 0.0 + }; + + let val_loss = model.validate(val_data)?; + + if val_loss < best_val_loss { + best_val_loss = val_loss; + patience_counter = 0; + } else { + patience_counter += 1; + } + + if epoch % 10 == 0 { + info!( + " Epoch {}/{}: train_loss={:.6}, val_loss={:.6}", + epoch + 1, self.epochs, last_train_loss, val_loss + ); + } + + if patience_counter >= self.early_stopping_patience { + info!("Early stopping at epoch {}", epoch + 1); + return Ok(DiffusionMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + epochs_completed: epoch + 1, + }); + } + } + + Ok(DiffusionMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + epochs_completed: self.epochs, + }) + })); + + let metrics = match training_result { + Ok(Ok(m)) => m, + Ok(Err(e)) => { + warn!("Diffusion training failed: {}", e); + DiffusionMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0 } + } + Err(_) => { + warn!("Diffusion training panicked (likely OOM)"); + DiffusionMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0 } + } + }; + + info!( + "Diffusion training completed: val_loss={:.6}, train_loss={:.6}, epochs={}", + metrics.val_loss, metrics.train_loss, metrics.epochs_completed + ); + + let duration_secs = trial_start.elapsed().as_secs_f64(); + 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(); + crate::hyperopt::shared_data::write_trial_result_json( + &self.training_paths.hyperopt_dir(), + &trial_result, + ).ok(); + + if self.device.is_cuda() { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + + Ok(metrics) + } + + fn extract_objective(metrics: &Self::Metrics) -> f64 { + metrics.val_loss + } +} + #[cfg(test)] mod tests { use super::*; @@ -177,7 +425,8 @@ mod tests { #[test] fn test_metrics_default() { let metrics = DiffusionMetrics::default(); - assert!(metrics.val_loss.is_nan()); + assert_eq!(metrics.val_loss, 0.0); + assert_eq!(metrics.train_loss, 0.0); assert_eq!(metrics.epochs_completed, 0); } diff --git a/crates/ml/src/hyperopt/adapters/kan.rs b/crates/ml/src/hyperopt/adapters/kan.rs index e0a36b5a0..edb07efe3 100644 --- a/crates/ml/src/hyperopt/adapters/kan.rs +++ b/crates/ml/src/hyperopt/adapters/kan.rs @@ -1,12 +1,21 @@ //! KAN (Kolmogorov-Arnold Network) Hyperparameter Optimization Adapter //! -//! Provides `KANParams` (implementing `ParameterSpace`) and `KANMetrics` -//! for hyperparameter optimization of the KAN model via the unified -//! optimization framework. +//! Provides `KANParams` (implementing `ParameterSpace`), `KANMetrics`, +//! and `KANTrainer` (implementing `HyperparameterOptimizable`) for +//! hyperparameter optimization of the KAN model via the unified framework. -use crate::hyperopt::traits::ParameterSpace; -use crate::MLError; +use candle_core::Device; use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use tracing::{info, warn}; + +use crate::features::extract_ml_features; +use crate::hyperopt::paths::TrainingPaths; +use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace}; +use crate::kan::config::KANConfig; +use crate::kan::trainable::KANTrainableAdapter; +use crate::training::unified_trainer::UnifiedTrainable; +use crate::MLError; /// Hyperparameters for KAN hyperopt tuning. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -117,6 +126,239 @@ pub struct KANMetrics { pub epochs_completed: usize, } +/// KAN trainer for hyperparameter optimization. +/// +/// Wraps the `KANTrainableAdapter` (which implements `UnifiedTrainable`) +/// and drives it through the `HyperparameterOptimizable` interface. +#[derive(Debug)] +pub struct KANTrainer { + data_dir: PathBuf, + epochs: usize, + device: Device, + training_paths: TrainingPaths, + early_stopping_patience: usize, + trial_counter: usize, +} + +impl KANTrainer { + /// Create a new KAN trainer. + /// + /// # Arguments + /// * `data_dir` - Directory containing .dbn/.dbn.zst files + /// * `epochs` - Number of training epochs per trial + pub fn new>(data_dir: P, epochs: usize) -> Result { + let data_dir = data_dir.into(); + if !data_dir.exists() { + return Err(MLError::ConfigError { + reason: format!("Data directory not found: {}", data_dir.display()), + }); + } + + let device = Device::new_cuda(0).unwrap_or_else(|e| { + warn!("CUDA unavailable ({}), falling back to CPU", e); + Device::Cpu + }); + + info!( + "KAN Trainer initialized: Device={:?}, Data={}, Epochs={}", + device, data_dir.display(), epochs + ); + + Ok(Self { + data_dir, + epochs, + device, + training_paths: TrainingPaths::new("/tmp/ml_training", "kan", "default"), + early_stopping_patience: 10, + trial_counter: 0, + }) + } + + /// Set training paths configuration. + pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self { + self.training_paths = paths; + self + } + + /// Configure early stopping patience. + pub fn with_early_stopping(mut self, patience: usize) -> Self { + self.early_stopping_patience = patience; + self + } +} + +impl HyperparameterOptimizable for KANTrainer { + type Params = KANParams; + type Metrics = KANMetrics; + + fn train_with_params(&mut self, params: Self::Params) -> Result { + let trial_start = std::time::Instant::now(); + let current_trial = self.trial_counter; + self.trial_counter += 1; + + info!( + "Training KAN: lr={:.6}, grid={}, spline_order={}, hidden_width={}, layers={}, batch={}", + params.learning_rate, params.grid_size, params.spline_order, + params.hidden_width, params.num_layers, params.batch_size + ); + + self.training_paths.create_all().map_err(|e| { + MLError::ModelError(format!("Failed to create training directories: {}", e)) + })?; + + let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir) + .map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?; + + let features = extract_ml_features(&bars) + .map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?; + + if features.len() < 2 { + return Err(MLError::ModelError("Insufficient features extracted".to_owned())); + } + + let feature_dim = 51; + let pairs = crate::hyperopt::shared_data::build_flat_pairs( + &features, &bars, feature_dim, &self.device, + )?; + + let split = (pairs.len() as f64 * 0.8) as usize; + let train_data = &pairs[..split]; + let val_data = &pairs[split..]; + + if val_data.len() < 2 { + return Err(MLError::ModelError("Validation set too small".to_owned())); + } + + // Build layer_widths: [input, hidden, ..., hidden, output] + let mut layer_widths = vec![feature_dim]; + for _ in 0..params.num_layers { + layer_widths.push(params.hidden_width); + } + layer_widths.push(1); // Single output + + let config = KANConfig { + grid_size: params.grid_size, + spline_order: params.spline_order, + layer_widths, + learning_rate: params.learning_rate, + weight_decay: params.weight_decay, + grad_clip: params.grad_clip, + }; + + let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result { + let mut model = KANTrainableAdapter::new(config, &self.device)?; + + let mut best_val_loss = f64::MAX; + let mut patience_counter = 0_usize; + let mut last_train_loss = 0.0_f64; + + for epoch in 0..self.epochs { + let mut epoch_loss = 0.0_f64; + let mut batch_count = 0_usize; + + for (input, target) in train_data { + let pred = model.forward(input)?; + let loss = model.compute_loss(&pred, target)?; + let loss_val = loss + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Loss scalar: {}", e)))? + as f64; + + if loss_val.is_nan() || loss_val.is_infinite() { + return Ok(KANMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + epochs_completed: epoch, + }); + } + + model.backward(&loss)?; + model.optimizer_step()?; + epoch_loss += loss_val; + batch_count += 1; + } + + last_train_loss = if batch_count > 0 { + epoch_loss / batch_count as f64 + } else { + 0.0 + }; + + let val_loss = model.validate(val_data)?; + + if val_loss < best_val_loss { + best_val_loss = val_loss; + patience_counter = 0; + } else { + patience_counter += 1; + } + + if epoch % 10 == 0 { + info!( + " Epoch {}/{}: train_loss={:.6}, val_loss={:.6}", + epoch + 1, self.epochs, last_train_loss, val_loss + ); + } + + if patience_counter >= self.early_stopping_patience { + info!("Early stopping at epoch {}", epoch + 1); + return Ok(KANMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + epochs_completed: epoch + 1, + }); + } + } + + Ok(KANMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + epochs_completed: self.epochs, + }) + })); + + let metrics = match training_result { + Ok(Ok(m)) => m, + Ok(Err(e)) => { + warn!("KAN training failed: {}", e); + KANMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0 } + } + Err(_) => { + warn!("KAN training panicked (likely OOM)"); + KANMetrics { val_loss: 1000.0, train_loss: 1000.0, epochs_completed: 0 } + } + }; + + info!( + "KAN training completed: val_loss={:.6}, train_loss={:.6}, epochs={}", + metrics.val_loss, metrics.train_loss, metrics.epochs_completed + ); + + let duration_secs = trial_start.elapsed().as_secs_f64(); + 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(); + crate::hyperopt::shared_data::write_trial_result_json( + &self.training_paths.hyperopt_dir(), + &trial_result, + ).ok(); + + if self.device.is_cuda() { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + + Ok(metrics) + } + + fn extract_objective(metrics: &Self::Metrics) -> f64 { + metrics.val_loss + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ml/src/hyperopt/adapters/liquid.rs b/crates/ml/src/hyperopt/adapters/liquid.rs index 910a8b5bd..8157fef37 100644 --- a/crates/ml/src/hyperopt/adapters/liquid.rs +++ b/crates/ml/src/hyperopt/adapters/liquid.rs @@ -15,9 +15,18 @@ //! | 5 | dropout | linear | [0.0, 0.5] | //! | 6 | gradient_clip | log | [0.1, 10.0] | +use candle_core::Device; use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use tracing::{info, warn}; -use crate::hyperopt::traits::ParameterSpace; +use crate::features::extract_ml_features; +use crate::gpu::DeviceConfig; +use crate::hyperopt::paths::TrainingPaths; +use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace}; +use crate::liquid::candle_cfc::CfCTrainConfig; +use crate::liquid::adapter::LiquidTrainableAdapter; +use crate::training::unified_trainer::UnifiedTrainable; use crate::MLError; /// Number of continuous dimensions in the Liquid CfC v2 parameter space. @@ -161,6 +170,275 @@ impl ParameterSpace for LiquidParams { } } +/// Metrics returned from Liquid CfC hyperopt training. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct LiquidMetrics { + /// Validation loss (primary objective for optimization) + pub val_loss: f64, + /// Training loss at end of run + pub train_loss: f64, + /// Number of epochs completed + pub epochs_completed: usize, +} + +/// Liquid CfC v2 trainer for hyperparameter optimization. +/// +/// Wraps the `LiquidTrainableAdapter` (which implements `UnifiedTrainable`) +/// and drives it through the `HyperparameterOptimizable` interface. +#[derive(Debug)] +pub struct LiquidTrainer { + data_dir: PathBuf, + epochs: usize, + device: Device, + training_paths: TrainingPaths, + early_stopping_patience: usize, + trial_counter: usize, +} + +impl LiquidTrainer { + /// Create a new Liquid CfC trainer. + /// + /// # Arguments + /// * `data_dir` - Directory containing .dbn/.dbn.zst files + /// * `epochs` - Number of training epochs per trial + pub fn new>(data_dir: P, epochs: usize) -> Result { + let data_dir = data_dir.into(); + if !data_dir.exists() { + return Err(MLError::ConfigError { + reason: format!("Data directory not found: {}", data_dir.display()), + }); + } + + let device = Device::new_cuda(0).unwrap_or_else(|e| { + warn!("CUDA unavailable ({}), falling back to CPU", e); + Device::Cpu + }); + + info!( + "Liquid Trainer initialized: Device={:?}, Data={}, Epochs={}", + device, + data_dir.display(), + epochs + ); + + Ok(Self { + data_dir, + epochs, + device, + training_paths: TrainingPaths::new("/tmp/ml_training", "liquid", "default"), + early_stopping_patience: 10, + trial_counter: 0, + }) + } + + /// Set training paths configuration. + pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self { + self.training_paths = paths; + self + } + + /// Configure early stopping patience. + pub fn with_early_stopping(mut self, patience: usize) -> Self { + self.early_stopping_patience = patience; + self + } +} + +impl HyperparameterOptimizable for LiquidTrainer { + type Params = LiquidParams; + type Metrics = LiquidMetrics; + + fn train_with_params(&mut self, params: Self::Params) -> Result { + let trial_start = std::time::Instant::now(); + let current_trial = self.trial_counter; + self.trial_counter += 1; + + info!( + "Training Liquid CfC: lr={:.6}, hidden={}, depth={}, tau=[{:.3},{:.3}], dropout={:.3}", + params.learning_rate, params.hidden_size, params.backbone_depth, + params.tau_min, params.tau_max, params.dropout + ); + + self.training_paths.create_all().map_err(|e| { + MLError::ModelError(format!("Failed to create training directories: {}", e)) + })?; + + // Load bars from DBN files + let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir) + .map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?; + + // Extract 51-dim features + let features = extract_ml_features(&bars) + .map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?; + + if features.len() < 2 { + return Err(MLError::ModelError("Insufficient features extracted".to_owned())); + } + + // Build (input, target) pairs: input = feature[i], target = close price change + let feature_dim = 51; + let pairs = crate::hyperopt::shared_data::build_flat_pairs(&features, &bars, feature_dim, &self.device)?; + + // Split train/val 80/20 + let split = (pairs.len() as f64 * 0.8) as usize; + let train_data = &pairs[..split]; + let val_data = &pairs[split..]; + + if val_data.len() < 2 { + return Err(MLError::ModelError("Validation set too small".to_owned())); + } + + // Build backbone hidden sizes from depth param + let backbone = vec![params.hidden_size; params.backbone_depth]; + + let config = CfCTrainConfig { + input_size: feature_dim, + hidden_size: params.hidden_size, + output_size: 1, + backbone_hidden_sizes: backbone, + tau_min: params.tau_min, + tau_max: params.tau_max, + learning_rate: params.learning_rate, + batch_size: 32, + seq_len: 1, + max_epochs: self.epochs, + early_stopping_patience: self.early_stopping_patience, + dropout: params.dropout, + gradient_clip: params.gradient_clip, + market_regime_adaptation: false, + device: if self.device.is_cuda() { + DeviceConfig::Cuda(0) + } else { + DeviceConfig::Cpu + }, + }; + + let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result { + let mut model = LiquidTrainableAdapter::new(config)?; + + let mut best_val_loss = f64::MAX; + let mut patience_counter = 0_usize; + let mut last_train_loss = 0.0_f64; + + for epoch in 0..self.epochs { + let mut epoch_loss = 0.0_f64; + let mut batch_count = 0_usize; + + for (input, target) in train_data { + let pred = model.forward(input)?; + let loss = model.compute_loss(&pred, target)?; + let loss_val = loss + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Loss scalar: {}", e)))? + as f64; + + if loss_val.is_nan() || loss_val.is_infinite() { + return Ok(LiquidMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + epochs_completed: epoch, + }); + } + + model.backward(&loss)?; + model.optimizer_step()?; + epoch_loss += loss_val; + batch_count += 1; + } + + last_train_loss = if batch_count > 0 { + epoch_loss / batch_count as f64 + } else { + 0.0 + }; + + // Validate + let val_loss = model.validate(val_data)?; + + if val_loss < best_val_loss { + best_val_loss = val_loss; + patience_counter = 0; + } else { + patience_counter += 1; + } + + if epoch % 10 == 0 { + info!( + " Epoch {}/{}: train_loss={:.6}, val_loss={:.6}", + epoch + 1, self.epochs, last_train_loss, val_loss + ); + } + + if patience_counter >= self.early_stopping_patience { + info!("Early stopping at epoch {}", epoch + 1); + return Ok(LiquidMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + epochs_completed: epoch + 1, + }); + } + } + + Ok(LiquidMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + epochs_completed: self.epochs, + }) + })); + + let metrics = match training_result { + Ok(Ok(m)) => m, + Ok(Err(e)) => { + warn!("Liquid training failed: {}", e); + LiquidMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + epochs_completed: 0, + } + } + Err(_) => { + warn!("Liquid training panicked (likely OOM)"); + LiquidMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + epochs_completed: 0, + } + } + }; + + info!( + "Liquid training completed: val_loss={:.6}, train_loss={:.6}, epochs={}", + metrics.val_loss, metrics.train_loss, metrics.epochs_completed + ); + + // Write trial result + let duration_secs = trial_start.elapsed().as_secs_f64(); + 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(); + crate::hyperopt::shared_data::write_trial_result_json( + &self.training_paths.hyperopt_dir(), + &trial_result, + ) + .ok(); + + // Cleanup + if self.device.is_cuda() { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + + Ok(metrics) + } + + fn extract_objective(metrics: &Self::Metrics) -> f64 { + metrics.val_loss + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ml/src/hyperopt/adapters/mod.rs b/crates/ml/src/hyperopt/adapters/mod.rs index 48e1e9497..ae48566ba 100644 --- a/crates/ml/src/hyperopt/adapters/mod.rs +++ b/crates/ml/src/hyperopt/adapters/mod.rs @@ -68,15 +68,15 @@ pub mod xlstm; // Re-export adapters for convenience pub use async_data_loader::AsyncDataLoader; -pub use kan::{KANMetrics, KANParams}; -pub use liquid::LiquidParams; +pub use kan::{KANMetrics, KANParams, KANTrainer}; +pub use liquid::{LiquidMetrics, LiquidParams, LiquidTrainer}; pub use continuous_ppo::{ContinuousPPOMetrics, ContinuousPPOParams, ContinuousPPOTrainer}; pub use dqn::{DQNMetrics, DQNParams, DQNTrainer}; pub use ensemble::{EnsembleParameterSpace, EnsembleSpaceConfig}; pub use mamba2::{Mamba2Metrics, Mamba2Params, Mamba2Trainer}; pub use ppo::{PPOMetrics, PPOParams, PPOTrainer}; pub use tft::{TFTMetrics, TFTParams, TFTTrainer as TFTHyperoptTrainer}; -pub use tggn::{TGGNMetrics, TGGNParams}; -pub use tlob::{TLOBMetrics, TLOBParams}; -pub use diffusion::{DiffusionMetrics, DiffusionParams}; -pub use xlstm::{XLSTMMetrics, XLSTMParams}; +pub use tggn::{TGGNMetrics, TGGNParams, TGGNTrainer}; +pub use tlob::{TLOBMetrics, TLOBParams, TLOBTrainer}; +pub use diffusion::{DiffusionMetrics, DiffusionParams, DiffusionTrainer}; +pub use xlstm::{XLSTMMetrics, XLSTMParams, XLSTMTrainer}; diff --git a/crates/ml/src/hyperopt/adapters/tggn.rs b/crates/ml/src/hyperopt/adapters/tggn.rs index 2b5075215..f35e6ef4e 100644 --- a/crates/ml/src/hyperopt/adapters/tggn.rs +++ b/crates/ml/src/hyperopt/adapters/tggn.rs @@ -1,12 +1,21 @@ //! TGGN (Temporal Graph Gated Network) Hyperparameter Optimization Adapter //! -//! Provides `TGGNParams` (implementing `ParameterSpace`) and `TGGNMetrics` -//! for hyperparameter optimization of the TGGN model via the unified -//! optimization framework. +//! Provides `TGGNParams` (implementing `ParameterSpace`), `TGGNMetrics`, +//! and `TGGNTrainer` (implementing `HyperparameterOptimizable`) for +//! hyperparameter optimization of the TGGN model via the unified framework. -use crate::hyperopt::traits::ParameterSpace; -use crate::MLError; +use candle_core::Device; use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use tracing::{info, warn}; + +use crate::features::extract_ml_features; +use crate::hyperopt::paths::TrainingPaths; +use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace}; +use crate::tgnn::TGGNConfig; +use crate::tgnn::trainable_adapter::TGGNTrainableAdapter; +use crate::training::unified_trainer::UnifiedTrainable; +use crate::MLError; /// Hyperparameters for TGGN hyperopt tuning. /// @@ -139,6 +148,242 @@ pub struct TGGNMetrics { pub epochs_completed: usize, } +/// TGGN trainer for hyperparameter optimization. +/// +/// Wraps the `TGGNTrainableAdapter` (which implements `UnifiedTrainable`) +/// and drives it through the `HyperparameterOptimizable` interface. +#[derive(Debug)] +pub struct TGGNTrainer { + data_dir: PathBuf, + epochs: usize, + device: Device, + training_paths: TrainingPaths, + early_stopping_patience: usize, + trial_counter: usize, +} + +impl TGGNTrainer { + /// Create a new TGGN trainer. + /// + /// # Arguments + /// * `data_dir` - Directory containing .dbn/.dbn.zst files + /// * `epochs` - Number of training epochs per trial + pub fn new>(data_dir: P, epochs: usize) -> Result { + let data_dir = data_dir.into(); + if !data_dir.exists() { + return Err(MLError::ConfigError { + reason: format!("Data directory not found: {}", data_dir.display()), + }); + } + + let device = Device::new_cuda(0).unwrap_or_else(|e| { + warn!("CUDA unavailable ({}), falling back to CPU", e); + Device::Cpu + }); + + info!( + "TGGN Trainer initialized: Device={:?}, Data={}, Epochs={}", + device, data_dir.display(), epochs + ); + + Ok(Self { + data_dir, + epochs, + device, + training_paths: TrainingPaths::new("/tmp/ml_training", "tggn", "default"), + early_stopping_patience: 10, + trial_counter: 0, + }) + } + + /// Set training paths configuration. + pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self { + self.training_paths = paths; + self + } + + /// Configure early stopping patience. + pub fn with_early_stopping(mut self, patience: usize) -> Self { + self.early_stopping_patience = patience; + self + } +} + +impl HyperparameterOptimizable for TGGNTrainer { + type Params = TGGNParams; + type Metrics = TGGNMetrics; + + fn train_with_params(&mut self, params: Self::Params) -> Result { + let trial_start = std::time::Instant::now(); + let current_trial = self.trial_counter; + self.trial_counter += 1; + + info!( + "Training TGGN: lr={:.6}, hidden={}, layers={}, node_dim={}, batch={}", + params.learning_rate, params.hidden_dim, params.num_layers, + params.node_dim, params.batch_size + ); + + self.training_paths.create_all().map_err(|e| { + MLError::ModelError(format!("Failed to create training directories: {}", e)) + })?; + + // Load bars and extract features + let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir) + .map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?; + + let features = extract_ml_features(&bars) + .map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?; + + if features.len() < 2 { + return Err(MLError::ModelError("Insufficient features extracted".to_owned())); + } + + let feature_dim = 51; + + // TGGN uses flat (non-sequence) input: node_dim is configurable + // We project 51-dim features into node_dim through the model's input projection + let pairs = crate::hyperopt::shared_data::build_flat_pairs( + &features, &bars, feature_dim, &self.device, + )?; + + let split = (pairs.len() as f64 * 0.8) as usize; + let train_data = &pairs[..split]; + let val_data = &pairs[split..]; + + if val_data.len() < 2 { + return Err(MLError::ModelError("Validation set too small".to_owned())); + } + + let config = TGGNConfig { + max_nodes: 64, + max_edges: 256, + node_dim: feature_dim, // Use raw feature dim as node_dim for the adapter + edge_dim: 4, + hidden_dim: params.hidden_dim, + num_layers: params.num_layers, + temporal_decay: params.temporal_decay, + update_frequency_ns: 1_000_000, + use_simd: false, + }; + + let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result { + let mut model = TGGNTrainableAdapter::new(config, &self.device)?; + + let mut best_val_loss = f64::MAX; + let mut patience_counter = 0_usize; + let mut last_train_loss = 0.0_f64; + + for epoch in 0..self.epochs { + let mut epoch_loss = 0.0_f64; + let mut batch_count = 0_usize; + + for (input, target) in train_data { + let pred = model.forward(input)?; + let loss = model.compute_loss(&pred, target)?; + let loss_val = loss + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Loss scalar: {}", e)))? + as f64; + + if loss_val.is_nan() || loss_val.is_infinite() { + return Ok(TGGNMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + directional_accuracy: 0.0, + epochs_completed: epoch, + }); + } + + model.backward(&loss)?; + model.optimizer_step()?; + epoch_loss += loss_val; + batch_count += 1; + } + + last_train_loss = if batch_count > 0 { + epoch_loss / batch_count as f64 + } else { + 0.0 + }; + + let val_loss = model.validate(val_data)?; + + if val_loss < best_val_loss { + best_val_loss = val_loss; + patience_counter = 0; + } else { + patience_counter += 1; + } + + if epoch % 10 == 0 { + info!( + " Epoch {}/{}: train_loss={:.6}, val_loss={:.6}", + epoch + 1, self.epochs, last_train_loss, val_loss + ); + } + + if patience_counter >= self.early_stopping_patience { + info!("Early stopping at epoch {}", epoch + 1); + return Ok(TGGNMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + directional_accuracy: 0.0, + epochs_completed: epoch + 1, + }); + } + } + + Ok(TGGNMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + directional_accuracy: 0.0, + epochs_completed: self.epochs, + }) + })); + + let metrics = match training_result { + Ok(Ok(m)) => m, + Ok(Err(e)) => { + warn!("TGGN training failed: {}", e); + TGGNMetrics { val_loss: 1000.0, train_loss: 1000.0, directional_accuracy: 0.0, epochs_completed: 0 } + } + Err(_) => { + warn!("TGGN training panicked (likely OOM)"); + TGGNMetrics { val_loss: 1000.0, train_loss: 1000.0, directional_accuracy: 0.0, epochs_completed: 0 } + } + }; + + info!( + "TGGN training completed: val_loss={:.6}, train_loss={:.6}, epochs={}", + metrics.val_loss, metrics.train_loss, metrics.epochs_completed + ); + + let duration_secs = trial_start.elapsed().as_secs_f64(); + 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(); + crate::hyperopt::shared_data::write_trial_result_json( + &self.training_paths.hyperopt_dir(), + &trial_result, + ).ok(); + + if self.device.is_cuda() { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + + Ok(metrics) + } + + fn extract_objective(metrics: &Self::Metrics) -> f64 { + metrics.val_loss + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ml/src/hyperopt/adapters/tlob.rs b/crates/ml/src/hyperopt/adapters/tlob.rs index dcc4bd67d..1ae46ae7e 100644 --- a/crates/ml/src/hyperopt/adapters/tlob.rs +++ b/crates/ml/src/hyperopt/adapters/tlob.rs @@ -1,12 +1,20 @@ //! TLOB (Time Limit Order Book) Hyperparameter Optimization Adapter //! -//! Provides `TLOBParams` (implementing `ParameterSpace`) and `TLOBMetrics` -//! for hyperparameter optimization of the TLOB model via the unified -//! optimization framework. +//! Provides `TLOBParams` (implementing `ParameterSpace`), `TLOBMetrics`, +//! and `TLOBTrainer` (implementing `HyperparameterOptimizable`) for +//! hyperparameter optimization of the TLOB model via the unified framework. -use crate::hyperopt::traits::ParameterSpace; -use crate::MLError; +use candle_core::Device; use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use tracing::{info, warn}; + +use crate::features::extract_ml_features; +use crate::hyperopt::paths::TrainingPaths; +use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace}; +use crate::tlob::trainable_adapter::{TLOBAdapterConfig, TLOBTrainableAdapter}; +use crate::training::unified_trainer::UnifiedTrainable; +use crate::MLError; /// Hyperparameters for TLOB hyperopt tuning. /// @@ -131,6 +139,234 @@ pub struct TLOBMetrics { pub epochs_completed: usize, } +/// TLOB trainer for hyperparameter optimization. +/// +/// Wraps the `TLOBTrainableAdapter` (which implements `UnifiedTrainable`) +/// and drives it through the `HyperparameterOptimizable` interface. +/// TLOB accepts sequence input of shape `(1, seq_len, feature_dim)`. +#[derive(Debug)] +pub struct TLOBTrainer { + data_dir: PathBuf, + epochs: usize, + device: Device, + training_paths: TrainingPaths, + early_stopping_patience: usize, + trial_counter: usize, +} + +impl TLOBTrainer { + /// Create a new TLOB trainer. + /// + /// # Arguments + /// * `data_dir` - Directory containing .dbn/.dbn.zst files + /// * `epochs` - Number of training epochs per trial + pub fn new>(data_dir: P, epochs: usize) -> Result { + let data_dir = data_dir.into(); + if !data_dir.exists() { + return Err(MLError::ConfigError { + reason: format!("Data directory not found: {}", data_dir.display()), + }); + } + + let device = Device::new_cuda(0).unwrap_or_else(|e| { + warn!("CUDA unavailable ({}), falling back to CPU", e); + Device::Cpu + }); + + info!( + "TLOB Trainer initialized: Device={:?}, Data={}, Epochs={}", + device, data_dir.display(), epochs + ); + + Ok(Self { + data_dir, + epochs, + device, + training_paths: TrainingPaths::new("/tmp/ml_training", "tlob", "default"), + early_stopping_patience: 10, + trial_counter: 0, + }) + } + + /// Set training paths configuration. + pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self { + self.training_paths = paths; + self + } + + /// Configure early stopping patience. + pub fn with_early_stopping(mut self, patience: usize) -> Self { + self.early_stopping_patience = patience; + self + } +} + +impl HyperparameterOptimizable for TLOBTrainer { + type Params = TLOBParams; + type Metrics = TLOBMetrics; + + fn train_with_params(&mut self, params: Self::Params) -> Result { + let trial_start = std::time::Instant::now(); + let current_trial = self.trial_counter; + self.trial_counter += 1; + + info!( + "Training TLOB: lr={:.6}, d_model={}, heads={}, layers={}, seq_len={}, batch={}", + params.learning_rate, params.d_model, params.num_heads, + params.num_layers, params.seq_len, params.batch_size + ); + + self.training_paths.create_all().map_err(|e| { + MLError::ModelError(format!("Failed to create training directories: {}", e)) + })?; + + let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir) + .map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?; + + let features = extract_ml_features(&bars) + .map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?; + + let feature_dim = 51; + let seq_len = params.seq_len; + + // TLOB uses sequence input: (1, seq_len, feature_dim) + let pairs = crate::hyperopt::shared_data::build_sequence_pairs( + &features, &bars, feature_dim, seq_len, &self.device, + )?; + + let split = (pairs.len() as f64 * 0.8) as usize; + let train_data = &pairs[..split]; + let val_data = &pairs[split..]; + + if val_data.len() < 2 { + return Err(MLError::ModelError("Validation set too small".to_owned())); + } + + let config = TLOBAdapterConfig { + d_model: params.d_model, + num_heads: params.num_heads, + num_layers: params.num_layers, + seq_len, + feature_dim, + }; + + let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result { + let mut model = TLOBTrainableAdapter::new(config, &self.device)?; + + let mut best_val_loss = f64::MAX; + let mut patience_counter = 0_usize; + let mut last_train_loss = 0.0_f64; + + for epoch in 0..self.epochs { + let mut epoch_loss = 0.0_f64; + let mut batch_count = 0_usize; + + for (input, target) in train_data { + let pred = model.forward(input)?; + let loss = model.compute_loss(&pred, target)?; + let loss_val = loss + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Loss scalar: {}", e)))? + as f64; + + if loss_val.is_nan() || loss_val.is_infinite() { + return Ok(TLOBMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + directional_accuracy: 0.0, + epochs_completed: epoch, + }); + } + + model.backward(&loss)?; + model.optimizer_step()?; + epoch_loss += loss_val; + batch_count += 1; + } + + last_train_loss = if batch_count > 0 { + epoch_loss / batch_count as f64 + } else { + 0.0 + }; + + let val_loss = model.validate(val_data)?; + + if val_loss < best_val_loss { + best_val_loss = val_loss; + patience_counter = 0; + } else { + patience_counter += 1; + } + + if epoch % 10 == 0 { + info!( + " Epoch {}/{}: train_loss={:.6}, val_loss={:.6}", + epoch + 1, self.epochs, last_train_loss, val_loss + ); + } + + if patience_counter >= self.early_stopping_patience { + info!("Early stopping at epoch {}", epoch + 1); + return Ok(TLOBMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + directional_accuracy: 0.0, + epochs_completed: epoch + 1, + }); + } + } + + Ok(TLOBMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + directional_accuracy: 0.0, + epochs_completed: self.epochs, + }) + })); + + let metrics = match training_result { + Ok(Ok(m)) => m, + Ok(Err(e)) => { + warn!("TLOB training failed: {}", e); + TLOBMetrics { val_loss: 1000.0, train_loss: 1000.0, directional_accuracy: 0.0, epochs_completed: 0 } + } + Err(_) => { + warn!("TLOB training panicked (likely OOM)"); + TLOBMetrics { val_loss: 1000.0, train_loss: 1000.0, directional_accuracy: 0.0, epochs_completed: 0 } + } + }; + + info!( + "TLOB training completed: val_loss={:.6}, train_loss={:.6}, epochs={}", + metrics.val_loss, metrics.train_loss, metrics.epochs_completed + ); + + let duration_secs = trial_start.elapsed().as_secs_f64(); + 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(); + crate::hyperopt::shared_data::write_trial_result_json( + &self.training_paths.hyperopt_dir(), + &trial_result, + ).ok(); + + if self.device.is_cuda() { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + + Ok(metrics) + } + + fn extract_objective(metrics: &Self::Metrics) -> f64 { + metrics.val_loss + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ml/src/hyperopt/adapters/xlstm.rs b/crates/ml/src/hyperopt/adapters/xlstm.rs index 5d1016cd5..aa8547841 100644 --- a/crates/ml/src/hyperopt/adapters/xlstm.rs +++ b/crates/ml/src/hyperopt/adapters/xlstm.rs @@ -1,10 +1,23 @@ //! Hyperopt adapter for xLSTM (Extended LSTM). +//! +//! Provides `XLSTMParams` (implementing `ParameterSpace`), `XLSTMMetrics`, +//! and `XLSTMTrainer` (implementing `HyperparameterOptimizable`) for +//! hyperparameter optimization of the xLSTM model via the unified framework. +use candle_core::Device; +use std::path::PathBuf; +use tracing::{info, warn}; + +use crate::features::extract_ml_features; +use crate::hyperopt::paths::TrainingPaths; +use crate::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace}; +use crate::training::unified_trainer::UnifiedTrainable; +use crate::xlstm::config::XLSTMConfig; +use crate::xlstm::trainable::XLSTMTrainableAdapter; use crate::MLError; -use crate::hyperopt::traits::ParameterSpace; /// Hyperparameters for xLSTM hyperopt tuning. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct XLSTMParams { pub learning_rate: f64, pub hidden_dim: usize, @@ -54,20 +67,20 @@ impl ParameterSpace for XLSTMParams { reason: format!("Expected 9 params, got {}", x.len()), }); } - let hidden_dim = x[1].round().max(32.0) as usize; - let num_heads = x[3].round().max(1.0) as usize; + let hidden_dim = x.get(1).copied().unwrap_or(64.0).round().max(32.0) as usize; + let num_heads = x.get(3).copied().unwrap_or(4.0).round().max(1.0) as usize; // Ensure hidden_dim is divisible by num_heads let hidden_dim = (hidden_dim / num_heads) * num_heads; Ok(Self { - learning_rate: x[0].exp(), + learning_rate: x.first().copied().unwrap_or(-6.9).exp(), hidden_dim: hidden_dim.max(num_heads), // at least 1 per head - num_blocks: x[2].round().max(2.0) as usize, + num_blocks: x.get(2).copied().unwrap_or(4.0).round().max(2.0) as usize, num_heads, - slstm_ratio: x[4].clamp(0.0, 1.0), - dropout: x[5].clamp(0.0, 0.5), - batch_size: x[6].round().max(8.0) as usize, - weight_decay: x[7].exp(), - grad_clip: x[8].exp(), + slstm_ratio: x.get(4).copied().unwrap_or(0.5).clamp(0.0, 1.0), + dropout: x.get(5).copied().unwrap_or(0.1).clamp(0.0, 0.5), + batch_size: x.get(6).copied().unwrap_or(32.0).round().max(8.0) as usize, + weight_decay: x.get(7).copied().unwrap_or(-9.21).exp(), + grad_clip: x.get(8).copied().unwrap_or(0.0).exp(), }) } @@ -94,7 +107,7 @@ impl ParameterSpace for XLSTMParams { } /// Metrics returned from xLSTM hyperopt training. -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] pub struct XLSTMMetrics { pub val_loss: f64, pub train_loss: f64, @@ -102,6 +115,241 @@ pub struct XLSTMMetrics { pub epochs_completed: usize, } +/// xLSTM trainer for hyperparameter optimization. +/// +/// Wraps the `XLSTMTrainableAdapter` (which implements `UnifiedTrainable`) +/// and drives it through the `HyperparameterOptimizable` interface. +#[derive(Debug)] +pub struct XLSTMTrainer { + data_dir: PathBuf, + epochs: usize, + device: Device, + training_paths: TrainingPaths, + early_stopping_patience: usize, + trial_counter: usize, +} + +impl XLSTMTrainer { + /// Create a new xLSTM trainer. + /// + /// # Arguments + /// * `data_dir` - Directory containing .dbn/.dbn.zst files + /// * `epochs` - Number of training epochs per trial + pub fn new>(data_dir: P, epochs: usize) -> Result { + let data_dir = data_dir.into(); + if !data_dir.exists() { + return Err(MLError::ConfigError { + reason: format!("Data directory not found: {}", data_dir.display()), + }); + } + + let device = Device::new_cuda(0).unwrap_or_else(|e| { + warn!("CUDA unavailable ({}), falling back to CPU", e); + Device::Cpu + }); + + info!( + "xLSTM Trainer initialized: Device={:?}, Data={}, Epochs={}", + device, data_dir.display(), epochs + ); + + Ok(Self { + data_dir, + epochs, + device, + training_paths: TrainingPaths::new("/tmp/ml_training", "xlstm", "default"), + early_stopping_patience: 10, + trial_counter: 0, + }) + } + + /// Set training paths configuration. + pub fn with_training_paths(mut self, paths: TrainingPaths) -> Self { + self.training_paths = paths; + self + } + + /// Configure early stopping patience. + pub fn with_early_stopping(mut self, patience: usize) -> Self { + self.early_stopping_patience = patience; + self + } +} + +impl HyperparameterOptimizable for XLSTMTrainer { + type Params = XLSTMParams; + type Metrics = XLSTMMetrics; + + fn train_with_params(&mut self, params: Self::Params) -> Result { + let trial_start = std::time::Instant::now(); + let current_trial = self.trial_counter; + self.trial_counter += 1; + + info!( + "Training xLSTM: lr={:.6}, hidden={}, blocks={}, heads={}, slstm_ratio={:.2}, batch={}", + params.learning_rate, params.hidden_dim, params.num_blocks, + params.num_heads, params.slstm_ratio, params.batch_size + ); + + self.training_paths.create_all().map_err(|e| { + MLError::ModelError(format!("Failed to create training directories: {}", e)) + })?; + + let bars = super::dbn_loader::load_bars_from_dbn_dir(&self.data_dir) + .map_err(|e| MLError::ModelError(format!("Failed to load DBN data: {}", e)))?; + + let features = extract_ml_features(&bars) + .map_err(|e| MLError::ModelError(format!("Failed to extract features: {}", e)))?; + + if features.len() < 2 { + return Err(MLError::ModelError("Insufficient features extracted".to_owned())); + } + + let feature_dim = 51; + + // xLSTM uses flat input (the adapter handles internal state) + let pairs = crate::hyperopt::shared_data::build_flat_pairs( + &features, &bars, feature_dim, &self.device, + )?; + + let split = (pairs.len() as f64 * 0.8) as usize; + let train_data = &pairs[..split]; + let val_data = &pairs[split..]; + + if val_data.len() < 2 { + return Err(MLError::ModelError("Validation set too small".to_owned())); + } + + let config = XLSTMConfig { + input_dim: feature_dim, + hidden_dim: params.hidden_dim, + num_blocks: params.num_blocks, + num_heads: params.num_heads, + slstm_ratio: params.slstm_ratio, + output_dim: 1, + dropout: params.dropout, + learning_rate: params.learning_rate, + weight_decay: params.weight_decay, + grad_clip: params.grad_clip, + }; + + let training_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| -> Result { + let mut model = XLSTMTrainableAdapter::new(config, &self.device)?; + + let mut best_val_loss = f64::MAX; + let mut patience_counter = 0_usize; + let mut last_train_loss = 0.0_f64; + + for epoch in 0..self.epochs { + let mut epoch_loss = 0.0_f64; + let mut batch_count = 0_usize; + + for (input, target) in train_data { + let pred = model.forward(input)?; + let loss = model.compute_loss(&pred, target)?; + let loss_val = loss + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Loss scalar: {}", e)))? + as f64; + + if loss_val.is_nan() || loss_val.is_infinite() { + return Ok(XLSTMMetrics { + val_loss: 1000.0, + train_loss: 1000.0, + directional_accuracy: 0.0, + epochs_completed: epoch, + }); + } + + model.backward(&loss)?; + model.optimizer_step()?; + epoch_loss += loss_val; + batch_count += 1; + } + + last_train_loss = if batch_count > 0 { + epoch_loss / batch_count as f64 + } else { + 0.0 + }; + + let val_loss = model.validate(val_data)?; + + if val_loss < best_val_loss { + best_val_loss = val_loss; + patience_counter = 0; + } else { + patience_counter += 1; + } + + if epoch % 10 == 0 { + info!( + " Epoch {}/{}: train_loss={:.6}, val_loss={:.6}", + epoch + 1, self.epochs, last_train_loss, val_loss + ); + } + + if patience_counter >= self.early_stopping_patience { + info!("Early stopping at epoch {}", epoch + 1); + return Ok(XLSTMMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + directional_accuracy: 0.0, + epochs_completed: epoch + 1, + }); + } + } + + Ok(XLSTMMetrics { + val_loss: best_val_loss, + train_loss: last_train_loss, + directional_accuracy: 0.0, + epochs_completed: self.epochs, + }) + })); + + let metrics = match training_result { + Ok(Ok(m)) => m, + Ok(Err(e)) => { + warn!("xLSTM training failed: {}", e); + XLSTMMetrics { val_loss: 1000.0, train_loss: 1000.0, directional_accuracy: 0.0, epochs_completed: 0 } + } + Err(_) => { + warn!("xLSTM training panicked (likely OOM)"); + XLSTMMetrics { val_loss: 1000.0, train_loss: 1000.0, directional_accuracy: 0.0, epochs_completed: 0 } + } + }; + + info!( + "xLSTM training completed: val_loss={:.6}, train_loss={:.6}, epochs={}", + metrics.val_loss, metrics.train_loss, metrics.epochs_completed + ); + + let duration_secs = trial_start.elapsed().as_secs_f64(); + 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(); + crate::hyperopt::shared_data::write_trial_result_json( + &self.training_paths.hyperopt_dir(), + &trial_result, + ).ok(); + + if self.device.is_cuda() { + std::thread::sleep(std::time::Duration::from_millis(100)); + } + + Ok(metrics) + } + + fn extract_objective(metrics: &Self::Metrics) -> f64 { + metrics.val_loss + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ml/src/hyperopt/mod.rs b/crates/ml/src/hyperopt/mod.rs index c9c966441..08eb11106 100644 --- a/crates/ml/src/hyperopt/mod.rs +++ b/crates/ml/src/hyperopt/mod.rs @@ -46,6 +46,7 @@ pub mod observer; pub mod optimizer; pub mod paths; pub mod sensitivity; +pub mod shared_data; pub mod traits; #[cfg(test)] diff --git a/crates/ml/src/hyperopt/shared_data.rs b/crates/ml/src/hyperopt/shared_data.rs new file mode 100644 index 000000000..97847f3dd --- /dev/null +++ b/crates/ml/src/hyperopt/shared_data.rs @@ -0,0 +1,300 @@ +//! Shared data preparation utilities for hyperopt trainers. +//! +//! Provides common functions for loading DBN data, extracting features, +//! normalizing, and building (input, target) tensor pairs for supervised +//! model training during hyperparameter optimization. + +use candle_core::{Device, Tensor}; +use serde::Serialize; +use std::fmt::Debug; +use std::path::Path; +use tracing::info; + +use crate::features::extraction::OHLCVBar; +use crate::MLError; + +/// Build flat (input, target) tensor pairs from extracted features and bars. +/// +/// Each pair maps a single feature vector (51-dim) to a normalized close-price +/// return target. Features are z-score normalized using training-set statistics. +/// +/// # Arguments +/// * `features` - Extracted feature vectors (51-dim each) +/// * `bars` - OHLCV bars aligned with features (offset by warmup period) +/// * `feature_dim` - Feature dimensionality (51) +/// * `device` - Device for tensor creation +pub fn build_flat_pairs( + features: &[[f64; 51]], + bars: &[OHLCVBar], + feature_dim: usize, + device: &Device, +) -> Result, MLError> { + if features.len() < 2 { + return Err(MLError::ModelError( + "Need at least 2 feature vectors to build pairs".to_owned(), + )); + } + + // The feature extractor has a warmup of 50 bars, so features[i] corresponds + // to bars[i + warmup]. We use the close-price return from the NEXT bar as target. + // Since we don't know the exact warmup offset here, we use price returns + // computed directly from consecutive features' corresponding bars. + // + // Simpler approach: target = normalized close-price return from consecutive features. + // features[i] -> target = (close[i+1] - close[i]) / close[i] where close comes + // from the bars that produced the features. + // + // Since we can't perfectly align bars<->features without the warmup constant, + // use the last (features.len()) bars for targets. + let n_bars = bars.len(); + let n_features = features.len(); + + // features correspond to bars[offset..offset+n_features] where offset = n_bars - n_features + let bar_offset = n_bars.saturating_sub(n_features); + + // Compute per-feature mean and std for z-score normalization + let mut means = vec![0.0_f64; feature_dim]; + let mut vars = vec![0.0_f64; feature_dim]; + let count = n_features as f64; + + for feat in features.iter() { + for (j, val) in feat.iter().enumerate() { + if let Some(m) = means.get_mut(j) { + *m += val; + } + } + } + for m in means.iter_mut() { + *m /= count; + } + + for feat in features.iter() { + for (j, val) in feat.iter().enumerate() { + if let (Some(v), Some(m)) = (vars.get_mut(j), means.get(j)) { + *v += (val - m) * (val - m); + } + } + } + let stds: Vec = vars + .iter() + .map(|v| (v / count).sqrt().max(1e-8)) + .collect(); + + // Compute target return statistics for normalization + let mut returns = Vec::with_capacity(n_features.saturating_sub(1)); + for i in 0..n_features.saturating_sub(1) { + let bar_idx = bar_offset + i; + let next_bar_idx = bar_offset + i + 1; + let close = bars.get(bar_idx).map(|b| b.close).unwrap_or(1.0); + let next_close = bars.get(next_bar_idx).map(|b| b.close).unwrap_or(1.0); + if close.abs() > 1e-12 { + returns.push((next_close - close) / close); + } else { + returns.push(0.0); + } + } + + let ret_mean = returns.iter().sum::() / returns.len().max(1) as f64; + let ret_std = (returns + .iter() + .map(|r| (r - ret_mean) * (r - ret_mean)) + .sum::() + / returns.len().max(1) as f64) + .sqrt() + .max(1e-8); + + let mut pairs = Vec::with_capacity(returns.len()); + for (i, ret) in returns.iter().enumerate() { + let feat = features.get(i).ok_or_else(|| { + MLError::ModelError(format!("Feature index {} out of bounds", i)) + })?; + + // Z-score normalize features + let normalized: Vec = feat + .iter() + .enumerate() + .map(|(j, val)| { + let m = means.get(j).copied().unwrap_or(0.0); + let s = stds.get(j).copied().unwrap_or(1.0); + ((val - m) / s) as f32 + }) + .collect(); + + let target_normalized = ((ret - ret_mean) / ret_std) as f32; + + let input = Tensor::new(normalized.as_slice(), device) + .map_err(|e| MLError::ModelError(format!("Input tensor: {}", e)))? + .reshape((1, feature_dim)) + .map_err(|e| MLError::ModelError(format!("Input reshape: {}", e)))?; + + let target = Tensor::new(&[target_normalized], device) + .map_err(|e| MLError::ModelError(format!("Target tensor: {}", e)))? + .reshape((1, 1)) + .map_err(|e| MLError::ModelError(format!("Target reshape: {}", e)))?; + + pairs.push((input, target)); + } + + info!( + "Built {} (input, target) pairs, feature_dim={}, ret_mean={:.6}, ret_std={:.6}", + pairs.len(), + feature_dim, + ret_mean, + ret_std + ); + + Ok(pairs) +} + +/// Build sequenced (input, target) tensor pairs for models requiring sequence input. +/// +/// Each input has shape `(1, seq_len, feature_dim)` and target is `(1, 1)`. +/// +/// # Arguments +/// * `features` - Extracted feature vectors +/// * `bars` - OHLCV bars +/// * `feature_dim` - Feature dimensionality (51) +/// * `seq_len` - Sequence length +/// * `device` - Device for tensor creation +pub fn build_sequence_pairs( + features: &[[f64; 51]], + bars: &[OHLCVBar], + feature_dim: usize, + seq_len: usize, + device: &Device, +) -> Result, MLError> { + if features.len() < seq_len + 1 { + return Err(MLError::ModelError(format!( + "Need at least {} features for seq_len={}, got {}", + seq_len + 1, + seq_len, + features.len() + ))); + } + + let n_bars = bars.len(); + let n_features = features.len(); + let bar_offset = n_bars.saturating_sub(n_features); + + // Compute per-feature stats for normalization + let mut means = vec![0.0_f64; feature_dim]; + let count = n_features as f64; + for feat in features.iter() { + for (j, val) in feat.iter().enumerate() { + if let Some(m) = means.get_mut(j) { + *m += val; + } + } + } + for m in means.iter_mut() { + *m /= count; + } + + let mut vars = vec![0.0_f64; feature_dim]; + for feat in features.iter() { + for (j, val) in feat.iter().enumerate() { + if let (Some(v), Some(m)) = (vars.get_mut(j), means.get(j)) { + *v += (val - m) * (val - m); + } + } + } + let stds: Vec = vars + .iter() + .map(|v| (v / count).sqrt().max(1e-8)) + .collect(); + + // Target: close-price returns + let mut returns = Vec::with_capacity(n_features.saturating_sub(1)); + for i in 0..n_features.saturating_sub(1) { + let bar_idx = bar_offset + i; + let next_bar_idx = bar_offset + i + 1; + let close = bars.get(bar_idx).map(|b| b.close).unwrap_or(1.0); + let next_close = bars.get(next_bar_idx).map(|b| b.close).unwrap_or(1.0); + if close.abs() > 1e-12 { + returns.push((next_close - close) / close); + } else { + returns.push(0.0); + } + } + + let ret_mean = returns.iter().sum::() / returns.len().max(1) as f64; + let ret_std = (returns + .iter() + .map(|r| (r - ret_mean) * (r - ret_mean)) + .sum::() + / returns.len().max(1) as f64) + .sqrt() + .max(1e-8); + + let num_samples = n_features.saturating_sub(seq_len); + let mut pairs = Vec::with_capacity(num_samples); + + for i in 0..num_samples { + // Build sequence of normalized features + let mut seq_data = Vec::with_capacity(seq_len * feature_dim); + for t in 0..seq_len { + let feat_idx = i + t; + let feat = features.get(feat_idx).ok_or_else(|| { + MLError::ModelError(format!("Feature index {} out of bounds", feat_idx)) + })?; + for (j, val) in feat.iter().enumerate() { + let m = means.get(j).copied().unwrap_or(0.0); + let s = stds.get(j).copied().unwrap_or(1.0); + seq_data.push(((val - m) / s) as f32); + } + } + + // Target: return at the step after the sequence + let ret_idx = i + seq_len - 1; + let target_val = returns + .get(ret_idx) + .map(|r| ((r - ret_mean) / ret_std) as f32) + .unwrap_or(0.0); + + let input = Tensor::new(seq_data.as_slice(), device) + .map_err(|e| MLError::ModelError(format!("Seq input tensor: {}", e)))? + .reshape((1, seq_len, feature_dim)) + .map_err(|e| MLError::ModelError(format!("Seq input reshape: {}", e)))?; + + let target = Tensor::new(&[target_val], device) + .map_err(|e| MLError::ModelError(format!("Seq target tensor: {}", e)))? + .reshape((1, 1)) + .map_err(|e| MLError::ModelError(format!("Seq target reshape: {}", e)))?; + + pairs.push((input, target)); + } + + info!( + "Built {} sequence pairs (seq_len={}, feature_dim={})", + pairs.len(), + seq_len, + feature_dim + ); + + Ok(pairs) +} + +/// Write a trial result to a JSON file (appending to existing trials). +pub fn write_trial_result_json( + hyperopt_dir: &Path, + trial_result: &crate::hyperopt::traits::TrialResult

, +) -> Result<(), std::io::Error> { + let trials_file = hyperopt_dir.join("trials.json"); + + let mut all_trials = if trials_file.exists() { + let content = std::fs::read_to_string(&trials_file)?; + serde_json::from_str::>(&content).unwrap_or_default() + } else { + Vec::new() + }; + + 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); + + 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(()) +}