617 lines
21 KiB
Rust
617 lines
21 KiB
Rust
//! 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
|
|
}
|
|
} |