rtx-hemodynamics

GPU-accelerated Physics-Informed Neural Network for hemodynamics simulation.

Overview

This crate provides the core PINN implementation for solving 2D incompressible Navier-Stokes equations in arterial hemodynamics simulation. It solves the inverse problem of determining pressure fields from velocity measurements (simulating data from 4D Flow MRI or ultrasound).

Modules

Module	Description
`navier_stokes`	2D incompressible Navier-Stokes residual computation
`network`	LFFN-MLP (Learnable Fourier Feature Network) architecture
`vessel`	Parametric vessel geometry using SDFs
`boundary`	Inlet/outlet/wall boundary condition handling
`wss`	Wall Shear Stress computation
`inference`	CUDA-optimized inference mode
`training`	PINN training loop with loss weighting
`config`	Configuration structures

Physics

Governing Equations

Momentum (2D incompressible Navier-Stokes):

ρ(∂u/∂t + u∂u/∂x + v∂u/∂y) = -∂p/∂x + μ(∂²u/∂x² + ∂²u/∂y²)
ρ(∂v/∂t + u∂v/∂x + v∂v/∂y) = -∂p/∂y + μ(∂²v/∂x² + ∂²v/∂y²)

Continuity:

∂u/∂x + ∂v/∂y = 0

Loss Function

L_total = λ_data * L_data
        + λ_momentum * L_momentum
        + λ_continuity * L_continuity
        + λ_boundary * L_boundary

Derived Quantities

Wall Shear Stress: τ_w = μ(∂u_t/∂n)|_wall
Velocity magnitude: |u| = √(u² + v²)
Pressure gradient: ∇p

Usage

use rtx_hemodynamics::{NavierStokesResidual, VesselPinn, PinnConfig};
use rtx_hemodynamics_shared::geometry::VesselGeometry;
use rtx_hemodynamics_shared::physics::FluidProperties;

// Create vessel geometry
let vessel = VesselGeometry::straight(0.1, 0.005).unwrap();

// Create PINN model
let config = PinnConfig::default();
let model = VesselPinn::new(config, &vessel).unwrap();

// Train on synthetic data
model.train(5000).await?;

// Infer pressure field
let points = vessel.sample_interior(1000);
let fields = model.infer(&points)?;

Network Architecture

Input: (x, y, t) coordinates
       ↓
Fourier Feature Encoding
  - Learnable frequency matrix B
  - Output: [sin(2π·xB), cos(2π·xB)]
       ↓
MLP: [2*ff_dim] → 256 → 256 → 256 → 256 → [3]
  - SiLU activation
  - Residual connections
       ↓
Output: (u, v, p) velocity + pressure

Vessel Geometry

Vessels are defined using Signed Distance Functions (SDF):

Straight: Simple cylindrical vessel
Stenosis: Smooth Gaussian narrowing
Aneurysm: Smooth Gaussian bulge
Bifurcation: Y-shaped branching

Tests

cargo test

Test cases include:

Poiseuille flow validation (analytical solution)
SDF correctness for all geometry types
Boundary condition enforcement
WSS computation accuracy

Performance

Optimizations:

CUDA Graph capture for repeated inference
Batched autograd operations
Fused kernel execution
Memory-efficient gradient computation

Dependencies

rtx-tensor: Tensor operations
rtx-autograd: Automatic differentiation
rtx-nn: Neural network layers
rtx-hemodynamics-shared: Shared type definitions

3.2 KiB Raw Blame History Unescape Escape