rustytorch/demos/rtx-hemodynamics

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