291 lines
9.6 KiB
Rust
291 lines
9.6 KiB
Rust
// Integration tests for physics solvers
|
||
// Following strict TDD - RED phase first
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use crate::solver::{Solver, SolverConfig};
|
||
use rtx_tensor::{Tensor, Device, DType};
|
||
use rtx_autograd::variable::Variable;
|
||
|
||
#[test]
|
||
fn test_harmonic_oscillator_solver() {
|
||
let device = Device::cpu();
|
||
|
||
// Simple harmonic oscillator: d²x/dt² = -ω²x
|
||
// Solution: x(t) = A*cos(ωt) + B*sin(ωt)
|
||
|
||
let omega = 2.0; // Angular frequency
|
||
let dt = 0.01; // Time step
|
||
let t_final = 10.0;
|
||
let steps = (t_final / dt) as usize;
|
||
|
||
// Initial conditions: x(0) = 1, dx/dt(0) = 0
|
||
let x0 = Tensor::scalar(1.0, DType::F32, &device).unwrap();
|
||
let v0 = Tensor::scalar(0.0, DType::F32, &device).unwrap();
|
||
|
||
// Create solver configuration
|
||
let config = SolverConfig {
|
||
method: "rk4".to_string(),
|
||
tolerance: 1e-6,
|
||
max_iterations: 1000,
|
||
adaptive_timestep: false,
|
||
};
|
||
|
||
let solver = Solver::new(config);
|
||
|
||
// Define the system dynamics
|
||
let dynamics = |x: &Tensor, _t: f32| -> Tensor {
|
||
x.mul_scalar(-omega * omega).unwrap()
|
||
};
|
||
|
||
// Solve the ODE
|
||
let mut x = x0.clone();
|
||
let mut v = v0.clone();
|
||
let mut trajectory = vec![x.to_vec().unwrap()[0]];
|
||
|
||
for i in 0..steps {
|
||
let t = i as f32 * dt;
|
||
|
||
// Update using velocity
|
||
x = x.add(&v.mul_scalar(dt).unwrap()).unwrap();
|
||
|
||
// Update velocity using acceleration
|
||
let accel = dynamics(&x, t);
|
||
v = v.add(&accel.mul_scalar(dt).unwrap()).unwrap();
|
||
|
||
trajectory.push(x.to_vec().unwrap()[0]);
|
||
}
|
||
|
||
// Verify solution is periodic
|
||
let period = 2.0 * std::f32::consts::PI / omega;
|
||
let cycles = t_final / period;
|
||
|
||
// Check that we completed full cycles
|
||
assert!(cycles > 1.0);
|
||
|
||
// Final position should be close to initial (after full periods)
|
||
let final_x = trajectory.last().unwrap();
|
||
let initial_x = trajectory.first().unwrap();
|
||
|
||
// Allow for numerical error accumulation
|
||
assert!((final_x - initial_x).abs() < 0.1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_heat_equation_solver() {
|
||
let device = Device::cpu();
|
||
|
||
// 1D heat equation: ∂u/∂t = α * ∂²u/∂x²
|
||
// Domain: [0, 1], initial condition: gaussian pulse
|
||
|
||
let nx = 50; // Spatial points
|
||
let dx = 1.0 / (nx as f32 - 1.0);
|
||
let alpha = 0.01; // Thermal diffusivity
|
||
let dt = 0.5 * dx * dx / alpha; // CFL condition
|
||
let t_final = 0.1;
|
||
let nt = (t_final / dt) as usize;
|
||
|
||
// Initial condition: Gaussian
|
||
let mut x_vals = Vec::new();
|
||
let mut u_vals = Vec::new();
|
||
for i in 0..nx {
|
||
let x = i as f32 * dx;
|
||
x_vals.push(x);
|
||
let u = (-50.0 * (x - 0.5) * (x - 0.5)).exp();
|
||
u_vals.push(u);
|
||
}
|
||
|
||
let u = Tensor::from_vec(u_vals.clone(), &[nx], &device).unwrap();
|
||
|
||
// Time evolution
|
||
let mut current_u = u;
|
||
for _ in 0..nt {
|
||
// Compute second derivative using finite differences
|
||
let u_vec = current_u.to_vec().unwrap();
|
||
let mut laplacian = vec![0.0; nx];
|
||
|
||
for i in 1..(nx - 1) {
|
||
laplacian[i] = (u_vec[i + 1] - 2.0 * u_vec[i] + u_vec[i - 1]) / (dx * dx);
|
||
}
|
||
|
||
// Boundary conditions (Dirichlet: u = 0 at boundaries)
|
||
laplacian[0] = 0.0;
|
||
laplacian[nx - 1] = 0.0;
|
||
|
||
let lap_tensor = Tensor::from_vec(laplacian, &[nx], &device).unwrap();
|
||
|
||
// Update using forward Euler
|
||
current_u = current_u.add(&lap_tensor.mul_scalar(alpha * dt).unwrap()).unwrap();
|
||
}
|
||
|
||
// Verify heat has diffused (peak should be lower)
|
||
let final_vals = current_u.to_vec().unwrap();
|
||
let initial_max = u_vals.iter().fold(0.0f32, |a, &b| a.max(b));
|
||
let final_max = final_vals.iter().fold(0.0f32, |a, &b| a.max(b));
|
||
|
||
assert!(final_max < initial_max);
|
||
assert!(final_max > 0.0); // Should still have some heat
|
||
}
|
||
|
||
#[test]
|
||
fn test_wave_equation_solver() {
|
||
let device = Device::cpu();
|
||
|
||
// 1D wave equation: ∂²u/∂t² = c² * ∂²u/∂x²
|
||
let nx = 100;
|
||
let dx = 1.0 / (nx as f32 - 1.0);
|
||
let c = 1.0; // Wave speed
|
||
let dt = 0.5 * dx / c; // CFL condition
|
||
let t_final = 2.0;
|
||
let nt = (t_final / dt) as usize;
|
||
|
||
// Initial conditions: Gaussian pulse, zero velocity
|
||
let mut u_vals = Vec::new();
|
||
for i in 0..nx {
|
||
let x = i as f32 * dx;
|
||
let u = (-100.0 * (x - 0.5) * (x - 0.5)).exp();
|
||
u_vals.push(u);
|
||
}
|
||
|
||
let u_prev = Tensor::from_vec(u_vals.clone(), &[nx], &device).unwrap();
|
||
let mut u_curr = u_prev.clone();
|
||
let mut u_next;
|
||
|
||
// Time evolution using finite differences
|
||
for _ in 0..nt {
|
||
let u_vec = u_curr.to_vec().unwrap();
|
||
let u_prev_vec = u_prev.to_vec().unwrap();
|
||
let mut new_u = vec![0.0; nx];
|
||
|
||
for i in 1..(nx - 1) {
|
||
let d2u_dx2 = (u_vec[i + 1] - 2.0 * u_vec[i] + u_vec[i - 1]) / (dx * dx);
|
||
new_u[i] = 2.0 * u_vec[i] - u_prev_vec[i] + c * c * dt * dt * d2u_dx2;
|
||
}
|
||
|
||
// Boundary conditions (fixed ends)
|
||
new_u[0] = 0.0;
|
||
new_u[nx - 1] = 0.0;
|
||
|
||
u_next = Tensor::from_vec(new_u, &[nx], &device).unwrap();
|
||
|
||
// Update for next iteration
|
||
u_prev = u_curr;
|
||
u_curr = u_next;
|
||
}
|
||
|
||
// Wave should have propagated
|
||
let final_vals = u_curr.to_vec().unwrap();
|
||
|
||
// Check energy is approximately conserved (with some numerical dissipation)
|
||
let initial_energy: f32 = u_vals.iter().map(|x| x * x).sum();
|
||
let final_energy: f32 = final_vals.iter().map(|x| x * x).sum();
|
||
|
||
// Allow up to 20% energy loss due to numerical dissipation
|
||
assert!(final_energy > 0.8 * initial_energy);
|
||
}
|
||
|
||
#[test]
|
||
fn test_gradient_flow_optimization() {
|
||
let device = Device::cpu();
|
||
|
||
// Minimize f(x, y) = x² + 2y² using gradient descent
|
||
let mut x = Variable::new(Tensor::scalar(3.0, DType::F32, &device).unwrap(), true);
|
||
let mut y = Variable::new(Tensor::scalar(4.0, DType::F32, &device).unwrap(), true);
|
||
|
||
let learning_rate = 0.1;
|
||
let iterations = 100;
|
||
|
||
for _ in 0..iterations {
|
||
// Compute f = x² + 2y²
|
||
let x_sq = x.multiply(&x).unwrap();
|
||
let y_sq = y.multiply(&y).unwrap();
|
||
let two_y_sq = y_sq.multiply_scalar(2.0).unwrap();
|
||
let f = x_sq.add(&two_y_sq).unwrap();
|
||
|
||
// Compute gradients
|
||
f.backward(None).unwrap();
|
||
|
||
// Get gradients
|
||
let x_grad = x.grad().unwrap();
|
||
let y_grad = y.grad().unwrap();
|
||
|
||
// Update parameters
|
||
let x_tensor = x.tensor();
|
||
let y_tensor = y.tensor();
|
||
|
||
let new_x_tensor = x_tensor.sub(&x_grad.mul_scalar(learning_rate).unwrap()).unwrap();
|
||
let new_y_tensor = y_tensor.sub(&y_grad.mul_scalar(learning_rate).unwrap()).unwrap();
|
||
|
||
// Create new variables
|
||
x = Variable::new(new_x_tensor, true);
|
||
y = Variable::new(new_y_tensor, true);
|
||
}
|
||
|
||
// Should converge to (0, 0)
|
||
let final_x = x.tensor().to_vec().unwrap()[0];
|
||
let final_y = y.tensor().to_vec().unwrap()[0];
|
||
|
||
assert!(final_x.abs() < 0.01);
|
||
assert!(final_y.abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_nonlinear_system_solver() {
|
||
let device = Device::cpu();
|
||
|
||
// Solve nonlinear system:
|
||
// x² + y² = 1 (circle)
|
||
// y = x² (parabola)
|
||
// Solutions: approximately (0.786, 0.618) and (-0.786, 0.618)
|
||
|
||
let mut x = Variable::new(Tensor::scalar(0.5, DType::F32, &device).unwrap(), true);
|
||
let mut y = Variable::new(Tensor::scalar(0.5, DType::F32, &device).unwrap(), true);
|
||
|
||
let iterations = 50;
|
||
let lr = 0.01;
|
||
|
||
for _ in 0..iterations {
|
||
// Compute residuals
|
||
let x_sq = x.multiply(&x).unwrap();
|
||
let y_sq = y.multiply(&y).unwrap();
|
||
|
||
// r1 = x² + y² - 1
|
||
let sum_sq = x_sq.add(&y_sq).unwrap();
|
||
let one = Variable::new(Tensor::scalar(1.0, DType::F32, &device).unwrap(), false);
|
||
let r1 = sum_sq.add(&one.multiply_scalar(-1.0).unwrap()).unwrap();
|
||
|
||
// r2 = y - x²
|
||
let r2 = y.add(&x_sq.multiply_scalar(-1.0).unwrap()).unwrap();
|
||
|
||
// Loss = r1² + r2²
|
||
let r1_sq = r1.multiply(&r1).unwrap();
|
||
let r2_sq = r2.multiply(&r2).unwrap();
|
||
let loss = r1_sq.add(&r2_sq).unwrap();
|
||
|
||
// Backward pass
|
||
loss.backward(None).unwrap();
|
||
|
||
// Update
|
||
let x_grad = x.grad().unwrap();
|
||
let y_grad = y.grad().unwrap();
|
||
|
||
let new_x = x.tensor().sub(&x_grad.mul_scalar(lr).unwrap()).unwrap();
|
||
let new_y = y.tensor().sub(&y_grad.mul_scalar(lr).unwrap()).unwrap();
|
||
|
||
x = Variable::new(new_x, true);
|
||
y = Variable::new(new_y, true);
|
||
}
|
||
|
||
// Check solution satisfies constraints
|
||
let x_val = x.tensor().to_vec().unwrap()[0];
|
||
let y_val = y.tensor().to_vec().unwrap()[0];
|
||
|
||
// Check x² + y² ≈ 1
|
||
let circle_error = (x_val * x_val + y_val * y_val - 1.0).abs();
|
||
assert!(circle_error < 0.1);
|
||
|
||
// Check y ≈ x²
|
||
let parabola_error = (y_val - x_val * x_val).abs();
|
||
assert!(parabola_error < 0.1);
|
||
}
|
||
} |