diff --git a/ml/src/ensemble/inference_ensemble.rs b/ml/src/ensemble/inference_ensemble.rs new file mode 100644 index 000000000..32a502c4c --- /dev/null +++ b/ml/src/ensemble/inference_ensemble.rs @@ -0,0 +1,297 @@ +//! Lightweight synchronous ensemble coordinator +//! +//! Wraps N [`ModelInferenceAdapter`] instances and aggregates their +//! predictions via confidence-weighted voting. Models that are not +//! ready or that fail inference are gracefully skipped with warnings. + +use std::collections::HashMap; + +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::{MLError, MLResult}; + +/// Synchronous coordinator that aggregates predictions from multiple +/// [`ModelInferenceAdapter`] instances using confidence-weighted voting. +/// +/// # Aggregation formula +/// +/// ```text +/// weighted_direction = sum(direction_i * weight_i * confidence_i) +/// / sum(weight_i * confidence_i) +/// ``` +/// +/// where `weight_i` defaults to 1.0 and can be overridden via +/// [`set_weight`](InferenceEnsemble::set_weight). +#[allow(missing_debug_implementations)] +pub struct InferenceEnsemble { + adapters: Vec>, + weights: HashMap, +} + +impl InferenceEnsemble { + /// Create a new ensemble from a vector of adapters. + /// + /// All adapters start with a default weight of 1.0. + pub fn new(adapters: Vec>) -> Self { + Self { + adapters, + weights: HashMap::new(), + } + } + + /// Set a custom weight for a model identified by name. + /// + /// Models without an explicit weight use 1.0. + pub fn set_weight(&mut self, model_name: &str, weight: f64) { + self.weights.insert(model_name.to_string(), weight); + } + + /// Return the number of adapters whose `is_ready()` returns true. + pub fn ready_count(&self) -> usize { + self.adapters.iter().filter(|a| a.is_ready()).count() + } + + /// Run inference across all ready adapters and aggregate via + /// confidence-weighted voting. + /// + /// Returns [`MLError::InferenceError`] if no models are ready. + pub fn predict(&self, features: &FeatureVector) -> MLResult { + let ready_adapters: Vec<&Box> = + self.adapters.iter().filter(|a| a.is_ready()).collect(); + + if ready_adapters.is_empty() { + return Err(MLError::InferenceError( + "No models are ready for inference".to_string(), + )); + } + + let mut weighted_direction_sum = 0.0_f64; + let mut weight_confidence_sum = 0.0_f64; + let mut confidence_sum = 0.0_f64; + let mut successful_count = 0_usize; + let mut model_names: Vec = Vec::new(); + + for adapter in &ready_adapters { + let model_name = adapter.model_name().to_string(); + match adapter.predict(features) { + Ok(pred) => { + let w = self + .weights + .get(&model_name) + .copied() + .unwrap_or(1.0); + let wc = w * pred.confidence; + weighted_direction_sum += pred.direction * wc; + weight_confidence_sum += wc; + confidence_sum += pred.confidence; + successful_count += 1; + model_names.push(model_name); + } + Err(e) => { + tracing::warn!( + model = %model_name, + error = %e, + "Model prediction failed, skipping" + ); + } + } + } + + if successful_count == 0 { + return Err(MLError::InferenceError( + "All ready models failed during inference".to_string(), + )); + } + + let direction = if weight_confidence_sum.abs() < f64::EPSILON { + 0.0 + } else { + weighted_direction_sum / weight_confidence_sum + }; + + let avg_confidence = confidence_sum / successful_count as f64; + + let ensemble_name = format!("ENSEMBLE({})", model_names.join("+")); + + Ok(EnsemblePrediction { + model_name: ensemble_name, + direction, + confidence: avg_confidence, + metadata: PredictionMeta::default(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A simple test adapter with configurable direction, confidence, and readiness. + struct DummyAdapter { + name: String, + direction: f64, + confidence: f64, + ready: bool, + } + + impl ModelInferenceAdapter for DummyAdapter { + fn model_name(&self) -> &str { + &self.name + } + + fn predict(&self, _features: &FeatureVector) -> MLResult { + Ok(EnsemblePrediction { + model_name: self.name.clone(), + direction: self.direction, + confidence: self.confidence, + metadata: PredictionMeta::default(), + }) + } + + fn is_ready(&self) -> bool { + self.ready + } + } + + fn make_features() -> FeatureVector { + FeatureVector { + values: vec![0.0; 51], + timestamp: 1_700_000_000_000_000, + } + } + + #[test] + fn test_ensemble_equal_weight_aggregation() { + // Model A: bullish (dir=1.0, conf=0.8) + // Model B: bearish (dir=-1.0, conf=0.6) + // weighted_direction = (1.0*1.0*0.8 + (-1.0)*1.0*0.6) / (1.0*0.8 + 1.0*0.6) + // = (0.8 - 0.6) / 1.4 + // = 0.2 / 1.4 + // ~ 0.1429 + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "A".to_string(), + direction: 1.0, + confidence: 0.8, + ready: true, + }), + Box::new(DummyAdapter { + name: "B".to_string(), + direction: -1.0, + confidence: 0.6, + ready: true, + }), + ]; + + let ensemble = InferenceEnsemble::new(adapters); + let features = make_features(); + let pred = ensemble.predict(&features).expect("predict should succeed"); + + // Net direction should be positive (bullish model has higher confidence) + assert!( + pred.direction > 0.0, + "direction should be positive, got {}", + pred.direction + ); + // But not too strong since the bearish model partially cancels + assert!( + pred.direction < 0.5, + "direction should be < 0.5, got {}", + pred.direction + ); + // Confidence = average = (0.8 + 0.6) / 2 = 0.7 + assert!( + (pred.confidence - 0.7).abs() < 1e-9, + "confidence should be 0.7, got {}", + pred.confidence + ); + } + + #[test] + fn test_ensemble_skips_unready_models() { + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "Ready".to_string(), + direction: 1.0, + confidence: 0.9, + ready: true, + }), + Box::new(DummyAdapter { + name: "NotReady".to_string(), + direction: -1.0, + confidence: 0.9, + ready: false, + }), + ]; + + let ensemble = InferenceEnsemble::new(adapters); + assert_eq!(ensemble.ready_count(), 1); + + let features = make_features(); + let pred = ensemble.predict(&features).expect("predict should succeed"); + + // Only the ready (bullish) model contributes + assert!( + pred.direction > 0.5, + "direction should be > 0.5 with only bullish model, got {}", + pred.direction + ); + // Model name should only include the ready model + assert!( + pred.model_name.contains("Ready"), + "model_name should contain 'Ready', got {}", + pred.model_name + ); + assert!( + !pred.model_name.contains("NotReady"), + "model_name should NOT contain 'NotReady', got {}", + pred.model_name + ); + } + + #[test] + fn test_ensemble_custom_weights() { + // Model A: bullish, weight 0.75 + // Model B: bearish, weight 0.25 + // weighted_direction = (1.0*0.75*0.8 + (-1.0)*0.25*0.8) / (0.75*0.8 + 0.25*0.8) + // = (0.6 - 0.2) / (0.6 + 0.2) + // = 0.4 / 0.8 + // = 0.5 + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "Heavy".to_string(), + direction: 1.0, + confidence: 0.8, + ready: true, + }), + Box::new(DummyAdapter { + name: "Light".to_string(), + direction: -1.0, + confidence: 0.8, + ready: true, + }), + ]; + + let mut ensemble = InferenceEnsemble::new(adapters); + ensemble.set_weight("Heavy", 0.75); + ensemble.set_weight("Light", 0.25); + + let features = make_features(); + let pred = ensemble.predict(&features).expect("predict should succeed"); + + // The heavier-weighted bullish model should dominate + assert!( + pred.direction > 0.0, + "direction should be positive (heavy bullish dominates), got {}", + pred.direction + ); + // With equal confidences, the 3:1 weight ratio strongly favors bullish + // Expected: 0.5, but any positive value > 0.3 confirms dominance + assert!( + pred.direction > 0.3, + "direction should be > 0.3 showing heavy model dominance, got {}", + pred.direction + ); + } +} diff --git a/ml/src/ensemble/mod.rs b/ml/src/ensemble/mod.rs index 18a087888..90f41afdc 100644 --- a/ml/src/ensemble/mod.rs +++ b/ml/src/ensemble/mod.rs @@ -18,6 +18,7 @@ pub mod training_integration; // Training integration for ML service pub mod voting; pub mod weights; pub mod inference_adapter; +pub mod inference_ensemble; pub mod adapters; // Re-export key types that are used across ensemble modules