Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,617 @@
|
||||
//! Augmented Neural ODEs implementation
|
||||
//!
|
||||
//! Augmented Neural ODEs increase the expressivity of Neural ODEs by adding
|
||||
//! extra dimensions to the state space. This helps overcome limitations of
|
||||
//! standard Neural ODEs in modeling complex dynamics.
|
||||
//!
|
||||
//! Based on "Augmented Neural ODEs" (Dupont et al., 2019)
|
||||
|
||||
use crate::prelude::*;
|
||||
use super::{ODEFunc, NeuralODE, NeuralODEConfig, Result, NeuralODEError};
|
||||
|
||||
/// Configuration for augmented Neural ODEs
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AugmentationConfig {
|
||||
/// Number of original dimensions
|
||||
pub original_dim: usize,
|
||||
/// Number of augmented dimensions to add
|
||||
pub augmented_dim: usize,
|
||||
/// Initialization strategy for augmented dimensions
|
||||
pub init_strategy: AugmentationInit,
|
||||
/// Whether to include augmented dimensions in output
|
||||
pub return_augmented: bool,
|
||||
}
|
||||
|
||||
/// Initialization strategies for augmented dimensions
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AugmentationInit {
|
||||
/// Initialize to zeros
|
||||
Zeros,
|
||||
/// Initialize to small random values
|
||||
SmallRandom { scale: f32 },
|
||||
/// Initialize using a linear transformation of the original state
|
||||
Linear { matrix: Tensor },
|
||||
/// Initialize to constant values
|
||||
Constant { value: f32 },
|
||||
}
|
||||
|
||||
impl Default for AugmentationConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
original_dim: 0, // Must be set explicitly
|
||||
augmented_dim: 1,
|
||||
init_strategy: AugmentationInit::SmallRandom { scale: 0.01 },
|
||||
return_augmented: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Augmented Neural ODE that adds extra dimensions to increase expressivity
|
||||
pub struct AugmentedNeuralODE {
|
||||
/// Underlying Neural ODE operating on augmented state
|
||||
neural_ode: NeuralODE,
|
||||
/// Augmentation configuration
|
||||
config: AugmentationConfig,
|
||||
/// Device for computations
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl AugmentedNeuralODE {
|
||||
/// Create a new Augmented Neural ODE
|
||||
///
|
||||
/// # Arguments
|
||||
/// - `dynamics`: ODE dynamics function for the augmented system
|
||||
/// - `ode_config`: Configuration for the underlying Neural ODE
|
||||
/// - `original_dim`: Number of original state dimensions
|
||||
/// - `augmented_dim`: Number of extra dimensions to add
|
||||
/// - `device`: Device for tensor operations
|
||||
pub fn new(
|
||||
dynamics: Box<dyn ODEFunc>,
|
||||
ode_config: NeuralODEConfig,
|
||||
original_dim: usize,
|
||||
augmented_dim: usize,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let augmentation_config = AugmentationConfig {
|
||||
original_dim,
|
||||
augmented_dim,
|
||||
init_strategy: AugmentationInit::SmallRandom { scale: 0.01 },
|
||||
return_augmented: false,
|
||||
};
|
||||
|
||||
let neural_ode = NeuralODE::new(dynamics, ode_config, device)?;
|
||||
|
||||
Ok(Self {
|
||||
neural_ode,
|
||||
config: augmentation_config,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Create with custom augmentation configuration
|
||||
pub fn with_config(
|
||||
dynamics: Box<dyn ODEFunc>,
|
||||
ode_config: NeuralODEConfig,
|
||||
augmentation_config: AugmentationConfig,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let neural_ode = NeuralODE::new(dynamics, ode_config, device)?;
|
||||
|
||||
Ok(Self {
|
||||
neural_ode,
|
||||
config: augmentation_config,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Augment the initial state with extra dimensions
|
||||
fn augment_state(&self, y0: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = if y0.ndim() == 2 {
|
||||
Some(y0.shape()[0])
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let original_shape = if let Some(bs) = batch_size {
|
||||
[bs, self.config.original_dim]
|
||||
} else {
|
||||
[self.config.original_dim]
|
||||
};
|
||||
|
||||
// Verify original state has correct dimensions
|
||||
let expected_shape = if batch_size.is_some() { &original_shape[..] } else { &original_shape[1..] };
|
||||
if y0.shape() != expected_shape {
|
||||
return Err(NeuralODEError::InvalidInput(format!(
|
||||
"Expected original state shape {:?}, got {:?}",
|
||||
expected_shape, y0.shape()
|
||||
)));
|
||||
}
|
||||
|
||||
// Create augmented dimensions
|
||||
let augmented_part = match &self.config.init_strategy {
|
||||
AugmentationInit::Zeros => {
|
||||
if let Some(bs) = batch_size {
|
||||
Tensor::zeros([bs, self.config.augmented_dim], &self.device)?
|
||||
} else {
|
||||
Tensor::zeros([self.config.augmented_dim], &self.device)?
|
||||
}
|
||||
}
|
||||
AugmentationInit::SmallRandom { scale } => {
|
||||
let shape = if let Some(bs) = batch_size {
|
||||
[bs, self.config.augmented_dim]
|
||||
} else {
|
||||
[self.config.augmented_dim]
|
||||
};
|
||||
Tensor::randn(shape, &self.device)?.mul_scalar(*scale)?
|
||||
}
|
||||
AugmentationInit::Linear { matrix } => {
|
||||
// Apply linear transformation to original state
|
||||
if let Some(_bs) = batch_size {
|
||||
y0.matmul(&matrix.t()?)?
|
||||
} else {
|
||||
matrix.matmul(y0)?
|
||||
}
|
||||
}
|
||||
AugmentationInit::Constant { value } => {
|
||||
let shape = if let Some(bs) = batch_size {
|
||||
[bs, self.config.augmented_dim]
|
||||
} else {
|
||||
[self.config.augmented_dim]
|
||||
};
|
||||
Tensor::full(shape, *value, &self.device)?
|
||||
}
|
||||
};
|
||||
|
||||
// Concatenate original and augmented parts
|
||||
let augmented_state = if batch_size.is_some() {
|
||||
Tensor::cat(&[y0.clone(), augmented_part], 1)?
|
||||
} else {
|
||||
Tensor::cat(&[y0.clone(), augmented_part], 0)?
|
||||
};
|
||||
|
||||
Ok(augmented_state)
|
||||
}
|
||||
|
||||
/// Extract original dimensions from augmented state
|
||||
fn extract_original(&self, augmented_state: &Tensor) -> Result<Tensor> {
|
||||
if augmented_state.ndim() == 1 {
|
||||
// Single trajectory: [total_dim] -> [original_dim]
|
||||
augmented_state.slice(0, 0..self.config.original_dim)
|
||||
} else if augmented_state.ndim() == 2 {
|
||||
if augmented_state.shape()[0] == self.config.original_dim + self.config.augmented_dim {
|
||||
// Single time series: [time_steps, total_dim] -> [time_steps, original_dim]
|
||||
augmented_state.slice(1, 0..self.config.original_dim)
|
||||
} else {
|
||||
// Batch: [batch_size, total_dim] -> [batch_size, original_dim]
|
||||
augmented_state.slice(1, 0..self.config.original_dim)
|
||||
}
|
||||
} else if augmented_state.ndim() == 3 {
|
||||
// Batch time series: [batch_size, time_steps, total_dim] -> [batch_size, time_steps, original_dim]
|
||||
augmented_state.slice(2, 0..self.config.original_dim)
|
||||
} else {
|
||||
Err(NeuralODEError::InvalidInput(format!(
|
||||
"Unsupported tensor dimensionality: {}",
|
||||
augmented_state.ndim()
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward pass with augmented state
|
||||
pub fn forward(&self, y0: &Tensor, t_span: &[f32]) -> Result<Tensor> {
|
||||
// Augment initial state
|
||||
let augmented_y0 = self.augment_state(y0)?;
|
||||
|
||||
// Solve ODE in augmented space
|
||||
let augmented_solution = self.neural_ode.forward(&augmented_y0, t_span)?;
|
||||
|
||||
// Return either full augmented state or just original dimensions
|
||||
if self.config.return_augmented {
|
||||
Ok(augmented_solution)
|
||||
} else {
|
||||
self.extract_original(&augmented_solution)
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward pass for batch of initial conditions
|
||||
pub fn forward_batch(&self, y0_batch: &Tensor, t_span: &[f32]) -> Result<Tensor> {
|
||||
if y0_batch.ndim() != 2 {
|
||||
return Err(NeuralODEError::InvalidInput(
|
||||
"Batch input must be 2D [batch_size, original_dim]".to_string()
|
||||
));
|
||||
}
|
||||
|
||||
// Augment batch
|
||||
let augmented_y0_batch = self.augment_state(y0_batch)?;
|
||||
|
||||
// Solve in augmented space
|
||||
let augmented_solution = self.neural_ode.forward_batch(&augmented_y0_batch, t_span)?;
|
||||
|
||||
// Extract original dimensions
|
||||
if self.config.return_augmented {
|
||||
Ok(augmented_solution)
|
||||
} else {
|
||||
self.extract_original(&augmented_solution)
|
||||
}
|
||||
}
|
||||
|
||||
/// Get all parameters from underlying Neural ODE
|
||||
pub fn parameters(&self) -> Vec<Tensor> {
|
||||
self.neural_ode.parameters()
|
||||
}
|
||||
|
||||
/// Get configuration
|
||||
pub fn config(&self) -> &AugmentationConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Set whether to return augmented dimensions in output
|
||||
pub fn set_return_augmented(&mut self, return_augmented: bool) {
|
||||
self.config.return_augmented = return_augmented;
|
||||
}
|
||||
|
||||
/// Get the underlying Neural ODE
|
||||
pub fn inner(&self) -> &NeuralODE {
|
||||
&self.neural_ode
|
||||
}
|
||||
|
||||
/// Evaluate dynamics in augmented space
|
||||
pub fn evaluate_augmented_dynamics(&self, t: f32, augmented_y: &Tensor) -> Result<Tensor> {
|
||||
self.neural_ode.evaluate_dynamics(t, augmented_y)
|
||||
}
|
||||
|
||||
/// Compare expressivity with standard Neural ODE
|
||||
///
|
||||
/// This runs both augmented and standard versions and returns metrics
|
||||
pub fn compare_with_standard(
|
||||
&self,
|
||||
standard_dynamics: Box<dyn ODEFunc>,
|
||||
y0: &Tensor,
|
||||
t_span: &[f32],
|
||||
) -> Result<ExpressionComparison> {
|
||||
// Run augmented version
|
||||
let augmented_result = self.forward(y0, t_span)?;
|
||||
|
||||
// Run standard version
|
||||
let standard_config = NeuralODEConfig::default(); // Use same config as augmented
|
||||
let standard_ode = NeuralODE::new(standard_dynamics, standard_config, &self.device)?;
|
||||
let standard_result = standard_ode.forward(y0, t_span)?;
|
||||
|
||||
// Compute comparison metrics
|
||||
let mse = augmented_result.sub(&standard_result)?.pow_scalar(2.0)?.mean()?;
|
||||
let max_diff = augmented_result.sub(&standard_result)?.abs()?.max()?;
|
||||
|
||||
let mse_val = mse.to_scalar::<f32>()?;
|
||||
let max_diff_val = max_diff.to_scalar::<f32>()?;
|
||||
|
||||
Ok(ExpressionComparison {
|
||||
mse: mse_val,
|
||||
max_absolute_diff: max_diff_val,
|
||||
augmented_trajectory: augmented_result,
|
||||
standard_trajectory: standard_result,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Results of comparing augmented vs standard Neural ODE
|
||||
#[derive(Debug)]
|
||||
pub struct ExpressionComparison {
|
||||
/// Mean squared error between trajectories
|
||||
pub mse: f32,
|
||||
/// Maximum absolute difference
|
||||
pub max_absolute_diff: f32,
|
||||
/// Trajectory from augmented Neural ODE
|
||||
pub augmented_trajectory: Tensor,
|
||||
/// Trajectory from standard Neural ODE
|
||||
pub standard_trajectory: Tensor,
|
||||
}
|
||||
|
||||
/// Wrapper that automatically augments any ODEFunc
|
||||
pub struct AugmentedODEFunc<F: ODEFunc> {
|
||||
inner: F,
|
||||
original_dim: usize,
|
||||
augmented_dim: usize,
|
||||
}
|
||||
|
||||
impl<F: ODEFunc> AugmentedODEFunc<F> {
|
||||
pub fn new(inner: F, original_dim: usize, augmented_dim: usize) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
original_dim,
|
||||
augmented_dim,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<F: ODEFunc> ODEFunc for AugmentedODEFunc<F> {
|
||||
fn forward(&self, t: f32, y: &Tensor) -> Result<Tensor> {
|
||||
// Extract original dimensions
|
||||
let y_original = if y.ndim() == 1 {
|
||||
y.slice(0, 0..self.original_dim)?
|
||||
} else {
|
||||
y.slice(1, 0..self.original_dim)?
|
||||
};
|
||||
|
||||
// Apply original dynamics to original dimensions
|
||||
let dy_original = self.inner.forward(t, &y_original)?;
|
||||
|
||||
// For augmented dimensions, we can use various strategies
|
||||
// Here we use a simple approach: augmented dimensions decay towards zero
|
||||
let y_augmented = if y.ndim() == 1 {
|
||||
y.slice(0, self.original_dim..)?
|
||||
} else {
|
||||
y.slice(1, self.original_dim..)?
|
||||
};
|
||||
|
||||
let dy_augmented = y_augmented.mul_scalar(-0.1)?; // Slow decay
|
||||
|
||||
// Combine derivatives
|
||||
let dy_combined = if y.ndim() == 1 {
|
||||
Tensor::cat(&[dy_original, dy_augmented], 0)?
|
||||
} else {
|
||||
Tensor::cat(&[dy_original, dy_augmented], 1)?
|
||||
};
|
||||
|
||||
Ok(dy_combined)
|
||||
}
|
||||
|
||||
fn parameters(&self) -> Vec<Tensor> {
|
||||
self.inner.parameters()
|
||||
}
|
||||
|
||||
fn is_autonomous(&self) -> bool {
|
||||
self.inner.is_autonomous()
|
||||
}
|
||||
|
||||
fn name(&self) -> String {
|
||||
format!("Augmented({})", self.inner.name())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "disabled_tests"))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::neural_ode::ode_func::LinearDynamics;
|
||||
|
||||
#[test]
|
||||
fn test_augmentation_config() {
|
||||
let config = AugmentationConfig::default();
|
||||
assert_eq!(config.original_dim, 0);
|
||||
assert_eq!(config.augmented_dim, 1);
|
||||
assert!(!config.return_augmented);
|
||||
|
||||
match config.init_strategy {
|
||||
AugmentationInit::SmallRandom { scale } => assert_eq!(scale, 0.01),
|
||||
_ => panic!("Expected SmallRandom initialization"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_augmented_neural_ode_creation() {
|
||||
let device = Device::cpu();
|
||||
let dynamics = LinearDynamics::decay(1.0, 4, device.clone()).unwrap(); // 4D total (2 original + 2 augmented)
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let augmented_ode = AugmentedNeuralODE::new(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
2, // original_dim
|
||||
2, // augmented_dim
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
assert_eq!(augmented_ode.config.original_dim, 2);
|
||||
assert_eq!(augmented_ode.config.augmented_dim, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_state_augmentation() {
|
||||
let device = Device::cpu();
|
||||
let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap(); // 3D total
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let augmented_ode = AugmentedNeuralODE::new(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
2, // original_dim
|
||||
1, // augmented_dim
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
let y0 = Tensor::ones([2], &device).unwrap();
|
||||
let augmented_y0 = augmented_ode.augment_state(&y0).unwrap();
|
||||
|
||||
assert_eq!(augmented_y0.shape(), &[3]); // Should be [original_dim + augmented_dim]
|
||||
|
||||
let aug_slice = augmented_y0.to_vec::<f32>().unwrap();
|
||||
// First two should be original state
|
||||
assert_eq!(aug_slice[0], 1.0);
|
||||
assert_eq!(aug_slice[1], 1.0);
|
||||
// Third should be small random value (approximately 0 for small scale)
|
||||
assert!(aug_slice[2].abs() < 0.1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_batch_augmentation() {
|
||||
let device = Device::cpu();
|
||||
let dynamics = LinearDynamics::decay(1.0, 4, device.clone()).unwrap(); // 4D total
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let augmented_ode = AugmentedNeuralODE::new(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
2, // original_dim
|
||||
2, // augmented_dim
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
let y0_batch = Tensor::ones([3, 2], &device).unwrap(); // Batch of 3
|
||||
let augmented_y0_batch = augmented_ode.augment_state(&y0_batch).unwrap();
|
||||
|
||||
assert_eq!(augmented_y0_batch.shape(), &[3, 4]); // [batch_size, total_dim]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_original_extraction() {
|
||||
let device = Device::cpu();
|
||||
let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let augmented_ode = AugmentedNeuralODE::new(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
2, // original_dim
|
||||
1, // augmented_dim
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
// Test extraction from time series: [time_steps, total_dim]
|
||||
let augmented_solution = Tensor::randn(&[5, 3], &device).unwrap();
|
||||
let original_part = augmented_ode.extract_original(&augmented_solution).unwrap();
|
||||
|
||||
assert_eq!(original_part.shape(), &[5, 2]); // [time_steps, original_dim]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_augmented_forward() {
|
||||
let device = Device::cpu();
|
||||
let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let augmented_ode = AugmentedNeuralODE::new(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
2, // original_dim
|
||||
1, // augmented_dim
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
let y0 = Tensor::ones([2], &device).unwrap();
|
||||
let t_span = vec![0.0, 0.5, 1.0];
|
||||
|
||||
let result = augmented_ode.forward(&y0, &t_span).unwrap();
|
||||
|
||||
// Should return only original dimensions
|
||||
assert_eq!(result.shape(), &[3, 2]); // [time_steps, original_dim]
|
||||
|
||||
let result_slice = result.to_vec::<f32>().unwrap();
|
||||
// Should start with initial conditions
|
||||
assert!((result_slice[0] - 1.0).abs() < 1e-6);
|
||||
assert!((result_slice[1] - 1.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_return_augmented_flag() {
|
||||
let device = Device::cpu();
|
||||
let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let mut augmented_ode = AugmentedNeuralODE::new(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
2, // original_dim
|
||||
1, // augmented_dim
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
let y0 = Tensor::ones([2], &device).unwrap();
|
||||
let t_span = vec![0.0, 1.0];
|
||||
|
||||
// Test with return_augmented = false (default)
|
||||
let result_original = augmented_ode.forward(&y0, &t_span).unwrap();
|
||||
assert_eq!(result_original.shape(), &[2, 2]); // Original dimensions only
|
||||
|
||||
// Test with return_augmented = true
|
||||
augmented_ode.set_return_augmented(true);
|
||||
let result_augmented = augmented_ode.forward(&y0, &t_span).unwrap();
|
||||
assert_eq!(result_augmented.shape(), &[2, 3]); // All dimensions
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_initialization_strategies() {
|
||||
let device = Device::cpu();
|
||||
|
||||
// Test zeros initialization
|
||||
let config_zeros = AugmentationConfig {
|
||||
original_dim: 2,
|
||||
augmented_dim: 2,
|
||||
init_strategy: AugmentationInit::Zeros,
|
||||
return_augmented: true,
|
||||
};
|
||||
|
||||
let dynamics = LinearDynamics::decay(1.0, 4, device.clone()).unwrap();
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let augmented_ode = AugmentedNeuralODE::with_config(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
config_zeros,
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
let y0 = Tensor::ones([2], &device).unwrap();
|
||||
let augmented_y0 = augmented_ode.augment_state(&y0).unwrap();
|
||||
|
||||
let aug_slice = augmented_y0.to_vec::<f32>().unwrap();
|
||||
assert_eq!(aug_slice[0], 1.0); // Original
|
||||
assert_eq!(aug_slice[1], 1.0); // Original
|
||||
assert_eq!(aug_slice[2], 0.0); // Augmented (zeros)
|
||||
assert_eq!(aug_slice[3], 0.0); // Augmented (zeros)
|
||||
|
||||
// Test constant initialization
|
||||
let config_const = AugmentationConfig {
|
||||
original_dim: 2,
|
||||
augmented_dim: 1,
|
||||
init_strategy: AugmentationInit::Constant { value: 0.5 },
|
||||
return_augmented: true,
|
||||
};
|
||||
|
||||
let dynamics2 = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
|
||||
let augmented_ode2 = AugmentedNeuralODE::with_config(
|
||||
Box::new(dynamics2),
|
||||
NeuralODEConfig::default(),
|
||||
config_const,
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
let augmented_y0_2 = augmented_ode2.augment_state(&y0).unwrap();
|
||||
let aug_slice_2 = augmented_y0_2.to_vec::<f32>().unwrap();
|
||||
assert_eq!(aug_slice_2[2], 0.5); // Should be constant value
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_augmented_ode_func_wrapper() {
|
||||
let device = Device::cpu();
|
||||
let inner_dynamics = LinearDynamics::decay(1.0, 2, device.clone()).unwrap();
|
||||
let augmented_dynamics = AugmentedODEFunc::new(inner_dynamics, 2, 1);
|
||||
|
||||
let y = Tensor::ones([3], &device).unwrap(); // [2 original + 1 augmented]
|
||||
let dy_dt = augmented_dynamics.forward(0.0, &y).unwrap();
|
||||
|
||||
assert_eq!(dy_dt.shape(), &[3]);
|
||||
|
||||
let dy_slice = dy_dt.to_vec::<f32>().unwrap();
|
||||
// Original dimensions should follow inner dynamics: -1 * y
|
||||
assert!((dy_slice[0] + 1.0).abs() < 1e-6);
|
||||
assert!((dy_slice[1] + 1.0).abs() < 1e-6);
|
||||
// Augmented dimension should decay slowly: -0.1 * y
|
||||
assert!((dy_slice[2] + 0.1).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parameters_passthrough() {
|
||||
let device = Device::cpu();
|
||||
let dynamics = LinearDynamics::decay(1.0, 3, device.clone()).unwrap();
|
||||
let ode_config = NeuralODEConfig::default();
|
||||
|
||||
let augmented_ode = AugmentedNeuralODE::new(
|
||||
Box::new(dynamics),
|
||||
ode_config,
|
||||
2,
|
||||
1,
|
||||
&device,
|
||||
).unwrap();
|
||||
|
||||
let params = augmented_ode.parameters();
|
||||
assert!(!params.is_empty()); // Should have parameters from underlying dynamics
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user