Initial commit

demos/rtx-hemodynamics/src/navier_stokes.rs

@@ -0,0 +1,538 @@
+//! 2D Incompressible Navier-Stokes residual computation
+//!
+//! This module implements the physics-informed loss terms for the 2D
+//! incompressible Navier-Stokes equations:
+//!
+//! **Momentum (x)**: ρ(∂u/∂t + u∂u/∂x + v∂u/∂y) = -∂p/∂x + μ(∂²u/∂x² + ∂²u/∂y²)
+//! **Momentum (y)**: ρ(∂v/∂t + u∂v/∂x + v∂v/∂y) = -∂p/∂y + μ(∂²v/∂x² + ∂²v/∂y²)
+//! **Continuity**: ∂u/∂x + ∂v/∂y = 0
+
+use rtx_hemodynamics_shared::physics::FluidProperties;
+
+/// Navier-Stokes residual computer for 2D incompressible flow
+///
+/// Computes the physics residuals for the momentum and continuity equations.
+#[derive(Debug, Clone)]
+pub struct NavierStokesResidual {
+    /// Fluid density (kg/m³)
+    rho: f64,
+    /// Dynamic viscosity (Pa.s)
+    mu: f64,
+    /// Whether flow is steady-state (no time derivatives)
+    steady_state: bool,
+}
+
+impl NavierStokesResidual {
+    /// Creates a new Navier-Stokes residual computer
+    #[must_use]
+    pub fn new(fluid: &FluidProperties, steady_state: bool) -> Self {
+        Self {
+            rho: fluid.density(),
+            mu: fluid.dynamic_viscosity(),
+            steady_state,
+        }
+    }
+
+    /// Creates residual computer for blood with steady-state assumption
+    #[must_use]
+    pub fn blood_steady() -> Self {
+        Self::new(&FluidProperties::blood(), true)
+    }
+
+    /// Creates residual computer for blood with transient flow
+    #[must_use]
+    pub fn blood_transient() -> Self {
+        Self::new(&FluidProperties::blood(), false)
+    }
+
+    /// Returns the fluid density
+    #[must_use]
+    pub const fn density(&self) -> f64 {
+        self.rho
+    }
+
+    /// Returns the dynamic viscosity
+    #[must_use]
+    pub const fn viscosity(&self) -> f64 {
+        self.mu
+    }
+
+    /// Returns whether using steady-state assumption
+    #[must_use]
+    pub const fn is_steady_state(&self) -> bool {
+        self.steady_state
+    }
+
+    /// Computes the x-momentum residual at a single point
+    ///
+    /// `R_x` = ρ(∂u/∂t + u∂u/∂x + v∂u/∂y) + ∂p/∂x - μ(∂²u/∂x² + ∂²u/∂y²)
+    ///
+    /// For steady-state: `R_x` = ρ(u∂u/∂x + v∂u/∂y) + ∂p/∂x - μ(∂²u/∂x² + ∂²u/∂y²)
+    ///
+    /// # Arguments
+    ///
+    /// * `u` - x-velocity
+    /// * `v` - y-velocity
+    /// * `du_dt` - time derivative of u (ignored if steady-state)
+    /// * `du_dx` - spatial derivative ∂u/∂x
+    /// * `du_dy` - spatial derivative ∂u/∂y
+    /// * `d2u_dx2` - second derivative ∂²u/∂x²
+    /// * `d2u_dy2` - second derivative ∂²u/∂y²
+    /// * `dp_dx` - pressure gradient ∂p/∂x
+    #[must_use]
+    #[allow(clippy::too_many_arguments)]
+    pub fn x_momentum_residual(
+        &self,
+        u: f64,
+        v: f64,
+        du_dt: f64,
+        du_dx: f64,
+        du_dy: f64,
+        d2u_dx2: f64,
+        d2u_dy2: f64,
+        dp_dx: f64,
+    ) -> f64 {
+        let time_term = if self.steady_state { 0.0 } else { du_dt };
+        let convection = u * du_dx + v * du_dy;
+        let diffusion = d2u_dx2 + d2u_dy2;
+
+        self.rho * (time_term + convection) + dp_dx - self.mu * diffusion
+    }
+
+    /// Computes the y-momentum residual at a single point
+    ///
+    /// `R_y` = ρ(∂v/∂t + u∂v/∂x + v∂v/∂y) + ∂p/∂y - μ(∂²v/∂x² + ∂²v/∂y²)
+    #[must_use]
+    #[allow(clippy::too_many_arguments)]
+    pub fn y_momentum_residual(
+        &self,
+        u: f64,
+        v: f64,
+        dv_dt: f64,
+        dv_dx: f64,
+        dv_dy: f64,
+        d2v_dx2: f64,
+        d2v_dy2: f64,
+        dp_dy: f64,
+    ) -> f64 {
+        let time_term = if self.steady_state { 0.0 } else { dv_dt };
+        let convection = u * dv_dx + v * dv_dy;
+        let diffusion = d2v_dx2 + d2v_dy2;
+
+        self.rho * (time_term + convection) + dp_dy - self.mu * diffusion
+    }
+
+    /// Computes the continuity residual at a single point
+    ///
+    /// `R_cont` = ∂u/∂x + ∂v/∂y
+    ///
+    /// For incompressible flow, this should be zero (divergence-free).
+    #[must_use]
+    pub fn continuity_residual(&self, du_dx: f64, dv_dy: f64) -> f64 {
+        du_dx + dv_dy
+    }
+
+    /// Computes all three residuals at a single point
+    ///
+    /// Returns (`R_x`, `R_y`, `R_cont`)
+    #[must_use]
+    #[allow(clippy::too_many_arguments)]
+    pub fn all_residuals(
+        &self,
+        u: f64,
+        v: f64,
+        du_dt: f64,
+        dv_dt: f64,
+        du_dx: f64,
+        du_dy: f64,
+        dv_dx: f64,
+        dv_dy: f64,
+        d2u_dx2: f64,
+        d2u_dy2: f64,
+        d2v_dx2: f64,
+        d2v_dy2: f64,
+        dp_dx: f64,
+        dp_dy: f64,
+    ) -> (f64, f64, f64) {
+        let r_x = self.x_momentum_residual(u, v, du_dt, du_dx, du_dy, d2u_dx2, d2u_dy2, dp_dx);
+        let r_y = self.y_momentum_residual(u, v, dv_dt, dv_dx, dv_dy, d2v_dx2, d2v_dy2, dp_dy);
+        let r_cont = self.continuity_residual(du_dx, dv_dy);
+
+        (r_x, r_y, r_cont)
+    }
+
+    /// Computes the total physics loss (MSE of all residuals)
+    ///
+    /// `L_physics` = `mean(R_x²` + `R_y²` + `R_cont²`)
+    #[must_use]
+    pub fn physics_loss(&self, residuals: &[(f64, f64, f64)]) -> f64 {
+        if residuals.is_empty() {
+            return 0.0;
+        }
+
+        let sum: f64 = residuals
+            .iter()
+            .map(|(rx, ry, rc)| rx.powi(2) + ry.powi(2) + rc.powi(2))
+            .sum();
+
+        sum / residuals.len() as f64
+    }
+
+    /// Computes individual loss components for monitoring
+    ///
+    /// Returns (`momentum_x_loss`, `momentum_y_loss`, `continuity_loss`)
+    #[must_use]
+    pub fn loss_components(&self, residuals: &[(f64, f64, f64)]) -> (f64, f64, f64) {
+        if residuals.is_empty() {
+            return (0.0, 0.0, 0.0);
+        }
+
+        let n = residuals.len() as f64;
+        let (sum_x, sum_y, sum_c) = residuals.iter().fold((0.0, 0.0, 0.0), |acc, (rx, ry, rc)| {
+            (acc.0 + rx.powi(2), acc.1 + ry.powi(2), acc.2 + rc.powi(2))
+        });
+
+        (sum_x / n, sum_y / n, sum_c / n)
+    }
+}
+
+/// Analytical solution for 2D planar Poiseuille flow (steady channel flow)
+///
+/// This represents flow between two parallel plates (2D channel flow),
+/// which exactly satisfies the 2D Cartesian Navier-Stokes equations.
+///
+/// For a channel of half-height h:
+/// - Velocity profile: u(y) = (1/2μ)(-dp/dx)(h² - y²)
+/// - Maximum velocity at centerline (y=0)
+/// - Zero velocity at walls (y=±h)
+#[derive(Debug, Clone)]
+pub struct PoiseuilleFlow {
+    /// Channel half-height (m) - distance from centerline to wall
+    half_height: f64,
+    /// Pressure gradient (Pa/m)
+    dp_dx: f64,
+    /// Dynamic viscosity (Pa.s)
+    mu: f64,
+}
+
+impl PoiseuilleFlow {
+    /// Creates a new 2D planar Poiseuille flow solution
+    ///
+    /// # Arguments
+    ///
+    /// * `half_height` - Channel half-height (distance from centerline to wall) in meters
+    /// * `pressure_drop` - Total pressure drop over channel length
+    /// * `length` - Channel length in meters
+    /// * `mu` - Dynamic viscosity in Pa.s
+    #[must_use]
+    pub fn new(half_height: f64, pressure_drop: f64, length: f64, mu: f64) -> Self {
+        Self {
+            half_height,
+            dp_dx: -pressure_drop / length, // Negative because pressure decreases
+            mu,
+        }
+    }
+
+    /// Creates Poiseuille flow with blood properties
+    #[must_use]
+    pub fn blood(half_height: f64, pressure_drop: f64, length: f64) -> Self {
+        Self::new(
+            half_height,
+            pressure_drop,
+            length,
+            FluidProperties::blood().dynamic_viscosity(),
+        )
+    }
+
+    /// Returns the channel half-height (also called radius for compatibility)
+    #[must_use]
+    pub const fn radius(&self) -> f64 {
+        self.half_height
+    }
+
+    /// Returns the pressure gradient
+    #[must_use]
+    pub const fn pressure_gradient(&self) -> f64 {
+        self.dp_dx
+    }
+
+    /// Computes the analytical velocity at position y from centerline
+    ///
+    /// u(y) = (1/2μ)(-dp/dx)(h² - y²)
+    ///
+    /// The velocity profile is parabolic with maximum at centerline.
+    #[must_use]
+    pub fn velocity(&self, y: f64) -> f64 {
+        if y.abs() > self.half_height {
+            return 0.0; // Outside channel
+        }
+        (-self.dp_dx / (2.0 * self.mu)) * (self.half_height.powi(2) - y.powi(2))
+    }
+
+    /// Returns the maximum (centerline) velocity
+    #[must_use]
+    pub fn max_velocity(&self) -> f64 {
+        self.velocity(0.0)
+    }
+
+    /// Returns the average velocity
+    ///
+    /// `u_avg` = (2/3) * `u_max` for 2D channel flow
+    #[must_use]
+    pub fn average_velocity(&self) -> f64 {
+        self.max_velocity() * 2.0 / 3.0
+    }
+
+    /// Computes the volumetric flow rate per unit depth
+    ///
+    /// Q = (2h³ / 3μ) * (-dp/dx)
+    #[must_use]
+    pub fn flow_rate(&self) -> f64 {
+        2.0 * self.half_height.powi(3) * (-self.dp_dx) / (3.0 * self.mu)
+    }
+
+    /// Computes the wall shear stress
+    ///
+    /// `τ_w` = μ|du/dy|_wall = h * (-dp/dx)
+    #[must_use]
+    pub fn wall_shear_stress(&self) -> f64 {
+        self.half_height * (-self.dp_dx)
+    }
+
+    /// Computes the velocity gradient at position y
+    ///
+    /// du/dy = (1/μ)(dp/dx) * y
+    #[must_use]
+    pub fn velocity_gradient(&self, y: f64) -> f64 {
+        (self.dp_dx / self.mu) * y
+    }
+
+    /// Computes the pressure at axial position x (relative to inlet)
+    ///
+    /// p(x) = `p_inlet` + dp/dx * x
+    #[must_use]
+    pub fn pressure(&self, x: f64, p_inlet: f64) -> f64 {
+        p_inlet + self.dp_dx * x
+    }
+
+    /// Validates that Navier-Stokes residuals are zero for this analytical solution
+    ///
+    /// For 2D planar Poiseuille flow:
+    /// - u = u(y), v = 0
+    /// - All time derivatives are zero
+    /// - du/dx = 0 (fully developed)
+    /// - dv/dx = dv/dy = 0
+    /// - d²u/dx² = 0
+    /// - d²u/dy² = (dp/dx) / μ
+    #[must_use]
+    pub fn validate_residuals(&self, ns: &NavierStokesResidual, y: f64) -> (f64, f64, f64) {
+        let u = self.velocity(y);
+        let v = 0.0;
+
+        // Time derivatives (steady-state)
+        let du_dt = 0.0;
+        let dv_dt = 0.0;
+
+        // Spatial derivatives
+        let du_dx = 0.0; // Fully developed
+        let du_dy = self.velocity_gradient(y);
+        let dv_dx = 0.0;
+        let dv_dy = 0.0;
+
+        // Second derivatives
+        let d2u_dx2 = 0.0;
+        // d²u/dy² = (dp/dx) / μ for 2D planar flow
+        let d2u_dy2 = self.dp_dx / self.mu;
+        let d2v_dx2 = 0.0;
+        let d2v_dy2 = 0.0;
+
+        // Pressure gradients
+        let dp_dx = self.dp_dx;
+        let dp_dy = 0.0;
+
+        ns.all_residuals(
+            u, v, du_dt, dv_dt, du_dx, du_dy, dv_dx, dv_dy, d2u_dx2, d2u_dy2, d2v_dx2, d2v_dy2,
+            dp_dx, dp_dy,
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_ns_residual_creation() {
+        let ns = NavierStokesResidual::blood_steady();
+        assert!((ns.density() - 1060.0).abs() < 1.0);
+        assert!((ns.viscosity() - 0.0035).abs() < 0.001);
+        assert!(ns.is_steady_state());
+    }
+
+    #[test]
+    fn test_continuity_residual() {
+        let ns = NavierStokesResidual::blood_steady();
+
+        // Divergence-free flow should have zero residual
+        let residual = ns.continuity_residual(0.5, -0.5);
+        assert!(residual.abs() < f64::EPSILON);
+
+        // Non-divergence-free should have non-zero residual
+        let residual = ns.continuity_residual(0.5, 0.3);
+        assert!((residual - 0.8).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn test_poiseuille_velocity_profile() {
+        // Create Poiseuille flow with known parameters
+        let radius = 0.005; // 5mm pipe
+        let pressure_drop = 100.0; // 100 Pa drop
+        let length = 0.1; // 10cm length
+        let mu = 0.0035; // Blood viscosity
+
+        let flow = PoiseuilleFlow::new(radius, pressure_drop, length, mu);
+
+        // Maximum velocity at centerline
+        let u_max = flow.max_velocity();
+        assert!(u_max > 0.0);
+
+        // Zero velocity at wall
+        let u_wall = flow.velocity(radius);
+        assert!(u_wall.abs() < 1e-10);
+
+        // Velocity at r=0 should be maximum
+        let u_center = flow.velocity(0.0);
+        assert!((u_center - u_max).abs() < f64::EPSILON);
+
+        // Parabolic profile: u(R/2) = 3/4 * u_max
+        let u_half = flow.velocity(radius / 2.0);
+        assert!((u_half - 0.75 * u_max).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_poiseuille_wall_shear_stress() {
+        let half_height = 0.005;
+        let pressure_drop = 100.0;
+        let length = 0.1;
+        let mu = 0.0035;
+
+        let flow = PoiseuilleFlow::new(half_height, pressure_drop, length, mu);
+        let wss = flow.wall_shear_stress();
+
+        // For 2D planar flow: τ_w = h * |dp/dx|
+        let expected = half_height * (pressure_drop / length);
+        assert!((wss - expected).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_poiseuille_validates_ns() {
+        // The Poiseuille solution should satisfy Navier-Stokes exactly
+        let radius = 0.005;
+        let pressure_drop = 100.0;
+        let length = 0.1;
+        let mu = 0.0035;
+
+        let flow = PoiseuilleFlow::new(radius, pressure_drop, length, mu);
+        let ns = NavierStokesResidual::new(&FluidProperties::new(1060.0, mu).unwrap(), true);
+
+        // Test at multiple radial positions
+        for i in 0..10 {
+            let r = radius * (i as f64) / 10.0;
+            let (rx, ry, rc) = flow.validate_residuals(&ns, r);
+
+            // All residuals should be near zero (within numerical precision)
+            assert!(
+                rx.abs() < 1e-6,
+                "X-momentum residual too large at r={}: {}",
+                r,
+                rx
+            );
+            assert!(
+                ry.abs() < 1e-10,
+                "Y-momentum residual too large at r={}: {}",
+                r,
+                ry
+            );
+            assert!(
+                rc.abs() < 1e-10,
+                "Continuity residual too large at r={}: {}",
+                r,
+                rc
+            );
+        }
+    }
+
+    #[test]
+    fn test_physics_loss_computation() {
+        let ns = NavierStokesResidual::blood_steady();
+
+        // Perfect solution should have zero loss
+        let perfect_residuals = vec![(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)];
+        let loss = ns.physics_loss(&perfect_residuals);
+        assert!(loss.abs() < f64::EPSILON);
+
+        // Non-zero residuals should give positive loss
+        let residuals = vec![(1.0, 0.0, 0.0), (0.0, 1.0, 0.0)];
+        let loss = ns.physics_loss(&residuals);
+        assert!((loss - 1.0).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn test_loss_components() {
+        let ns = NavierStokesResidual::blood_steady();
+
+        let residuals = vec![(1.0, 2.0, 3.0), (1.0, 2.0, 3.0)];
+
+        let (lx, ly, lc) = ns.loss_components(&residuals);
+        assert!((lx - 1.0).abs() < f64::EPSILON);
+        assert!((ly - 4.0).abs() < f64::EPSILON);
+        assert!((lc - 9.0).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn test_poiseuille_flow_rate() {
+        let half_height = 0.005;
+        let pressure_drop = 100.0;
+        let length = 0.1;
+        let mu = 0.0035;
+
+        let flow = PoiseuilleFlow::new(half_height, pressure_drop, length, mu);
+
+        // 2D planar flow rate per unit depth: Q = (2h³ / 3μ) * |dp/dx|
+        let dp_dx = pressure_drop / length;
+        let expected_q = 2.0 * half_height.powi(3) * dp_dx / (3.0 * mu);
+        let actual_q = flow.flow_rate();
+
+        assert!((actual_q - expected_q).abs() < 1e-15);
+    }
+
+    #[test]
+    fn test_steady_vs_transient() {
+        let fluid = FluidProperties::blood();
+        let steady = NavierStokesResidual::new(&fluid, true);
+        let transient = NavierStokesResidual::new(&fluid, false);
+
+        // With non-zero du/dt, steady should ignore it
+        let r_steady = steady.x_momentum_residual(
+            1.0, 0.0,  // u, v
+            10.0, // du_dt (should be ignored)
+            0.0, 0.0, // du_dx, du_dy
+            0.0, 0.0, // d2u_dx2, d2u_dy2
+            0.0, // dp_dx
+        );
+
+        let r_transient = transient.x_momentum_residual(
+            1.0, 0.0,  // u, v
+            10.0, // du_dt (should be included)
+            0.0, 0.0, // du_dx, du_dy
+            0.0, 0.0, // d2u_dx2, d2u_dy2
+            0.0, // dp_dx
+        );
+
+        // Steady should have zero (no time term)
+        assert!(r_steady.abs() < f64::EPSILON);
+
+        // Transient should have ρ * du/dt = 1060 * 10 = 10600
+        assert!((r_transient - 1060.0 * 10.0).abs() < 0.1);
+    }
+}