mirror of
https://github.com/ruvnet/RuView
synced 2026-08-09 20:21:43 +00:00
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:
+484
@@ -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
|
||||
}
|
||||
}
|
||||
+322
@@ -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
|
||||
}
|
||||
}
|
||||
Vendored
+191
@@ -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);
|
||||
}
|
||||
}
|
||||
+549
@@ -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(¤t_output)?;
|
||||
}
|
||||
|
||||
// Apply pooling to get fixed-size representation
|
||||
let pooled = Self::apply_pooling(¤t_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);
|
||||
}
|
||||
}
|
||||
+480
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user