305 lines
10 KiB
Rust
305 lines
10 KiB
Rust
//! Tests for PINN training implementation
|
|
//!
|
|
//! These tests verify the correctness and robustness of the Physics-Informed
|
|
//! Neural Network training algorithms.
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::super::{PINNTrainer, TrainingConfig, LossComponents};
|
|
use crate::physics::{PINN, BoundaryConditions, HeatEquation};
|
|
use crate::error::Result;
|
|
|
|
#[tokio::test]
|
|
async fn test_training_config_defaults() {
|
|
let config = TrainingConfig::default();
|
|
|
|
assert_eq!(config.learning_rate, 1e-3);
|
|
assert_eq!(config.physics_weight, 1.0);
|
|
assert_eq!(config.boundary_weight, 100.0);
|
|
assert_eq!(config.convergence_tolerance, 1e-6);
|
|
assert_eq!(config.patience, 100);
|
|
assert_eq!(config.conservation_weights, vec![1.0]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_trainer_creation() {
|
|
let config = TrainingConfig::default();
|
|
let trainer = PINNTrainer::new(config);
|
|
|
|
// Verify trainer is created successfully
|
|
assert!(trainer.config.learning_rate > 0.0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_latin_hypercube_sampling() -> Result<()> {
|
|
let config = TrainingConfig::default();
|
|
let trainer = PINNTrainer::new(config);
|
|
|
|
// Create mock boundary conditions
|
|
let boundary_data = create_test_boundary_conditions();
|
|
|
|
// Generate training points
|
|
let points = trainer.generate_training_points(&boundary_data, 100)?;
|
|
|
|
// Verify correct number of points
|
|
assert_eq!(points.len(), 100);
|
|
|
|
// Verify points are within domain bounds
|
|
let (x_min, x_max) = boundary_data.spatial_bounds();
|
|
let (t_min, t_max) = boundary_data.temporal_bounds();
|
|
|
|
for (x, t) in &points {
|
|
assert!(*x >= x_min && *x <= x_max, "x coordinate {} out of bounds [{}, {}]", x, x_min, x_max);
|
|
assert!(*t >= t_min && *t <= t_max, "t coordinate {} out of bounds [{}, {}]", t, t_min, t_max);
|
|
}
|
|
|
|
// Verify good space coverage (no clustering)
|
|
let mut x_coords: Vec<f64> = points.iter().map(|(x, _)| *x).collect();
|
|
x_coords.sort_by(|a, b| a.total_cmp(b));
|
|
|
|
// Check that points are reasonably spread out
|
|
let min_spacing = (x_max - x_min) / 200.0; // Expect at least this spacing
|
|
for i in 1..x_coords.len() {
|
|
let spacing = x_coords[i] - x_coords[i-1];
|
|
if spacing < min_spacing {
|
|
// Allow some clustering but not excessive
|
|
let clustered_count = x_coords.windows(2)
|
|
.filter(|w| w[1] - w[0] < min_spacing)
|
|
.count();
|
|
assert!(clustered_count < points.len() / 4, "Too much point clustering detected");
|
|
break;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_boundary_point_sampling() -> Result<()> {
|
|
let config = TrainingConfig::default();
|
|
let trainer = PINNTrainer::new(config);
|
|
|
|
let boundary_data = create_test_boundary_conditions();
|
|
let points = trainer.sample_boundary_points(&boundary_data, 100)?;
|
|
|
|
// Verify correct number of points
|
|
assert_eq!(points.len(), 100);
|
|
|
|
let (x_min, x_max) = boundary_data.spatial_bounds();
|
|
let (t_min, t_max) = boundary_data.temporal_bounds();
|
|
|
|
// Count points on each boundary
|
|
let mut left_boundary = 0;
|
|
let mut right_boundary = 0;
|
|
let mut bottom_boundary = 0;
|
|
let mut top_boundary = 0;
|
|
|
|
for (x, t) in &points {
|
|
if (*x - x_min).abs() < 1e-10 { left_boundary += 1; }
|
|
if (*x - x_max).abs() < 1e-10 { right_boundary += 1; }
|
|
if (*t - t_min).abs() < 1e-10 { bottom_boundary += 1; }
|
|
if (*t - t_max).abs() < 1e-10 { top_boundary += 1; }
|
|
}
|
|
|
|
// Verify roughly equal distribution on boundaries
|
|
let expected_per_boundary = 100 / 4;
|
|
assert!(left_boundary >= expected_per_boundary - 2);
|
|
assert!(right_boundary >= expected_per_boundary - 2);
|
|
assert!(bottom_boundary >= expected_per_boundary - 2);
|
|
assert!(top_boundary >= expected_per_boundary - 2);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_convergence_detection() -> Result<()> {
|
|
let mut config = TrainingConfig::default();
|
|
config.convergence_tolerance = 1e-4;
|
|
config.patience = 10;
|
|
|
|
let trainer = PINNTrainer::new(config);
|
|
|
|
// Test converged case
|
|
let mut loss_history = std::collections::VecDeque::new();
|
|
for _ in 0..15 {
|
|
loss_history.push_back(1e-5); // Loss below tolerance
|
|
}
|
|
|
|
let converged = trainer.check_convergence(&loss_history, 1e-5)?;
|
|
assert!(converged, "Should detect convergence when loss is below tolerance");
|
|
|
|
// Test non-converged case
|
|
loss_history.clear();
|
|
for i in 0..15 {
|
|
loss_history.push_back(0.1 - (i as f64) * 0.001); // Steadily decreasing but above tolerance
|
|
}
|
|
|
|
let converged = trainer.check_convergence(&loss_history, 0.085)?;
|
|
assert!(!converged, "Should not detect convergence when loss is above tolerance");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_adam_optimizer() -> Result<()> {
|
|
use super::super::AdamOptimizer;
|
|
|
|
let mut optimizer = AdamOptimizer::new(0.001);
|
|
|
|
// Test initial state
|
|
assert!((optimizer.learning_rate - 0.001).abs() < 1e-10);
|
|
|
|
// Test learning rate update
|
|
optimizer.set_learning_rate(0.01);
|
|
assert!((optimizer.learning_rate - 0.01).abs() < 1e-10);
|
|
|
|
// Test optimizer step (mock parameters)
|
|
let mock_params = vec![];
|
|
let result = optimizer.step(mock_params).await;
|
|
assert!(result.is_ok(), "Optimizer step should succeed");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_cosine_annealing_scheduler() -> Result<()> {
|
|
use super::super::CosineAnnealingScheduler;
|
|
|
|
let scheduler = CosineAnnealingScheduler::new(0.01, 100, 0.001);
|
|
|
|
// Test initial learning rate
|
|
let lr_0 = scheduler.step(0);
|
|
assert!((lr_0 - 0.01).abs() < 1e-6, "Initial LR should be close to max");
|
|
|
|
// Test middle learning rate
|
|
let lr_50 = scheduler.step(50);
|
|
assert!(lr_50 < 0.01 && lr_50 > 0.001, "Middle LR should be between min and max");
|
|
|
|
// Test final learning rate
|
|
let lr_100 = scheduler.step(100);
|
|
assert!((lr_100 - 0.001).abs() < 1e-6, "Final LR should be close to min");
|
|
|
|
// Test monotonic decrease in first half
|
|
let lr_25 = scheduler.step(25);
|
|
assert!(lr_25 < lr_0, "LR should decrease in first quarter");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_loss_components_creation() {
|
|
let loss = LossComponents {
|
|
physics: 0.1,
|
|
boundary: 0.05,
|
|
initial: 0.02,
|
|
conservation: vec![0.01, 0.015],
|
|
data: 0.0,
|
|
total: 0.185,
|
|
iteration: 42,
|
|
};
|
|
|
|
assert!((loss.physics - 0.1).abs() < 1e-10);
|
|
assert!((loss.boundary - 0.05).abs() < 1e-10);
|
|
assert_eq!(loss.conservation.len(), 2);
|
|
assert_eq!(loss.iteration, 42);
|
|
|
|
// Test total loss calculation
|
|
let expected_total = loss.physics + loss.boundary + loss.initial +
|
|
loss.conservation.iter().sum::<f64>() + loss.data;
|
|
assert!((loss.total - expected_total).abs() < 1e-10);
|
|
}
|
|
|
|
// Helper function to create test boundary conditions
|
|
fn create_test_boundary_conditions() -> BoundaryConditions {
|
|
// This would be a mock implementation
|
|
// In the real codebase, would use actual BoundaryConditions
|
|
use std::collections::HashMap;
|
|
|
|
// Mock implementation that provides the interface needed for testing
|
|
BoundaryConditions::new_for_testing(0.0, 1.0, 0.0, 2.0)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_physics_gradients_structure() {
|
|
use super::super::PhysicsGradients;
|
|
use rtx_autograd::Variable;
|
|
use crate::Tensor;
|
|
use crate::Device;
|
|
|
|
// Create mock gradients for testing
|
|
let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0], &[3], Device::cuda(0).unwrap_or(Device::default())).unwrap();
|
|
let variable = Variable::from_tensor(tensor, true);
|
|
|
|
let gradients = PhysicsGradients {
|
|
du_dx: variable.clone(),
|
|
du_dt: variable.clone(),
|
|
d2u_dx2: variable.clone(),
|
|
d2u_dt2: variable.clone(),
|
|
d2u_dxdt: variable.clone(),
|
|
};
|
|
|
|
// Test that all gradient components are properly initialized
|
|
// In a real test, would verify gradient computations
|
|
assert!(true); // Placeholder assertion for structure test
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_training_integration_mock() -> Result<()> {
|
|
// This would be a comprehensive integration test
|
|
// For now, test that the training interface is properly structured
|
|
|
|
let config = TrainingConfig {
|
|
learning_rate: 0.001,
|
|
physics_weight: 1.0,
|
|
boundary_weight: 10.0,
|
|
conservation_weights: vec![1.0],
|
|
convergence_tolerance: 1e-6,
|
|
patience: 50,
|
|
};
|
|
|
|
let trainer = PINNTrainer::new(config);
|
|
|
|
// Test that trainer has all necessary components
|
|
assert!(trainer.config.learning_rate > 0.0);
|
|
assert!(trainer.config.physics_weight > 0.0);
|
|
assert!(trainer.config.boundary_weight > 0.0);
|
|
assert!(!trainer.config.conservation_weights.is_empty());
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// Mock BoundaryConditions for testing (would be in actual module)
|
|
impl BoundaryConditions {
|
|
pub fn new_for_testing(x_min: f64, x_max: f64, t_min: f64, t_max: f64) -> Self {
|
|
// Mock implementation for testing
|
|
Self {
|
|
spatial_domain: (x_min, x_max),
|
|
temporal_domain: (t_min, t_max),
|
|
// ... other fields would be properly initialized
|
|
}
|
|
}
|
|
|
|
pub fn spatial_bounds(&self) -> (f64, f64) {
|
|
self.spatial_domain
|
|
}
|
|
|
|
pub fn temporal_bounds(&self) -> (f64, f64) {
|
|
self.temporal_domain
|
|
}
|
|
|
|
pub fn get_value_at(&self, _x: f64, _t: f64) -> Option<f64> {
|
|
Some(0.0) // Mock boundary value
|
|
}
|
|
|
|
pub fn get_initial_value_at(&self, _x: f64) -> Option<f64> {
|
|
Some(1.0) // Mock initial value
|
|
}
|
|
}
|
|
|
|
// Mock BoundaryConditions structure for testing
|
|
#[derive(Debug)]
|
|
pub struct BoundaryConditions {
|
|
spatial_domain: (f64, f64),
|
|
temporal_domain: (f64, f64),
|
|
} |