515 lines
16 KiB
Rust
515 lines
16 KiB
Rust
//! Classical PDE solvers for comparison.
|
|
//!
|
|
//! Provides finite difference solvers to compare against neural operators.
|
|
|
|
use neuralop_studio_shared::{BCType, PDEDefinition, PDEType, SolutionField};
|
|
|
|
/// Classical PDE solver.
|
|
#[derive(Debug)]
|
|
pub struct ClassicalSolver {
|
|
/// Maximum iterations.
|
|
max_iterations: usize,
|
|
/// Convergence tolerance.
|
|
tolerance: f64,
|
|
}
|
|
|
|
impl Default for ClassicalSolver {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl ClassicalSolver {
|
|
/// Create a new classical solver.
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
max_iterations: 10000,
|
|
tolerance: 1e-6,
|
|
}
|
|
}
|
|
|
|
/// Set maximum iterations.
|
|
#[must_use]
|
|
pub fn with_max_iterations(mut self, max_iter: usize) -> Self {
|
|
self.max_iterations = max_iter;
|
|
self
|
|
}
|
|
|
|
/// Set convergence tolerance.
|
|
#[must_use]
|
|
pub fn with_tolerance(mut self, tol: f64) -> Self {
|
|
self.tolerance = tol;
|
|
self
|
|
}
|
|
|
|
/// Solve the PDE.
|
|
pub fn solve(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
|
|
match pde.pde_type {
|
|
PDEType::Poisson => self.solve_poisson(pde),
|
|
PDEType::Heat => self.solve_heat(pde),
|
|
PDEType::Wave => self.solve_wave(pde),
|
|
PDEType::Burgers => self.solve_burgers(pde),
|
|
_ => Err(format!("Unsupported PDE type: {:?}", pde.pde_type)),
|
|
}
|
|
}
|
|
|
|
/// Solve Poisson equation using Gauss-Seidel iteration.
|
|
fn solve_poisson(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
|
|
let nx = pde.domain.resolution[0];
|
|
let ny = if pde.domain.dimensions > 1 {
|
|
pde.domain.resolution.get(1).copied().unwrap_or(1)
|
|
} else {
|
|
1
|
|
};
|
|
|
|
let x_range = pde.domain.bounds[0];
|
|
let y_range = if pde.domain.dimensions > 1 {
|
|
pde.domain.bounds.get(1).copied().unwrap_or((0.0, 1.0))
|
|
} else {
|
|
(0.0, 1.0)
|
|
};
|
|
|
|
let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
|
|
let dy = if ny > 1 {
|
|
(y_range.1 - y_range.0) / (ny - 1) as f64
|
|
} else {
|
|
dx
|
|
};
|
|
|
|
let mut u = vec![0.0; nx * ny];
|
|
let mut u_new = vec![0.0; nx * ny];
|
|
|
|
// Apply boundary conditions
|
|
self.apply_boundary_conditions(&mut u, pde, nx, ny);
|
|
|
|
// Source term
|
|
let source: Vec<f64> = (0..nx * ny)
|
|
.map(|idx| {
|
|
let i = idx % nx;
|
|
let j = idx / nx;
|
|
let x = x_range.0 + i as f64 * dx;
|
|
let y = y_range.0 + j as f64 * dy;
|
|
self.evaluate_source(pde, x, y)
|
|
})
|
|
.collect();
|
|
|
|
// Gauss-Seidel iteration
|
|
for iteration in 0..self.max_iterations {
|
|
u_new.copy_from_slice(&u);
|
|
|
|
for j in 1..ny.saturating_sub(1).max(1) {
|
|
for i in 1..nx.saturating_sub(1) {
|
|
let idx = j * nx + i;
|
|
let idx_left = j * nx + i - 1;
|
|
let idx_right = j * nx + i + 1;
|
|
let idx_down = (j.saturating_sub(1)) * nx + i;
|
|
let idx_up = (j + 1).min(ny - 1) * nx + i;
|
|
|
|
if ny == 1 {
|
|
// 1D case
|
|
u_new[idx] = 0.5 * (u[idx_left] + u[idx_right] + dx * dx * source[idx]);
|
|
} else {
|
|
// 2D case
|
|
u_new[idx] = 0.25
|
|
* (u[idx_left]
|
|
+ u_new[idx_right.min(nx * ny - 1)]
|
|
+ u[idx_down]
|
|
+ u_new[idx_up.min(nx * ny - 1)]
|
|
+ dx * dy * source[idx]);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check convergence
|
|
let diff: f64 = u
|
|
.iter()
|
|
.zip(u_new.iter())
|
|
.map(|(a, b)| (a - b).abs())
|
|
.sum::<f64>()
|
|
/ (nx * ny) as f64;
|
|
|
|
u.copy_from_slice(&u_new);
|
|
|
|
if diff < self.tolerance {
|
|
break;
|
|
}
|
|
|
|
if iteration == self.max_iterations - 1 {
|
|
// Warning: did not converge, but continue
|
|
}
|
|
}
|
|
|
|
// Generate coordinates
|
|
let coordinates: Vec<Vec<f64>> = (0..nx * ny)
|
|
.map(|idx| {
|
|
let i = idx % nx;
|
|
let j = idx / nx;
|
|
let x = x_range.0 + i as f64 * dx;
|
|
let y = y_range.0 + j as f64 * dy;
|
|
vec![x, y]
|
|
})
|
|
.collect();
|
|
|
|
Ok(SolutionField {
|
|
coordinates,
|
|
values: u,
|
|
name: "u".to_string(),
|
|
time: None,
|
|
})
|
|
}
|
|
|
|
/// Solve heat equation using explicit finite differences.
|
|
fn solve_heat(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
|
|
let nx = pde.domain.resolution[0];
|
|
let (t_start, t_end) = pde.domain.time_bounds.unwrap_or((0.0, 1.0));
|
|
let nt = pde.domain.time_resolution.unwrap_or(100);
|
|
|
|
let x_range = pde.domain.bounds[0];
|
|
let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
|
|
let dt = (t_end - t_start) / nt as f64;
|
|
|
|
let alpha = pde.parameters.diffusion.unwrap_or(0.01);
|
|
|
|
// CFL condition check
|
|
let cfl = alpha * dt / (dx * dx);
|
|
if cfl > 0.5 {
|
|
// Adjust dt if unstable
|
|
let _dt = 0.4 * dx * dx / alpha;
|
|
}
|
|
|
|
let mut u = vec![0.0; nx];
|
|
let mut u_new = vec![0.0; nx];
|
|
|
|
// Apply initial condition
|
|
if let Some(ref ic) = pde.initial_condition {
|
|
for i in 0..nx {
|
|
let x = x_range.0 + i as f64 * dx;
|
|
u[i] = self.evaluate_initial(ic, x);
|
|
}
|
|
}
|
|
|
|
// Time stepping
|
|
for _ in 0..nt {
|
|
u_new.copy_from_slice(&u);
|
|
|
|
for i in 1..nx - 1 {
|
|
let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
|
|
u_new[i] = u[i] + dt * alpha * laplacian;
|
|
}
|
|
|
|
// Apply boundary conditions
|
|
u_new[0] = 0.0;
|
|
u_new[nx - 1] = 0.0;
|
|
|
|
u.copy_from_slice(&u_new);
|
|
}
|
|
|
|
let coordinates: Vec<Vec<f64>> = (0..nx)
|
|
.map(|i| {
|
|
let x = x_range.0 + i as f64 * dx;
|
|
vec![x]
|
|
})
|
|
.collect();
|
|
|
|
Ok(SolutionField {
|
|
coordinates,
|
|
values: u,
|
|
name: "u".to_string(),
|
|
time: Some(t_end),
|
|
})
|
|
}
|
|
|
|
/// Solve wave equation using explicit finite differences.
|
|
fn solve_wave(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
|
|
let nx = pde.domain.resolution[0];
|
|
let (t_start, t_end) = pde.domain.time_bounds.unwrap_or((0.0, 1.0));
|
|
let nt = pde.domain.time_resolution.unwrap_or(100);
|
|
|
|
let x_range = pde.domain.bounds[0];
|
|
let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
|
|
let dt = (t_end - t_start) / nt as f64;
|
|
|
|
let c = pde.parameters.wave_speed.unwrap_or(1.0);
|
|
let courant = c * dt / dx;
|
|
|
|
if courant > 1.0 {
|
|
// Warn about stability
|
|
}
|
|
|
|
let mut u = vec![0.0; nx];
|
|
let mut u_prev = vec![0.0; nx];
|
|
let mut u_new = vec![0.0; nx];
|
|
|
|
// Initial conditions
|
|
if let Some(ref ic) = pde.initial_condition {
|
|
for i in 0..nx {
|
|
let x = x_range.0 + i as f64 * dx;
|
|
u[i] = self.evaluate_initial(ic, x);
|
|
u_prev[i] = u[i]; // Zero initial velocity assumed
|
|
}
|
|
}
|
|
|
|
// Time stepping
|
|
for _ in 0..nt {
|
|
for i in 1..nx - 1 {
|
|
let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
|
|
u_new[i] = 2.0 * u[i] - u_prev[i] + dt * dt * c * c * laplacian;
|
|
}
|
|
|
|
u_new[0] = 0.0;
|
|
u_new[nx - 1] = 0.0;
|
|
|
|
u_prev.copy_from_slice(&u);
|
|
u.copy_from_slice(&u_new);
|
|
}
|
|
|
|
let coordinates: Vec<Vec<f64>> = (0..nx).map(|i| vec![x_range.0 + i as f64 * dx]).collect();
|
|
|
|
Ok(SolutionField {
|
|
coordinates,
|
|
values: u,
|
|
name: "u".to_string(),
|
|
time: Some(t_end),
|
|
})
|
|
}
|
|
|
|
/// Solve Burgers' equation using Lax-Wendroff scheme.
|
|
fn solve_burgers(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
|
|
let nx = pde.domain.resolution[0];
|
|
let (t_start, t_end) = pde.domain.time_bounds.unwrap_or((0.0, 1.0));
|
|
let nt = pde.domain.time_resolution.unwrap_or(100);
|
|
|
|
let x_range = pde.domain.bounds[0];
|
|
let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
|
|
let dt = (t_end - t_start) / nt as f64;
|
|
|
|
let nu = pde.parameters.diffusion.unwrap_or(0.01);
|
|
|
|
let mut u = vec![0.0; nx];
|
|
let mut u_new = vec![0.0; nx];
|
|
|
|
// Initial condition
|
|
if let Some(ref ic) = pde.initial_condition {
|
|
for i in 0..nx {
|
|
let x = x_range.0 + i as f64 * dx;
|
|
u[i] = self.evaluate_initial(ic, x);
|
|
}
|
|
}
|
|
|
|
// Time stepping
|
|
for _ in 0..nt {
|
|
u_new.copy_from_slice(&u);
|
|
|
|
for i in 1..nx - 1 {
|
|
let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
|
|
let du_dx = (u[i + 1] - u[i - 1]) / (2.0 * dx);
|
|
let convection = u[i] * du_dx;
|
|
|
|
u_new[i] = u[i] + dt * (nu * laplacian - convection);
|
|
}
|
|
|
|
// Periodic boundary conditions
|
|
u_new[0] = u_new[nx - 2];
|
|
u_new[nx - 1] = u_new[1];
|
|
|
|
u.copy_from_slice(&u_new);
|
|
}
|
|
|
|
let coordinates: Vec<Vec<f64>> = (0..nx).map(|i| vec![x_range.0 + i as f64 * dx]).collect();
|
|
|
|
Ok(SolutionField {
|
|
coordinates,
|
|
values: u,
|
|
name: "u".to_string(),
|
|
time: Some(t_end),
|
|
})
|
|
}
|
|
|
|
/// Apply boundary conditions.
|
|
fn apply_boundary_conditions(&self, u: &mut [f64], pde: &PDEDefinition, nx: usize, ny: usize) {
|
|
for bc in &pde.boundary_conditions {
|
|
match bc.bc_type {
|
|
BCType::Dirichlet => {
|
|
let value: f64 = bc.value.parse().unwrap_or(0.0);
|
|
|
|
match bc.boundary.as_str() {
|
|
"left" => {
|
|
for j in 0..ny {
|
|
u[j * nx] = value;
|
|
}
|
|
}
|
|
"right" => {
|
|
for j in 0..ny {
|
|
u[j * nx + nx - 1] = value;
|
|
}
|
|
}
|
|
"bottom" => {
|
|
for i in 0..nx {
|
|
u[i] = value;
|
|
}
|
|
}
|
|
"top" => {
|
|
for i in 0..nx {
|
|
u[(ny - 1) * nx + i] = value;
|
|
}
|
|
}
|
|
"all" | _ => {
|
|
// All boundaries
|
|
for i in 0..nx {
|
|
u[i] = value;
|
|
u[(ny - 1) * nx + i] = value;
|
|
}
|
|
for j in 0..ny {
|
|
u[j * nx] = value;
|
|
u[j * nx + nx - 1] = value;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
BCType::Periodic => {
|
|
// Handled in time stepping
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Evaluate source term.
|
|
fn evaluate_source(&self, pde: &PDEDefinition, x: f64, y: f64) -> f64 {
|
|
if let Some(ref source) = pde.parameters.source_term {
|
|
// Simple pattern matching for common sources
|
|
if source.contains("sin") {
|
|
(std::f64::consts::PI * x).sin() * (std::f64::consts::PI * y).sin()
|
|
} else {
|
|
source.parse().unwrap_or(0.0)
|
|
}
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
|
|
/// Evaluate initial condition.
|
|
fn evaluate_initial(&self, ic: &neuralop_studio_shared::InitialCondition, x: f64) -> f64 {
|
|
if ic.function.contains("sin") {
|
|
(std::f64::consts::PI * x).sin()
|
|
} else if ic.function.contains("gauss") {
|
|
(-((x - 0.5) * (x - 0.5)) / 0.01).exp()
|
|
} else {
|
|
ic.function.parse().unwrap_or(0.0)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compare neural operator solution with classical solution.
|
|
#[must_use]
|
|
pub fn compare_solutions(neural_solution: &[f64], classical_solution: &[f64]) -> (f64, f64, f64) {
|
|
if neural_solution.is_empty() || classical_solution.is_empty() {
|
|
return (0.0, 0.0, 0.0);
|
|
}
|
|
|
|
let min_len = neural_solution.len().min(classical_solution.len());
|
|
|
|
let mse: f64 = neural_solution[..min_len]
|
|
.iter()
|
|
.zip(classical_solution[..min_len].iter())
|
|
.map(|(n, c)| (n - c).powi(2))
|
|
.sum::<f64>()
|
|
/ min_len as f64;
|
|
|
|
let classical_norm: f64 = classical_solution[..min_len]
|
|
.iter()
|
|
.map(|c| c.powi(2))
|
|
.sum::<f64>()
|
|
.sqrt();
|
|
|
|
let error_norm: f64 = neural_solution[..min_len]
|
|
.iter()
|
|
.zip(classical_solution[..min_len].iter())
|
|
.map(|(n, c)| (n - c).powi(2))
|
|
.sum::<f64>()
|
|
.sqrt();
|
|
|
|
let relative_l2 = if classical_norm > 1e-10 {
|
|
error_norm / classical_norm
|
|
} else {
|
|
error_norm
|
|
};
|
|
|
|
let max_error: f64 = neural_solution[..min_len]
|
|
.iter()
|
|
.zip(classical_solution[..min_len].iter())
|
|
.map(|(n, c)| (n - c).abs())
|
|
.fold(0.0, f64::max);
|
|
|
|
(mse, relative_l2, max_error)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use neuralop_studio_shared::{
|
|
sample_burgers_problem, sample_heat_problem, sample_poisson_problem,
|
|
};
|
|
|
|
#[test]
|
|
fn test_solver_creation() {
|
|
let solver = ClassicalSolver::new();
|
|
assert_eq!(solver.max_iterations, 10000);
|
|
}
|
|
|
|
#[test]
|
|
fn test_solve_poisson() {
|
|
let solver = ClassicalSolver::new().with_max_iterations(100);
|
|
let pde = sample_poisson_problem();
|
|
|
|
let result = solver.solve(&pde);
|
|
assert!(result.is_ok());
|
|
|
|
let solution = result.unwrap();
|
|
assert!(!solution.values.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_solve_heat() {
|
|
let solver = ClassicalSolver::new();
|
|
let pde = sample_heat_problem();
|
|
|
|
let result = solver.solve(&pde);
|
|
assert!(result.is_ok());
|
|
|
|
let solution = result.unwrap();
|
|
assert!(solution.time.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_solve_burgers() {
|
|
let solver = ClassicalSolver::new();
|
|
let pde = sample_burgers_problem();
|
|
|
|
let result = solver.solve(&pde);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_compare_solutions() {
|
|
let neural = vec![0.0, 0.1, 0.2, 0.1, 0.0];
|
|
let classical = vec![0.0, 0.1, 0.2, 0.1, 0.0];
|
|
|
|
let (mse, rel_l2, max_err) = compare_solutions(&neural, &classical);
|
|
assert!(mse < 1e-10);
|
|
assert!(rel_l2 < 1e-10);
|
|
assert!(max_err < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compare_solutions_with_error() {
|
|
let neural = vec![0.0, 0.11, 0.2, 0.1, 0.0]; // Small error
|
|
let classical = vec![0.0, 0.1, 0.2, 0.1, 0.0];
|
|
|
|
let (mse, _rel_l2, max_err) = compare_solutions(&neural, &classical);
|
|
assert!(mse > 0.0);
|
|
assert!((max_err - 0.01).abs() < 1e-10);
|
|
}
|
|
}
|