feat: vendor midstream and sublinear-time-solver libraries (#109)

Add ruvnet/midstream (AIMDS real-time inference) and
ruvnet/sublinear-time-solver (sublinear optimization algorithms)
as vendored dependencies under vendor/.
This commit is contained in:
rUv
2026-03-02 23:34:05 -05:00
committed by GitHub
parent 14902e6b4e
commit 407b46b206
1600 changed files with 1852646 additions and 0 deletions
@@ -0,0 +1,484 @@
//! Neural network layer implementations optimized for temporal prediction
use crate::error::{Result, TemporalNeuralError};
use nalgebra::{DMatrix, DVector};
use serde::{Deserialize, Serialize};
use rand::{Rng, distributions::{Distribution, Uniform}};
use std::f64::consts::PI;
/// GRU (Gated Recurrent Unit) layer optimized for micro-networks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GruLayer {
/// Input size
pub input_size: usize,
/// Hidden state size
pub hidden_size: usize,
/// Reset gate weights (input)
pub weight_ir: DMatrix<f64>,
/// Reset gate weights (hidden)
pub weight_hr: DMatrix<f64>,
/// Reset gate bias
pub bias_r: DVector<f64>,
/// Update gate weights (input)
pub weight_iz: DMatrix<f64>,
/// Update gate weights (hidden)
pub weight_hz: DMatrix<f64>,
/// Update gate bias
pub bias_z: DVector<f64>,
/// New gate weights (input)
pub weight_in: DMatrix<f64>,
/// New gate weights (hidden)
pub weight_hn: DMatrix<f64>,
/// New gate bias
pub bias_n: DVector<f64>,
/// Current hidden state
hidden_state: Option<DVector<f64>>,
}
impl GruLayer {
/// Create a new GRU layer
pub fn new(input_size: usize, hidden_size: usize) -> Self {
let mut layer = Self {
input_size,
hidden_size,
weight_ir: DMatrix::zeros(hidden_size, input_size),
weight_hr: DMatrix::zeros(hidden_size, hidden_size),
bias_r: DVector::zeros(hidden_size),
weight_iz: DMatrix::zeros(hidden_size, input_size),
weight_hz: DMatrix::zeros(hidden_size, hidden_size),
bias_z: DVector::zeros(hidden_size),
weight_in: DMatrix::zeros(hidden_size, input_size),
weight_hn: DMatrix::zeros(hidden_size, hidden_size),
bias_n: DVector::zeros(hidden_size),
hidden_state: None,
};
layer.initialize_weights();
layer
}
/// Initialize weights using Xavier/Glorot initialization
pub fn initialize_weights(&mut self) {
let mut rng = rand::thread_rng();
// Xavier initialization scale
let input_scale = (6.0 / (self.input_size + self.hidden_size) as f64).sqrt();
let hidden_scale = (6.0 / (2.0 * self.hidden_size) as f64).sqrt();
// Initialize input weights
let uniform = Uniform::new(-input_scale, input_scale);
self.weight_ir = DMatrix::from_fn(self.hidden_size, self.input_size, |_, _| {
uniform.sample(&mut rng)
});
self.weight_iz = DMatrix::from_fn(self.hidden_size, self.input_size, |_, _| {
uniform.sample(&mut rng)
});
self.weight_in = DMatrix::from_fn(self.hidden_size, self.input_size, |_, _| {
uniform.sample(&mut rng)
});
// Initialize hidden weights
let hidden_uniform = Uniform::new(-hidden_scale, hidden_scale);
self.weight_hr = DMatrix::from_fn(self.hidden_size, self.hidden_size, |_, _| {
hidden_uniform.sample(&mut rng)
});
self.weight_hz = DMatrix::from_fn(self.hidden_size, self.hidden_size, |_, _| {
hidden_uniform.sample(&mut rng)
});
self.weight_hn = DMatrix::from_fn(self.hidden_size, self.hidden_size, |_, _| {
hidden_uniform.sample(&mut rng)
});
// Initialize biases to small positive values for forget gates
self.bias_z.fill(1.0); // Update gate bias - helps with gradient flow
}
/// Forward pass through GRU layer
pub fn forward(&mut self, input: &DVector<f64>) -> Result<DVector<f64>> {
if input.len() != self.input_size {
return Err(TemporalNeuralError::ModelError {
component: "GruLayer".to_string(),
message: format!(
"Input size mismatch: expected {}, got {}",
self.input_size, input.len()
),
context: vec![],
});
}
// Initialize hidden state if needed
if self.hidden_state.is_none() {
self.hidden_state = Some(DVector::zeros(self.hidden_size));
}
let h_prev = self.hidden_state.as_ref().unwrap().clone();
// Reset gate: r = sigmoid(W_ir @ x + W_hr @ h + b_r)
let r = sigmoid(&(&self.weight_ir * input + &self.weight_hr * &h_prev + &self.bias_r));
// Update gate: z = sigmoid(W_iz @ x + W_hz @ h + b_z)
let z = sigmoid(&(&self.weight_iz * input + &self.weight_hz * &h_prev + &self.bias_z));
// New gate: n = tanh(W_in @ x + W_hn @ (r ⊙ h) + b_n)
let r_h = r.component_mul(&h_prev);
let n = tanh(&(&self.weight_in * input + &self.weight_hn * &r_h + &self.bias_n));
// Hidden state: h = (1 - z) ⊙ n + z ⊙ h_prev
let one_minus_z = DVector::from_element(self.hidden_size, 1.0) - &z;
let h_new = one_minus_z.component_mul(&n) + z.component_mul(&h_prev);
self.hidden_state = Some(h_new.clone());
Ok(h_new)
}
/// Process a sequence of inputs
pub fn forward_sequence(&mut self, inputs: &DMatrix<f64>) -> Result<DMatrix<f64>> {
let seq_len = inputs.ncols();
let mut outputs = DMatrix::zeros(self.hidden_size, seq_len);
for t in 0..seq_len {
let input = inputs.column(t);
let output = self.forward(&input.into())?;
outputs.set_column(t, &output);
}
Ok(outputs)
}
/// Reset hidden state
pub fn reset_state(&mut self) {
self.hidden_state = None;
}
/// Get current hidden state
pub fn get_state(&self) -> Option<&DVector<f64>> {
self.hidden_state.as_ref()
}
/// Set hidden state (for initialization or transfer)
pub fn set_state(&mut self, state: DVector<f64>) -> Result<()> {
if state.len() != self.hidden_size {
return Err(TemporalNeuralError::ModelError {
component: "GruLayer".to_string(),
message: format!(
"State size mismatch: expected {}, got {}",
self.hidden_size, state.len()
),
context: vec![],
});
}
self.hidden_state = Some(state);
Ok(())
}
/// Get parameter count
pub fn parameter_count(&self) -> usize {
self.weight_ir.len() + self.weight_hr.len() + self.bias_r.len() +
self.weight_iz.len() + self.weight_hz.len() + self.bias_z.len() +
self.weight_in.len() + self.weight_hn.len() + self.bias_n.len()
}
}
/// Temporal Convolutional Network (TCN) layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TcnLayer {
/// Number of input channels
pub input_channels: usize,
/// Number of output channels
pub output_channels: usize,
/// Kernel size
pub kernel_size: usize,
/// Dilation factor
pub dilation: usize,
/// Convolution weights
pub weight: DMatrix<f64>,
/// Bias terms
pub bias: DVector<f64>,
/// Whether to use residual connections
pub residual: bool,
/// Residual projection weights (if needed)
pub residual_weight: Option<DMatrix<f64>>,
}
impl TcnLayer {
/// Create a new TCN layer
pub fn new(
input_channels: usize,
output_channels: usize,
kernel_size: usize,
dilation: usize,
residual: bool,
) -> Self {
let weight_size = output_channels * input_channels * kernel_size;
let mut layer = Self {
input_channels,
output_channels,
kernel_size,
dilation,
weight: DMatrix::zeros(output_channels, input_channels * kernel_size),
bias: DVector::zeros(output_channels),
residual,
residual_weight: None,
};
// Create residual projection if channel sizes don't match
if residual && input_channels != output_channels {
layer.residual_weight = Some(DMatrix::zeros(output_channels, input_channels));
}
layer.initialize_weights();
layer
}
/// Initialize weights
pub fn initialize_weights(&mut self) {
let mut rng = rand::thread_rng();
let fan_in = self.input_channels * self.kernel_size;
let fan_out = self.output_channels * self.kernel_size;
let scale = (2.0 / (fan_in + fan_out) as f64).sqrt();
let uniform = Uniform::new(-scale, scale);
self.weight = DMatrix::from_fn(self.output_channels, self.input_channels * self.kernel_size, |_, _| {
uniform.sample(&mut rng)
});
if let Some(ref mut res_weight) = self.residual_weight {
let res_scale = (2.0 / (self.input_channels + self.output_channels) as f64).sqrt();
let res_uniform = Uniform::new(-res_scale, res_scale);
*res_weight = DMatrix::from_fn(self.output_channels, self.input_channels, |_, _| {
res_uniform.sample(&mut rng)
});
}
}
/// Forward pass with causal convolution
pub fn forward(&self, input: &DMatrix<f64>) -> Result<DMatrix<f64>> {
let (channels, seq_len) = (input.nrows(), input.ncols());
if channels != self.input_channels {
return Err(TemporalNeuralError::ModelError {
component: "TcnLayer".to_string(),
message: format!(
"Input channel mismatch: expected {}, got {}",
self.input_channels, channels
),
context: vec![],
});
}
// Calculate output sequence length (causal convolution doesn't reduce length)
let output_len = seq_len;
let mut output = DMatrix::zeros(self.output_channels, output_len);
// Perform causal dilated convolution
for t in 0..output_len {
for out_ch in 0..self.output_channels {
let mut sum = self.bias[out_ch];
for k in 0..self.kernel_size {
let input_t = t as i64 - (k * self.dilation) as i64;
if input_t >= 0 {
let input_t = input_t as usize;
for in_ch in 0..self.input_channels {
let weight_idx = in_ch * self.kernel_size + k;
sum += self.weight[(out_ch, weight_idx)] * input[(in_ch, input_t)];
}
}
}
output[(out_ch, t)] = sum;
}
}
// Apply residual connection if enabled
if self.residual {
if let Some(ref res_weight) = self.residual_weight {
// Project input to match output channels
let residual = res_weight * input;
output += residual;
} else if self.input_channels == self.output_channels {
// Direct residual connection
output += input;
}
}
Ok(output)
}
/// Get parameter count
pub fn parameter_count(&self) -> usize {
let mut count = self.weight.len() + self.bias.len();
if let Some(ref res_weight) = self.residual_weight {
count += res_weight.len();
}
count
}
}
/// Dense (fully connected) layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DenseLayer {
/// Input dimension
pub input_dim: usize,
/// Output dimension
pub output_dim: usize,
/// Weight matrix
pub weight: DMatrix<f64>,
/// Bias vector
pub bias: DVector<f64>,
/// Activation function
pub activation: ActivationFunction,
}
/// Activation function types
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum ActivationFunction {
/// Linear (no activation)
Linear,
/// ReLU activation
Relu,
/// Tanh activation
Tanh,
/// Sigmoid activation
Sigmoid,
/// GELU activation
Gelu,
}
impl DenseLayer {
/// Create a new dense layer
pub fn new(input_dim: usize, output_dim: usize, activation: ActivationFunction) -> Self {
let mut layer = Self {
input_dim,
output_dim,
weight: DMatrix::zeros(output_dim, input_dim),
bias: DVector::zeros(output_dim),
activation,
};
layer.initialize_weights();
layer
}
/// Initialize weights
pub fn initialize_weights(&mut self) {
let mut rng = rand::thread_rng();
let scale = (2.0 / (self.input_dim + self.output_dim) as f64).sqrt();
let uniform = Uniform::new(-scale, scale);
self.weight = DMatrix::from_fn(self.output_dim, self.input_dim, |_, _| {
uniform.sample(&mut rng)
});
}
/// Forward pass
pub fn forward(&self, input: &DVector<f64>) -> Result<DVector<f64>> {
if input.len() != self.input_dim {
return Err(TemporalNeuralError::ModelError {
component: "DenseLayer".to_string(),
message: format!(
"Input dimension mismatch: expected {}, got {}",
self.input_dim, input.len()
),
context: vec![],
});
}
let linear_output = &self.weight * input + &self.bias;
let activated_output = apply_activation(&linear_output, self.activation);
Ok(activated_output)
}
/// Get parameter count
pub fn parameter_count(&self) -> usize {
self.weight.len() + self.bias.len()
}
}
// Activation functions
fn sigmoid(x: &DVector<f64>) -> DVector<f64> {
x.map(|val| 1.0 / (1.0 + (-val).exp()))
}
fn tanh(x: &DVector<f64>) -> DVector<f64> {
x.map(|val| val.tanh())
}
fn relu(x: &DVector<f64>) -> DVector<f64> {
x.map(|val| val.max(0.0))
}
fn gelu(x: &DVector<f64>) -> DVector<f64> {
x.map(|val| 0.5 * val * (1.0 + (val * (2.0 / PI).sqrt()).tanh()))
}
fn apply_activation(x: &DVector<f64>, activation: ActivationFunction) -> DVector<f64> {
match activation {
ActivationFunction::Linear => x.clone(),
ActivationFunction::Relu => relu(x),
ActivationFunction::Tanh => tanh(x),
ActivationFunction::Sigmoid => sigmoid(x),
ActivationFunction::Gelu => gelu(x),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_gru_forward() {
let mut gru = GruLayer::new(4, 8);
let input = DVector::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
let output = gru.forward(&input).unwrap();
assert_eq!(output.len(), 8);
// Test state persistence
let output2 = gru.forward(&input).unwrap();
assert_ne!(output, output2); // Should be different due to state
}
#[test]
fn test_gru_sequence() {
let mut gru = GruLayer::new(2, 4);
let inputs = DMatrix::from_row_slice(2, 3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
let outputs = gru.forward_sequence(&inputs).unwrap();
assert_eq!(outputs.shape(), (4, 3));
}
#[test]
fn test_tcn_forward() {
let tcn = TcnLayer::new(2, 4, 3, 1, true);
let input = DMatrix::from_row_slice(2, 5, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]);
let output = tcn.forward(&input).unwrap();
assert_eq!(output.shape(), (4, 5));
}
#[test]
fn test_dense_forward() {
let dense = DenseLayer::new(4, 2, ActivationFunction::Relu);
let input = DVector::from_vec(vec![1.0, -2.0, 3.0, -4.0]);
let output = dense.forward(&input).unwrap();
assert_eq!(output.len(), 2);
}
#[test]
fn test_activation_functions() {
let x = DVector::from_vec(vec![-2.0, -1.0, 0.0, 1.0, 2.0]);
let relu_out = apply_activation(&x, ActivationFunction::Relu);
assert_eq!(relu_out[0], 0.0); // ReLU of negative should be 0
assert_eq!(relu_out[3], 1.0); // ReLU of positive should be unchanged
let tanh_out = apply_activation(&x, ActivationFunction::Tanh);
assert!(tanh_out[2] == 0.0); // tanh(0) = 0
let sigmoid_out = apply_activation(&x, ActivationFunction::Sigmoid);
assert!(sigmoid_out[2] == 0.5); // sigmoid(0) = 0.5
}
}
@@ -0,0 +1,322 @@
//! Neural network model implementations for temporal prediction
use crate::{
config::{Config, ModelConfig},
error::{Result, TemporalNeuralError},
};
use nalgebra::{DVector, DMatrix};
use serde::{Deserialize, Serialize};
pub mod system_a;
pub mod system_b;
pub mod layers;
pub mod quantization;
pub use system_a::SystemA;
pub use system_b::SystemB;
pub use layers::{GruLayer, TcnLayer, DenseLayer};
pub use quantization::{QuantizedModel, QuantizationScheme};
/// Common trait for all neural network models
pub trait ModelTrait: Send + Sync {
/// Model parameter type
type Params: ModelParams;
/// Create a new model with the given configuration
fn new(config: &ModelConfig) -> Result<Self> where Self: Sized;
/// Forward pass through the network
fn forward(&self, input: &DMatrix<f64>) -> Result<DVector<f64>>;
/// Get model parameters for training/serialization
fn parameters(&self) -> &Self::Params;
/// Get mutable model parameters for training
fn parameters_mut(&mut self) -> &mut Self::Params;
/// Load parameters from another model (for transfer learning)
fn load_parameters(&mut self, params: Self::Params) -> Result<()>;
/// Get the number of parameters in the model
fn parameter_count(&self) -> usize;
/// Get model memory usage in bytes
fn memory_usage(&self) -> usize;
/// Get the expected input shape
fn input_shape(&self) -> (usize, usize);
/// Get the output dimension
fn output_dim(&self) -> usize;
/// Model name for identification
fn model_name(&self) -> &'static str;
/// Prepare model for inference (e.g., quantization)
fn prepare_for_inference(&mut self) -> Result<()> {
Ok(()) // Default implementation does nothing
}
/// Check if model is ready for inference
fn is_inference_ready(&self) -> bool {
true // Default implementation always ready
}
/// Get model configuration
fn config(&self) -> &ModelConfig;
/// Validate input dimensions
fn validate_input(&self, input: &DMatrix<f64>) -> Result<()> {
let expected_shape = self.input_shape();
let actual_shape = (input.nrows(), input.ncols());
if actual_shape != expected_shape {
return Err(TemporalNeuralError::ModelError {
component: self.model_name().to_string(),
message: format!(
"Input shape mismatch: expected {:?}, got {:?}",
expected_shape, actual_shape
),
context: vec![
("expected_rows".to_string(), expected_shape.0.to_string()),
("expected_cols".to_string(), expected_shape.1.to_string()),
("actual_rows".to_string(), actual_shape.0.to_string()),
("actual_cols".to_string(), actual_shape.1.to_string()),
],
});
}
Ok(())
}
}
/// Trait for model parameters
pub trait ModelParams: Clone + Send + Sync + Serialize + for<'de> Deserialize<'de> {
/// Initialize parameters with the given configuration
fn initialize(config: &ModelConfig, rng: &mut impl rand::Rng) -> Self;
/// Get the number of parameters
fn parameter_count(&self) -> usize;
/// Get parameter memory usage in bytes
fn memory_usage(&self) -> usize {
self.parameter_count() * std::mem::size_of::<f64>()
}
/// Apply L2 regularization to parameters
fn apply_l2_regularization(&mut self, weight_decay: f64);
/// Clip gradients to prevent explosion
fn clip_gradients(&mut self, max_norm: f64);
/// Zero out gradients
fn zero_gradients(&mut self);
/// Update parameters using gradients
fn update_parameters(&mut self, learning_rate: f64);
/// Get parameter statistics for monitoring
fn parameter_stats(&self) -> ParameterStats;
}
/// Statistics about model parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParameterStats {
/// Mean absolute value of parameters
pub mean_abs_value: f64,
/// Standard deviation of parameters
pub std_dev: f64,
/// Minimum parameter value
pub min_value: f64,
/// Maximum parameter value
pub max_value: f64,
/// Mean absolute gradient (if available)
pub mean_abs_gradient: Option<f64>,
/// Gradient norm (if available)
pub gradient_norm: Option<f64>,
}
/// Model metadata for serialization and identification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMetadata {
/// Model name and version
pub name: String,
/// Configuration used to create the model
pub config: ModelConfig,
/// Training information
pub training_info: Option<TrainingInfo>,
/// Performance metrics
pub performance_metrics: Option<PerformanceMetrics>,
/// Creation timestamp
pub created_at: chrono::DateTime<chrono::Utc>,
/// Last modified timestamp
pub modified_at: chrono::DateTime<chrono::Utc>,
}
/// Training information for model provenance
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingInfo {
/// Number of training epochs completed
pub epochs_trained: u32,
/// Training dataset size
pub training_samples: usize,
/// Validation dataset size
pub validation_samples: usize,
/// Final training loss
pub final_train_loss: f64,
/// Final validation loss
pub final_val_loss: f64,
/// Training time in seconds
pub training_time_sec: f64,
/// Optimizer used
pub optimizer: String,
/// Learning rate schedule
pub learning_rate_schedule: Vec<(u32, f64)>,
}
/// Performance metrics for model evaluation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PerformanceMetrics {
/// Mean squared error on test set
pub mse: f64,
/// Mean absolute error on test set
pub mae: f64,
/// P90 absolute error
pub p90_error: f64,
/// P99 absolute error
pub p99_error: f64,
/// Average inference latency in milliseconds
pub avg_latency_ms: f64,
/// P50 inference latency in milliseconds
pub p50_latency_ms: f64,
/// P99.9 inference latency in milliseconds
pub p99_9_latency_ms: f64,
/// Model memory usage in bytes
pub memory_usage_bytes: usize,
/// Throughput in predictions per second
pub throughput_pred_per_sec: f64,
}
/// Factory function to create models from configuration
pub fn create_model(config: &Config) -> Result<Box<dyn ModelTrait<Params = Box<dyn ModelParams>>>> {
match config.system {
crate::config::SystemConfig::Traditional(_) => {
let system_a = SystemA::new(&config.model)?;
Ok(Box::new(system_a) as Box<dyn ModelTrait<Params = Box<dyn ModelParams>>>)
}
crate::config::SystemConfig::TemporalSolver(_) => {
let system_b = SystemB::new(&config.model)?;
Ok(Box::new(system_b) as Box<dyn ModelTrait<Params = Box<dyn ModelParams>>>)
}
}
}
/// Load a model from file
pub fn load_model(path: &std::path::Path) -> Result<(Box<dyn ModelTrait<Params = Box<dyn ModelParams>>>, ModelMetadata)> {
let content = std::fs::read_to_string(path)?;
let data: serde_json::Value = serde_json::from_str(&content)?;
let metadata: ModelMetadata = serde_json::from_value(data["metadata"].clone())?;
let model = match metadata.name.as_str() {
"SystemA" => {
let system_a = SystemA::new(&metadata.config)?;
Box::new(system_a) as Box<dyn ModelTrait<Params = Box<dyn ModelParams>>>
}
"SystemB" => {
let system_b = SystemB::new(&metadata.config)?;
Box::new(system_b) as Box<dyn ModelTrait<Params = Box<dyn ModelParams>>>
}
_ => {
return Err(TemporalNeuralError::ModelError {
component: "model_loader".to_string(),
message: format!("Unknown model type: {}", metadata.name),
context: vec![],
});
}
};
Ok((model, metadata))
}
/// Save a model to file
pub fn save_model(
model: &dyn ModelTrait<Params = Box<dyn ModelParams>>,
metadata: &ModelMetadata,
path: &std::path::Path,
) -> Result<()> {
let data = serde_json::json!({
"metadata": metadata,
"parameters": model.parameters(),
});
let content = serde_json::to_string_pretty(&data)?;
std::fs::write(path, content)?;
Ok(())
}
/// Model comparison utilities
pub mod comparison {
use super::*;
/// Compare two models and return similarity metrics
pub fn compare_models(
model1: &dyn ModelTrait<Params = Box<dyn ModelParams>>,
model2: &dyn ModelTrait<Params = Box<dyn ModelParams>>,
) -> ModelComparison {
let params1 = model1.parameters();
let params2 = model2.parameters();
ModelComparison {
parameter_count_diff: (model1.parameter_count() as i64 - model2.parameter_count() as i64).abs() as usize,
memory_usage_diff: (model1.memory_usage() as i64 - model2.memory_usage() as i64).abs() as usize,
architecture_match: model1.model_name() == model2.model_name(),
config_match: model1.config() == model2.config(),
}
}
/// Result of model comparison
#[derive(Debug, Clone)]
pub struct ModelComparison {
/// Difference in parameter count
pub parameter_count_diff: usize,
/// Difference in memory usage
pub memory_usage_diff: usize,
/// Whether architectures match
pub architecture_match: bool,
/// Whether configurations match
pub config_match: bool,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
#[test]
fn test_model_creation() {
let config = Config::default();
let model = create_model(&config);
assert!(model.is_ok());
}
#[test]
fn test_model_validation() {
let config = Config::default();
let model = create_model(&config).unwrap();
let input = DMatrix::zeros(10, 4); // Wrong size
assert!(model.validate_input(&input).is_err());
let correct_input = DMatrix::zeros(256, 4); // Correct size
assert!(model.validate_input(&correct_input).is_ok());
}
#[test]
fn test_parameter_stats() {
// This would be implemented once we have concrete parameter types
// For now, just test that the trait compiles
}
}
@@ -0,0 +1,191 @@
//! INT8 quantization implementation for neural network optimization
use crate::error::{Result, TemporalNeuralError};
use nalgebra::{DMatrix, DVector};
use serde::{Deserialize, Serialize};
/// Quantization scheme for model optimization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QuantizationScheme {
/// 8-bit integer quantization
Int8,
/// 4-bit integer quantization
Int4,
/// Binary quantization
Binary,
}
/// Quantized model wrapper
#[derive(Debug, Clone)]
pub struct QuantizedModel {
/// Quantization scheme used
pub scheme: QuantizationScheme,
/// Quantized weights
pub quantized_weights: Vec<i8>,
/// Scale factors for dequantization
pub scales: Vec<f32>,
/// Zero points for symmetric quantization
pub zero_points: Vec<i8>,
/// Original model shape information
pub shape_info: ModelShapeInfo,
}
/// Information about model structure for quantization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelShapeInfo {
/// Layer dimensions
pub layer_shapes: Vec<(usize, usize)>,
/// Total parameter count
pub total_params: usize,
/// Memory savings achieved
pub memory_savings_ratio: f64,
}
impl QuantizedModel {
/// Create quantized model from floating point weights
pub fn quantize_int8(weights: &[f64]) -> Result<Self> {
if weights.is_empty() {
return Err(TemporalNeuralError::QuantizationError {
message: "Cannot quantize empty weight vector".to_string(),
scheme: Some("int8".to_string()),
accuracy_loss: None,
});
}
let mut quantized_weights = Vec::with_capacity(weights.len());
let mut scales = Vec::new();
let mut zero_points = Vec::new();
// For simplicity, quantize the entire weight vector with single scale/zero_point
// In practice, you'd quantize per-layer or per-channel
let min_val = weights.iter().fold(f64::INFINITY, |a, &b| a.min(b));
let max_val = weights.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
let scale = (max_val - min_val) / 255.0; // 8-bit range
let zero_point = (-min_val / scale).round() as i8;
scales.push(scale as f32);
zero_points.push(zero_point);
for &weight in weights {
let quantized = ((weight / scale) + zero_point as f64).round().clamp(-128.0, 127.0) as i8;
quantized_weights.push(quantized);
}
let shape_info = ModelShapeInfo {
layer_shapes: vec![(weights.len(), 1)], // Simplified
total_params: weights.len(),
memory_savings_ratio: 4.0, // 32-bit -> 8-bit
};
Ok(Self {
scheme: QuantizationScheme::Int8,
quantized_weights,
scales,
zero_points,
shape_info,
})
}
/// Dequantize weights back to floating point
pub fn dequantize(&self) -> Result<Vec<f64>> {
if self.quantized_weights.len() != self.shape_info.total_params {
return Err(TemporalNeuralError::QuantizationError {
message: "Weight count mismatch during dequantization".to_string(),
scheme: Some(format!("{:?}", self.scheme)),
accuracy_loss: None,
});
}
let mut dequantized = Vec::with_capacity(self.quantized_weights.len());
let scale = self.scales[0] as f64; // Using first scale for simplicity
let zero_point = self.zero_points[0] as f64;
for &quantized_weight in &self.quantized_weights {
let dequantized_weight = (quantized_weight as f64 - zero_point) * scale;
dequantized.push(dequantized_weight);
}
Ok(dequantized)
}
/// Get memory usage in bytes
pub fn memory_usage(&self) -> usize {
self.quantized_weights.len() * std::mem::size_of::<i8>() +
self.scales.len() * std::mem::size_of::<f32>() +
self.zero_points.len() * std::mem::size_of::<i8>()
}
/// Estimate accuracy loss from quantization
pub fn estimate_accuracy_loss(&self, original_weights: &[f64]) -> Result<f64> {
let dequantized = self.dequantize()?;
if dequantized.len() != original_weights.len() {
return Err(TemporalNeuralError::QuantizationError {
message: "Length mismatch for accuracy estimation".to_string(),
scheme: Some(format!("{:?}", self.scheme)),
accuracy_loss: None,
});
}
// Compute MSE between original and dequantized weights
let mse = original_weights.iter()
.zip(dequantized.iter())
.map(|(&orig, &deq)| (orig - deq).powi(2))
.sum::<f64>() / original_weights.len() as f64;
Ok(mse.sqrt()) // Return RMSE
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_int8_quantization() {
let weights = vec![1.0, 2.5, -1.5, 0.0, 3.2, -2.1];
let quantized = QuantizedModel::quantize_int8(&weights).unwrap();
assert_eq!(quantized.quantized_weights.len(), weights.len());
assert_eq!(quantized.scales.len(), 1);
assert_eq!(quantized.zero_points.len(), 1);
}
#[test]
fn test_dequantization() {
let weights = vec![1.0, 2.5, -1.5, 0.0, 3.2, -2.1];
let quantized = QuantizedModel::quantize_int8(&weights).unwrap();
let dequantized = quantized.dequantize().unwrap();
assert_eq!(dequantized.len(), weights.len());
// Check that dequantized values are reasonably close to originals
for (orig, deq) in weights.iter().zip(dequantized.iter()) {
assert!((orig - deq).abs() < 0.1); // Should be within quantization error
}
}
#[test]
fn test_accuracy_loss_estimation() {
let weights = vec![1.0, 2.5, -1.5, 0.0, 3.2, -2.1];
let quantized = QuantizedModel::quantize_int8(&weights).unwrap();
let accuracy_loss = quantized.estimate_accuracy_loss(&weights).unwrap();
assert!(accuracy_loss >= 0.0);
assert!(accuracy_loss < 1.0); // Should be reasonable for this simple case
}
#[test]
fn test_memory_usage() {
let weights = vec![1.0; 1000];
let quantized = QuantizedModel::quantize_int8(&weights).unwrap();
let memory_usage = quantized.memory_usage();
assert!(memory_usage > 0);
// Should use roughly 1/4 the memory (32-bit -> 8-bit)
let original_memory = weights.len() * std::mem::size_of::<f64>();
assert!(memory_usage < original_memory / 2);
}
}
@@ -0,0 +1,549 @@
//! System A: Traditional micro-neural network implementation
//!
//! This module implements the baseline traditional micro-network for comparison
//! against the temporal solver approach. It uses standard neural architectures
//! without mathematical solver integration.
use crate::{
config::ModelConfig,
error::{Result, TemporalNeuralError},
models::{
layers::{GruLayer, TcnLayer, DenseLayer, ActivationFunction},
ModelTrait, ModelParams, ParameterStats,
},
};
use nalgebra::{DMatrix, DVector};
use serde::{Deserialize, Serialize};
/// System A: Traditional micro-neural network
///
/// This system implements a standard approach using either GRU or TCN
/// for sequence modeling followed by a dense output layer.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemA {
/// Model configuration
config: ModelConfig,
/// Parameters for the system
params: SystemAParams,
/// Current architecture type
architecture: ArchitectureType,
}
/// Architecture variants for System A
#[derive(Debug, Clone, Serialize, Deserialize)]
enum ArchitectureType {
/// GRU-based architecture
Gru {
/// GRU layers
layers: Vec<GruLayer>,
/// Output projection layer
output_layer: DenseLayer,
},
/// TCN-based architecture
Tcn {
/// TCN layers
layers: Vec<TcnLayer>,
/// Global pooling type
pooling: PoolingType,
/// Output projection layer
output_layer: DenseLayer,
},
}
/// Pooling types for TCN architecture
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
enum PoolingType {
/// Take the last timestep
Last,
/// Global average pooling
Average,
/// Global max pooling
Max,
/// Attention-based pooling
Attention,
}
/// Parameters for System A
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemAParams {
/// All parameter values flattened
values: Vec<f64>,
/// Gradients for backpropagation
gradients: Vec<f64>,
/// Parameter structure metadata
structure: ParameterStructure,
}
/// Metadata about parameter structure for serialization
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ParameterStructure {
/// Total parameter count
total_params: usize,
/// Parameter layout (layer_name -> (start_idx, count))
layout: Vec<(String, usize, usize)>,
}
impl SystemA {
/// Create a new System A model
pub fn new(config: &ModelConfig) -> Result<Self> {
let architecture = match config.model_type.as_str() {
"micro_gru" => Self::create_gru_architecture(config)?,
"micro_tcn" => Self::create_tcn_architecture(config)?,
_ => {
return Err(TemporalNeuralError::ConfigurationError {
message: format!("Unsupported model type: {}", config.model_type),
field: Some("model_type".to_string()),
});
}
};
let params = SystemAParams::initialize(config, &mut rand::thread_rng());
Ok(Self {
config: config.clone(),
params,
architecture,
})
}
fn create_gru_architecture(config: &ModelConfig) -> Result<ArchitectureType> {
let mut layers = Vec::new();
// First layer takes input features
let first_layer = GruLayer::new(config.hidden_size, config.hidden_size);
layers.push(first_layer);
// Additional layers if specified
for _ in 1..config.num_layers {
let layer = GruLayer::new(config.hidden_size, config.hidden_size);
layers.push(layer);
}
// Output layer - predict 2D position (x, y) at horizon
let activation = match config.activation.as_str() {
"relu" => ActivationFunction::Relu,
"tanh" => ActivationFunction::Tanh,
"gelu" => ActivationFunction::Gelu,
"linear" => ActivationFunction::Linear,
_ => ActivationFunction::Tanh, // Default
};
let output_layer = DenseLayer::new(config.hidden_size, 2, activation);
Ok(ArchitectureType::Gru {
layers,
output_layer,
})
}
fn create_tcn_architecture(config: &ModelConfig) -> Result<ArchitectureType> {
let mut layers = Vec::new();
let mut channels = config.hidden_size as usize;
// Create dilated TCN layers
for layer_idx in 0..config.num_layers as usize {
let dilation = 2_usize.pow(layer_idx as u32);
let layer = TcnLayer::new(
channels,
config.hidden_size as usize,
3, // kernel size
dilation,
config.residual,
);
layers.push(layer);
channels = config.hidden_size as usize;
}
// Global pooling strategy
let pooling = PoolingType::Last; // Simple approach - take last timestep
// Output layer
let activation = match config.activation.as_str() {
"relu" => ActivationFunction::Relu,
"tanh" => ActivationFunction::Tanh,
"gelu" => ActivationFunction::Gelu,
"linear" => ActivationFunction::Linear,
_ => ActivationFunction::Tanh,
};
let output_layer = DenseLayer::new(config.hidden_size as usize, 2, activation);
Ok(ArchitectureType::Tcn {
layers,
pooling,
output_layer,
})
}
/// Apply pooling to TCN output
fn apply_pooling(output: &DMatrix<f64>, pooling: PoolingType) -> DVector<f64> {
match pooling {
PoolingType::Last => {
// Take the last timestep
output.column(output.ncols() - 1).into()
}
PoolingType::Average => {
// Global average pooling
let mut result = DVector::zeros(output.nrows());
for i in 0..output.nrows() {
result[i] = output.row(i).mean();
}
result
}
PoolingType::Max => {
// Global max pooling
let mut result = DVector::zeros(output.nrows());
for i in 0..output.nrows() {
result[i] = output.row(i).max();
}
result
}
PoolingType::Attention => {
// Simple attention: uniform weights for now
// In a full implementation, this would be learned
let seq_len = output.ncols();
let weights = DVector::from_element(seq_len, 1.0 / seq_len as f64);
output * weights
}
}
}
}
impl ModelTrait for SystemA {
type Params = SystemAParams;
fn new(config: &ModelConfig) -> Result<Self> {
Self::new(config)
}
fn forward(&self, input: &DMatrix<f64>) -> Result<DVector<f64>> {
self.validate_input(input)?;
match &self.architecture {
ArchitectureType::Gru { layers, output_layer } => {
// Process through GRU layers
let mut current_input = input.clone();
// For GRU, we need to process the sequence through each layer
// Input shape: (features, sequence_length)
// We need to process timestep by timestep
let seq_len = input.ncols();
let mut hidden_states = Vec::new();
// Create temporary GRU layers for forward pass (clone to avoid mutation issues)
let mut temp_layers: Vec<GruLayer> = layers.iter().cloned().collect();
// Process sequence through all GRU layers
for t in 0..seq_len {
let timestep_input = current_input.column(t).into();
// Process through each layer for this timestep
let mut layer_input = timestep_input;
for layer in temp_layers.iter_mut() {
layer_input = layer.forward(&layer_input)?;
}
hidden_states.push(layer_input);
}
// Take the final hidden state
let final_hidden = hidden_states.last().unwrap();
// Project to output
output_layer.forward(final_hidden)
}
ArchitectureType::Tcn { layers, pooling, output_layer } => {
// Process through TCN layers
let mut current_output = input.clone();
for layer in layers {
current_output = layer.forward(&current_output)?;
}
// Apply pooling to get fixed-size representation
let pooled = Self::apply_pooling(&current_output, *pooling);
// Project to output
output_layer.forward(&pooled)
}
}
}
fn parameters(&self) -> &Self::Params {
&self.params
}
fn parameters_mut(&mut self) -> &mut Self::Params {
&mut self.params
}
fn load_parameters(&mut self, params: Self::Params) -> Result<()> {
if params.values.len() != self.params.values.len() {
return Err(TemporalNeuralError::ModelError {
component: "SystemA".to_string(),
message: format!(
"Parameter count mismatch: expected {}, got {}",
self.params.values.len(),
params.values.len()
),
context: vec![],
});
}
self.params = params;
Ok(())
}
fn parameter_count(&self) -> usize {
self.params.parameter_count()
}
fn memory_usage(&self) -> usize {
self.params.memory_usage()
}
fn input_shape(&self) -> (usize, usize) {
// Returns (sequence_length, features)
// This will be set based on configuration in a real implementation
(256, 4) // Default for 128ms window at 2kHz with 4 features
}
fn output_dim(&self) -> usize {
2 // Predicting 2D position (x, y)
}
fn model_name(&self) -> &'static str {
"SystemA"
}
fn config(&self) -> &ModelConfig {
&self.config
}
fn prepare_for_inference(&mut self) -> Result<()> {
// For System A, we might apply quantization here
// For now, just mark as ready
Ok(())
}
}
impl ModelParams for SystemAParams {
fn initialize(config: &ModelConfig, rng: &mut impl rand::Rng) -> Self {
// Calculate total parameter count based on architecture
let total_params = Self::calculate_param_count(config);
// Initialize with small random values
let values: Vec<f64> = (0..total_params)
.map(|_| (rng.sample(rand::distributions::Standard) as f64 - 0.5) * 0.02)
.collect();
let gradients = vec![0.0; total_params];
let structure = Self::build_structure(config);
Self {
values,
gradients,
structure,
}
}
fn parameter_count(&self) -> usize {
self.values.len()
}
fn apply_l2_regularization(&mut self, weight_decay: f64) {
for (param, grad) in self.values.iter().zip(self.gradients.iter_mut()) {
*grad += weight_decay * param;
}
}
fn clip_gradients(&mut self, max_norm: f64) {
let grad_norm: f64 = self.gradients.iter().map(|g| g * g).sum::<f64>().sqrt();
if grad_norm > max_norm {
let scale = max_norm / grad_norm;
for grad in self.gradients.iter_mut() {
*grad *= scale;
}
}
}
fn zero_gradients(&mut self) {
self.gradients.fill(0.0);
}
fn update_parameters(&mut self, learning_rate: f64) {
for (param, grad) in self.values.iter_mut().zip(self.gradients.iter()) {
*param -= learning_rate * grad;
}
}
fn parameter_stats(&self) -> ParameterStats {
let n = self.values.len() as f64;
let mean = self.values.iter().sum::<f64>() / n;
let variance = self.values.iter()
.map(|x| (x - mean).powi(2))
.sum::<f64>() / n;
let mean_abs_value = self.values.iter().map(|x| x.abs()).sum::<f64>() / n;
let min_value = self.values.iter().fold(f64::INFINITY, |a, &b| a.min(b));
let max_value = self.values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
let mean_abs_gradient = if !self.gradients.is_empty() {
Some(self.gradients.iter().map(|g| g.abs()).sum::<f64>() / self.gradients.len() as f64)
} else {
None
};
let gradient_norm = if !self.gradients.is_empty() {
Some(self.gradients.iter().map(|g| g * g).sum::<f64>().sqrt())
} else {
None
};
ParameterStats {
mean_abs_value,
std_dev: variance.sqrt(),
min_value,
max_value,
mean_abs_gradient,
gradient_norm,
}
}
}
impl SystemAParams {
fn calculate_param_count(config: &ModelConfig) -> usize {
match config.model_type.as_str() {
"micro_gru" => {
// GRU parameters: 3 * (input_size * hidden + hidden * hidden + hidden) per layer
// Plus output layer: hidden * output + output
let hidden = config.hidden_size as usize;
let input_size = hidden; // Simplified assumption
let gru_params_per_layer = 3 * (input_size * hidden + hidden * hidden + hidden);
let gru_params = gru_params_per_layer * config.num_layers as usize;
let output_params = hidden * 2 + 2; // 2D output
gru_params + output_params
}
"micro_tcn" => {
// TCN parameters: kernel_size * input_channels * output_channels + output_channels per layer
// Plus output layer
let channels = config.hidden_size as usize;
let kernel_size = 3;
let tcn_params_per_layer = kernel_size * channels * channels + channels;
let tcn_params = tcn_params_per_layer * config.num_layers as usize;
let output_params = channels * 2 + 2; // 2D output
tcn_params + output_params
}
_ => 1000, // Default fallback
}
}
fn build_structure(config: &ModelConfig) -> ParameterStructure {
let mut layout = Vec::new();
let mut offset = 0;
match config.model_type.as_str() {
"micro_gru" => {
for i in 0..config.num_layers {
let layer_params = Self::gru_layer_param_count(config.hidden_size as usize);
layout.push((format!("gru_layer_{}", i), offset, layer_params));
offset += layer_params;
}
}
"micro_tcn" => {
for i in 0..config.num_layers {
let layer_params = Self::tcn_layer_param_count(config.hidden_size as usize);
layout.push((format!("tcn_layer_{}", i), offset, layer_params));
offset += layer_params;
}
}
_ => {}
}
// Output layer
let output_params = config.hidden_size as usize * 2 + 2;
layout.push(("output_layer".to_string(), offset, output_params));
ParameterStructure {
total_params: Self::calculate_param_count(config),
layout,
}
}
fn gru_layer_param_count(hidden_size: usize) -> usize {
3 * (hidden_size * hidden_size + hidden_size * hidden_size + hidden_size)
}
fn tcn_layer_param_count(channels: usize) -> usize {
3 * channels * channels + channels // kernel_size=3
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_system_a_creation() {
let mut config = ModelConfig {
model_type: "micro_gru".to_string(),
hidden_size: 16,
num_layers: 2,
dropout: 0.1,
residual: true,
activation: "tanh".to_string(),
layer_norm: false,
};
let system = SystemA::new(&config).unwrap();
assert_eq!(system.model_name(), "SystemA");
assert_eq!(system.output_dim(), 2);
// Test TCN variant
config.model_type = "micro_tcn".to_string();
let system_tcn = SystemA::new(&config).unwrap();
assert_eq!(system_tcn.model_name(), "SystemA");
}
#[test]
fn test_forward_pass() {
let config = ModelConfig {
model_type: "micro_gru".to_string(),
hidden_size: 8,
num_layers: 1,
dropout: 0.0,
residual: false,
activation: "tanh".to_string(),
layer_norm: false,
};
let system = SystemA::new(&config).unwrap();
let input = DMatrix::from_element(4, 10, 1.0); // 4 features, 10 timesteps
let output = system.forward(&input).unwrap();
assert_eq!(output.len(), 2);
}
#[test]
fn test_parameter_operations() {
let config = ModelConfig {
model_type: "micro_gru".to_string(),
hidden_size: 4,
num_layers: 1,
dropout: 0.0,
residual: false,
activation: "tanh".to_string(),
layer_norm: false,
};
let system = SystemA::new(&config).unwrap();
let param_count = system.parameter_count();
assert!(param_count > 0);
let stats = system.parameters().parameter_stats();
assert!(stats.mean_abs_value >= 0.0);
}
}
@@ -0,0 +1,480 @@
//! System B: Temporal solver neural network implementation
//!
//! This module implements the novel temporal solver approach that combines
//! neural networks with Kalman filter priors and sublinear solver gating
//! for improved latency and mathematical verification.
use crate::{
config::{ModelConfig, TemporalSolverConfig},
error::{Result, TemporalNeuralError},
models::{
layers::{GruLayer, TcnLayer, DenseLayer, ActivationFunction},
ModelTrait, ModelParams, ParameterStats,
system_a::{SystemA, SystemAParams}, // Reuse SystemA architecture
},
solvers::{KalmanFilter, SolverGate, PageRankSelector},
};
use nalgebra::{DMatrix, DVector};
use serde::{Deserialize, Serialize};
/// System B: Temporal solver neural network
///
/// This system combines a neural network (same architecture as System A)
/// with Kalman filter priors and sublinear solver verification.
/// The key innovation is residual learning: the network predicts the
/// residual between the Kalman prior and the true target.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemB {
/// Base neural network (same as System A)
base_network: SystemA,
/// Kalman filter for prior predictions
kalman_filter: KalmanFilter,
/// Sublinear solver gate for verification
solver_gate: SolverGate,
/// PageRank-based active selector for training
active_selector: Option<PageRankSelector>,
/// Temporal solver configuration
solver_config: TemporalSolverConfig,
/// Whether the system is in training or inference mode
inference_mode: bool,
}
/// Prediction result from System B
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemBPrediction {
/// Kalman filter prior prediction
pub prior: DVector<f64>,
/// Neural network residual prediction
pub residual: DVector<f64>,
/// Final combined prediction (prior + residual)
pub prediction: DVector<f64>,
/// Solver gate verification result
pub gate_result: GateResult,
/// Computation timing breakdown
pub timing: PredictionTiming,
}
/// Result from the solver gate verification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GateResult {
/// Whether the gate passed (prediction is verified)
pub passed: bool,
/// Certificate error from sublinear solver
pub certificate_error: f64,
/// Computational work performed
pub work_performed: u64,
/// Fallback strategy used if gate failed
pub fallback_used: Option<String>,
}
/// Timing breakdown for performance analysis
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionTiming {
/// Time spent in Kalman filter update (microseconds)
pub kalman_us: f64,
/// Time spent in neural network forward pass (microseconds)
pub network_us: f64,
/// Time spent in solver gate verification (microseconds)
pub gate_us: f64,
/// Total prediction time (microseconds)
pub total_us: f64,
}
impl SystemB {
/// Create a new System B model
pub fn new(config: &ModelConfig, solver_config: &TemporalSolverConfig) -> Result<Self> {
// Create base network (same architecture as System A)
let base_network = SystemA::new(config)?;
// Initialize Kalman filter
let kalman_filter = KalmanFilter::new(&solver_config.prior)?;
// Initialize solver gate
let solver_gate = SolverGate::new(&solver_config.solver_gate)?;
// Initialize active selector for training
let active_selector = Some(PageRankSelector::new(&solver_config.active_selection)?);
Ok(Self {
base_network,
kalman_filter,
solver_gate,
active_selector,
solver_config: solver_config.clone(),
inference_mode: false,
})
}
/// Set inference mode (disables active selector)
pub fn set_inference_mode(&mut self, inference_mode: bool) {
self.inference_mode = inference_mode;
if inference_mode {
self.active_selector = None; // Save memory in inference
}
}
/// Predict with full temporal solver pipeline
pub fn predict_with_solver(&mut self, input: &DMatrix<f64>) -> Result<SystemBPrediction> {
let start_time = std::time::Instant::now();
// Step 1: Update Kalman filter and get prior
let kalman_start = std::time::Instant::now();
let prior = self.kalman_filter.predict(input)?;
let kalman_time = kalman_start.elapsed().as_micros() as f64;
// Step 2: Neural network predicts residual
let network_start = std::time::Instant::now();
let residual = self.base_network.forward(input)?;
let network_time = network_start.elapsed().as_micros() as f64;
// Step 3: Combine prior and residual
let prediction = &prior + &residual;
// Step 4: Verify with solver gate
let gate_start = std::time::Instant::now();
let gate_result = self.solver_gate.verify(&prior, &residual, &prediction)?;
let gate_time = gate_start.elapsed().as_micros() as f64;
let total_time = start_time.elapsed().as_micros() as f64;
// Apply fallback if gate failed
let final_prediction = if gate_result.passed {
prediction
} else {
self.apply_fallback(&prior, &residual, &gate_result)?
};
Ok(SystemBPrediction {
prior,
residual,
prediction: final_prediction,
gate_result,
timing: PredictionTiming {
kalman_us: kalman_time,
network_us: network_time,
gate_us: gate_time,
total_us: total_time,
},
})
}
/// Apply fallback strategy when solver gate fails
fn apply_fallback(
&self,
prior: &DVector<f64>,
residual: &DVector<f64>,
gate_result: &GateResult,
) -> Result<DVector<f64>> {
match self.solver_config.solver_gate.fallback_strategy.as_str() {
"kalman_only" => Ok(prior.clone()),
"hold_last" => {
// In a real implementation, this would hold the last verified prediction
// For now, use the prior
Ok(prior.clone())
}
"disable_gate" => {
// Use the combined prediction anyway
Ok(prior + residual)
}
"weighted_blend" => {
// Blend based on certificate error
let cert_error = gate_result.certificate_error;
let weight = 1.0 / (1.0 + cert_error * 10.0); // Sigmoid-like weighting
Ok(prior * (1.0 - weight) + (prior + residual) * weight)
}
_ => {
return Err(TemporalNeuralError::ConfigurationError {
message: format!(
"Unknown fallback strategy: {}",
self.solver_config.solver_gate.fallback_strategy
),
field: Some("fallback_strategy".to_string()),
});
}
}
}
/// Update Kalman filter state with ground truth (for training)
pub fn update_kalman_state(&mut self, measurement: &DVector<f64>) -> Result<()> {
self.kalman_filter.update(measurement)
}
/// Get active sample selector for training
pub fn active_selector(&mut self) -> Option<&mut PageRankSelector> {
self.active_selector.as_mut()
}
/// Reset all internal states
pub fn reset_states(&mut self) -> Result<()> {
self.kalman_filter.reset()?;
self.solver_gate.reset()?;
if let Some(ref mut selector) = self.active_selector {
selector.reset()?;
}
Ok(())
}
/// Get solver statistics for monitoring
pub fn get_solver_stats(&self) -> SolverStats {
SolverStats {
gate_pass_rate: self.solver_gate.get_pass_rate(),
avg_certificate_error: self.solver_gate.get_avg_certificate_error(),
avg_computational_work: self.solver_gate.get_avg_work(),
kalman_prediction_error: self.kalman_filter.get_prediction_error(),
total_predictions: self.solver_gate.get_prediction_count(),
}
}
/// Configure solver parameters dynamically
pub fn configure_solver(&mut self, epsilon: Option<f64>, budget: Option<u64>) -> Result<()> {
if let Some(eps) = epsilon {
self.solver_gate.set_epsilon(eps)?;
}
if let Some(b) = budget {
self.solver_gate.set_budget(b)?;
}
Ok(())
}
}
/// Statistics from the solver components
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SolverStats {
/// Percentage of predictions that passed the gate
pub gate_pass_rate: f64,
/// Average certificate error across all predictions
pub avg_certificate_error: f64,
/// Average computational work performed
pub avg_computational_work: f64,
/// Kalman filter prediction error
pub kalman_prediction_error: f64,
/// Total number of predictions made
pub total_predictions: u64,
}
impl ModelTrait for SystemB {
type Params = SystemAParams; // Reuse SystemA parameters
fn new(config: &ModelConfig) -> Result<Self> {
// This is a simplified constructor - in practice would need full solver config
let default_solver_config = TemporalSolverConfig::default();
Self::new(config, &default_solver_config)
}
fn forward(&self, input: &DMatrix<f64>) -> Result<DVector<f64>> {
// For the ModelTrait interface, we provide a simplified forward pass
// without solver verification (for compatibility)
let prior = self.kalman_filter.predict_const(input)?;
let residual = self.base_network.forward(input)?;
Ok(&prior + &residual)
}
fn parameters(&self) -> &Self::Params {
self.base_network.parameters()
}
fn parameters_mut(&mut self) -> &mut Self::Params {
self.base_network.parameters_mut()
}
fn load_parameters(&mut self, params: Self::Params) -> Result<()> {
self.base_network.load_parameters(params)
}
fn parameter_count(&self) -> usize {
self.base_network.parameter_count()
}
fn memory_usage(&self) -> usize {
self.base_network.memory_usage() +
self.kalman_filter.memory_usage() +
self.solver_gate.memory_usage() +
self.active_selector.as_ref().map_or(0, |s| s.memory_usage())
}
fn input_shape(&self) -> (usize, usize) {
self.base_network.input_shape()
}
fn output_dim(&self) -> usize {
self.base_network.output_dim()
}
fn model_name(&self) -> &'static str {
"SystemB"
}
fn config(&self) -> &ModelConfig {
self.base_network.config()
}
fn prepare_for_inference(&mut self) -> Result<()> {
self.set_inference_mode(true);
self.base_network.prepare_for_inference()?;
self.kalman_filter.prepare_for_inference()?;
self.solver_gate.prepare_for_inference()?;
Ok(())
}
fn is_inference_ready(&self) -> bool {
self.inference_mode &&
self.base_network.is_inference_ready() &&
self.kalman_filter.is_inference_ready() &&
self.solver_gate.is_inference_ready()
}
}
/// Training utilities specific to System B
impl SystemB {
/// Compute residual learning loss
pub fn compute_residual_loss(
&self,
predictions: &[SystemBPrediction],
targets: &[DVector<f64>],
) -> Result<f64> {
if predictions.len() != targets.len() {
return Err(TemporalNeuralError::TrainingError {
epoch: 0,
message: "Predictions and targets length mismatch".to_string(),
metrics: None,
});
}
let mut total_loss = 0.0;
let mut valid_samples = 0;
for (pred, target) in predictions.iter().zip(targets.iter()) {
// Residual learning: target - prior should equal residual
let expected_residual = target - &pred.prior;
let residual_error = &pred.residual - &expected_residual;
// MSE loss
let sample_loss: f64 = residual_error.iter().map(|x| x * x).sum();
total_loss += sample_loss;
valid_samples += 1;
}
if valid_samples == 0 {
return Err(TemporalNeuralError::TrainingError {
epoch: 0,
message: "No valid samples for loss computation".to_string(),
metrics: None,
});
}
Ok(total_loss / valid_samples as f64)
}
/// Compute additional regularization terms
pub fn compute_regularization_loss(&self, predictions: &[SystemBPrediction]) -> f64 {
let mut reg_loss = 0.0;
let smoothness_weight = 0.1; // From config
// Smoothness penalty on velocity predictions
for pred in predictions {
if pred.prediction.len() >= 2 {
// Assume prediction is [x, y] - penalize large velocities
let velocity_mag = pred.prediction[0].powi(2) + pred.prediction[1].powi(2);
reg_loss += smoothness_weight * velocity_mag;
}
}
// Gate verification penalty
let gate_penalty = 1.0;
for pred in predictions {
if !pred.gate_result.passed {
reg_loss += gate_penalty * pred.gate_result.certificate_error;
}
}
reg_loss / predictions.len() as f64
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_system_b_creation() {
let config = ModelConfig {
model_type: "micro_gru".to_string(),
hidden_size: 8,
num_layers: 1,
dropout: 0.0,
residual: false,
activation: "tanh".to_string(),
layer_norm: false,
};
let solver_config = TemporalSolverConfig::default();
let system = SystemB::new(&config, &solver_config).unwrap();
assert_eq!(system.model_name(), "SystemB");
assert_eq!(system.output_dim(), 2);
}
#[test]
fn test_inference_mode() {
let config = ModelConfig {
model_type: "micro_gru".to_string(),
hidden_size: 4,
num_layers: 1,
dropout: 0.0,
residual: false,
activation: "tanh".to_string(),
layer_norm: false,
};
let solver_config = TemporalSolverConfig::default();
let mut system = SystemB::new(&config, &solver_config).unwrap();
assert!(!system.inference_mode);
assert!(system.active_selector.is_some());
system.set_inference_mode(true);
assert!(system.inference_mode);
assert!(system.active_selector.is_none());
}
#[test]
fn test_forward_pass() {
let config = ModelConfig {
model_type: "micro_gru".to_string(),
hidden_size: 4,
num_layers: 1,
dropout: 0.0,
residual: false,
activation: "tanh".to_string(),
layer_norm: false,
};
let solver_config = TemporalSolverConfig::default();
let system = SystemB::new(&config, &solver_config).unwrap();
let input = DMatrix::from_element(4, 10, 1.0);
let output = system.forward(&input).unwrap();
assert_eq!(output.len(), 2);
}
#[test]
fn test_solver_stats() {
let config = ModelConfig {
model_type: "micro_gru".to_string(),
hidden_size: 4,
num_layers: 1,
dropout: 0.0,
residual: false,
activation: "tanh".to_string(),
layer_norm: false,
};
let solver_config = TemporalSolverConfig::default();
let system = SystemB::new(&config, &solver_config).unwrap();
let stats = system.get_solver_stats();
assert!(stats.gate_pass_rate >= 0.0 && stats.gate_pass_rate <= 1.0);
assert!(stats.avg_certificate_error >= 0.0);
}
}