From 2dd3326aecbccedb4ca0a6f4ea830767e45b388f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 21 Feb 2026 14:58:36 +0100 Subject: [PATCH] docs: ensemble real inference design and implementation plan Co-Authored-By: Claude Opus 4.6 --- ...26-02-21-ensemble-real-inference-design.md | 217 +++ ...-ensemble-real-inference-implementation.md | 1694 +++++++++++++++++ 2 files changed, 1911 insertions(+) create mode 100644 docs/plans/2026-02-21-ensemble-real-inference-design.md create mode 100644 docs/plans/2026-02-21-ensemble-real-inference-implementation.md diff --git a/docs/plans/2026-02-21-ensemble-real-inference-design.md b/docs/plans/2026-02-21-ensemble-real-inference-design.md new file mode 100644 index 000000000..c7cb4fb8c --- /dev/null +++ b/docs/plans/2026-02-21-ensemble-real-inference-design.md @@ -0,0 +1,217 @@ +# Ensemble Real Inference Design + +**Date**: 2026-02-21 +**Status**: Approved +**Scope**: Replace mock predictions in EnsembleCoordinator with actual model inference, modernize Mamba2, validate all models individually. + +## Problem Statement + +The `EnsembleCoordinator` in `ml/src/ensemble/coordinator.rs` uses `simulate_trained_model_prediction()` (lines 157-195) which computes arithmetic on raw features instead of calling `model.forward()`. The trading service (`services/trading_service/src/services/ml.rs`, lines 42-47) also hardcodes predictions. This means no trained model ever produces a real inference in production. + +## Design Decisions (User-Approved) + +1. **Per-model adapter pattern** — each model gets its own `ModelInferenceAdapter` implementation +2. **Hybrid feature reconciliation** — one canonical 51-dim `FeatureVector`, adapters pad/extend/buffer for sequence models +3. **Device placement** — `Auto` (GPU if available, CPU fallback). All 4 models fit in 4GB VRAM for inference (3.7MB total: DQN 0.67MB + PPO 0.31MB + TFT 1.98MB + Mamba2 0.74MB) + +## Architecture + +### 1. ModelInferenceAdapter Trait + +```rust +// ml/src/ensemble/inference_adapter.rs + +pub trait ModelInferenceAdapter: Send + Sync { + /// Human-readable model name for logging + fn model_name(&self) -> &str; + + /// Run inference on a single feature vector, return ensemble-compatible prediction + fn predict(&self, features: &FeatureVector) -> Result; + + /// Whether this adapter has a loaded model ready for inference + fn is_ready(&self) -> bool; +} + +pub struct FeatureVector { + pub values: Vec, // 51-dim canonical features + pub timestamp: i64, // epoch micros +} + +pub struct EnsemblePrediction { + pub model_name: String, + pub direction: f64, // [-1.0, 1.0] bearish→bullish + pub confidence: f64, // [0.0, 1.0] + pub metadata: PredictionMeta, // model-specific extras +} + +pub struct PredictionMeta { + pub latency_us: u64, + pub quantiles: Option>, // TFT only + pub attention_weights: Option>, // TFT only + pub q_values: Option>, // DQN only +} +``` + +### 2. Per-Model Adapters + +#### DQN Adapter +- **Input**: `FeatureVector.values` → `Tensor([1, 51], f32)` +- **Forward**: `dqn.forward(&tensor)` → `Tensor([1, 45])` Q-values +- **Output**: `direction = (best_action - 22) / 22.0`, `confidence = softmax(q_values).max()` +- **No buffering needed** — single-step model + +#### PPO Adapter +- **Input**: `FeatureVector.values` → pad to 64-dim → `Tensor([1, 64], f32)` +- **Forward**: `ppo.action_probabilities(&tensor)` → `Tensor([1, 45])` action probs +- **Output**: `direction = weighted_sum(probs * action_values)`, `confidence = probs.max()` +- **LSTM state**: Maintain `(h_t, c_t)` across calls, reset on new episode + +#### TFT Adapter +- **Input**: Buffer last `sequence_length` feature vectors (default 10) +- **Static**: Extract from first feature vector (asset type, session) +- **Historical**: `Tensor([1, seq, num_unknown])` from buffered features +- **Future**: `Tensor([1, horizon, num_known])` — calendar features (hour, day, etc.) +- **Forward**: `tft.forward(static, historical, future)` → `Tensor([1, 10, 9])` quantile forecasts +- **Output**: `direction = sign(median_quantile[0])`, `confidence = 1 - (q75 - q25) / q50.abs()` + +#### Mamba2 Adapter +- **Input**: Buffer last `max_seq_len` feature vectors +- **Pad**: Each 51-dim vector → 225-dim (zero-pad right) +- **Forward**: `mamba.forward(&tensor)` → `Tensor([1, seq, 1])` probability +- **Output**: `direction = 2 * last_prob - 1`, `confidence = (last_prob - 0.5).abs() * 2` + +### 3. Model Loading + +```rust +pub struct ModelLoader { + device: Device, // Auto-detected +} + +impl ModelLoader { + /// Load a model from a safetensors checkpoint path + pub fn load_dqn(&self, path: &Path, config: &DQNConfig) -> Result; + pub fn load_ppo(&self, path: &Path, config: &PPOConfig) -> Result; + pub fn load_tft(&self, path: &Path, config: &TFTConfig) -> Result; + pub fn load_mamba2(&self, path: &Path, config: &Mamba2Config) -> Result; +} +``` + +Each loader: +1. Creates a `VarMap` and loads weights from safetensors +2. Constructs the model with the provided config +3. Sets model to eval mode (no dropout) +4. Wraps in the adapter with pre-allocated buffers + +### 4. Coordinator Integration + +Replace `simulate_trained_model_prediction()` in `EnsembleCoordinator`: + +```rust +pub struct EnsembleCoordinator { + adapters: Vec>, + // ... existing fields +} + +impl EnsembleCoordinator { + pub fn predict(&self, features: &FeatureVector) -> Result { + let predictions: Vec = self.adapters + .iter() + .filter(|a| a.is_ready()) + .map(|a| a.predict(features)) + .collect::>>()?; + + self.aggregate(predictions) // existing weighted averaging logic + } +} +``` + +### 5. Feature-to-Tensor Mapping + +Canonical 51-dim `FeatureVector` layout (matches DQN state_dim): + +| Index | Feature | Description | +|-------|---------|-------------| +| 0-9 | Price features | OHLCV, VWAP, spread, tick | +| 10-19 | Technical indicators | RSI, MACD, Bollinger, ATR | +| 20-29 | Order book | Depth levels, imbalance | +| 30-39 | Microstructure | Trade flow, toxicity | +| 40-50 | Position/risk | PnL, exposure, drawdown | + +Per-model mapping: +- **DQN**: Direct pass-through (51-dim) +- **PPO**: Pad indices 51-63 with zeros (→ 64-dim) +- **TFT**: Split into static (indices 40-45), unknown (indices 0-39), known (calendar from timestamp) +- **Mamba2**: Zero-pad to 225-dim + +## Mamba2 Modernization + +### 1. Directional Loss Function + +Add `DirectionalMSE` to `ml/src/mamba/mod.rs`: + +```rust +pub fn directional_mse_loss( + predictions: &Tensor, + targets: &Tensor, + direction_weight: f64, // default 2.0 — penalize wrong-sign 2x +) -> Result { + let mse = (predictions - targets)?.sqr()?; + let pred_sign = predictions.ge(&Tensor::zeros_like(predictions)?)?; + let target_sign = targets.ge(&Tensor::zeros_like(targets)?)?; + let wrong_direction = pred_sign.ne(&target_sign)?.to_dtype(DType::F64)?; + let weights = (&wrong_direction * (direction_weight - 1.0) + 1.0)?; + (&mse * &weights)?.mean_all() +} +``` + +### 2. Standardize Discretization + +In `ml/src/mamba/mod.rs`, replace the ZOH discretization (lines ~913-929) with bilinear (Tustin) to match `ssd_layer.rs`: + +```rust +// Before (ZOH): A_disc = I + A_cont * dt +// After (Bilinear): A_disc = (I + A*dt/2) * (I - A*dt/2)^{-1} +// Approximation: A_disc = I + A*dt + (A*dt)^2/2 (matches ssd_layer.rs) +``` + +## Testing Strategy + +### Individual Model Validation (before ensemble wiring) + +Each model gets a standalone inference test that: +1. Creates model with known config +2. Saves checkpoint to temp file +3. Loads via `ModelLoader` +4. Runs `predict()` on synthetic `FeatureVector` +5. Validates output shape, range, and determinism + +### Ensemble Integration Tests + +1. Load all 4 models → predict → verify aggregation +2. Partial availability (2 of 4 models) → graceful degradation +3. Latency measurement — all 4 models inference < 1ms total + +## Files to Create/Modify + +**New files:** +- `ml/src/ensemble/inference_adapter.rs` — trait + types +- `ml/src/ensemble/adapters/mod.rs` — adapter module +- `ml/src/ensemble/adapters/dqn.rs` — DQN adapter +- `ml/src/ensemble/adapters/ppo.rs` — PPO adapter +- `ml/src/ensemble/adapters/tft.rs` — TFT adapter +- `ml/src/ensemble/adapters/mamba2.rs` — Mamba2 adapter +- `ml/src/ensemble/model_loader.rs` — checkpoint loading +- `ml/src/mamba/loss.rs` — DirectionalMSE loss + +**Modified files:** +- `ml/src/ensemble/coordinator.rs` — replace mock predictions +- `ml/src/ensemble/mod.rs` — add module declarations +- `ml/src/mamba/mod.rs` — fix discretization, add loss module +- `services/trading_service/src/services/ml.rs` — remove hardcoded predictions (future PR) + +## Non-Goals + +- **Complex-valued state space** — deferred, requires Candle complex tensor infrastructure +- **MIMO Mamba** — deferred, single-asset focus first +- **Trading service integration** — separate PR after ensemble works +- **Quantized inference** — not needed at 3.7MB total model size diff --git a/docs/plans/2026-02-21-ensemble-real-inference-implementation.md b/docs/plans/2026-02-21-ensemble-real-inference-implementation.md new file mode 100644 index 000000000..bd8fabed3 --- /dev/null +++ b/docs/plans/2026-02-21-ensemble-real-inference-implementation.md @@ -0,0 +1,1694 @@ +# Ensemble Real Inference Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Replace mock predictions in EnsembleCoordinator with actual model forward() calls for DQN, PPO, TFT, and Mamba2, plus modernize Mamba2's loss/discretization. + +**Architecture:** Per-model adapter pattern with a shared `ModelInferenceAdapter` trait. Each adapter handles its own feature-to-tensor mapping (padding, buffering, splitting). Models loaded from safetensors checkpoints via `candle_core::safetensors::load()`. Coordinator iterates ready adapters and aggregates results. + +**Tech Stack:** Rust, Candle v0.9.1 (`candle_core`, `candle_nn`), safetensors, `DType::F32` (DQN/PPO), `DType::F64` (Mamba2), tokio (async coordinator) + +**Build command:** `SQLX_OFFLINE=true cargo check -p ml` +**Test command:** `SQLX_OFFLINE=true cargo test -p ml --lib` +**Clippy rules:** `#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::indexing_slicing)]` — must use `.get()`, `?`, `.ok_or()` everywhere. + +--- + +### Task 1: Inference Adapter Trait & Types + +**Files:** +- Create: `ml/src/ensemble/inference_adapter.rs` +- Modify: `ml/src/ensemble/mod.rs` (add module declaration + re-exports) + +**Step 1: Write the failing test** + +Add to the bottom of `ml/src/ensemble/inference_adapter.rs`: + +```rust +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_feature_vector_creation() { + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000_000_000, + }; + assert_eq!(fv.values.len(), 51); + assert_eq!(fv.timestamp, 1700000000_000_000); + } + + #[test] + fn test_ensemble_prediction_direction_bounds() { + let pred = EnsemblePrediction { + model_name: "test".to_string(), + direction: 0.75, + confidence: 0.9, + metadata: PredictionMeta::default(), + }; + assert!(pred.direction >= -1.0 && pred.direction <= 1.0); + assert!(pred.confidence >= 0.0 && pred.confidence <= 1.0); + } + + #[test] + fn test_prediction_meta_default() { + let meta = PredictionMeta::default(); + assert_eq!(meta.latency_us, 0); + assert!(meta.quantiles.is_none()); + assert!(meta.attention_weights.is_none()); + assert!(meta.q_values.is_none()); + } +} +``` + +**Step 2: Write the implementation** + +Create `ml/src/ensemble/inference_adapter.rs`: + +```rust +//! Model inference adapter trait for ensemble prediction +//! +//! Defines the contract that each model adapter must implement +//! to participate in ensemble prediction. + +use crate::MLResult; + +/// Canonical feature vector for ensemble inference. +/// 51-dim layout matching DQN state_dim: +/// 0-9: price features (OHLCV, VWAP, spread, tick) +/// 10-19: technical indicators (RSI, MACD, Bollinger, ATR) +/// 20-29: order book (depth levels, imbalance) +/// 30-39: microstructure (trade flow, toxicity) +/// 40-50: position/risk (PnL, exposure, drawdown) +#[derive(Debug, Clone)] +pub struct FeatureVector { + pub values: Vec, + pub timestamp: i64, +} + +/// Prediction output from a single model adapter, normalized for ensemble aggregation. +#[derive(Debug, Clone)] +pub struct EnsemblePrediction { + pub model_name: String, + /// Directional signal: -1.0 (bearish) to 1.0 (bullish) + pub direction: f64, + /// Model confidence: 0.0 to 1.0 + pub confidence: f64, + /// Model-specific metadata + pub metadata: PredictionMeta, +} + +/// Optional model-specific metadata attached to predictions. +#[derive(Debug, Clone, Default)] +pub struct PredictionMeta { + /// Inference latency in microseconds + pub latency_us: u64, + /// Quantile forecasts (TFT only) + pub quantiles: Option>, + /// Attention weights (TFT only) + pub attention_weights: Option>, + /// Raw Q-values (DQN only) + pub q_values: Option>, +} + +/// Adapter trait for running inference on a loaded model. +/// +/// Each model type (DQN, PPO, TFT, Mamba2) implements this trait +/// to handle its own feature-to-tensor mapping and output normalization. +pub trait ModelInferenceAdapter: Send + Sync { + /// Human-readable model name for logging + fn model_name(&self) -> &str; + + /// Run inference on a canonical feature vector. + /// Returns a normalized ensemble prediction. + fn predict(&self, features: &FeatureVector) -> MLResult; + + /// Whether this adapter has a loaded model ready for inference + fn is_ready(&self) -> bool; +} +``` + +**Step 3: Register the module** + +In `ml/src/ensemble/mod.rs`, add after the existing module declarations: + +```rust +pub mod inference_adapter; +``` + +And add re-exports: + +```rust +pub use inference_adapter::{EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta}; +``` + +**Step 4: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::inference_adapter::tests -v` +Expected: 3 tests pass + +**Step 5: Commit** + +```bash +git add ml/src/ensemble/inference_adapter.rs ml/src/ensemble/mod.rs +git commit -m "feat(ensemble): add ModelInferenceAdapter trait and types" +``` + +--- + +### Task 2: DQN Inference Adapter + +**Files:** +- Create: `ml/src/ensemble/adapters/mod.rs` +- Create: `ml/src/ensemble/adapters/dqn.rs` +- Modify: `ml/src/ensemble/mod.rs` (add `pub mod adapters;`) + +**Context:** +- `DQN::new(config: DQNConfig) -> Result` in `ml/src/dqn/dqn.rs:793` +- `DQN::forward(&self, state: &Tensor) -> Result` at line 1011 +- Input: `Tensor([batch, 51], F32)`, Output: `Tensor([batch, num_actions], F32)` Q-values +- `DQN::load_from_safetensors(&mut self, path: &str) -> Result<(), MLError>` at line 2532 +- `DQNConfig` at `ml/src/dqn/dqn.rs:33` — needs `state_dim`, `num_actions`, `hidden_dims` + +**Step 1: Write the test** + +Add to bottom of `ml/src/ensemble/adapters/dqn.rs`: + +```rust +#[cfg(test)] +mod tests { + use super::*; + use crate::ensemble::inference_adapter::FeatureVector; + + fn test_config() -> crate::dqn::dqn::DQNConfig { + crate::dqn::dqn::DQNConfig { + state_dim: 51, + num_actions: 45, + hidden_dims: vec![64, 64], + ..Default::default() + } + } + + #[test] + fn test_dqn_adapter_creation() { + let config = test_config(); + let adapter = DqnInferenceAdapter::new(config); + assert!(adapter.is_ok()); + let adapter = adapter.unwrap(); + assert_eq!(adapter.model_name(), "DQN"); + assert!(adapter.is_ready()); // model created with random weights + } + + #[test] + fn test_dqn_adapter_predict_direction_range() { + let config = test_config(); + let adapter = DqnInferenceAdapter::new(config).unwrap(); + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000, + }; + let pred = adapter.predict(&fv).unwrap(); + assert!(pred.direction >= -1.0 && pred.direction <= 1.0, + "direction {} out of [-1,1]", pred.direction); + assert!(pred.confidence >= 0.0 && pred.confidence <= 1.0, + "confidence {} out of [0,1]", pred.confidence); + assert_eq!(pred.model_name, "DQN"); + assert!(pred.metadata.q_values.is_some()); + } + + #[test] + fn test_dqn_adapter_deterministic() { + let config = test_config(); + let adapter = DqnInferenceAdapter::new(config).unwrap(); + let fv = FeatureVector { + values: vec![0.5; 51], + timestamp: 1700000000, + }; + let p1 = adapter.predict(&fv).unwrap(); + let p2 = adapter.predict(&fv).unwrap(); + assert!((p1.direction - p2.direction).abs() < 1e-6, + "Non-deterministic: {} vs {}", p1.direction, p2.direction); + } +} +``` + +**Step 2: Write the implementation** + +Create `ml/src/ensemble/adapters/dqn.rs`: + +```rust +//! DQN inference adapter for ensemble prediction + +use std::sync::Mutex; +use candle_core::{DType, Device, Tensor}; +use crate::dqn::dqn::{DQN, DQNConfig}; +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::{MLError, MLResult}; + +/// DQN inference adapter. +/// +/// Converts 51-dim FeatureVector to Tensor, runs DQN forward pass, +/// and maps Q-values to direction/confidence. +pub struct DqnInferenceAdapter { + model: Mutex, + device: Device, +} + +impl DqnInferenceAdapter { + /// Create adapter with a new DQN model (random weights). + pub fn new(config: DQNConfig) -> MLResult { + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let model = DQN::new(config)?; + Ok(Self { + model: Mutex::new(model), + device, + }) + } + + /// Create adapter and load weights from safetensors checkpoint. + pub fn from_checkpoint(config: DQNConfig, path: &str) -> MLResult { + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let mut model = DQN::new(config)?; + model.load_from_safetensors(path)?; + Ok(Self { + model: Mutex::new(model), + device, + }) + } +} + +impl ModelInferenceAdapter for DqnInferenceAdapter { + fn model_name(&self) -> &str { + "DQN" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + // Convert f64 features to f32 tensor [1, 51] + let f32_values: Vec = features.values.iter().map(|&v| v as f32).collect(); + let input = Tensor::from_vec(f32_values, (1, features.values.len()), &self.device) + .map_err(|e| MLError::TensorCreationError { + operation: "DQN feature vector to tensor".to_string(), + reason: e.to_string(), + })?; + + // Forward pass — returns [1, num_actions] Q-values + let model = self.model.lock().map_err(|e| { + MLError::LockError(format!("DQN model lock failed: {}", e)) + })?; + let q_values = model.forward(&input)?; + + // Extract Q-values as Vec + let q_vec: Vec = q_values + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("squeeze failed: {}", e)))? + .to_vec1::() + .map_err(|e| MLError::ModelError(format!("to_vec1 failed: {}", e)))?; + + let num_actions = q_vec.len(); + + // Find best action + let (best_idx, best_q) = q_vec + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .ok_or_else(|| MLError::InferenceError("empty Q-values".to_string()))?; + + // Direction: map action index to [-1, 1] + // Actions centered at num_actions/2 = hold + let center = num_actions as f64 / 2.0; + let direction = ((best_idx as f64 - center) / center).clamp(-1.0, 1.0); + + // Confidence: softmax probability of best action + let max_q = q_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let exp_sum: f32 = q_vec.iter().map(|&q| (q - max_q).exp()).sum(); + let confidence = ((best_q - max_q).exp() / exp_sum) as f64; + + let latency = start.elapsed().as_micros() as u64; + + Ok(EnsemblePrediction { + model_name: "DQN".to_string(), + direction, + confidence: confidence.clamp(0.0, 1.0), + metadata: PredictionMeta { + latency_us: latency, + q_values: Some(q_vec.iter().map(|&v| v as f64).collect()), + ..Default::default() + }, + }) + } + + fn is_ready(&self) -> bool { + self.model.lock().is_ok() + } +} + +// Send + Sync: Mutex is Send+Sync, Device is Send+Sync +unsafe impl Send for DqnInferenceAdapter {} +unsafe impl Sync for DqnInferenceAdapter {} +``` + +Create `ml/src/ensemble/adapters/mod.rs`: + +```rust +//! Per-model inference adapters for ensemble prediction + +pub mod dqn; + +pub use dqn::DqnInferenceAdapter; +``` + +**Step 3: Register the adapters module** + +In `ml/src/ensemble/mod.rs`, add: + +```rust +pub mod adapters; +``` + +**Step 4: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::adapters::dqn::tests -v` +Expected: 3 tests pass + +**Step 5: Commit** + +```bash +git add ml/src/ensemble/adapters/mod.rs ml/src/ensemble/adapters/dqn.rs ml/src/ensemble/mod.rs +git commit -m "feat(ensemble): add DQN inference adapter" +``` + +--- + +### Task 3: PPO Inference Adapter + +**Files:** +- Create: `ml/src/ensemble/adapters/ppo.rs` +- Modify: `ml/src/ensemble/adapters/mod.rs` (add `pub mod ppo;`) + +**Context:** +- `PPO::new(config: PPOConfig) -> Result` at `ml/src/ppo/ppo.rs:705` +- `ActorNetwork::MLP` → `action_probabilities(&Tensor) -> Result` at line 427 +- `ActorNetwork::LSTM` → requires hidden state (not supported for ensemble initially, use MLP) +- Input: `Tensor([batch, state_dim], F32)`, Output: `Tensor([batch, num_actions], F32)` action probs +- PPO state_dim = 64 (needs 13-dim zero-pad from 51-dim FeatureVector) + +**Step 1: Write the test** + +Add to `ml/src/ensemble/adapters/ppo.rs`: + +```rust +#[cfg(test)] +mod tests { + use super::*; + use crate::ensemble::inference_adapter::FeatureVector; + + fn test_config() -> crate::ppo::ppo::PPOConfig { + crate::ppo::ppo::PPOConfig { + state_dim: 64, + num_actions: 45, + policy_hidden_dims: vec![64, 64], + value_hidden_dims: vec![64, 64], + ..Default::default() + } + } + + #[test] + fn test_ppo_adapter_creation() { + let config = test_config(); + let adapter = PpoInferenceAdapter::new(config); + assert!(adapter.is_ok()); + let adapter = adapter.unwrap(); + assert_eq!(adapter.model_name(), "PPO"); + assert!(adapter.is_ready()); + } + + #[test] + fn test_ppo_adapter_pads_input_to_64() { + let config = test_config(); + let adapter = PpoInferenceAdapter::new(config).unwrap(); + // 51-dim input should be zero-padded to 64 + let fv = FeatureVector { + values: vec![0.3; 51], + timestamp: 1700000000, + }; + let pred = adapter.predict(&fv).unwrap(); + assert!(pred.direction >= -1.0 && pred.direction <= 1.0); + assert!(pred.confidence >= 0.0 && pred.confidence <= 1.0); + assert_eq!(pred.model_name, "PPO"); + } + + #[test] + fn test_ppo_adapter_deterministic() { + let config = test_config(); + let adapter = PpoInferenceAdapter::new(config).unwrap(); + let fv = FeatureVector { + values: vec![0.5; 51], + timestamp: 1700000000, + }; + let p1 = adapter.predict(&fv).unwrap(); + let p2 = adapter.predict(&fv).unwrap(); + assert!((p1.direction - p2.direction).abs() < 1e-6); + } +} +``` + +**Step 2: Write the implementation** + +```rust +//! PPO inference adapter for ensemble prediction + +use std::sync::Mutex; +use candle_core::{Device, Tensor}; +use crate::ppo::ppo::{PPO, PPOConfig}; +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::{MLError, MLResult}; + +/// PPO inference adapter (MLP actor only, LSTM not supported). +/// +/// Zero-pads 51-dim FeatureVector to 64-dim for PPO state_dim. +/// Converts action probabilities to direction/confidence. +pub struct PpoInferenceAdapter { + model: Mutex, + state_dim: usize, + device: Device, +} + +impl PpoInferenceAdapter { + pub fn new(config: PPOConfig) -> MLResult { + let state_dim = config.state_dim; + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let model = PPO::new(config)?; + Ok(Self { + model: Mutex::new(model), + state_dim, + device, + }) + } + + /// Pad feature vector to match PPO state_dim (zero-pad right) + fn pad_features(&self, values: &[f64]) -> Vec { + let mut padded = Vec::with_capacity(self.state_dim); + for &v in values.iter().take(self.state_dim) { + padded.push(v as f32); + } + // Zero-pad remaining dimensions + padded.resize(self.state_dim, 0.0); + padded + } +} + +impl ModelInferenceAdapter for PpoInferenceAdapter { + fn model_name(&self) -> &str { + "PPO" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + let padded = self.pad_features(&features.values); + let input = Tensor::from_vec(padded, (1, self.state_dim), &self.device) + .map_err(|e| MLError::TensorCreationError { + operation: "PPO feature vector to tensor".to_string(), + reason: e.to_string(), + })?; + + let model = self.model.lock().map_err(|e| { + MLError::LockError(format!("PPO model lock failed: {}", e)) + })?; + + // MLP actor: action_probabilities returns softmax probs [1, num_actions] + let probs = model.actor.action_probabilities(&input)?; + + let probs_vec: Vec = probs + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("squeeze failed: {}", e)))? + .to_vec1::() + .map_err(|e| MLError::ModelError(format!("to_vec1 failed: {}", e)))?; + + let num_actions = probs_vec.len(); + let center = num_actions as f64 / 2.0; + + // Weighted sum of action probabilities × action values → direction + let direction: f64 = probs_vec + .iter() + .enumerate() + .map(|(i, &p)| { + let action_val = (i as f64 - center) / center; + p as f64 * action_val + }) + .sum(); + + // Confidence: max probability + let confidence = probs_vec + .iter() + .cloned() + .fold(f32::NEG_INFINITY, f32::max) as f64; + + let latency = start.elapsed().as_micros() as u64; + + Ok(EnsemblePrediction { + model_name: "PPO".to_string(), + direction: direction.clamp(-1.0, 1.0), + confidence: confidence.clamp(0.0, 1.0), + metadata: PredictionMeta { + latency_us: latency, + ..Default::default() + }, + }) + } + + fn is_ready(&self) -> bool { + self.model.lock().is_ok() + } +} + +unsafe impl Send for PpoInferenceAdapter {} +unsafe impl Sync for PpoInferenceAdapter {} +``` + +**Step 3: Register module** + +In `ml/src/ensemble/adapters/mod.rs`, add: + +```rust +pub mod ppo; +pub use ppo::PpoInferenceAdapter; +``` + +**Step 4: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::adapters::ppo::tests -v` +Expected: 3 tests pass + +**Step 5: Commit** + +```bash +git add ml/src/ensemble/adapters/ppo.rs ml/src/ensemble/adapters/mod.rs +git commit -m "feat(ensemble): add PPO inference adapter" +``` + +--- + +### Task 4: TFT Inference Adapter + +**Files:** +- Create: `ml/src/ensemble/adapters/tft.rs` +- Modify: `ml/src/ensemble/adapters/mod.rs` + +**Context:** +- `TemporalFusionTransformer::new(config: TFTConfig) -> Result` at `ml/src/tft/mod.rs:303` +- `forward(&mut self, static, historical, future) -> Result` at line 514 +- TFT needs `&mut self` (writes attention weights to RwLock internally) +- Input: 3 tensors — static `[1, num_static]`, historical `[1, seq, num_unknown]`, future `[1, horizon, num_known]` +- Output: `[1, horizon, num_quantiles]` quantile forecasts +- Config: `TFTConfig` at `ml/src/tft/mod.rs:108` +- Must buffer multiple feature vectors for sequence input + +**Step 1: Write the test** + +```rust +#[cfg(test)] +mod tests { + use super::*; + use crate::ensemble::inference_adapter::FeatureVector; + + fn test_config() -> crate::tft::TFTConfig { + crate::tft::TFTConfig { + input_dim: 20, // 6 static + 6 known + 8 unknown = 20 + hidden_dim: 32, + num_heads: 2, + num_layers: 1, + prediction_horizon: 5, + sequence_length: 4, + num_quantiles: 9, + num_static_features: 6, + num_known_features: 6, + num_unknown_features: 8, + learning_rate: 0.001, + batch_size: 1, + dropout_rate: 0.0, // Disable for deterministic test + ..Default::default() + } + } + + #[test] + fn test_tft_adapter_creation() { + let config = test_config(); + let adapter = TftInferenceAdapter::new(config, 4); + assert!(adapter.is_ok()); + let adapter = adapter.unwrap(); + assert_eq!(adapter.model_name(), "TFT"); + assert!(!adapter.is_ready()); // not ready until buffer filled + } + + #[test] + fn test_tft_adapter_buffers_and_predicts() { + let config = test_config(); + let mut adapter = TftInferenceAdapter::new(config, 4).unwrap(); + + // Feed 4 feature vectors to fill the buffer + for i in 0..4 { + let fv = FeatureVector { + values: vec![0.1 * (i as f64 + 1.0); 51], + timestamp: 1700000000 + i * 1000, + }; + let _ = adapter.predict(&fv); + } + assert!(adapter.is_ready()); + + // 5th prediction should work with full buffer + let fv = FeatureVector { + values: vec![0.5; 51], + timestamp: 1700004000, + }; + let pred = adapter.predict(&fv); + assert!(pred.is_ok(), "prediction failed: {:?}", pred.err()); + let pred = pred.unwrap(); + assert!(pred.direction >= -1.0 && pred.direction <= 1.0); + assert!(pred.confidence >= 0.0 && pred.confidence <= 1.0); + assert!(pred.metadata.quantiles.is_some()); + } +} +``` + +**Step 2: Write the implementation** + +```rust +//! TFT inference adapter for ensemble prediction + +use std::collections::VecDeque; +use std::sync::Mutex; +use candle_core::{DType, Device, Tensor}; +use crate::tft::{TFTConfig, TemporalFusionTransformer}; +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::{MLError, MLResult}; + +/// TFT inference adapter with sequence buffering. +/// +/// Buffers `sequence_length` feature vectors before producing predictions. +/// Splits 51-dim FeatureVector into static/historical/future tensors. +pub struct TftInferenceAdapter { + model: Mutex, + buffer: Mutex>, + sequence_length: usize, + num_static: usize, + num_known: usize, + num_unknown: usize, + prediction_horizon: usize, + num_quantiles: usize, + device: Device, +} + +impl TftInferenceAdapter { + pub fn new(config: TFTConfig, sequence_length: usize) -> MLResult { + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let num_static = config.num_static_features; + let num_known = config.num_known_features; + let num_unknown = config.num_unknown_features; + let prediction_horizon = config.prediction_horizon; + let num_quantiles = config.num_quantiles; + let model = TemporalFusionTransformer::new_with_device(config, device.clone())?; + + Ok(Self { + model: Mutex::new(model), + buffer: Mutex::new(VecDeque::with_capacity(sequence_length + 1)), + sequence_length, + num_static, + num_known, + num_unknown, + prediction_horizon, + num_quantiles, + device, + }) + } + + /// Extract static features from first feature vector (indices 40-45). + fn extract_static(&self, fv: &FeatureVector) -> Vec { + let mut static_feats = Vec::with_capacity(self.num_static); + for i in 40..(40 + self.num_static) { + static_feats.push(fv.values.get(i).copied().unwrap_or(0.0) as f32); + } + static_feats + } + + /// Extract unknown features from a feature vector (indices 0..num_unknown). + fn extract_unknown(&self, fv: &FeatureVector) -> Vec { + let mut unknown = Vec::with_capacity(self.num_unknown); + for i in 0..self.num_unknown { + unknown.push(fv.values.get(i).copied().unwrap_or(0.0) as f32); + } + unknown + } + + /// Generate known future features from timestamps (calendar features). + fn generate_future(&self, base_timestamp: i64) -> Vec { + let mut futures = Vec::with_capacity(self.prediction_horizon * self.num_known); + for step in 0..self.prediction_horizon { + let ts = base_timestamp + (step as i64 * 60_000_000); // 1-min bars + // Simple calendar: hour_sin, hour_cos, day_sin, day_cos, minute_sin, minute_cos + let hour = ((ts / 3_600_000_000) % 24) as f64; + let day = ((ts / 86_400_000_000) % 7) as f64; + let minute = ((ts / 60_000_000) % 60) as f64; + let mut step_feats = vec![ + (hour * std::f64::consts::PI / 12.0).sin() as f32, + (hour * std::f64::consts::PI / 12.0).cos() as f32, + (day * std::f64::consts::PI / 3.5).sin() as f32, + (day * std::f64::consts::PI / 3.5).cos() as f32, + (minute * std::f64::consts::PI / 30.0).sin() as f32, + (minute * std::f64::consts::PI / 30.0).cos() as f32, + ]; + step_feats.resize(self.num_known, 0.0); + futures.extend(step_feats); + } + futures + } +} + +impl ModelInferenceAdapter for TftInferenceAdapter { + fn model_name(&self) -> &str { + "TFT" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + // Add to buffer + let mut buffer = self.buffer.lock().map_err(|e| { + MLError::LockError(format!("TFT buffer lock failed: {}", e)) + })?; + buffer.push_back(features.clone()); + if buffer.len() > self.sequence_length { + buffer.pop_front(); + } + + // Not enough history yet — return neutral prediction + if buffer.len() < self.sequence_length { + return Ok(EnsemblePrediction { + model_name: "TFT".to_string(), + direction: 0.0, + confidence: 0.0, + metadata: PredictionMeta { + latency_us: start.elapsed().as_micros() as u64, + ..Default::default() + }, + }); + } + + // Build tensors from buffer + let first = buffer.front().ok_or_else(|| { + MLError::InferenceError("empty buffer".to_string()) + })?; + + // Static: [1, num_static] + let static_data = self.extract_static(first); + let static_tensor = Tensor::from_vec( + static_data, (1, self.num_static), &self.device, + ).map_err(|e| MLError::TensorCreationError { + operation: "TFT static tensor".to_string(), + reason: e.to_string(), + })?; + + // Historical: [1, seq_len, num_unknown] + let mut hist_data = Vec::with_capacity(self.sequence_length * self.num_unknown); + for fv in buffer.iter() { + hist_data.extend(self.extract_unknown(fv)); + } + let hist_tensor = Tensor::from_vec( + hist_data, (1, self.sequence_length, self.num_unknown), &self.device, + ).map_err(|e| MLError::TensorCreationError { + operation: "TFT historical tensor".to_string(), + reason: e.to_string(), + })?; + + // Future: [1, horizon, num_known] + let future_data = self.generate_future(features.timestamp); + let future_tensor = Tensor::from_vec( + future_data, (1, self.prediction_horizon, self.num_known), &self.device, + ).map_err(|e| MLError::TensorCreationError { + operation: "TFT future tensor".to_string(), + reason: e.to_string(), + })?; + + drop(buffer); // Release buffer lock before model lock + + // Forward pass + let mut model = self.model.lock().map_err(|e| { + MLError::LockError(format!("TFT model lock failed: {}", e)) + })?; + let output = model.forward(&static_tensor, &hist_tensor, &future_tensor)?; + + // Output: [1, horizon, num_quantiles] + // Extract first timestep quantiles + let quantiles_tensor = output + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("squeeze batch: {}", e)))?; + let first_step = quantiles_tensor + .get(0) + .map_err(|e| MLError::ModelError(format!("get first step: {}", e)))?; + let quantiles: Vec = first_step + .to_vec1::() + .map_err(|e| MLError::ModelError(format!("quantiles to_vec1: {}", e)))?; + + // Direction from median quantile (index num_quantiles/2) + let median_idx = self.num_quantiles / 2; + let median = quantiles.get(median_idx).copied().unwrap_or(0.0) as f64; + let direction = median.signum() * (1.0 - (-median.abs()).exp()); // smooth mapping to [-1,1] + + // Confidence from IQR: narrow spread = high confidence + let q25_idx = self.num_quantiles / 4; + let q75_idx = 3 * self.num_quantiles / 4; + let q25 = quantiles.get(q25_idx).copied().unwrap_or(0.0) as f64; + let q75 = quantiles.get(q75_idx).copied().unwrap_or(0.0) as f64; + let iqr = (q75 - q25).abs(); + let confidence = if median.abs() > 1e-8 { + (1.0 - iqr / median.abs()).clamp(0.0, 1.0) + } else { + 0.5 + }; + + let latency = start.elapsed().as_micros() as u64; + + Ok(EnsemblePrediction { + model_name: "TFT".to_string(), + direction: direction.clamp(-1.0, 1.0), + confidence: confidence.clamp(0.0, 1.0), + metadata: PredictionMeta { + latency_us: latency, + quantiles: Some(quantiles.iter().map(|&q| q as f64).collect()), + ..Default::default() + }, + }) + } + + fn is_ready(&self) -> bool { + self.buffer + .lock() + .map(|b| b.len() >= self.sequence_length) + .unwrap_or(false) + } +} + +unsafe impl Send for TftInferenceAdapter {} +unsafe impl Sync for TftInferenceAdapter {} +``` + +**Step 3: Register module** + +In `ml/src/ensemble/adapters/mod.rs`, add: + +```rust +pub mod tft; +pub use tft::TftInferenceAdapter; +``` + +**Step 4: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::adapters::tft::tests -v` +Expected: 3 tests pass + +**Step 5: Commit** + +```bash +git add ml/src/ensemble/adapters/tft.rs ml/src/ensemble/adapters/mod.rs +git commit -m "feat(ensemble): add TFT inference adapter with sequence buffering" +``` + +--- + +### Task 5: Mamba2 Inference Adapter + +**Files:** +- Create: `ml/src/ensemble/adapters/mamba2.rs` +- Modify: `ml/src/ensemble/adapters/mod.rs` + +**Context:** +- `Mamba2SSM::new(config: Mamba2Config, device: &Device) -> Result` at `ml/src/mamba/mod.rs:640` +- `forward(&mut self, input: &Tensor) -> Result` at line 784 +- Input: `Tensor([batch, seq, d_model], F64)`, Output: `Tensor([batch, seq, 1], F64)` +- d_model = 225 — zero-pad 51-dim feature to 225 +- Must buffer sequence of feature vectors +- Uses DType::F64 + +**Step 1: Write the test** + +```rust +#[cfg(test)] +mod tests { + use super::*; + use crate::ensemble::inference_adapter::FeatureVector; + + fn test_config() -> crate::mamba::Mamba2Config { + crate::mamba::Mamba2Config { + d_model: 32, // Small for testing + d_state: 8, + d_head: 8, + num_heads: 2, + expand: 2, + num_layers: 1, + max_seq_len: 4, + dropout: 0.0, + ..Default::default() + } + } + + #[test] + fn test_mamba2_adapter_creation() { + let config = test_config(); + let adapter = Mamba2InferenceAdapter::new(config, 4); + assert!(adapter.is_ok()); + let adapter = adapter.unwrap(); + assert_eq!(adapter.model_name(), "MAMBA-2"); + assert!(!adapter.is_ready()); // not ready until buffer filled + } + + #[test] + fn test_mamba2_adapter_buffers_and_predicts() { + let config = test_config(); + let mut adapter = Mamba2InferenceAdapter::new(config, 4).unwrap(); + + // Fill buffer + for i in 0..4 { + let fv = FeatureVector { + values: vec![0.1 * (i as f64 + 1.0); 51], + timestamp: 1700000000 + i * 1000, + }; + let _ = adapter.predict(&fv); + } + + // Predict with full buffer + let fv = FeatureVector { + values: vec![0.5; 51], + timestamp: 1700004000, + }; + let pred = adapter.predict(&fv); + assert!(pred.is_ok(), "prediction failed: {:?}", pred.err()); + let pred = pred.unwrap(); + assert!(pred.direction >= -1.0 && pred.direction <= 1.0); + assert!(pred.confidence >= 0.0 && pred.confidence <= 1.0); + } +} +``` + +**Step 2: Write the implementation** + +```rust +//! Mamba2 SSM inference adapter for ensemble prediction + +use std::collections::VecDeque; +use std::sync::Mutex; +use candle_core::{DType, Device, Tensor}; +use crate::mamba::{Mamba2Config, Mamba2SSM}; +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::{MLError, MLResult}; + +/// Mamba2 SSM inference adapter with sequence buffering. +/// +/// Buffers feature vectors, zero-pads each to d_model dimensions, +/// and runs the SSM forward pass. +pub struct Mamba2InferenceAdapter { + model: Mutex, + buffer: Mutex>>, + sequence_length: usize, + d_model: usize, + device: Device, +} + +impl Mamba2InferenceAdapter { + pub fn new(config: Mamba2Config, sequence_length: usize) -> MLResult { + let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu); + let d_model = config.d_model; + let model = Mamba2SSM::new(config, &device)?; + Ok(Self { + model: Mutex::new(model), + buffer: Mutex::new(VecDeque::with_capacity(sequence_length + 1)), + sequence_length, + d_model, + device, + }) + } + + /// Zero-pad a 51-dim feature vector to d_model dimensions. + fn pad_to_d_model(&self, values: &[f64]) -> Vec { + let mut padded = Vec::with_capacity(self.d_model); + for &v in values.iter().take(self.d_model) { + padded.push(v); + } + padded.resize(self.d_model, 0.0); + padded + } +} + +impl ModelInferenceAdapter for Mamba2InferenceAdapter { + fn model_name(&self) -> &str { + "MAMBA-2" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + let padded = self.pad_to_d_model(&features.values); + + let mut buffer = self.buffer.lock().map_err(|e| { + MLError::LockError(format!("Mamba2 buffer lock failed: {}", e)) + })?; + buffer.push_back(padded); + if buffer.len() > self.sequence_length { + buffer.pop_front(); + } + + if buffer.len() < self.sequence_length { + return Ok(EnsemblePrediction { + model_name: "MAMBA-2".to_string(), + direction: 0.0, + confidence: 0.0, + metadata: PredictionMeta { + latency_us: start.elapsed().as_micros() as u64, + ..Default::default() + }, + }); + } + + // Build tensor [1, seq_len, d_model] from buffer (F64) + let flat: Vec = buffer.iter().flat_map(|v| v.iter().copied()).collect(); + let seq_len = buffer.len(); + drop(buffer); // Release buffer lock + + let input = Tensor::from_vec( + flat, (1, seq_len, self.d_model), &self.device, + ).map_err(|e| MLError::TensorCreationError { + operation: "Mamba2 sequence tensor".to_string(), + reason: e.to_string(), + })?; + + // Forward pass → [1, seq_len, 1] + let mut model = self.model.lock().map_err(|e| { + MLError::LockError(format!("Mamba2 model lock failed: {}", e)) + })?; + let output = model.forward(&input)?; + + // Extract last timestep's prediction + let last_step = output + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("squeeze batch: {}", e)))? + .get(seq_len - 1) + .map_err(|e| MLError::ModelError(format!("get last step: {}", e)))? + .to_vec1::() + .map_err(|e| MLError::ModelError(format!("to_vec1: {}", e)))?; + + let prob = last_step.first().copied().unwrap_or(0.5); + + // Direction: map [0,1] probability to [-1,1] + let direction = 2.0 * prob - 1.0; + // Confidence: distance from 0.5 (max uncertainty) + let confidence = (prob - 0.5).abs() * 2.0; + + let latency = start.elapsed().as_micros() as u64; + + Ok(EnsemblePrediction { + model_name: "MAMBA-2".to_string(), + direction: direction.clamp(-1.0, 1.0), + confidence: confidence.clamp(0.0, 1.0), + metadata: PredictionMeta { + latency_us: latency, + ..Default::default() + }, + }) + } + + fn is_ready(&self) -> bool { + self.buffer + .lock() + .map(|b| b.len() >= self.sequence_length) + .unwrap_or(false) + } +} + +unsafe impl Send for Mamba2InferenceAdapter {} +unsafe impl Sync for Mamba2InferenceAdapter {} +``` + +**Step 3: Register module** + +In `ml/src/ensemble/adapters/mod.rs`, add: + +```rust +pub mod mamba2; +pub use mamba2::Mamba2InferenceAdapter; +``` + +**Step 4: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::adapters::mamba2::tests -v` +Expected: 2 tests pass + +**Step 5: Commit** + +```bash +git add ml/src/ensemble/adapters/mamba2.rs ml/src/ensemble/adapters/mod.rs +git commit -m "feat(ensemble): add Mamba2 inference adapter with sequence buffering" +``` + +--- + +### Task 6: Mamba2 Directional Loss Function + +**Files:** +- Create: `ml/src/mamba/loss.rs` +- Modify: `ml/src/mamba/mod.rs` (add `pub mod loss;` declaration) + +**Context:** +- Current Mamba2 uses MSE only (`ml/src/mamba/mod.rs`) +- Need `directional_mse_loss` that penalizes wrong-direction predictions more heavily +- All Mamba2 tensors use DType::F64 + +**Step 1: Write the test** + +```rust +#[cfg(test)] +mod tests { + use super::*; + use candle_core::{Device, Tensor}; + + #[test] + fn test_directional_loss_same_as_mse_when_correct_direction() { + let device = Device::Cpu; + let preds = Tensor::from_vec(vec![0.3, 0.5, 0.8], 3, &device).unwrap(); + let targets = Tensor::from_vec(vec![0.2, 0.6, 0.9], 3, &device).unwrap(); + + let dir_loss = directional_mse_loss(&preds, &targets, 2.0).unwrap(); + let mse = mse_loss(&preds, &targets).unwrap(); + + // All same direction (positive), so directional loss == MSE + let dir_val: f64 = dir_loss.to_vec0().unwrap(); + let mse_val: f64 = mse.to_vec0().unwrap(); + assert!((dir_val - mse_val).abs() < 1e-6, + "same-direction: dir={} vs mse={}", dir_val, mse_val); + } + + #[test] + fn test_directional_loss_higher_for_wrong_direction() { + let device = Device::Cpu; + // Prediction positive, target negative → wrong direction + let preds = Tensor::from_vec(vec![0.3_f64], 1, &device).unwrap(); + let targets = Tensor::from_vec(vec![-0.2_f64], 1, &device).unwrap(); + + let dir_loss = directional_mse_loss(&preds, &targets, 2.0).unwrap(); + let mse = mse_loss(&preds, &targets).unwrap(); + + let dir_val: f64 = dir_loss.to_vec0().unwrap(); + let mse_val: f64 = mse.to_vec0().unwrap(); + assert!(dir_val > mse_val, + "wrong-direction loss {} should be > mse {}", dir_val, mse_val); + } + + #[test] + fn test_directional_loss_weight_1_equals_mse() { + let device = Device::Cpu; + let preds = Tensor::from_vec(vec![0.3_f64, -0.1], 2, &device).unwrap(); + let targets = Tensor::from_vec(vec![-0.2_f64, 0.5], 2, &device).unwrap(); + + let dir_loss = directional_mse_loss(&preds, &targets, 1.0).unwrap(); + let mse = mse_loss(&preds, &targets).unwrap(); + + let dir_val: f64 = dir_loss.to_vec0().unwrap(); + let mse_val: f64 = mse.to_vec0().unwrap(); + assert!((dir_val - mse_val).abs() < 1e-6, + "weight=1.0: dir={} should == mse={}", dir_val, mse_val); + } +} +``` + +**Step 2: Write the implementation** + +Create `ml/src/mamba/loss.rs`: + +```rust +//! Loss functions for Mamba2 SSM training +//! +//! Includes directional MSE that penalizes wrong-sign predictions +//! more heavily — critical for trading where direction matters more +//! than magnitude. + +use candle_core::{DType, Tensor}; +use crate::{MLError, MLResult}; + +/// Standard MSE loss (baseline for comparison) +pub fn mse_loss(predictions: &Tensor, targets: &Tensor) -> MLResult { + let diff = predictions.sub(targets) + .map_err(|e| MLError::ModelError(format!("mse sub: {}", e)))?; + let sq = diff.sqr() + .map_err(|e| MLError::ModelError(format!("mse sqr: {}", e)))?; + sq.mean_all() + .map_err(|e| MLError::ModelError(format!("mse mean: {}", e))) +} + +/// Directional MSE loss. +/// +/// When prediction and target have opposite signs (wrong direction), +/// the squared error is multiplied by `direction_weight` (default: 2.0). +/// When signs agree, normal MSE is used. +/// +/// `direction_weight = 1.0` produces identical results to standard MSE. +/// `direction_weight = 2.0` penalizes wrong-direction errors 2x. +pub fn directional_mse_loss( + predictions: &Tensor, + targets: &Tensor, + direction_weight: f64, +) -> MLResult { + let diff = predictions.sub(targets) + .map_err(|e| MLError::ModelError(format!("dir_mse sub: {}", e)))?; + let sq = diff.sqr() + .map_err(|e| MLError::ModelError(format!("dir_mse sqr: {}", e)))?; + + // Detect wrong direction: sign(pred) != sign(target) + let zeros = Tensor::zeros_like(predictions) + .map_err(|e| MLError::ModelError(format!("dir_mse zeros: {}", e)))?; + let pred_positive = predictions.ge(&zeros) + .map_err(|e| MLError::ModelError(format!("dir_mse pred_ge: {}", e)))?; + let target_positive = targets.ge(&zeros) + .map_err(|e| MLError::ModelError(format!("dir_mse target_ge: {}", e)))?; + + // wrong_direction = (pred_positive XOR target_positive) as float + let wrong = pred_positive.ne(&target_positive) + .map_err(|e| MLError::ModelError(format!("dir_mse ne: {}", e)))? + .to_dtype(predictions.dtype()) + .map_err(|e| MLError::ModelError(format!("dir_mse to_dtype: {}", e)))?; + + // weight = 1.0 + wrong * (direction_weight - 1.0) + let extra = wrong.affine(direction_weight - 1.0, 1.0) + .map_err(|e| MLError::ModelError(format!("dir_mse affine: {}", e)))?; + + let weighted = sq.mul(&extra) + .map_err(|e| MLError::ModelError(format!("dir_mse mul: {}", e)))?; + + weighted.mean_all() + .map_err(|e| MLError::ModelError(format!("dir_mse mean: {}", e))) +} +``` + +**Step 3: Register module** + +Find the module declarations section near the top of `ml/src/mamba/mod.rs` and add: + +```rust +pub mod loss; +``` + +**Step 4: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib mamba::loss::tests -v` +Expected: 3 tests pass + +**Step 5: Commit** + +```bash +git add ml/src/mamba/loss.rs ml/src/mamba/mod.rs +git commit -m "feat(mamba): add directional MSE loss function" +``` + +--- + +### Task 7: Mamba2 Discretization Fix + +**Files:** +- Modify: `ml/src/mamba/mod.rs` (lines ~913-929) + +**Context:** +- Current `discretize_ssm()` at line 913 uses ZOH: `A_disc = I + A_cont * dt` +- `ssd_layer.rs` uses bilinear: `A_disc = I + A*dt + (A*dt)^2/2` +- Need to standardize on bilinear in `mod.rs` for consistency + +**Step 1: Write the test** + +Add to `ml/src/mamba/mod.rs` tests section (find `#[cfg(test)]`): + +```rust + #[test] + fn test_bilinear_discretization_more_accurate_than_zoh() { + // For a simple scalar A=-1, dt=0.1: + // Exact: exp(-0.1) = 0.9048 + // ZOH: 1 + (-1)*0.1 = 0.9000 + // Bilinear: 1 + (-1)*0.1 + ((-1)*0.1)^2/2 = 0.9050 + // Bilinear is closer to exact + let zoh = 1.0 + (-1.0) * 0.1; + let bilinear = 1.0 + (-1.0) * 0.1 + ((-1.0) * 0.1_f64).powi(2) / 2.0; + let exact = (-0.1_f64).exp(); + + assert!((bilinear - exact).abs() < (zoh - exact).abs(), + "bilinear {} should be closer to exact {} than zoh {}", + bilinear, exact, zoh); + } +``` + +**Step 2: Fix the discretization** + +In `ml/src/mamba/mod.rs`, replace the body of `discretize_ssm` (around lines 913-929): + +```rust + fn discretize_ssm(&self, A_cont: &Tensor, dt: &Tensor) -> Result { + // dt is [d_model] but A_cont is [d_state, d_state] + // Use mean of dt as a scalar for discretization + let dt_mean = dt.mean_all()?; + let dt_scalar = dt_mean.to_vec0::()?; + + let dt_tensor = Tensor::from_slice(&[dt_scalar], &[1], A_cont.device())?.reshape(&[])?; + + // Bilinear (Tustin) approximation: A_disc = I + A*dt + (A*dt)^2 / 2 + // More accurate than ZOH (I + A*dt), matches ssd_layer.rs + let A_dt = A_cont.broadcast_mul(&dt_tensor)?; + let A_dt_sq = A_dt.matmul(&A_dt)?; + let half = Tensor::from_slice(&[0.5_f64], &[1], A_cont.device())?.reshape(&[])?; + let second_order = A_dt_sq.broadcast_mul(&half)?; + + let identity = Tensor::eye(A_cont.dim(0)?, DType::F64, A_cont.device())?; + let A_discrete = ((&identity + &A_dt)? + &second_order)?; + + Ok(A_discrete) + } +``` + +**Step 3: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib mamba -v` (run all mamba tests) +Expected: All existing + new test pass + +**Step 4: Commit** + +```bash +git add ml/src/mamba/mod.rs +git commit -m "fix(mamba): standardize discretization to bilinear (Tustin)" +``` + +--- + +### Task 8: Coordinator Integration + +**Files:** +- Modify: `ml/src/ensemble/coordinator.rs` + +**Context:** +- Replace `generate_mock_predictions` with real adapter calls +- Keep backward compatibility: if no adapters registered, fall back to mock +- `EnsembleCoordinator` already has `active_models: Arc>` +- New field: `adapters: Vec>` +- Must convert adapter `EnsemblePrediction` → existing `ModelPrediction` + +**Step 1: Write the test** + +Add to `ml/src/ensemble/coordinator.rs` tests: + +```rust + #[tokio::test] + async fn test_ensemble_with_real_adapter() { + use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, + }; + + // Create a simple mock adapter for testing + struct TestAdapter; + impl ModelInferenceAdapter for TestAdapter { + fn model_name(&self) -> &str { "test" } + fn predict(&self, _features: &FeatureVector) -> crate::MLResult { + Ok(EnsemblePrediction { + model_name: "test".to_string(), + direction: 0.6, + confidence: 0.85, + metadata: PredictionMeta::default(), + }) + } + fn is_ready(&self) -> bool { true } + } + + let mut coordinator = EnsembleCoordinator::new(); + coordinator.add_adapter(Box::new(TestAdapter)); + coordinator.register_model("test".to_string(), 1.0).await.unwrap(); + + let features = Features::new( + vec![0.5; 10], + vec!["f1","f2","f3","f4","f5","f6","f7","f8","f9","f10"] + .into_iter().map(String::from).collect(), + ); + let decision = coordinator.predict(&features).await.unwrap(); + assert!(decision.confidence > 0.0); + assert_eq!(decision.model_count(), 1); + } + + #[tokio::test] + async fn test_ensemble_graceful_no_adapters() { + // Without adapters, falls back to mock predictions + let coordinator = EnsembleCoordinator::new(); + coordinator.register_model("DQN".to_string(), 0.5).await.unwrap(); + coordinator.register_model("PPO".to_string(), 0.5).await.unwrap(); + + let features = Features::new( + vec![0.5; 5], + vec!["f1","f2","f3","f4","f5"] + .into_iter().map(String::from).collect(), + ); + let decision = coordinator.predict(&features).await.unwrap(); + assert!(decision.confidence >= 0.0); + } +``` + +**Step 2: Modify the coordinator** + +In `ml/src/ensemble/coordinator.rs`: + +1. Add the `adapters` field to `EnsembleCoordinator`: + +```rust +pub struct EnsembleCoordinator { + active_models: Arc>, + aggregator: Arc, + model_weights: Arc>>, + config: EnsembleConfig, + /// Real model inference adapters (when available) + adapters: Vec>, +} +``` + +2. Add `adapters: Vec::new()` to `new()`. + +3. Add `add_adapter` method: + +```rust + /// Add a real model inference adapter + pub fn add_adapter(&mut self, adapter: Box) { + info!("Added inference adapter: {}", adapter.model_name()); + self.adapters.push(adapter); + } +``` + +4. Update `predict()` to use adapters when available: + +Replace `generate_mock_predictions` call with logic that tries adapters first, falls back to mocks for models without adapters. + +In `generate_mock_predictions`, add at the top of the loop body: + +```rust + for (model_id, checkpoint_opt) in model_info { + // Try real adapter first + if let Some(adapter) = self.adapters.iter().find(|a| a.model_name() == model_id && a.is_ready()) { + let fv = FeatureVector { + values: features.values.clone(), + timestamp: features.timestamp as i64, + }; + match adapter.predict(&fv) { + Ok(ensemble_pred) => { + let prediction = ModelPrediction::new( + model_id.clone(), + ensemble_pred.direction, + ensemble_pred.confidence, + ); + predictions.push(prediction); + continue; + } + Err(e) => { + warn!("Adapter {} inference failed, falling back to mock: {}", model_id, e); + } + } + } + // ... existing mock fallback below +``` + +5. Add the import at the top: + +```rust +use crate::ensemble::inference_adapter::{FeatureVector, ModelInferenceAdapter}; +``` + +**Step 3: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::coordinator::tests -v` +Expected: All existing + 2 new tests pass + +**Step 4: Commit** + +```bash +git add ml/src/ensemble/coordinator.rs +git commit -m "feat(ensemble): integrate real inference adapters into coordinator" +``` + +--- + +### Task 9: Individual Model Validation Tests + +**Files:** +- Add tests to each adapter file (inline `#[cfg(test)]` modules) + +**Purpose:** Validate each model can do a round-trip: create → save → load → predict → verify output. + +**Step 1: Add round-trip test to DQN adapter** + +In `ml/src/ensemble/adapters/dqn.rs` tests: + +```rust + #[test] + fn test_dqn_checkpoint_round_trip() { + let config = test_config(); + let adapter = DqnInferenceAdapter::new(config.clone()).unwrap(); + let fv = FeatureVector { + values: vec![0.3; 51], + timestamp: 1700000000, + }; + // Predict with random weights + let pred1 = adapter.predict(&fv).unwrap(); + + // Save checkpoint + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("dqn_test.safetensors"); + { + let model = adapter.model.lock().unwrap(); + model.get_q_network_vars() + .save(&path) + .unwrap(); + } + + // Load into new adapter + let adapter2 = DqnInferenceAdapter::from_checkpoint( + config, + path.to_str().unwrap(), + ).unwrap(); + let pred2 = adapter2.predict(&fv).unwrap(); + + // Same weights → same prediction + assert!((pred1.direction - pred2.direction).abs() < 1e-4, + "round-trip mismatch: {} vs {}", pred1.direction, pred2.direction); + } +``` + +**Step 2: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::adapters -v` +Expected: All adapter tests pass + +**Step 3: Commit** + +```bash +git add ml/src/ensemble/adapters/dqn.rs +git commit -m "test(ensemble): add model checkpoint round-trip validation" +``` + +--- + +### Task 10: Full Ensemble Integration Test + +**Files:** +- Add to `ml/src/ensemble/coordinator.rs` tests + +**Step 1: Write the integration test** + +```rust + #[tokio::test] + async fn test_full_ensemble_with_dqn_adapter() { + use crate::ensemble::adapters::DqnInferenceAdapter; + use crate::dqn::dqn::DQNConfig; + + let dqn_config = DQNConfig { + state_dim: 51, + num_actions: 45, + hidden_dims: vec![64, 64], + ..Default::default() + }; + + let adapter = DqnInferenceAdapter::new(dqn_config).unwrap(); + + let mut coordinator = EnsembleCoordinator::new(); + coordinator.add_adapter(Box::new(adapter)); + coordinator.register_model("DQN".to_string(), 1.0).await.unwrap(); + + let features = Features::new( + vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, + 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, + 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, + 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, + 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.55], + (0..51).map(|i| format!("f{}", i)).collect(), + ); + + let decision = coordinator.predict(&features).await.unwrap(); + + // Real adapter should have produced a real prediction + assert!(decision.confidence > 0.0); + assert!(decision.signal >= -1.0 && decision.signal <= 1.0); + assert_eq!(decision.model_count(), 1); + } +``` + +**Step 2: Verify** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble::coordinator::tests -v` +Expected: All tests pass including the new one + +**Step 3: Commit** + +```bash +git add ml/src/ensemble/coordinator.rs +git commit -m "test(ensemble): add full DQN adapter integration test" +``` + +--- + +### Task 11: Final Compilation Check & Cleanup + +**Step 1: Full crate check** + +Run: `SQLX_OFFLINE=true cargo check -p ml` +Expected: No errors (warnings ok) + +**Step 2: Run all ensemble tests** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib ensemble -v` +Expected: All tests pass + +**Step 3: Run mamba tests** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib mamba -v` +Expected: All tests pass including loss and discretization + +**Step 4: Run full ml test suite** + +Run: `SQLX_OFFLINE=true cargo test -p ml --lib` +Expected: No regressions + +**Step 5: Fix any warnings** + +Address unused variable warnings or Clippy violations found during compilation. + +**Step 6: Commit** + +```bash +git add -A +git commit -m "chore(ensemble): cleanup warnings and verify full test suite" +```