428 lines
12 KiB
Rust
428 lines
12 KiB
Rust
//! Optimized inference engine for hemodynamics PINN
|
|
//!
|
|
//! Provides efficient batched inference with optional GPU acceleration.
|
|
|
|
use crate::network::VesselPinn;
|
|
use crate::vessel::VesselSdf;
|
|
use crate::wss::WssComputer;
|
|
use rtx_hemodynamics_shared::fields::{FieldResponse, PressureField, VelocityField, WssField};
|
|
use rtx_hemodynamics_shared::geometry::Point2D;
|
|
|
|
/// Inference engine for field evaluation
|
|
#[derive(Debug)]
|
|
pub struct InferenceEngine {
|
|
/// The trained PINN model
|
|
model: VesselPinn,
|
|
/// Vessel geometry
|
|
vessel: VesselSdf,
|
|
/// WSS computer
|
|
wss_computer: WssComputer,
|
|
/// Inference batch size
|
|
batch_size: usize,
|
|
}
|
|
|
|
impl InferenceEngine {
|
|
/// Creates a new inference engine
|
|
#[must_use]
|
|
pub fn new(model: VesselPinn, vessel: VesselSdf) -> Self {
|
|
Self {
|
|
model,
|
|
vessel,
|
|
wss_computer: WssComputer::blood(),
|
|
batch_size: 1000,
|
|
}
|
|
}
|
|
|
|
/// Sets the batch size for inference
|
|
#[must_use]
|
|
pub const fn with_batch_size(mut self, size: usize) -> Self {
|
|
self.batch_size = size;
|
|
self
|
|
}
|
|
|
|
/// Returns the model
|
|
#[must_use]
|
|
pub const fn model(&self) -> &VesselPinn {
|
|
&self.model
|
|
}
|
|
|
|
/// Returns the vessel geometry
|
|
#[must_use]
|
|
pub const fn vessel(&self) -> &VesselSdf {
|
|
&self.vessel
|
|
}
|
|
|
|
/// Evaluates velocity and pressure at given points
|
|
#[must_use]
|
|
pub fn evaluate(&self, points: &[Point2D], time: f64) -> Vec<(f64, f64, f64)> {
|
|
self.model.forward_batch(points, time)
|
|
}
|
|
|
|
/// Evaluates and returns full field response
|
|
#[must_use]
|
|
pub fn evaluate_fields(&self, points: &[Point2D], time: f64) -> FieldResponse {
|
|
let predictions = self.evaluate(points, time);
|
|
|
|
let u: Vec<f64> = predictions.iter().map(|(u, _, _)| *u).collect();
|
|
let v: Vec<f64> = predictions.iter().map(|(_, v, _)| *v).collect();
|
|
let p: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();
|
|
|
|
let velocity = VelocityField::new(points.to_vec(), u, v).unwrap();
|
|
let pressure = PressureField::new(points.to_vec(), p).unwrap();
|
|
|
|
FieldResponse::new(velocity, pressure, None)
|
|
}
|
|
|
|
/// Evaluates fields including WSS at wall points
|
|
#[must_use]
|
|
pub fn evaluate_with_wss(
|
|
&self,
|
|
interior_points: &[Point2D],
|
|
wall_points: &[Point2D],
|
|
time: f64,
|
|
) -> FieldResponse {
|
|
// Evaluate interior fields
|
|
let interior_preds = self.evaluate(interior_points, time);
|
|
|
|
let u: Vec<f64> = interior_preds.iter().map(|(u, _, _)| *u).collect();
|
|
let v: Vec<f64> = interior_preds.iter().map(|(_, v, _)| *v).collect();
|
|
let p: Vec<f64> = interior_preds.iter().map(|(_, _, p)| *p).collect();
|
|
|
|
let velocity = VelocityField::new(interior_points.to_vec(), u, v).unwrap();
|
|
let pressure = PressureField::new(interior_points.to_vec(), p).unwrap();
|
|
|
|
// Compute WSS at wall points
|
|
let wss_values = self.compute_wss(wall_points, time);
|
|
let wss = WssField::new(wall_points.to_vec(), wss_values).unwrap();
|
|
|
|
FieldResponse::new(velocity, pressure, Some(wss))
|
|
}
|
|
|
|
/// Computes WSS at wall points
|
|
#[must_use]
|
|
pub fn compute_wss(&self, wall_points: &[Point2D], time: f64) -> Vec<f64> {
|
|
let velocity_fn = |p: &Point2D| {
|
|
let (u, v, _) = self.model.forward(p.x, p.y, time);
|
|
(u, v)
|
|
};
|
|
|
|
self.wss_computer
|
|
.compute_wss_batch(wall_points, &self.vessel, velocity_fn)
|
|
}
|
|
|
|
/// Evaluates on a regular grid for visualization
|
|
#[must_use]
|
|
pub fn evaluate_grid(&self, nx: usize, ny: usize, time: f64) -> GridFieldData {
|
|
let length = self.vessel.geometry().length();
|
|
let (bbox_min, bbox_max) = self.vessel.geometry().bounding_box();
|
|
|
|
let dx = length / (nx - 1) as f64;
|
|
let dy = (bbox_max.y - bbox_min.y) / (ny - 1) as f64;
|
|
|
|
let mut points = Vec::with_capacity(nx * ny);
|
|
let mut mask = Vec::with_capacity(nx * ny);
|
|
|
|
for j in 0..ny {
|
|
for i in 0..nx {
|
|
let x = i as f64 * dx;
|
|
let y = bbox_min.y + j as f64 * dy;
|
|
let p = Point2D::new(x, y);
|
|
|
|
points.push(p);
|
|
mask.push(self.vessel.is_inside(&p));
|
|
}
|
|
}
|
|
|
|
let predictions = self.evaluate(&points, time);
|
|
|
|
let u: Vec<f64> = predictions.iter().map(|(u, _, _)| *u).collect();
|
|
let v: Vec<f64> = predictions.iter().map(|(_, v, _)| *v).collect();
|
|
let p: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();
|
|
|
|
GridFieldData {
|
|
nx,
|
|
ny,
|
|
dx,
|
|
dy,
|
|
x_min: 0.0,
|
|
y_min: bbox_min.y,
|
|
u,
|
|
v,
|
|
p,
|
|
mask,
|
|
}
|
|
}
|
|
|
|
/// Samples streamlines for visualization
|
|
#[must_use]
|
|
pub fn compute_streamlines(
|
|
&self,
|
|
seed_points: &[Point2D],
|
|
time: f64,
|
|
max_steps: usize,
|
|
step_size: f64,
|
|
) -> Vec<Vec<Point2D>> {
|
|
seed_points
|
|
.iter()
|
|
.map(|seed| self.trace_streamline(*seed, time, max_steps, step_size))
|
|
.collect()
|
|
}
|
|
|
|
/// Traces a single streamline from a seed point
|
|
fn trace_streamline(
|
|
&self,
|
|
seed: Point2D,
|
|
time: f64,
|
|
max_steps: usize,
|
|
step_size: f64,
|
|
) -> Vec<Point2D> {
|
|
let mut streamline = vec![seed];
|
|
let mut current = seed;
|
|
|
|
for _ in 0..max_steps {
|
|
// Get velocity at current point
|
|
let (u, v, _) = self.model.forward(current.x, current.y, time);
|
|
let vel_mag = (u * u + v * v).sqrt();
|
|
|
|
if vel_mag < 1e-10 {
|
|
break; // Stagnation point
|
|
}
|
|
|
|
// Normalize and step
|
|
let next = Point2D::new(
|
|
current.x + step_size * u / vel_mag,
|
|
current.y + step_size * v / vel_mag,
|
|
);
|
|
|
|
// Check if still inside vessel
|
|
if !self.vessel.is_inside(&next) {
|
|
break;
|
|
}
|
|
|
|
streamline.push(next);
|
|
current = next;
|
|
}
|
|
|
|
streamline
|
|
}
|
|
|
|
/// Computes flow statistics for the current solution
|
|
#[must_use]
|
|
pub fn compute_statistics(&self, time: f64) -> FlowStatistics {
|
|
let interior = self.vessel.sample_interior(1000, 42);
|
|
let wall = self.vessel.sample_boundary(200);
|
|
|
|
let predictions = self.evaluate(&interior, time);
|
|
let wss_values = self.compute_wss(&wall, time);
|
|
|
|
let velocities: Vec<f64> = predictions
|
|
.iter()
|
|
.map(|(u, v, _)| (u * u + v * v).sqrt())
|
|
.collect();
|
|
let pressures: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();
|
|
|
|
FlowStatistics {
|
|
max_velocity: velocities.iter().copied().fold(0.0_f64, f64::max),
|
|
mean_velocity: velocities.iter().sum::<f64>() / velocities.len() as f64,
|
|
max_pressure: pressures.iter().copied().fold(f64::NEG_INFINITY, f64::max),
|
|
min_pressure: pressures.iter().copied().fold(f64::INFINITY, f64::min),
|
|
pressure_drop: pressures[0] - pressures[pressures.len() - 1],
|
|
max_wss: WssComputer::max_wss(&wss_values),
|
|
mean_wss: WssComputer::mean_wss(&wss_values),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Grid-based field data for visualization
|
|
#[derive(Debug, Clone)]
|
|
pub struct GridFieldData {
|
|
/// Number of grid points in x
|
|
pub nx: usize,
|
|
/// Number of grid points in y
|
|
pub ny: usize,
|
|
/// Grid spacing in x
|
|
pub dx: f64,
|
|
/// Grid spacing in y
|
|
pub dy: f64,
|
|
/// Minimum x coordinate
|
|
pub x_min: f64,
|
|
/// Minimum y coordinate
|
|
pub y_min: f64,
|
|
/// X-velocity component (row-major: [j * nx + i])
|
|
pub u: Vec<f64>,
|
|
/// Y-velocity component
|
|
pub v: Vec<f64>,
|
|
/// Pressure
|
|
pub p: Vec<f64>,
|
|
/// Interior mask (true if inside vessel)
|
|
pub mask: Vec<bool>,
|
|
}
|
|
|
|
impl GridFieldData {
|
|
/// Returns the value at grid index (i, j)
|
|
///
|
|
/// Returns None if outside grid or masked.
|
|
#[must_use]
|
|
pub fn get(&self, i: usize, j: usize) -> Option<(f64, f64, f64)> {
|
|
if i >= self.nx || j >= self.ny {
|
|
return None;
|
|
}
|
|
|
|
let idx = j * self.nx + i;
|
|
if !self.mask[idx] {
|
|
return None;
|
|
}
|
|
|
|
Some((self.u[idx], self.v[idx], self.p[idx]))
|
|
}
|
|
|
|
/// Converts grid index to physical coordinates
|
|
#[must_use]
|
|
pub fn grid_to_physical(&self, i: usize, j: usize) -> Point2D {
|
|
Point2D::new(
|
|
self.x_min + i as f64 * self.dx,
|
|
self.y_min + j as f64 * self.dy,
|
|
)
|
|
}
|
|
}
|
|
|
|
/// Flow statistics summary
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct FlowStatistics {
|
|
/// Maximum velocity magnitude
|
|
pub max_velocity: f64,
|
|
/// Mean velocity magnitude
|
|
pub mean_velocity: f64,
|
|
/// Maximum pressure
|
|
pub max_pressure: f64,
|
|
/// Minimum pressure
|
|
pub min_pressure: f64,
|
|
/// Pressure drop across vessel
|
|
pub pressure_drop: f64,
|
|
/// Maximum wall shear stress
|
|
pub max_wss: f64,
|
|
/// Mean wall shear stress
|
|
pub mean_wss: f64,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::PinnConfig;
|
|
|
|
fn create_test_engine() -> InferenceEngine {
|
|
let config = PinnConfig::default().with_layers(2).with_hidden_dim(16);
|
|
let model = VesselPinn::new(config);
|
|
let vessel = VesselSdf::straight(0.1, 0.005).unwrap();
|
|
|
|
InferenceEngine::new(model, vessel)
|
|
}
|
|
|
|
#[test]
|
|
fn test_inference_engine_creation() {
|
|
let engine = create_test_engine();
|
|
assert!(engine.batch_size > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_evaluate_points() {
|
|
let engine = create_test_engine();
|
|
let points = vec![Point2D::new(0.01, 0.0), Point2D::new(0.05, 0.002)];
|
|
|
|
let results = engine.evaluate(&points, 0.0);
|
|
|
|
assert_eq!(results.len(), 2);
|
|
for (u, v, p) in &results {
|
|
assert!(u.is_finite());
|
|
assert!(v.is_finite());
|
|
assert!(p.is_finite());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_evaluate_fields() {
|
|
let engine = create_test_engine();
|
|
let points = engine.vessel.sample_interior(50, 42);
|
|
|
|
let response = engine.evaluate_fields(&points, 0.0);
|
|
|
|
assert_eq!(response.velocity().len(), 50);
|
|
assert_eq!(response.pressure().len(), 50);
|
|
assert!(!response.has_wss());
|
|
}
|
|
|
|
#[test]
|
|
fn test_evaluate_with_wss() {
|
|
let engine = create_test_engine();
|
|
let interior = engine.vessel.sample_interior(50, 42);
|
|
let wall = engine.vessel.sample_boundary(20);
|
|
|
|
let response = engine.evaluate_with_wss(&interior, &wall, 0.0);
|
|
|
|
assert_eq!(response.velocity().len(), 50);
|
|
assert!(response.has_wss());
|
|
assert_eq!(response.wss().as_ref().unwrap().len(), 20);
|
|
}
|
|
|
|
#[test]
|
|
fn test_evaluate_grid() {
|
|
let engine = create_test_engine();
|
|
let grid = engine.evaluate_grid(10, 5, 0.0);
|
|
|
|
assert_eq!(grid.nx, 10);
|
|
assert_eq!(grid.ny, 5);
|
|
assert_eq!(grid.u.len(), 50);
|
|
assert_eq!(grid.mask.len(), 50);
|
|
}
|
|
|
|
#[test]
|
|
fn test_streamline_tracing() {
|
|
let engine = create_test_engine();
|
|
let seeds = vec![Point2D::new(0.01, 0.0)];
|
|
|
|
let streamlines = engine.compute_streamlines(&seeds, 0.0, 100, 0.001);
|
|
|
|
assert_eq!(streamlines.len(), 1);
|
|
assert!(!streamlines[0].is_empty());
|
|
|
|
// First point should be seed
|
|
let first = &streamlines[0][0];
|
|
assert!((first.x - 0.01).abs() < f64::EPSILON);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_statistics() {
|
|
let engine = create_test_engine();
|
|
let stats = engine.compute_statistics(0.0);
|
|
|
|
assert!(stats.max_velocity >= 0.0);
|
|
assert!(stats.mean_velocity >= 0.0);
|
|
assert!(stats.max_wss >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_grid_field_data_access() {
|
|
let grid = GridFieldData {
|
|
nx: 3,
|
|
ny: 2,
|
|
dx: 0.05,
|
|
dy: 0.005,
|
|
x_min: 0.0,
|
|
y_min: -0.005,
|
|
u: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
|
|
v: vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
|
|
p: vec![100.0, 90.0, 80.0, 100.0, 90.0, 80.0],
|
|
mask: vec![true, true, true, true, true, true],
|
|
};
|
|
|
|
let (u, v, p) = grid.get(1, 0).unwrap();
|
|
assert!((u - 2.0).abs() < f64::EPSILON);
|
|
assert!((v - 0.2).abs() < f64::EPSILON);
|
|
assert!((p - 90.0).abs() < f64::EPSILON);
|
|
|
|
let point = grid.grid_to_physical(1, 1);
|
|
assert!((point.x - 0.05).abs() < f64::EPSILON);
|
|
assert!((point.y - 0.0).abs() < f64::EPSILON);
|
|
}
|
|
}
|