431 lines
13 KiB
Rust
431 lines
13 KiB
Rust
//! Physics-informed losses for CFD neural operators.
|
||
|
||
use aeroflow_shared::{FlowConditions, FlowField2D};
|
||
|
||
/// Navier-Stokes residual calculator for physics-informed losses.
|
||
#[derive(Debug)]
|
||
pub struct NavierStokesLoss {
|
||
/// Kinematic viscosity (m²/s).
|
||
nu: f32,
|
||
/// Density (kg/m³).
|
||
rho: f32,
|
||
}
|
||
|
||
impl Default for NavierStokesLoss {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
impl NavierStokesLoss {
|
||
/// Create a new Navier-Stokes loss calculator.
|
||
pub fn new() -> Self {
|
||
Self {
|
||
nu: 1.5e-5, // Air at sea level
|
||
rho: 1.225,
|
||
}
|
||
}
|
||
|
||
/// Create with specific fluid properties.
|
||
pub fn with_properties(nu: f32, rho: f32) -> Self {
|
||
Self { nu, rho }
|
||
}
|
||
|
||
/// Evaluate physics residuals on the flow field.
|
||
pub fn evaluate(&self, flow_field: &FlowField2D, conditions: &FlowConditions) -> f32 {
|
||
let continuity = self.continuity_residual(flow_field);
|
||
let momentum_x = self.momentum_x_residual(flow_field, conditions);
|
||
let momentum_y = self.momentum_y_residual(flow_field, conditions);
|
||
|
||
(continuity.powi(2) + momentum_x.powi(2) + momentum_y.powi(2)).sqrt()
|
||
}
|
||
|
||
/// Compute all residual components.
|
||
pub fn compute_residuals(
|
||
&self,
|
||
flow_field: &FlowField2D,
|
||
conditions: &FlowConditions,
|
||
) -> PhysicsResiduals {
|
||
PhysicsResiduals {
|
||
continuity: self.continuity_residual(flow_field),
|
||
momentum_x: self.momentum_x_residual(flow_field, conditions),
|
||
momentum_y: self.momentum_y_residual(flow_field, conditions),
|
||
boundary: self.boundary_residual(flow_field, conditions),
|
||
}
|
||
}
|
||
|
||
/// Continuity equation residual: ∂u/∂x + ∂v/∂y = 0
|
||
fn continuity_residual(&self, flow_field: &FlowField2D) -> f32 {
|
||
let nx = flow_field.velocity_x.len();
|
||
let ny = flow_field.velocity_x[0].len();
|
||
|
||
if nx < 3 || ny < 3 {
|
||
return 0.0;
|
||
}
|
||
|
||
let dx = (flow_field.bounds.x_max - flow_field.bounds.x_min) / (nx - 1) as f32;
|
||
let dy = (flow_field.bounds.y_max - flow_field.bounds.y_min) / (ny - 1) as f32;
|
||
|
||
let mut total_residual = 0.0;
|
||
let mut count = 0;
|
||
|
||
for i in 1..nx - 1 {
|
||
for j in 1..ny - 1 {
|
||
let du_dx = (flow_field.velocity_x[i + 1][j] - flow_field.velocity_x[i - 1][j])
|
||
/ (2.0 * dx);
|
||
let dv_dy = (flow_field.velocity_y[i][j + 1] - flow_field.velocity_y[i][j - 1])
|
||
/ (2.0 * dy);
|
||
|
||
total_residual += (du_dx + dv_dy).abs();
|
||
count += 1;
|
||
}
|
||
}
|
||
|
||
if count > 0 {
|
||
total_residual / count as f32
|
||
} else {
|
||
0.0
|
||
}
|
||
}
|
||
|
||
/// X-momentum equation residual.
|
||
/// u·∂u/∂x + v·∂u/∂y = -1/ρ·∂p/∂x + ν·∇²u
|
||
fn momentum_x_residual(&self, flow_field: &FlowField2D, _conditions: &FlowConditions) -> f32 {
|
||
let nx = flow_field.velocity_x.len();
|
||
let ny = flow_field.velocity_x[0].len();
|
||
|
||
if nx < 3 || ny < 3 {
|
||
return 0.0;
|
||
}
|
||
|
||
let dx = (flow_field.bounds.x_max - flow_field.bounds.x_min) / (nx - 1) as f32;
|
||
let dy = (flow_field.bounds.y_max - flow_field.bounds.y_min) / (ny - 1) as f32;
|
||
|
||
let mut total_residual = 0.0;
|
||
let mut count = 0;
|
||
|
||
for i in 2..nx - 2 {
|
||
for j in 2..ny - 2 {
|
||
let u = flow_field.velocity_x[i][j];
|
||
let v = flow_field.velocity_y[i][j];
|
||
|
||
// First derivatives
|
||
let du_dx = (flow_field.velocity_x[i + 1][j] - flow_field.velocity_x[i - 1][j])
|
||
/ (2.0 * dx);
|
||
let du_dy = (flow_field.velocity_x[i][j + 1] - flow_field.velocity_x[i][j - 1])
|
||
/ (2.0 * dy);
|
||
|
||
// Second derivatives (Laplacian)
|
||
let d2u_dx2 = (flow_field.velocity_x[i + 1][j] - 2.0 * flow_field.velocity_x[i][j]
|
||
+ flow_field.velocity_x[i - 1][j])
|
||
/ (dx * dx);
|
||
let d2u_dy2 = (flow_field.velocity_x[i][j + 1] - 2.0 * flow_field.velocity_x[i][j]
|
||
+ flow_field.velocity_x[i][j - 1])
|
||
/ (dy * dy);
|
||
|
||
// Pressure gradient
|
||
let dp_dx =
|
||
(flow_field.pressure[i + 1][j] - flow_field.pressure[i - 1][j]) / (2.0 * dx);
|
||
|
||
// Momentum residual
|
||
let convection = u * du_dx + v * du_dy;
|
||
let pressure_term = dp_dx / self.rho;
|
||
let diffusion = self.nu * (d2u_dx2 + d2u_dy2);
|
||
|
||
let residual = convection + pressure_term - diffusion;
|
||
total_residual += residual.abs();
|
||
count += 1;
|
||
}
|
||
}
|
||
|
||
if count > 0 {
|
||
total_residual / count as f32
|
||
} else {
|
||
0.0
|
||
}
|
||
}
|
||
|
||
/// Y-momentum equation residual.
|
||
fn momentum_y_residual(&self, flow_field: &FlowField2D, _conditions: &FlowConditions) -> f32 {
|
||
let nx = flow_field.velocity_y.len();
|
||
let ny = flow_field.velocity_y[0].len();
|
||
|
||
if nx < 3 || ny < 3 {
|
||
return 0.0;
|
||
}
|
||
|
||
let dx = (flow_field.bounds.x_max - flow_field.bounds.x_min) / (nx - 1) as f32;
|
||
let dy = (flow_field.bounds.y_max - flow_field.bounds.y_min) / (ny - 1) as f32;
|
||
|
||
let mut total_residual = 0.0;
|
||
let mut count = 0;
|
||
|
||
for i in 2..nx - 2 {
|
||
for j in 2..ny - 2 {
|
||
let u = flow_field.velocity_x[i][j];
|
||
let v = flow_field.velocity_y[i][j];
|
||
|
||
// First derivatives
|
||
let dv_dx = (flow_field.velocity_y[i + 1][j] - flow_field.velocity_y[i - 1][j])
|
||
/ (2.0 * dx);
|
||
let dv_dy = (flow_field.velocity_y[i][j + 1] - flow_field.velocity_y[i][j - 1])
|
||
/ (2.0 * dy);
|
||
|
||
// Second derivatives
|
||
let d2v_dx2 = (flow_field.velocity_y[i + 1][j] - 2.0 * flow_field.velocity_y[i][j]
|
||
+ flow_field.velocity_y[i - 1][j])
|
||
/ (dx * dx);
|
||
let d2v_dy2 = (flow_field.velocity_y[i][j + 1] - 2.0 * flow_field.velocity_y[i][j]
|
||
+ flow_field.velocity_y[i][j - 1])
|
||
/ (dy * dy);
|
||
|
||
// Pressure gradient
|
||
let dp_dy =
|
||
(flow_field.pressure[i][j + 1] - flow_field.pressure[i][j - 1]) / (2.0 * dy);
|
||
|
||
// Momentum residual
|
||
let convection = u * dv_dx + v * dv_dy;
|
||
let pressure_term = dp_dy / self.rho;
|
||
let diffusion = self.nu * (d2v_dx2 + d2v_dy2);
|
||
|
||
let residual = convection + pressure_term - diffusion;
|
||
total_residual += residual.abs();
|
||
count += 1;
|
||
}
|
||
}
|
||
|
||
if count > 0 {
|
||
total_residual / count as f32
|
||
} else {
|
||
0.0
|
||
}
|
||
}
|
||
|
||
/// Boundary condition residual.
|
||
fn boundary_residual(&self, flow_field: &FlowField2D, conditions: &FlowConditions) -> f32 {
|
||
let nx = flow_field.velocity_x.len();
|
||
let ny = flow_field.velocity_x[0].len();
|
||
|
||
let mut residual = 0.0;
|
||
let mut count = 0;
|
||
|
||
// Inlet boundary (left edge)
|
||
for j in 0..ny {
|
||
let u_inlet = conditions.velocity * conditions.angle_of_attack.to_radians().cos();
|
||
let v_inlet = conditions.velocity * conditions.angle_of_attack.to_radians().sin();
|
||
|
||
residual += (flow_field.velocity_x[0][j] - u_inlet).abs();
|
||
residual += (flow_field.velocity_y[0][j] - v_inlet).abs();
|
||
count += 2;
|
||
}
|
||
|
||
// Farfield boundaries (top and bottom)
|
||
for i in 0..nx {
|
||
// Top
|
||
residual += (flow_field.velocity_x[i][ny - 1] - conditions.velocity).abs();
|
||
// Bottom
|
||
residual += (flow_field.velocity_x[i][0] - conditions.velocity).abs();
|
||
count += 2;
|
||
}
|
||
|
||
if count > 0 {
|
||
residual / count as f32
|
||
} else {
|
||
0.0
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Physics residuals for monitoring.
|
||
#[derive(Debug, Clone, Copy, Default)]
|
||
pub struct PhysicsResiduals {
|
||
/// Continuity equation residual.
|
||
pub continuity: f32,
|
||
/// X-momentum residual.
|
||
pub momentum_x: f32,
|
||
/// Y-momentum residual.
|
||
pub momentum_y: f32,
|
||
/// Boundary condition residual.
|
||
pub boundary: f32,
|
||
}
|
||
|
||
impl PhysicsResiduals {
|
||
/// Total L2 norm of residuals.
|
||
pub fn total(&self) -> f32 {
|
||
(self.continuity.powi(2)
|
||
+ self.momentum_x.powi(2)
|
||
+ self.momentum_y.powi(2)
|
||
+ self.boundary.powi(2))
|
||
.sqrt()
|
||
}
|
||
}
|
||
|
||
/// Euler equations residual (inviscid).
|
||
#[derive(Debug)]
|
||
pub struct EulerLoss {
|
||
/// Ratio of specific heats.
|
||
#[allow(dead_code)]
|
||
gamma: f32,
|
||
}
|
||
|
||
impl Default for EulerLoss {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
impl EulerLoss {
|
||
/// Create a new Euler loss calculator.
|
||
pub fn new() -> Self {
|
||
Self { gamma: 1.4 }
|
||
}
|
||
|
||
/// Evaluate Euler equation residuals.
|
||
pub fn evaluate(&self, _flow_field: &FlowField2D, conditions: &FlowConditions) -> f32 {
|
||
// For compressible flow, would compute:
|
||
// ∂ρ/∂t + ∇·(ρu) = 0 (mass)
|
||
// ∂(ρu)/∂t + ∇·(ρuu + pI) = 0 (momentum)
|
||
// ∂E/∂t + ∇·((E + p)u) = 0 (energy)
|
||
|
||
// Simplified: just check for low Mach incompressible behavior
|
||
if conditions.mach < 0.3 {
|
||
0.0
|
||
} else {
|
||
0.1 * (conditions.mach - 0.3)
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Wall function for near-wall treatment.
|
||
#[derive(Debug)]
|
||
pub struct WallFunction {
|
||
/// Von Karman constant.
|
||
kappa: f32,
|
||
/// Integration constant.
|
||
b: f32,
|
||
}
|
||
|
||
impl Default for WallFunction {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
impl WallFunction {
|
||
/// Create a new wall function.
|
||
pub fn new() -> Self {
|
||
Self {
|
||
kappa: 0.41,
|
||
b: 5.0,
|
||
}
|
||
}
|
||
|
||
/// Law of the wall: u+ = (1/κ)ln(y+) + B
|
||
pub fn u_plus(&self, y_plus: f32) -> f32 {
|
||
if y_plus < 11.6 {
|
||
// Viscous sublayer
|
||
y_plus
|
||
} else {
|
||
// Log layer
|
||
(1.0 / self.kappa) * y_plus.ln() + self.b
|
||
}
|
||
}
|
||
|
||
/// Compute wall shear stress.
|
||
pub fn wall_shear(&self, u_tau: f32, rho: f32) -> f32 {
|
||
rho * u_tau.powi(2)
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use aeroflow_shared::DomainBounds;
|
||
|
||
fn create_test_flow_field() -> FlowField2D {
|
||
let nx = 10;
|
||
let ny = 8;
|
||
|
||
FlowField2D {
|
||
pressure: vec![vec![101325.0; ny]; nx],
|
||
velocity_x: vec![vec![100.0; ny]; nx],
|
||
velocity_y: vec![vec![0.0; ny]; nx],
|
||
velocity_magnitude: vec![vec![100.0; ny]; nx],
|
||
cp: vec![vec![0.0; ny]; nx],
|
||
vorticity: vec![vec![0.0; ny]; nx],
|
||
tke: None,
|
||
grid_x: vec![vec![0.0; ny]; nx],
|
||
grid_y: vec![vec![0.0; ny]; nx],
|
||
bounds: DomainBounds::default(),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_ns_loss_creation() {
|
||
let loss = NavierStokesLoss::new();
|
||
assert!((loss.nu - 1.5e-5).abs() < 1e-7);
|
||
}
|
||
|
||
#[test]
|
||
fn test_continuity_residual() {
|
||
let loss = NavierStokesLoss::new();
|
||
let flow_field = create_test_flow_field();
|
||
let residual = loss.continuity_residual(&flow_field);
|
||
// Uniform flow should have zero continuity residual
|
||
assert!(residual.abs() < 1e-3);
|
||
}
|
||
|
||
#[test]
|
||
fn test_evaluate() {
|
||
let loss = NavierStokesLoss::new();
|
||
let flow_field = create_test_flow_field();
|
||
let conditions = FlowConditions::default();
|
||
let total = loss.evaluate(&flow_field, &conditions);
|
||
assert!(total.is_finite());
|
||
}
|
||
|
||
#[test]
|
||
fn test_physics_residuals() {
|
||
let loss = NavierStokesLoss::new();
|
||
let flow_field = create_test_flow_field();
|
||
let conditions = FlowConditions::default();
|
||
let residuals = loss.compute_residuals(&flow_field, &conditions);
|
||
|
||
assert!(residuals.continuity.is_finite());
|
||
assert!(residuals.momentum_x.is_finite());
|
||
assert!(residuals.momentum_y.is_finite());
|
||
assert!(residuals.total().is_finite());
|
||
}
|
||
|
||
#[test]
|
||
fn test_euler_loss() {
|
||
let loss = EulerLoss::new();
|
||
let flow_field = create_test_flow_field();
|
||
|
||
let low_mach = FlowConditions {
|
||
mach: 0.2,
|
||
..Default::default()
|
||
};
|
||
assert_eq!(loss.evaluate(&flow_field, &low_mach), 0.0);
|
||
|
||
let high_mach = FlowConditions {
|
||
mach: 0.5,
|
||
..Default::default()
|
||
};
|
||
assert!(loss.evaluate(&flow_field, &high_mach) > 0.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_wall_function() {
|
||
let wf = WallFunction::new();
|
||
|
||
// Viscous sublayer: u+ = y+
|
||
assert!((wf.u_plus(5.0) - 5.0).abs() < 0.01);
|
||
|
||
// Log layer: u+ = (1/κ)ln(y+) + B
|
||
let y_plus: f32 = 100.0;
|
||
let expected = (1.0 / 0.41) * y_plus.ln() + 5.0;
|
||
assert!((wf.u_plus(y_plus) - expected).abs() < 0.01);
|
||
}
|
||
}
|