125 lines
3.2 KiB
Markdown
125 lines
3.2 KiB
Markdown
# 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
|
||
|
||
```rust
|
||
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
|
||
|
||
```bash
|
||
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
|