Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,514 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user