From 4eacd4e22f37883a103d4c624e6ffa6672b6b080 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 23 Feb 2026 01:32:25 +0100 Subject: [PATCH] feat(ml): add KAN architecture + TLOB/KAN trainable/hyperopt adapters Phase 2-3 of ensemble expansion: - KAN (Kolmogorov-Arnold Network): B-spline basis, layer, network, trainable adapter - TLOB UnifiedTrainable adapter with 3D input support (batch, seq, features) - Hyperopt adapters for both KAN and TLOB (ParameterSpace + metrics) - ModelType::KAN variant registered in common, coordinator, lib.rs - 44 new tests, all passing, zero warnings Co-Authored-By: Claude Opus 4.6 --- common/src/model_types.rs | 7 + ml/src/hyperopt/adapters/kan.rs | 217 ++++++++++ ml/src/hyperopt/adapters/mod.rs | 4 + ml/src/hyperopt/adapters/tlob.rs | 229 ++++++++++ ml/src/integration/coordinator.rs | 5 + ml/src/kan/config.rs | 42 ++ ml/src/kan/layer.rs | 151 +++++++ ml/src/kan/mod.rs | 15 + ml/src/kan/network.rs | 121 ++++++ ml/src/kan/spline.rs | 188 +++++++++ ml/src/kan/trainable.rs | 493 ++++++++++++++++++++++ ml/src/lib.rs | 1 + ml/src/tlob/mod.rs | 2 + ml/src/tlob/trainable_adapter.rs | 671 ++++++++++++++++++++++++++++++ 14 files changed, 2146 insertions(+) create mode 100644 ml/src/hyperopt/adapters/kan.rs create mode 100644 ml/src/hyperopt/adapters/tlob.rs create mode 100644 ml/src/kan/config.rs create mode 100644 ml/src/kan/layer.rs create mode 100644 ml/src/kan/mod.rs create mode 100644 ml/src/kan/network.rs create mode 100644 ml/src/kan/spline.rs create mode 100644 ml/src/kan/trainable.rs create mode 100644 ml/src/tlob/trainable_adapter.rs diff --git a/common/src/model_types.rs b/common/src/model_types.rs index 663188ba3..8ecc458dc 100644 --- a/common/src/model_types.rs +++ b/common/src/model_types.rs @@ -42,6 +42,8 @@ pub enum ModelType { TGNN, /// Ensemble of multiple models Ensemble, + /// Kolmogorov-Arnold Network + KAN, } impl ModelType { @@ -60,6 +62,7 @@ impl ModelType { Self::PPO => "ppo", Self::Transformer => "transformer", Self::Ensemble => "ensemble", + Self::KAN => "kan", } } @@ -76,6 +79,7 @@ impl ModelType { Self::DistilledMicroNet => "distilled", Self::Transformer => "transformer", Self::Ensemble => "ensemble", + Self::KAN => "kan", } } @@ -92,6 +96,7 @@ impl ModelType { Self::TGGN | Self::TGNN => "tggn", Self::DistilledMicroNet => "distilled", Self::Transformer => "transformer", + Self::KAN => "kan", } } @@ -110,6 +115,7 @@ impl ModelType { Self::TLOB => "TLOB", Self::Transformer => "TRANSFORMER", Self::Ensemble => "ENSEMBLE", + Self::KAN => "KAN", } } @@ -134,6 +140,7 @@ impl ModelType { "ppo" => Some(Self::PPO), "transformer" => Some(Self::Transformer), "ensemble" => Some(Self::Ensemble), + "kan" => Some(Self::KAN), _ => None, } } diff --git a/ml/src/hyperopt/adapters/kan.rs b/ml/src/hyperopt/adapters/kan.rs new file mode 100644 index 000000000..e0a36b5a0 --- /dev/null +++ b/ml/src/hyperopt/adapters/kan.rs @@ -0,0 +1,217 @@ +//! 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. + +use crate::hyperopt::traits::ParameterSpace; +use crate::MLError; +use serde::{Deserialize, Serialize}; + +/// Hyperparameters for KAN hyperopt tuning. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct KANParams { + /// Learning rate for AdamW optimizer (log-scale) + pub learning_rate: f64, + /// Number of B-spline grid points per activation + pub grid_size: usize, + /// B-spline order (degree + 1) + pub spline_order: usize, + /// Hidden layer width (used for all hidden layers) + pub hidden_width: usize, + /// Number of hidden layers (total layers = num_layers + 1 for output) + pub num_layers: usize, + /// L2 weight decay (log-scale) + pub weight_decay: f64, + /// Maximum gradient norm for clipping (log-scale) + pub grad_clip: f64, + /// Training batch size + pub batch_size: usize, +} + +impl Default for KANParams { + fn default() -> Self { + Self { + learning_rate: 1e-3, + grid_size: 5, + spline_order: 4, + hidden_width: 32, + num_layers: 2, + weight_decay: 1e-4, + grad_clip: 1.0, + batch_size: 32, + } + } +} + +/// Number of continuous parameters in the KAN parameter space. +const NUM_PARAMS: usize = 8; + +impl ParameterSpace for KANParams { + fn continuous_bounds() -> Vec<(f64, f64)> { + vec![ + (1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log) + (3.0, 12.0), // grid_size (linear) + (2.0, 5.0), // spline_order (linear) + (16.0, 64.0), // hidden_width (linear) + (2.0, 4.0), // num_layers (linear) + (1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log) + (0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log) + (8.0, 64.0), // batch_size (linear) + ] + } + + fn from_continuous(x: &[f64]) -> Result { + if x.len() != NUM_PARAMS { + return Err(MLError::ConfigError { + reason: format!("Expected {} continuous parameters, got {}", NUM_PARAMS, x.len()), + }); + } + Ok(Self { + learning_rate: x.first().copied().unwrap_or(0.0).exp(), + grid_size: x.get(1).copied().unwrap_or(5.0).round().max(3.0) as usize, + spline_order: x.get(2).copied().unwrap_or(4.0).round().max(2.0) as usize, + hidden_width: x.get(3).copied().unwrap_or(32.0).round().max(16.0) as usize, + num_layers: x.get(4).copied().unwrap_or(2.0).round().max(2.0) as usize, + weight_decay: x.get(5).copied().unwrap_or(0.0).exp(), + grad_clip: x.get(6).copied().unwrap_or(0.0).exp(), + batch_size: x.get(7).copied().unwrap_or(32.0).round().max(8.0) as usize, + }) + } + + fn to_continuous(&self) -> Vec { + vec![ + self.learning_rate.ln(), + self.grid_size as f64, + self.spline_order as f64, + self.hidden_width as f64, + self.num_layers as f64, + self.weight_decay.ln(), + self.grad_clip.ln(), + self.batch_size as f64, + ] + } + + fn param_names() -> Vec<&'static str> { + vec![ + "learning_rate", + "grid_size", + "spline_order", + "hidden_width", + "num_layers", + "weight_decay", + "grad_clip", + "batch_size", + ] + } +} + +/// Metrics returned from KAN hyperopt training. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct KANMetrics { + /// 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, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bounds_count_matches_param_names() { + let bounds = KANParams::continuous_bounds(); + let names = KANParams::param_names(); + assert_eq!(bounds.len(), names.len()); + assert_eq!(bounds.len(), NUM_PARAMS); + } + + #[test] + fn test_roundtrip_continuous() { + let params = KANParams::default(); + let continuous = params.to_continuous(); + let restored = KANParams::from_continuous(&continuous).unwrap(); + assert!( + (params.learning_rate - restored.learning_rate).abs() < 1e-6, + "lr mismatch: {} vs {}", + params.learning_rate, + restored.learning_rate + ); + assert_eq!(params.grid_size, restored.grid_size); + assert_eq!(params.spline_order, restored.spline_order); + assert_eq!(params.hidden_width, restored.hidden_width); + assert_eq!(params.num_layers, restored.num_layers); + assert!( + (params.weight_decay - restored.weight_decay).abs() / params.weight_decay < 1e-6, + "weight_decay mismatch" + ); + assert!( + (params.grad_clip - restored.grad_clip).abs() < 1e-6, + "grad_clip mismatch" + ); + assert_eq!(params.batch_size, restored.batch_size); + } + + #[test] + fn test_from_continuous_wrong_length_errors() { + let result = KANParams::from_continuous(&[0.1, 0.2]); + assert!(result.is_err()); + let err_msg = format!("{}", result.unwrap_err()); + assert!( + err_msg.contains("Expected 8"), + "Error should mention expected count: {}", + err_msg + ); + } + + #[test] + fn test_bounds_are_valid() { + for (i, (min, max)) in KANParams::continuous_bounds().iter().enumerate() { + assert!( + min < max, + "Invalid bounds for param {} ({}): {} >= {}", + i, + KANParams::param_names().get(i).copied().unwrap_or("?"), + min, + max + ); + } + } + + #[test] + fn test_default_within_bounds() { + let params = KANParams::default(); + let continuous = params.to_continuous(); + let bounds = KANParams::continuous_bounds(); + let names = KANParams::param_names(); + for (i, (val, (min, max))) in continuous.iter().zip(bounds.iter()).enumerate() { + assert!( + *val >= *min && *val <= *max, + "Param {} ({}) = {} outside [{}, {}]", + i, + names.get(i).copied().unwrap_or("?"), + val, + min, + max + ); + } + } + + #[test] + fn test_metrics_default() { + let metrics = KANMetrics::default(); + assert_eq!(metrics.val_loss, 0.0); + assert_eq!(metrics.train_loss, 0.0); + assert_eq!(metrics.epochs_completed, 0); + } + + #[test] + fn test_continuous_length() { + let params = KANParams::default(); + let continuous = params.to_continuous(); + assert_eq!(continuous.len(), NUM_PARAMS); + } +} diff --git a/ml/src/hyperopt/adapters/mod.rs b/ml/src/hyperopt/adapters/mod.rs index d5dab1ba0..741ec0680 100644 --- a/ml/src/hyperopt/adapters/mod.rs +++ b/ml/src/hyperopt/adapters/mod.rs @@ -52,14 +52,17 @@ pub mod async_data_loader; pub mod continuous_ppo; pub mod dqn; +pub mod kan; pub mod liquid; pub mod mamba2; pub mod ppo; pub mod tft; pub mod tggn; +pub mod tlob; // Re-export adapters for convenience pub use async_data_loader::AsyncDataLoader; +pub use kan::{KANMetrics, KANParams}; pub use liquid::LiquidParams; pub use continuous_ppo::{ContinuousPPOMetrics, ContinuousPPOParams, ContinuousPPOTrainer}; pub use dqn::{DQNMetrics, DQNParams, DQNTrainer}; @@ -67,3 +70,4 @@ 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}; diff --git a/ml/src/hyperopt/adapters/tlob.rs b/ml/src/hyperopt/adapters/tlob.rs new file mode 100644 index 000000000..dcc4bd67d --- /dev/null +++ b/ml/src/hyperopt/adapters/tlob.rs @@ -0,0 +1,229 @@ +//! 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. + +use crate::hyperopt::traits::ParameterSpace; +use crate::MLError; +use serde::{Deserialize, Serialize}; + +/// Hyperparameters for TLOB hyperopt tuning. +/// +/// Controls the TLOB transformer's architecture and training: +/// - Architecture: `d_model`, `num_heads`, `num_layers`, `seq_len` +/// - Regularization: `dropout`, `weight_decay`, `grad_clip` +/// - Training: `learning_rate`, `batch_size` +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct TLOBParams { + /// Learning rate for AdamW optimizer (log-scale) + pub learning_rate: f64, + /// Model / hidden dimension for projection layers + pub d_model: usize, + /// Number of attention heads + pub num_heads: usize, + /// Number of transformer/projection layers + pub num_layers: usize, + /// Sequence length (number of time steps in LOB snapshot) + pub seq_len: usize, + /// Dropout rate for regularization + pub dropout: f64, + /// Training batch size + pub batch_size: usize, + /// L2 weight decay (log-scale) + pub weight_decay: f64, + /// Maximum gradient norm for clipping (log-scale) + pub grad_clip: f64, +} + +impl Default for TLOBParams { + fn default() -> Self { + Self { + learning_rate: 1e-3, + d_model: 128, + num_heads: 4, + num_layers: 2, + seq_len: 128, + dropout: 0.1, + batch_size: 32, + weight_decay: 1e-4, + grad_clip: 1.0, + } + } +} + +/// Number of continuous parameters in the TLOB parameter space. +const NUM_PARAMS: usize = 9; + +impl ParameterSpace for TLOBParams { + fn continuous_bounds() -> Vec<(f64, f64)> { + vec![ + (1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log) + (64.0, 512.0), // d_model (linear) + (2.0, 16.0), // num_heads (linear) + (1.0, 8.0), // num_layers (linear) + (32.0, 256.0), // seq_len (linear) + (0.0, 0.5), // dropout (linear) + (8.0, 64.0), // batch_size (linear) + (1e-6_f64.ln(), 1e-2_f64.ln()), // weight_decay (log) + (0.5_f64.ln(), 5.0_f64.ln()), // grad_clip (log) + ] + } + + fn from_continuous(x: &[f64]) -> Result { + if x.len() != NUM_PARAMS { + return Err(MLError::ConfigError { + reason: format!("Expected {} continuous parameters, got {}", NUM_PARAMS, x.len()), + }); + } + Ok(Self { + learning_rate: x.first().copied().unwrap_or(0.0).exp(), + d_model: x.get(1).copied().unwrap_or(128.0).round().max(64.0) as usize, + num_heads: x.get(2).copied().unwrap_or(4.0).round().max(2.0) as usize, + num_layers: x.get(3).copied().unwrap_or(2.0).round().max(1.0) as usize, + seq_len: x.get(4).copied().unwrap_or(128.0).round().max(32.0) as usize, + 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(0.0).exp(), + grad_clip: x.get(8).copied().unwrap_or(0.0).exp(), + }) + } + + fn to_continuous(&self) -> Vec { + vec![ + self.learning_rate.ln(), + self.d_model as f64, + self.num_heads as f64, + self.num_layers as f64, + self.seq_len as f64, + self.dropout, + self.batch_size as f64, + self.weight_decay.ln(), + self.grad_clip.ln(), + ] + } + + fn param_names() -> Vec<&'static str> { + vec![ + "learning_rate", + "d_model", + "num_heads", + "num_layers", + "seq_len", + "dropout", + "batch_size", + "weight_decay", + "grad_clip", + ] + } +} + +/// Metrics returned from TLOB hyperopt training. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct TLOBMetrics { + /// Validation loss (primary objective for optimization) + pub val_loss: f64, + /// Training loss at end of run + pub train_loss: f64, + /// Directional accuracy on validation set + pub directional_accuracy: f64, + /// Number of epochs completed + pub epochs_completed: usize, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bounds_count_matches_param_names() { + let bounds = TLOBParams::continuous_bounds(); + let names = TLOBParams::param_names(); + assert_eq!(bounds.len(), names.len()); + assert_eq!(bounds.len(), NUM_PARAMS); + } + + #[test] + fn test_roundtrip_continuous() { + let params = TLOBParams::default(); + let continuous = params.to_continuous(); + let restored = TLOBParams::from_continuous(&continuous).unwrap(); + assert!( + (params.learning_rate - restored.learning_rate).abs() < 1e-6, + "lr mismatch: {} vs {}", + params.learning_rate, + restored.learning_rate + ); + assert_eq!(params.d_model, restored.d_model); + assert_eq!(params.num_heads, restored.num_heads); + assert_eq!(params.num_layers, restored.num_layers); + assert_eq!(params.seq_len, restored.seq_len); + assert!((params.dropout - restored.dropout).abs() < 1e-6, "dropout mismatch"); + assert_eq!(params.batch_size, restored.batch_size); + assert!( + (params.weight_decay - restored.weight_decay).abs() / params.weight_decay < 1e-6, + "weight_decay mismatch" + ); + assert!( + (params.grad_clip - restored.grad_clip).abs() < 1e-6, + "grad_clip mismatch" + ); + } + + #[test] + fn test_from_continuous_wrong_length_errors() { + let result = TLOBParams::from_continuous(&[0.1, 0.2]); + assert!(result.is_err()); + let err_msg = format!("{}", result.unwrap_err()); + assert!(err_msg.contains("Expected 9"), "Error should mention expected count: {}", err_msg); + } + + #[test] + fn test_bounds_are_valid() { + for (i, (min, max)) in TLOBParams::continuous_bounds().iter().enumerate() { + assert!( + min < max, + "Invalid bounds for param {} ({}): {} >= {}", + i, + TLOBParams::param_names().get(i).copied().unwrap_or("?"), + min, + max + ); + } + } + + #[test] + fn test_default_within_bounds() { + let params = TLOBParams::default(); + let continuous = params.to_continuous(); + let bounds = TLOBParams::continuous_bounds(); + let names = TLOBParams::param_names(); + for (i, (val, (min, max))) in continuous.iter().zip(bounds.iter()).enumerate() { + assert!( + *val >= *min && *val <= *max, + "Param {} ({}) = {} outside [{}, {}]", + i, + names.get(i).copied().unwrap_or("?"), + val, + min, + max + ); + } + } + + #[test] + fn test_metrics_default() { + let metrics = TLOBMetrics::default(); + assert_eq!(metrics.val_loss, 0.0); + assert_eq!(metrics.train_loss, 0.0); + assert_eq!(metrics.directional_accuracy, 0.0); + assert_eq!(metrics.epochs_completed, 0); + } + + #[test] + fn test_continuous_length() { + let params = TLOBParams::default(); + let continuous = params.to_continuous(); + assert_eq!(continuous.len(), NUM_PARAMS); + } +} diff --git a/ml/src/integration/coordinator.rs b/ml/src/integration/coordinator.rs index b70e98f82..1dab9df61 100644 --- a/ml/src/integration/coordinator.rs +++ b/ml/src/integration/coordinator.rs @@ -471,6 +471,10 @@ impl EnsembleCoordinator { // Ensemble methods - use weighted combination approach self.deep_q_prediction(input_features, model.weight * 1.2) // Enhanced prediction for ensemble }, + ModelType::KAN => { + // Kolmogorov-Arnold Network prediction + self.deep_q_prediction(input_features, model.weight) + }, }; let latency_us = start_time.elapsed().as_micros() as u64; @@ -882,6 +886,7 @@ impl EnsembleCoordinator { ModelType::PPO => 0.83, // Good confidence for policy optimization ModelType::Transformer => 0.88, // High confidence for sequence modeling ModelType::Ensemble => 0.92, // Highest confidence for ensemble methods + ModelType::KAN => 0.86, // Good confidence for learned activation functions }; // Adjust confidence based on feature quality diff --git a/ml/src/kan/config.rs b/ml/src/kan/config.rs new file mode 100644 index 000000000..8c2c7b61e --- /dev/null +++ b/ml/src/kan/config.rs @@ -0,0 +1,42 @@ +//! KAN (Kolmogorov-Arnold Network) configuration. +//! +//! Defines hyperparameters for KAN models including B-spline grid size, +//! spline order, layer widths, and training parameters. + +use serde::{Deserialize, Serialize}; + +/// Configuration for a Kolmogorov-Arnold Network. +/// +/// KAN replaces fixed activation functions (ReLU) with learnable B-spline +/// activations on each edge, enabling the network to discover arbitrary +/// non-linear relationships. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KANConfig { + /// Number of B-spline grid points per activation function. + /// More points = finer approximation but more parameters. + pub grid_size: usize, + /// B-spline order (degree + 1). 4 = cubic splines. + pub spline_order: usize, + /// Layer widths including input and output dimensions. + /// Example: `[51, 32, 16, 1]` means input=51, two hidden layers, output=1. + pub layer_widths: Vec, + /// Learning rate for AdamW optimizer. + pub learning_rate: f64, + /// L2 weight decay for regularization. + pub weight_decay: f64, + /// Maximum gradient norm for clipping. + pub grad_clip: f64, +} + +impl Default for KANConfig { + fn default() -> Self { + Self { + grid_size: 5, + spline_order: 4, + layer_widths: vec![51, 32, 16, 1], + learning_rate: 1e-3, + weight_decay: 1e-4, + grad_clip: 1.0, + } + } +} diff --git a/ml/src/kan/layer.rs b/ml/src/kan/layer.rs new file mode 100644 index 000000000..24eb5625e --- /dev/null +++ b/ml/src/kan/layer.rs @@ -0,0 +1,151 @@ +//! KAN layer implementation. +//! +//! A single KAN layer maps `(batch, in_dim)` to `(batch, out_dim)` by evaluating +//! learned B-spline activations on each input dimension and combining them with +//! learnable coefficients plus a linear residual connection. + +use candle_core::Tensor; +use candle_nn::VarBuilder; +use crate::MLError; + +use super::spline::BSplineBasis; + +/// A single Kolmogorov-Arnold Network layer. +/// +/// Forward pass: +/// 1. Expand each input dim through B-spline basis: `(batch, in_dim)` -> `(batch, in_dim * num_bases)` +/// 2. Multiply by coefficients: `(batch, in_dim * num_bases) @ coefficients(in_dim * num_bases, out_dim)` -> `(batch, out_dim)` +/// 3. Add residual: `input @ residual_weight(in_dim, out_dim)` -> `(batch, out_dim)` +#[derive(Debug)] +pub struct KANLayer { + /// B-spline basis evaluator (shared across all input dims) + spline_basis: BSplineBasis, + /// Learnable spline coefficients: (in_dim * num_bases, out_dim) + coefficients: Tensor, + /// Learnable residual weights: (in_dim, out_dim) + residual_weight: Tensor, + /// Input dimension + in_dim: usize, + /// Number of basis functions + num_bases: usize, +} + +impl KANLayer { + /// Create a new KAN layer. + /// + /// # Arguments + /// * `in_dim` - Input dimension + /// * `out_dim` - Output dimension + /// * `grid_size` - B-spline grid points + /// * `spline_order` - B-spline order (degree + 1) + /// * `vb` - VarBuilder for creating learnable parameters + pub fn new( + in_dim: usize, + out_dim: usize, + grid_size: usize, + spline_order: usize, + vb: VarBuilder<'_>, + ) -> Result { + let spline_basis = BSplineBasis::new(grid_size, spline_order, -2.0, 2.0)?; + let num_bases = spline_basis.num_bases(); + + // Xavier-like init scale + let scale = (2.0 / (in_dim * num_bases + out_dim) as f64).sqrt(); + + let coefficients = vb + .get_with_hints( + (in_dim * num_bases, out_dim), + "coefficients", + candle_nn::Init::Randn { mean: 0.0, stdev: scale }, + ) + .map_err(|e| MLError::ModelError(format!("Failed to create coefficients: {}", e)))?; + + let residual_weight = vb + .get_with_hints( + (in_dim, out_dim), + "residual", + candle_nn::Init::Randn { + mean: 0.0, + stdev: (1.0 / in_dim as f64).sqrt(), + }, + ) + .map_err(|e| MLError::ModelError(format!("Failed to create residual weight: {}", e)))?; + + Ok(Self { + spline_basis, + coefficients, + residual_weight, + in_dim, + num_bases, + }) + } + + /// Forward pass: `(batch, in_dim)` -> `(batch, out_dim)`. + pub fn forward(&self, input: &Tensor) -> Result { + let device = input.device(); + let dims = input.shape().dims(); + let batch = *dims.first().ok_or_else(|| { + MLError::ModelError("Input tensor has no dimensions".to_string()) + })?; + + // Flatten input to (batch * in_dim,) for basis evaluation + let flat_input = input.reshape(&[batch * self.in_dim]).map_err(|e| { + MLError::ModelError(format!("KAN layer reshape failed: {}", e)) + })?; + + // Evaluate B-spline basis: (batch * in_dim, num_bases) + let basis_vals = self.spline_basis.evaluate(&flat_input, device)?; + + // Reshape to (batch, in_dim * num_bases) + let basis_flat = basis_vals + .reshape(&[batch, self.in_dim * self.num_bases]) + .map_err(|e| { + MLError::ModelError(format!("KAN basis reshape failed: {}", e)) + })?; + + // Spline contribution: (batch, in_dim * num_bases) @ (in_dim * num_bases, out_dim) + let spline_out = basis_flat.matmul(&self.coefficients).map_err(|e| { + MLError::ModelError(format!("KAN spline matmul failed: {}", e)) + })?; + + // Residual contribution: (batch, in_dim) @ (in_dim, out_dim) + let residual_out = input.matmul(&self.residual_weight).map_err(|e| { + MLError::ModelError(format!("KAN residual matmul failed: {}", e)) + })?; + + // Combine + spline_out.add(&residual_out).map_err(|e| { + MLError::ModelError(format!("KAN layer add failed: {}", e)) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::{DType, Device}; + use candle_nn::VarMap; + + #[test] + fn test_kan_layer_output_shape() { + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, &Device::Cpu); + let layer = KANLayer::new(4, 8, 5, 4, vb.pp("layer0")).unwrap(); + + let input = Tensor::randn(0.0f32, 0.5, &[2, 4], &Device::Cpu).unwrap(); + let output = layer.forward(&input).unwrap(); + let dims = output.shape().dims(); + assert_eq!(dims, &[2, 8]); + } + + #[test] + fn test_kan_layer_learnable_params() { + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, &Device::Cpu); + let _layer = KANLayer::new(4, 8, 5, 4, vb.pp("layer0")).unwrap(); + + let all_vars = var_map.all_vars(); + // Should have coefficients + residual = 2 tensors + assert_eq!(all_vars.len(), 2, "Expected 2 learnable tensors, got {}", all_vars.len()); + } +} diff --git a/ml/src/kan/mod.rs b/ml/src/kan/mod.rs new file mode 100644 index 000000000..3bb8168bf --- /dev/null +++ b/ml/src/kan/mod.rs @@ -0,0 +1,15 @@ +//! KAN (Kolmogorov-Arnold Network) module. +//! +//! Implements KAN with learnable B-spline activation functions on each edge +//! of the network, replacing fixed activations (ReLU) with data-driven +//! non-linearities. + +pub mod config; +pub mod layer; +pub mod network; +pub mod spline; +pub mod trainable; + +pub use config::KANConfig; +pub use network::KANNetwork; +pub use trainable::KANTrainableAdapter; diff --git a/ml/src/kan/network.rs b/ml/src/kan/network.rs new file mode 100644 index 000000000..b3bb90146 --- /dev/null +++ b/ml/src/kan/network.rs @@ -0,0 +1,121 @@ +//! KAN network: a stack of KAN layers. +//! +//! Chains multiple KAN layers according to `layer_widths` from the config. +//! No explicit activation between layers because each KAN layer contains +//! its own learned activation functions (B-splines). + +use candle_core::Tensor; +use candle_nn::VarBuilder; +use crate::MLError; + +use super::config::KANConfig; +use super::layer::KANLayer; + +/// A multi-layer Kolmogorov-Arnold Network. +#[derive(Debug)] +pub struct KANNetwork { + layers: Vec, +} + +impl KANNetwork { + /// Build a KAN network from config. + /// + /// Creates one KAN layer per consecutive pair of widths in `config.layer_widths`. + /// For example, widths `[51, 32, 16, 1]` creates 3 layers: + /// 51->32, 32->16, 16->1. + pub fn new(config: &KANConfig, vb: VarBuilder<'_>) -> Result { + if config.layer_widths.len() < 2 { + return Err(MLError::ConfigError { + reason: "layer_widths must have at least 2 entries".to_string(), + }); + } + + let mut layers = Vec::with_capacity(config.layer_widths.len() - 1); + for i in 0..config.layer_widths.len() - 1 { + let in_dim = *config.layer_widths.get(i).ok_or_else(|| MLError::ConfigError { + reason: format!("Missing layer_widths index {}", i), + })?; + let out_dim = *config.layer_widths.get(i + 1).ok_or_else(|| MLError::ConfigError { + reason: format!("Missing layer_widths index {}", i + 1), + })?; + + let layer = KANLayer::new( + in_dim, + out_dim, + config.grid_size, + config.spline_order, + vb.pp(format!("layer_{}", i)), + )?; + layers.push(layer); + } + + Ok(Self { layers }) + } + + /// Forward pass: chains all layers sequentially. + /// + /// Input shape: `(batch, layer_widths[0])` + /// Output shape: `(batch, layer_widths[last])` + pub fn forward(&self, input: &Tensor) -> Result { + let mut x = input.clone(); + for layer in &self.layers { + x = layer.forward(&x)?; + } + Ok(x) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::{DType, Device}; + use candle_nn::VarMap; + + fn make_config() -> KANConfig { + KANConfig { + grid_size: 3, + spline_order: 3, + layer_widths: vec![8, 4, 1], + learning_rate: 1e-3, + weight_decay: 1e-4, + grad_clip: 1.0, + } + } + + #[test] + fn test_kan_network_forward() { + let config = make_config(); + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, &Device::Cpu); + let net = KANNetwork::new(&config, vb).unwrap(); + + let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap(); + let output = net.forward(&input).unwrap(); + let dims = output.shape().dims(); + assert_eq!(dims, &[2, 1], "Expected (2, 1), got {:?}", dims); + } + + #[test] + fn test_kan_network_produces_gradients() { + let config = make_config(); + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, &Device::Cpu); + let net = KANNetwork::new(&config, vb).unwrap(); + + let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap(); + let output = net.forward(&input).unwrap(); + let loss = output.sqr().unwrap().mean_all().unwrap(); + let grads = loss.backward().unwrap(); + + // Check that at least some variables received gradients + let vars_lock = var_map.data().lock().unwrap(); + let mut has_grad = false; + for (_name, var) in vars_lock.iter() { + if grads.get(var.as_tensor()).is_some() { + has_grad = true; + break; + } + } + assert!(has_grad, "No gradients found for any variable"); + } +} diff --git a/ml/src/kan/spline.rs b/ml/src/kan/spline.rs new file mode 100644 index 000000000..5ffa59446 --- /dev/null +++ b/ml/src/kan/spline.rs @@ -0,0 +1,188 @@ +//! B-spline basis function evaluator for KAN. +//! +//! Implements Cox-de Boor recursion on CPU to compute B-spline basis values. +//! Knots are stored as `Vec` (no gradient needed for the knot vector). +//! The evaluator produces a `(batch, num_bases)` tensor for use in KAN layers. + +use candle_core::{DType, Device, Tensor}; +use crate::MLError; + +/// B-spline basis evaluator. +/// +/// Stores a uniform knot vector and evaluates all basis functions for a +/// batch of scalar inputs using the Cox-de Boor recursion. +#[derive(Debug, Clone)] +pub struct BSplineBasis { + /// Uniform knot vector (length = grid_size + 2 * order). + knots: Vec, + /// Spline order (degree + 1). + order: usize, + /// Number of basis functions = grid_size + order - 1. + num_bases: usize, +} + +impl BSplineBasis { + /// Create a new B-spline basis with uniform knots over `[x_min, x_max]`. + /// + /// # Arguments + /// * `grid_size` - Number of interior grid points + /// * `order` - Spline order (degree + 1), e.g. 4 for cubic + /// * `x_min` - Left boundary of the domain + /// * `x_max` - Right boundary of the domain + pub fn new(grid_size: usize, order: usize, x_min: f64, x_max: f64) -> Result { + if grid_size == 0 || order == 0 { + return Err(MLError::ConfigError { + reason: "grid_size and order must be > 0".to_string(), + }); + } + if x_max <= x_min { + return Err(MLError::ConfigError { + reason: format!("x_max ({}) must be > x_min ({})", x_max, x_min), + }); + } + + let num_bases = grid_size + order - 1; + // Total knots: num_bases + order = grid_size + 2*order - 1 + let n_knots = num_bases + order; + let mut knots = Vec::with_capacity(n_knots); + + // Extend domain slightly beyond [x_min, x_max] for boundary padding. + let step = (x_max - x_min) / grid_size as f64; + let pad = (order as f64 - 1.0) * step; + let start = x_min - pad; + + for i in 0..n_knots { + knots.push((start + i as f64 * step) as f32); + } + + Ok(Self { + knots, + order, + num_bases, + }) + } + + /// Number of basis functions produced. + pub fn num_bases(&self) -> usize { + self.num_bases + } + + /// Evaluate all basis functions for a batch of scalars. + /// + /// # Arguments + /// * `x` - Tensor of shape `(n,)` containing scalar inputs + /// * `device` - Device to place the output on + /// + /// # Returns + /// Tensor of shape `(n, num_bases)` with basis values. + pub fn evaluate(&self, x: &Tensor, device: &Device) -> Result { + // Flatten to 1-D regardless of input shape + let x_flat = x.flatten_all().map_err(|e| { + MLError::ModelError(format!("BSpline flatten failed: {}", e)) + })?; + let n = x_flat.dims1().map_err(|e| { + MLError::ModelError(format!("BSpline dims failed: {}", e)) + })?; + let x_vec: Vec = x_flat.to_vec1().map_err(|e| { + MLError::ModelError(format!("BSpline to_vec1 failed: {}", e)) + })?; + + // Compute basis values on CPU via Cox-de Boor + let mut result = vec![0.0f32; n * self.num_bases]; + for (i, &xi) in x_vec.iter().enumerate() { + for j in 0..self.num_bases { + let val = self.cox_de_boor(xi, j, self.order); + // Safe index: i * num_bases + j is always within bounds + if let Some(slot) = result.get_mut(i * self.num_bases + j) { + *slot = val; + } + } + } + + Tensor::from_vec(result, &[n, self.num_bases], device) + .map_err(|e| MLError::ModelError(format!("BSpline tensor creation failed: {}", e)))? + .to_dtype(DType::F32) + .map_err(|e| MLError::ModelError(format!("BSpline dtype conversion failed: {}", e))) + } + + /// Cox-de Boor recursion for basis function `j` of given `order` at point `x`. + fn cox_de_boor(&self, x: f32, j: usize, k: usize) -> f32 { + if k == 1 { + let t_j = self.knots.get(j).copied().unwrap_or(0.0); + let t_j1 = self.knots.get(j + 1).copied().unwrap_or(0.0); + return if x >= t_j && x < t_j1 { 1.0 } else { 0.0 }; + } + + let t_j = self.knots.get(j).copied().unwrap_or(0.0); + let t_jk = self.knots.get(j + k).copied().unwrap_or(0.0); + let t_jk1 = self.knots.get(j + k - 1).copied().unwrap_or(0.0); + let t_j1 = self.knots.get(j + 1).copied().unwrap_or(0.0); + + let denom_left = t_jk1 - t_j; + let left = if denom_left.abs() > 1e-10 { + ((x - t_j) / denom_left) * self.cox_de_boor(x, j, k - 1) + } else { + 0.0 + }; + + let denom_right = t_jk - t_j1; + let right = if denom_right.abs() > 1e-10 { + ((t_jk - x) / denom_right) * self.cox_de_boor(x, j + 1, k - 1) + } else { + 0.0 + }; + + left + right + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bspline_basis_shape() { + let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap(); + assert_eq!(basis.num_bases(), 8); // 5 + 4 - 1 + + let x = Tensor::new(&[0.0f32, 0.5, -0.5, 0.1], &Device::Cpu).unwrap(); + let out = basis.evaluate(&x, &Device::Cpu).unwrap(); + let dims = out.shape().dims(); + assert_eq!(dims.len(), 2); + assert_eq!(dims[0], 4); + assert_eq!(dims[1], 8); + } + + #[test] + fn test_bspline_partition_of_unity() { + let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap(); + // Points well inside the domain should sum close to 1.0 + let x = Tensor::new(&[0.0f32, 0.3, -0.3], &Device::Cpu).unwrap(); + let out = basis.evaluate(&x, &Device::Cpu).unwrap(); + let sums = out.sum(1).unwrap(); + let sums_vec: Vec = sums.to_vec1().unwrap(); + for (i, &s) in sums_vec.iter().enumerate() { + assert!( + (s - 1.0).abs() < 0.15, + "Basis sum at index {} = {} (expected ~1.0)", + i, + s + ); + } + } + + #[test] + fn test_bspline_non_negative() { + let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap(); + let x = Tensor::new(&[0.0f32, 0.5, -0.5, 0.9], &Device::Cpu).unwrap(); + let out = basis.evaluate(&x, &Device::Cpu).unwrap(); + let flat: Vec = out.flatten_all().unwrap().to_vec1().unwrap(); + for &v in &flat { + assert!( + v >= -1e-6, + "Basis value {} is negative (< -epsilon)", + v + ); + } + } +} diff --git a/ml/src/kan/trainable.rs b/ml/src/kan/trainable.rs new file mode 100644 index 000000000..813fea0f3 --- /dev/null +++ b/ml/src/kan/trainable.rs @@ -0,0 +1,493 @@ +//! UnifiedTrainable adapter for KAN (Kolmogorov-Arnold Network). +//! +//! Wraps a KANNetwork with VarMap + AdamW optimizer to provide the +//! UnifiedTrainable interface. Follows the same pattern as TGGNTrainableAdapter. + +use candle_core::{backprop::GradStore, DType, Device, Tensor}; +use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap}; +use std::collections::HashMap; + +use super::config::KANConfig; +use super::network::KANNetwork; +use crate::training::unified_trainer::{checkpoint, CheckpointMetadata, TrainingMetrics, UnifiedTrainable}; +use crate::MLError; + +/// KAN trainable adapter implementing UnifiedTrainable. +pub struct KANTrainableAdapter { + config: KANConfig, + var_map: VarMap, + network: KANNetwork, + optimizer: AdamW, + grads: Option, + device: Device, + learning_rate: f64, + step: usize, + latest_metrics: TrainingMetrics, + loss_history: Vec, + last_grad_norm: f64, +} + +impl std::fmt::Debug for KANTrainableAdapter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KANTrainableAdapter") + .field("config", &self.config) + .field("device", &format!("{:?}", self.device)) + .field("learning_rate", &self.learning_rate) + .field("step", &self.step) + .field("loss_history_len", &self.loss_history.len()) + .field("last_grad_norm", &self.last_grad_norm) + .finish_non_exhaustive() + } +} + +impl KANTrainableAdapter { + /// Create a new KAN trainable adapter. + pub fn new(config: KANConfig, device: &Device) -> Result { + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, device); + + let network = KANNetwork::new(&config, vb)?; + + let learning_rate = config.learning_rate; + let weight_decay = config.weight_decay; + let all_vars = var_map.all_vars(); + let optimizer = AdamW::new( + all_vars, + ParamsAdamW { + lr: learning_rate, + beta1: 0.9, + beta2: 0.999, + eps: 1e-8, + weight_decay, + }, + ) + .map_err(|e| MLError::ModelError(format!("Failed to create AdamW optimizer: {}", e)))?; + + Ok(Self { + config, + var_map, + network, + optimizer, + grads: None, + device: device.clone(), + learning_rate, + step: 0, + latest_metrics: TrainingMetrics::default(), + loss_history: Vec::new(), + last_grad_norm: 0.0, + }) + } + + /// Access the KAN configuration. + pub fn kan_config(&self) -> &KANConfig { + &self.config + } +} + +impl UnifiedTrainable for KANTrainableAdapter { + fn model_type(&self) -> &str { + "KAN" + } + + fn device(&self) -> &Device { + &self.device + } + + fn forward(&mut self, input: &Tensor) -> Result { + self.network.forward(input) + } + + fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result { + let diff = predictions.sub(targets).map_err(|e| { + MLError::ModelError(format!("Loss subtraction failed: {}", e)) + })?; + let squared = diff.powf(2.0).map_err(|e| { + MLError::ModelError(format!("Loss squaring failed: {}", e)) + })?; + squared.mean_all().map_err(|e| { + MLError::ModelError(format!("Loss mean failed: {}", e)) + }) + } + + fn backward(&mut self, loss: &Tensor) -> Result { + let grads = loss.backward().map_err(|e| { + MLError::TrainingError(format!("Backward pass failed: {}", e)) + })?; + + let mut grad_norm_sq = 0.0; + let vars_lock = self + .var_map + .data() + .lock() + .map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?; + + for (_name, var) in vars_lock.iter() { + if let Some(grad) = grads.get(var.as_tensor()) { + let norm = grad + .sqr() + .and_then(|s| s.sum_all()) + .and_then(|s| s.to_scalar::()) + .map_err(|e| { + MLError::ModelError(format!("Failed to compute grad norm: {}", e)) + })?; + grad_norm_sq += norm as f64; + } + } + drop(vars_lock); + + let grad_norm = grad_norm_sq.sqrt(); + self.last_grad_norm = grad_norm; + self.latest_metrics.grad_norm = Some(grad_norm); + self.grads = Some(grads); + + Ok(grad_norm) + } + + fn optimizer_step(&mut self) -> Result<(), MLError> { + if let Some(grads) = self.grads.take() { + self.optimizer.step(&grads).map_err(|e| { + MLError::TrainingError(format!("Optimizer step failed: {}", e)) + })?; + } + self.step += 1; + Ok(()) + } + + fn zero_grad(&mut self) -> Result<(), MLError> { + Ok(()) + } + + fn get_learning_rate(&self) -> f64 { + self.learning_rate + } + + fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> { + self.learning_rate = lr; + let all_vars = self.var_map.all_vars(); + self.optimizer = AdamW::new( + all_vars, + ParamsAdamW { + lr, + beta1: 0.9, + beta2: 0.999, + eps: 1e-8, + weight_decay: self.config.weight_decay, + }, + ) + .map_err(|e| { + MLError::ModelError(format!("Failed to recreate optimizer with new lr: {}", e)) + })?; + Ok(()) + } + + fn get_step(&self) -> usize { + self.step + } + + fn collect_metrics(&self) -> TrainingMetrics { + let mut metrics = self.latest_metrics.clone(); + + if !self.loss_history.is_empty() { + let recent: Vec = self.loss_history.iter().rev().take(100).copied().collect(); + metrics.loss = recent.iter().sum::() / recent.len() as f64; + } + + metrics.learning_rate = self.learning_rate; + metrics.custom_metrics.insert("training_steps".to_string(), self.step as f64); + metrics.custom_metrics.insert("loss_history_len".to_string(), self.loss_history.len() as f64); + metrics.custom_metrics.insert("grid_size".to_string(), self.config.grid_size as f64); + metrics.custom_metrics.insert("spline_order".to_string(), self.config.spline_order as f64); + metrics.custom_metrics.insert("num_layers".to_string(), (self.config.layer_widths.len().saturating_sub(1)) as f64); + + metrics + } + + fn save_checkpoint(&self, checkpoint_path: &str) -> Result { + let metadata = CheckpointMetadata { + model_type: "KAN".to_string(), + version: "1.0.0".to_string(), + epoch: 0, + step: self.step, + timestamp: std::time::SystemTime::now(), + config: serde_json::to_value(&self.config).map_err(|e| { + MLError::SerializationError { + reason: format!("Failed to serialize config: {}", e), + } + })?, + metrics: self.collect_metrics(), + }; + + checkpoint::save_metadata(&metadata, checkpoint_path)?; + + let safetensors_path = format!("{}.safetensors", checkpoint_path); + let vars_lock = self + .var_map + .data() + .lock() + .map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?; + + let mut tensors: HashMap = HashMap::new(); + for (name, var) in vars_lock.iter() { + tensors.insert(name.clone(), var.as_tensor().clone()); + } + drop(vars_lock); + + candle_core::safetensors::save(&tensors, &safetensors_path).map_err(|e| { + MLError::CheckpointError(format!("Failed to save safetensors: {}", e)) + })?; + + tracing::info!("Saved KAN checkpoint to {} (step {})", checkpoint_path, self.step); + + Ok(checkpoint_path.to_string()) + } + + fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result { + let metadata = checkpoint::load_metadata(checkpoint_path)?; + + if metadata.model_type != "KAN" { + return Err(MLError::CheckpointError(format!( + "Invalid model type in checkpoint: expected KAN, got {}", + metadata.model_type + ))); + } + + let safetensors_path = format!("{}.safetensors", checkpoint_path); + let tensors = candle_core::safetensors::load(&safetensors_path, &self.device).map_err( + |e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)), + )?; + + let vars_lock = self + .var_map + .data() + .lock() + .map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?; + + for (name, tensor) in &tensors { + if let Some(var) = vars_lock.get(name) { + var.set(tensor).map_err(|e| { + MLError::CheckpointError(format!("Failed to set var {}: {}", name, e)) + })?; + } else { + tracing::warn!("Checkpoint contains unknown variable: {}", name); + } + } + drop(vars_lock); + + self.step = metadata.step; + self.latest_metrics = metadata.metrics.clone(); + + tracing::info!("Loaded KAN checkpoint from {} (step {})", checkpoint_path, metadata.step); + + Ok(metadata) + } + + fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result { + if val_data.is_empty() { + return Err(MLError::ValidationError { + message: "Empty validation dataset".to_string(), + }); + } + + let mut total_loss = 0.0; + let mut count = 0usize; + + for (input, target) in val_data { + let prediction = self.forward(input)?; + let loss = self.compute_loss(&prediction, target)?; + let loss_val = loss.to_scalar::().map_err(|e| MLError::ValidationError { + message: format!("Failed to extract loss value: {}", e), + })?; + total_loss += loss_val as f64; + count += 1; + } + + let avg_loss = if count > 0 { + total_loss / count as f64 + } else { + 0.0 + }; + + self.latest_metrics.val_loss = Some(avg_loss); + tracing::debug!("KAN validation loss: {:.6}", avg_loss); + + Ok(avg_loss) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_config() -> KANConfig { + KANConfig { + grid_size: 3, + spline_order: 3, + layer_widths: vec![8, 4, 1], + learning_rate: 1e-3, + weight_decay: 1e-4, + grad_clip: 1.0, + } + } + + #[test] + fn test_model_type() { + let cfg = make_config(); + let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + assert_eq!(adapter.model_type(), "KAN"); + } + + #[test] + fn test_device() { + let cfg = make_config(); + let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + assert!(matches!(adapter.device(), &Device::Cpu)); + } + + #[test] + fn test_forward_shape() { + let cfg = make_config(); + let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap(); + let output = adapter.forward(&input).unwrap(); + + let dims = output.shape().dims(); + assert_eq!(dims.len(), 2); + assert_eq!(dims[0], 2); + assert_eq!(dims[1], 1); + } + + #[test] + fn test_compute_loss() { + let cfg = make_config(); + let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let preds = Tensor::new(&[[1.0f32], [2.0]], &Device::Cpu).unwrap(); + let targets = Tensor::new(&[[1.5f32], [2.5]], &Device::Cpu).unwrap(); + + let loss = adapter.compute_loss(&preds, &targets).unwrap(); + let loss_val: f32 = loss.to_scalar().unwrap(); + assert!((loss_val - 0.25).abs() < 1e-5, "Expected ~0.25, got {}", loss_val); + } + + #[test] + fn test_backward_returns_grad_norm() { + let cfg = make_config(); + let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap(); + let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + + let output = adapter.forward(&input).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + let grad_norm = adapter.backward(&loss).unwrap(); + + assert!(grad_norm >= 0.0, "Gradient norm should be non-negative"); + assert!(adapter.grads.is_some(), "Grads should be stored after backward"); + } + + #[test] + fn test_train_step_cycle() { + let cfg = make_config(); + let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + assert_eq!(adapter.get_step(), 0); + + let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap(); + let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + + let output = adapter.forward(&input).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + let loss_val: f64 = loss.to_scalar::().unwrap() as f64; + + adapter.backward(&loss).unwrap(); + adapter.optimizer_step().unwrap(); + + assert_eq!(adapter.get_step(), 1); + + adapter.loss_history.push(loss_val); + let metrics = adapter.collect_metrics(); + assert!(metrics.loss >= 0.0); + assert_eq!(metrics.custom_metrics.get("training_steps").copied(), Some(1.0)); + } + + #[test] + fn test_learning_rate_get_set() { + let cfg = make_config(); + let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let original_lr = adapter.get_learning_rate(); + assert!((original_lr - 1e-3).abs() < 1e-10); + + adapter.set_learning_rate(5e-4).unwrap(); + assert!((adapter.get_learning_rate() - 5e-4).abs() < 1e-10); + } + + #[test] + fn test_collect_metrics() { + let cfg = make_config(); + let adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let metrics = adapter.collect_metrics(); + assert!(metrics.custom_metrics.contains_key("training_steps")); + assert!(metrics.custom_metrics.contains_key("grid_size")); + assert!(metrics.custom_metrics.contains_key("spline_order")); + assert!(metrics.custom_metrics.contains_key("num_layers")); + assert_eq!(metrics.custom_metrics.get("grid_size").copied(), Some(3.0)); + } + + #[test] + fn test_checkpoint_roundtrip() { + let cfg = make_config(); + let mut adapter = KANTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap(); + + let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap(); + let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + let output = adapter.forward(&input).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + adapter.backward(&loss).unwrap(); + adapter.optimizer_step().unwrap(); + + let tmp_dir = std::env::temp_dir(); + let checkpoint_path = tmp_dir.join("kan_test_ckpt"); + let path_str = checkpoint_path.to_str().unwrap(); + adapter.save_checkpoint(path_str).unwrap(); + + let mut adapter2 = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + let metadata = adapter2.load_checkpoint(path_str).unwrap(); + + assert_eq!(metadata.model_type, "KAN"); + assert_eq!(metadata.step, 1); + assert_eq!(adapter2.get_step(), 1); + + // Clean up + let _ = std::fs::remove_file(format!("{}.json", path_str)); + let _ = std::fs::remove_file(format!("{}.safetensors", path_str)); + } + + #[test] + fn test_validate() { + let cfg = make_config(); + let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let val_data: Vec<(Tensor, Tensor)> = (0..3) + .map(|_| { + let input = Tensor::randn(0.0f32, 0.5, &[2, 8], &Device::Cpu).unwrap(); + let target = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + (input, target) + }) + .collect(); + + let val_loss = adapter.validate(&val_data).unwrap(); + assert!(val_loss >= 0.0, "Validation loss should be non-negative"); + assert!(adapter.latest_metrics.val_loss.is_some(), "val_loss should be set"); + } + + #[test] + fn test_validate_empty_errors() { + let cfg = make_config(); + let mut adapter = KANTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + let result = adapter.validate(&[]); + assert!(result.is_err(), "Validating empty data should error"); + } +} diff --git a/ml/src/lib.rs b/ml/src/lib.rs index f24ab35b5..65f6bc223 100644 --- a/ml/src/lib.rs +++ b/ml/src/lib.rs @@ -692,6 +692,7 @@ pub mod evaluation; // DQN evaluation engine (backtest metrics, Sharpe ratio) pub mod flash_attention; pub mod hyperopt; // Bayesian hyperparameter optimization (egobox) pub mod integration; +pub mod kan; pub mod labeling; pub mod liquid; pub mod mamba; diff --git a/ml/src/tlob/mod.rs b/ml/src/tlob/mod.rs index 96d342064..59495bce4 100644 --- a/ml/src/tlob/mod.rs +++ b/ml/src/tlob/mod.rs @@ -7,6 +7,7 @@ pub mod analytics; pub mod features; pub mod mbp10_feature_extractor; // MBP-10 to TLOB feature extraction pub mod performance; +pub mod trainable_adapter; pub mod transformer; // Re-export key types for external use @@ -17,6 +18,7 @@ pub use features::{ TLOBFeatures as TLOBInputFeatures, TLOB_FEATURE_COUNT, }; +pub use trainable_adapter::{TLOBAdapterConfig, TLOBTrainableAdapter}; pub use transformer::{TLOBConfig, TLOBMetrics, TLOBTransformer}; // Re-export transformer-specific TLOBFeatures with a different name to avoid conflicts diff --git a/ml/src/tlob/trainable_adapter.rs b/ml/src/tlob/trainable_adapter.rs new file mode 100644 index 000000000..4b0790445 --- /dev/null +++ b/ml/src/tlob/trainable_adapter.rs @@ -0,0 +1,671 @@ +//! UnifiedTrainable adapter for TLOB (Time Limit Order Book) +//! +//! Wraps a candle-based projection network to provide the UnifiedTrainable +//! interface for TLOB. The projection network maps flattened order book features +//! through hidden layers, enabling gradient-based training via the unified +//! training orchestrator. +//! +//! Architecture: input_linear(seq_len*feature_dim -> d_model) -> ReLU -> output_linear(d_model -> 1) + +use candle_core::{backprop::GradStore, DType, Device, Module, Tensor}; +use candle_nn::{linear, AdamW, Linear, Optimizer, ParamsAdamW, VarBuilder, VarMap}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +use crate::training::unified_trainer::{checkpoint, CheckpointMetadata, TrainingMetrics, UnifiedTrainable}; +use crate::MLError; + +/// Configuration for the TLOB trainable adapter projection network. +/// +/// This is separate from the ONNX-oriented `TLOBConfig` in `transformer.rs`. +/// It defines the architecture for the candle-based projection network used +/// by the unified training orchestrator. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TLOBAdapterConfig { + /// Model / hidden dimension for projection layers + pub d_model: usize, + /// Number of attention heads (reserved for future transformer layers) + pub num_heads: usize, + /// Number of projection layers (reserved for future deeper networks) + pub num_layers: usize, + /// Sequence length (number of time steps in LOB snapshot) + pub seq_len: usize, + /// Feature dimension per time step (e.g. 51 for TLOB features) + pub feature_dim: usize, +} + +impl Default for TLOBAdapterConfig { + fn default() -> Self { + Self { + d_model: 128, + num_heads: 4, + num_layers: 2, + seq_len: 128, + feature_dim: 51, + } + } +} + +/// Adapter wrapping TLOB with a candle-based projection network for unified training. +/// +/// The projection network has two linear layers: +/// - `input_linear`: projects from `seq_len * feature_dim` to `d_model` +/// - `output_linear`: projects from `d_model` to 1 (scalar prediction) +/// +/// Training uses AdamW with gradient tracking via `GradStore`. +pub struct TLOBTrainableAdapter { + /// TLOB adapter configuration + config: TLOBAdapterConfig, + /// Candle variable map holding learnable parameters + var_map: VarMap, + /// Input projection layer (seq_len*feature_dim -> d_model) + input_linear: Linear, + /// Output projection layer (d_model -> 1) + output_linear: Linear, + /// AdamW optimizer + optimizer: AdamW, + /// Gradient store from last backward pass (consumed by optimizer_step) + grads: Option, + /// Device (CPU or CUDA) + device: Device, + /// Current learning rate + learning_rate: f64, + /// Current training step + step: usize, + /// Latest training metrics + latest_metrics: TrainingMetrics, + /// Loss history for rolling average + loss_history: Vec, + /// Last computed gradient norm + last_grad_norm: f64, +} + +impl std::fmt::Debug for TLOBTrainableAdapter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TLOBTrainableAdapter") + .field("config", &self.config) + .field("device", &format!("{:?}", self.device)) + .field("learning_rate", &self.learning_rate) + .field("step", &self.step) + .field("loss_history_len", &self.loss_history.len()) + .field("last_grad_norm", &self.last_grad_norm) + .finish_non_exhaustive() + } +} + +impl TLOBTrainableAdapter { + /// Create a new TLOB trainable adapter with projection network. + /// + /// # Arguments + /// * `config` - TLOB adapter configuration specifying dimensions + /// * `device` - Device to create tensors on (CPU or CUDA) + /// + /// # Returns + /// Initialized adapter ready for training + pub fn new(config: TLOBAdapterConfig, device: &Device) -> Result { + let var_map = VarMap::new(); + let vb = VarBuilder::from_varmap(&var_map, DType::F32, device); + + let input_dim = config.seq_len * config.feature_dim; + + let input_linear = linear(input_dim, config.d_model, vb.pp("input")) + .map_err(|e| MLError::ModelError(format!("Failed to create input linear: {}", e)))?; + + let output_linear = linear(config.d_model, 1, vb.pp("output")) + .map_err(|e| MLError::ModelError(format!("Failed to create output linear: {}", e)))?; + + let learning_rate = 1e-3; + let all_vars = var_map.all_vars(); + let optimizer = AdamW::new( + all_vars, + ParamsAdamW { + lr: learning_rate, + beta1: 0.9, + beta2: 0.999, + eps: 1e-8, + weight_decay: 1e-4, + }, + ) + .map_err(|e| MLError::ModelError(format!("Failed to create AdamW optimizer: {}", e)))?; + + Ok(Self { + config, + var_map, + input_linear, + output_linear, + optimizer, + grads: None, + device: device.clone(), + learning_rate, + step: 0, + latest_metrics: TrainingMetrics::default(), + loss_history: Vec::new(), + last_grad_norm: 0.0, + }) + } + + /// Access the underlying TLOB adapter configuration. + pub fn tlob_config(&self) -> &TLOBAdapterConfig { + &self.config + } +} + +impl UnifiedTrainable for TLOBTrainableAdapter { + fn model_type(&self) -> &str { + "TLOB" + } + + fn device(&self) -> &Device { + &self.device + } + + fn forward(&mut self, input: &Tensor) -> Result { + // input shape: [batch, seq_len, feature_dim] or [batch, seq_len*feature_dim] + // Flatten to [batch, seq_len*feature_dim] if 3D + let dims = input.shape().dims(); + let flat_input = if dims.len() == 3 { + let batch = dims.first().copied().ok_or_else(|| { + MLError::ModelError("Empty input dimensions".to_string()) + })?; + let flat_dim = self.config.seq_len * self.config.feature_dim; + input.reshape(&[batch, flat_dim]).map_err(|e| { + MLError::ModelError(format!("Failed to reshape input to 2D: {}", e)) + })? + } else { + input.clone() + }; + + // input_linear: seq_len*feature_dim -> d_model + let hidden = self.input_linear.forward(&flat_input).map_err(|e| { + MLError::ModelError(format!("Input linear forward failed: {}", e)) + })?; + + // ReLU activation + let activated = hidden.relu().map_err(|e| { + MLError::ModelError(format!("ReLU activation failed: {}", e)) + })?; + + // output_linear: d_model -> 1 + let output = self.output_linear.forward(&activated).map_err(|e| { + MLError::ModelError(format!("Output linear forward failed: {}", e)) + })?; + + Ok(output) + } + + fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result { + // MSE loss: mean((predictions - targets)^2) + let diff = predictions.sub(targets).map_err(|e| { + MLError::ModelError(format!("Loss subtraction failed: {}", e)) + })?; + let squared = diff.powf(2.0).map_err(|e| { + MLError::ModelError(format!("Loss squaring failed: {}", e)) + })?; + let loss = squared.mean_all().map_err(|e| { + MLError::ModelError(format!("Loss mean failed: {}", e)) + })?; + Ok(loss) + } + + fn backward(&mut self, loss: &Tensor) -> Result { + let grads = loss.backward().map_err(|e| { + MLError::TrainingError(format!("Backward pass failed: {}", e)) + })?; + + // Compute gradient norm across all variables for monitoring + let mut grad_norm_sq = 0.0; + let vars_lock = self + .var_map + .data() + .lock() + .map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?; + + for (_name, var) in vars_lock.iter() { + if let Some(grad) = grads.get(var.as_tensor()) { + let norm = grad + .sqr() + .and_then(|s| s.sum_all()) + .and_then(|s| s.to_scalar::()) + .map_err(|e| { + MLError::ModelError(format!("Failed to compute grad norm: {}", e)) + })?; + grad_norm_sq += norm as f64; + } + } + drop(vars_lock); + + let grad_norm = grad_norm_sq.sqrt(); + self.last_grad_norm = grad_norm; + self.latest_metrics.grad_norm = Some(grad_norm); + self.grads = Some(grads); + + Ok(grad_norm) + } + + fn optimizer_step(&mut self) -> Result<(), MLError> { + if let Some(grads) = self.grads.take() { + self.optimizer.step(&grads).map_err(|e| { + MLError::TrainingError(format!("Optimizer step failed: {}", e)) + })?; + } + self.step += 1; + Ok(()) + } + + fn zero_grad(&mut self) -> Result<(), MLError> { + // Gradients are implicitly zeroed by creating a new GradStore in backward(). + // No explicit zeroing needed with candle's approach. + Ok(()) + } + + fn get_learning_rate(&self) -> f64 { + self.learning_rate + } + + fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> { + self.learning_rate = lr; + // Recreate optimizer with new learning rate + let all_vars = self.var_map.all_vars(); + self.optimizer = AdamW::new( + all_vars, + ParamsAdamW { + lr, + beta1: 0.9, + beta2: 0.999, + eps: 1e-8, + weight_decay: 1e-4, + }, + ) + .map_err(|e| { + MLError::ModelError(format!("Failed to recreate optimizer with new lr: {}", e)) + })?; + Ok(()) + } + + fn get_step(&self) -> usize { + self.step + } + + fn collect_metrics(&self) -> TrainingMetrics { + let mut metrics = self.latest_metrics.clone(); + + // Rolling average of recent losses + if !self.loss_history.is_empty() { + let recent: Vec = self.loss_history.iter().rev().take(100).copied().collect(); + metrics.loss = recent.iter().sum::() / recent.len() as f64; + } + + metrics.learning_rate = self.learning_rate; + + // TLOB-specific custom metrics + metrics + .custom_metrics + .insert("training_steps".to_string(), self.step as f64); + metrics.custom_metrics.insert( + "loss_history_len".to_string(), + self.loss_history.len() as f64, + ); + metrics.custom_metrics.insert( + "d_model".to_string(), + self.config.d_model as f64, + ); + metrics.custom_metrics.insert( + "seq_len".to_string(), + self.config.seq_len as f64, + ); + metrics.custom_metrics.insert( + "feature_dim".to_string(), + self.config.feature_dim as f64, + ); + metrics.custom_metrics.insert( + "num_layers".to_string(), + self.config.num_layers as f64, + ); + + metrics + } + + fn save_checkpoint(&self, checkpoint_path: &str) -> Result { + let metadata = CheckpointMetadata { + model_type: "TLOB".to_string(), + version: "1.0.0".to_string(), + epoch: 0, + step: self.step, + timestamp: std::time::SystemTime::now(), + config: serde_json::to_value(&self.config).map_err(|e| { + MLError::SerializationError { + reason: format!("Failed to serialize config: {}", e), + } + })?, + metrics: self.collect_metrics(), + }; + + // Save JSON metadata + checkpoint::save_metadata(&metadata, checkpoint_path)?; + + // Save model weights via safetensors + let safetensors_path = format!("{}.safetensors", checkpoint_path); + let vars_lock = self + .var_map + .data() + .lock() + .map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?; + + let mut tensors: HashMap = HashMap::new(); + for (name, var) in vars_lock.iter() { + tensors.insert(name.clone(), var.as_tensor().clone()); + } + drop(vars_lock); + + candle_core::safetensors::save(&tensors, &safetensors_path).map_err(|e| { + MLError::CheckpointError(format!("Failed to save safetensors: {}", e)) + })?; + + tracing::info!( + "Saved TLOB checkpoint to {} (step {})", + checkpoint_path, + self.step + ); + + Ok(checkpoint_path.to_string()) + } + + fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result { + let metadata = checkpoint::load_metadata(checkpoint_path)?; + + if metadata.model_type != "TLOB" { + return Err(MLError::CheckpointError(format!( + "Invalid model type in checkpoint: expected TLOB, got {}", + metadata.model_type + ))); + } + + // Load weights from safetensors + let safetensors_path = format!("{}.safetensors", checkpoint_path); + let tensors = candle_core::safetensors::load(&safetensors_path, &self.device).map_err( + |e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)), + )?; + + // Set loaded tensors into VarMap + let vars_lock = self + .var_map + .data() + .lock() + .map_err(|e| MLError::LockError(format!("Failed to lock var_map: {}", e)))?; + + for (name, tensor) in &tensors { + if let Some(var) = vars_lock.get(name) { + var.set(tensor).map_err(|e| { + MLError::CheckpointError(format!("Failed to set var {}: {}", name, e)) + })?; + } else { + tracing::warn!("Checkpoint contains unknown variable: {}", name); + } + } + drop(vars_lock); + + // Restore training state + self.step = metadata.step; + self.latest_metrics = metadata.metrics.clone(); + + tracing::info!( + "Loaded TLOB checkpoint from {} (step {})", + checkpoint_path, + metadata.step + ); + + Ok(metadata) + } + + fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result { + if val_data.is_empty() { + return Err(MLError::ValidationError { + message: "Empty validation dataset".to_string(), + }); + } + + let mut total_loss = 0.0; + let mut count = 0usize; + + for (input, target) in val_data { + let prediction = self.forward(input)?; + let loss = self.compute_loss(&prediction, target)?; + let loss_val = loss.to_scalar::().map_err(|e| MLError::ValidationError { + message: format!("Failed to extract loss value: {}", e), + })?; + total_loss += loss_val as f64; + count += 1; + } + + let avg_loss = if count > 0 { + total_loss / count as f64 + } else { + 0.0 + }; + + self.latest_metrics.val_loss = Some(avg_loss); + + tracing::debug!("TLOB validation loss: {:.6}", avg_loss); + + Ok(avg_loss) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_config() -> TLOBAdapterConfig { + TLOBAdapterConfig { + d_model: 32, + num_heads: 2, + num_layers: 2, + seq_len: 32, + feature_dim: 51, + } + } + + #[test] + fn test_model_type_returns_tlob() { + let cfg = make_config(); + let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + assert_eq!(adapter.model_type(), "TLOB"); + } + + #[test] + fn test_device_returns_cpu() { + let cfg = make_config(); + let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + assert!(matches!(adapter.device(), &Device::Cpu)); + } + + #[test] + fn test_forward_produces_output() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap(); + + // 3D input: [batch=2, seq_len=32, feature_dim=51] + let input = + Tensor::zeros(&[2, cfg.seq_len, cfg.feature_dim], DType::F32, &Device::Cpu).unwrap(); + let output = adapter.forward(&input).unwrap(); + + let dims = output.shape().dims(); + assert_eq!(dims.len(), 2); + assert_eq!(dims[0], 2); // batch + assert_eq!(dims[1], 1); // scalar prediction + } + + #[test] + fn test_forward_accepts_2d_input() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap(); + + // 2D input: [batch=2, seq_len*feature_dim] + let flat_dim = cfg.seq_len * cfg.feature_dim; + let input = Tensor::zeros(&[2, flat_dim], DType::F32, &Device::Cpu).unwrap(); + let output = adapter.forward(&input).unwrap(); + + let dims = output.shape().dims(); + assert_eq!(dims.len(), 2); + assert_eq!(dims[0], 2); + assert_eq!(dims[1], 1); + } + + #[test] + fn test_compute_loss_returns_scalar() { + let cfg = make_config(); + let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let preds = Tensor::new(&[[1.0f32], [2.0]], &Device::Cpu).unwrap(); + let targets = Tensor::new(&[[1.5f32], [2.5]], &Device::Cpu).unwrap(); + + let loss = adapter.compute_loss(&preds, &targets).unwrap(); + let loss_val: f32 = loss.to_scalar().unwrap(); + + // MSE of (0.5^2 + 0.5^2)/2 = 0.25 + assert!((loss_val - 0.25).abs() < 1e-5, "Expected ~0.25, got {}", loss_val); + } + + #[test] + fn test_backward_returns_grad_norm() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap(); + + let flat_dim = cfg.seq_len * cfg.feature_dim; + let input = Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap(); + let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + + let output = adapter.forward(&input).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + let grad_norm = adapter.backward(&loss).unwrap(); + + assert!(grad_norm >= 0.0, "Gradient norm should be non-negative"); + assert!(adapter.grads.is_some(), "Grads should be stored after backward"); + } + + #[test] + fn test_train_step_cycle() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap(); + + assert_eq!(adapter.get_step(), 0); + + // Full train cycle: zero_grad -> forward -> loss -> backward -> optimizer_step + adapter.zero_grad().unwrap(); + + let flat_dim = cfg.seq_len * cfg.feature_dim; + let input = Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap(); + let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + + let output = adapter.forward(&input).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + let loss_val: f64 = loss.to_scalar::().unwrap() as f64; + + adapter.backward(&loss).unwrap(); + adapter.optimizer_step().unwrap(); + + assert!(adapter.get_step() >= 1); + + // Record loss and check metrics + adapter.loss_history.push(loss_val); + let metrics = adapter.collect_metrics(); + assert!(metrics.loss >= 0.0); + assert_eq!( + metrics.custom_metrics.get("training_steps").copied(), + Some(1.0) + ); + } + + #[test] + fn test_learning_rate_get_set() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let original_lr = adapter.get_learning_rate(); + assert!((original_lr - 1e-3).abs() < 1e-10); + + adapter.set_learning_rate(5e-4).unwrap(); + assert!((adapter.get_learning_rate() - 5e-4).abs() < 1e-10); + } + + #[test] + fn test_collect_metrics() { + let cfg = make_config(); + let adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let metrics = adapter.collect_metrics(); + assert!(metrics.custom_metrics.contains_key("training_steps")); + assert!(metrics.custom_metrics.contains_key("d_model")); + assert!(metrics.custom_metrics.contains_key("seq_len")); + assert!(metrics.custom_metrics.contains_key("feature_dim")); + assert!(metrics.custom_metrics.contains_key("num_layers")); + assert_eq!(metrics.custom_metrics.get("d_model").copied(), Some(32.0)); + assert_eq!(metrics.custom_metrics.get("seq_len").copied(), Some(32.0)); + assert_eq!(metrics.custom_metrics.get("feature_dim").copied(), Some(51.0)); + } + + #[test] + fn test_checkpoint_roundtrip() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap(); + + // Run a training step to have non-zero state + let flat_dim = cfg.seq_len * cfg.feature_dim; + let input = Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap(); + let targets = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + let output = adapter.forward(&input).unwrap(); + let loss = adapter.compute_loss(&output, &targets).unwrap(); + adapter.backward(&loss).unwrap(); + adapter.optimizer_step().unwrap(); + + // Save checkpoint + let tmp_dir = std::env::temp_dir(); + let checkpoint_path = tmp_dir.join("tlob_test_ckpt"); + let path_str = checkpoint_path.to_str().unwrap(); + adapter.save_checkpoint(path_str).unwrap(); + + // Load into fresh adapter + let mut adapter2 = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + let metadata = adapter2.load_checkpoint(path_str).unwrap(); + + assert_eq!(metadata.model_type, "TLOB"); + assert_eq!(metadata.step, 1); + assert_eq!(adapter2.get_step(), 1); + + // Clean up checkpoint files + let _ = std::fs::remove_file(format!("{}.json", path_str)); + let _ = std::fs::remove_file(format!("{}.safetensors", path_str)); + } + + #[test] + fn test_validate_returns_loss() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg.clone(), &Device::Cpu).unwrap(); + + let flat_dim = cfg.seq_len * cfg.feature_dim; + let val_data: Vec<(Tensor, Tensor)> = (0..3) + .map(|_| { + let input = + Tensor::randn(0.0f32, 1.0, &[2, flat_dim], &Device::Cpu).unwrap(); + let target = Tensor::randn(0.0f32, 1.0, &[2, 1], &Device::Cpu).unwrap(); + (input, target) + }) + .collect(); + + let val_loss = adapter.validate(&val_data).unwrap(); + assert!(val_loss >= 0.0, "Validation loss should be non-negative"); + assert!( + adapter.latest_metrics.val_loss.is_some(), + "val_loss should be set after validate" + ); + } + + #[test] + fn test_validate_empty_errors() { + let cfg = make_config(); + let mut adapter = TLOBTrainableAdapter::new(cfg, &Device::Cpu).unwrap(); + + let result = adapter.validate(&[]); + assert!(result.is_err(), "Validating empty data should error"); + } +}