429 lines
13 KiB
Rust
429 lines
13 KiB
Rust
//! Vessel geometry with signed distance functions for PINN training
|
||
//!
|
||
//! This module provides GPU-compatible vessel geometry representations
|
||
//! with efficient signed distance function (SDF) evaluation for use
|
||
//! in physics-informed neural network training.
|
||
|
||
use rtx_hemodynamics_shared::geometry::{Point2D, StenosisParams, VesselGeometry};
|
||
|
||
/// Vessel SDF evaluator for efficient batched SDF computation
|
||
///
|
||
/// Wraps a `VesselGeometry` and provides methods optimized for
|
||
/// PINN training, including batched evaluation and gradient computation.
|
||
#[derive(Debug, Clone)]
|
||
pub struct VesselSdf {
|
||
/// Underlying vessel geometry
|
||
geometry: VesselGeometry,
|
||
/// Precomputed bounding box
|
||
bbox_min: Point2D,
|
||
bbox_max: Point2D,
|
||
}
|
||
|
||
impl VesselSdf {
|
||
/// Creates a new vessel SDF from geometry
|
||
#[must_use]
|
||
pub fn new(geometry: VesselGeometry) -> Self {
|
||
let (bbox_min, bbox_max) = geometry.bounding_box();
|
||
Self {
|
||
geometry,
|
||
bbox_min,
|
||
bbox_max,
|
||
}
|
||
}
|
||
|
||
/// Creates a straight vessel SDF
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns an error if length or radius is not positive.
|
||
pub fn straight(length: f64, radius: f64) -> Result<Self, String> {
|
||
VesselGeometry::straight(length, radius)
|
||
.map(Self::new)
|
||
.map_err(|e| e.to_string())
|
||
}
|
||
|
||
/// Returns the underlying geometry
|
||
#[must_use]
|
||
pub const fn geometry(&self) -> &VesselGeometry {
|
||
&self.geometry
|
||
}
|
||
|
||
/// Returns the bounding box minimum corner
|
||
#[must_use]
|
||
pub const fn bbox_min(&self) -> &Point2D {
|
||
&self.bbox_min
|
||
}
|
||
|
||
/// Returns the bounding box maximum corner
|
||
#[must_use]
|
||
pub const fn bbox_max(&self) -> &Point2D {
|
||
&self.bbox_max
|
||
}
|
||
|
||
/// Evaluates the signed distance at a single point
|
||
#[must_use]
|
||
pub fn sdf(&self, point: &Point2D) -> f64 {
|
||
self.geometry.signed_distance(point)
|
||
}
|
||
|
||
/// Evaluates the signed distance at multiple points
|
||
#[must_use]
|
||
pub fn sdf_batch(&self, points: &[Point2D]) -> Vec<f64> {
|
||
points.iter().map(|p| self.sdf(p)).collect()
|
||
}
|
||
|
||
/// Computes the SDF gradient (normal direction) at a point
|
||
///
|
||
/// Uses finite differences with specified epsilon.
|
||
#[must_use]
|
||
pub fn sdf_gradient(&self, point: &Point2D, eps: f64) -> Point2D {
|
||
let sdf_center = self.sdf(point);
|
||
let sdf_x = self.sdf(&Point2D::new(point.x + eps, point.y));
|
||
let sdf_y = self.sdf(&Point2D::new(point.x, point.y + eps));
|
||
|
||
Point2D::new((sdf_x - sdf_center) / eps, (sdf_y - sdf_center) / eps)
|
||
}
|
||
|
||
/// Computes the outward normal at a boundary point
|
||
#[must_use]
|
||
pub fn normal(&self, point: &Point2D) -> Point2D {
|
||
self.sdf_gradient(point, 1e-6).normalize()
|
||
}
|
||
|
||
/// Checks if a point is inside the vessel (SDF < 0)
|
||
#[must_use]
|
||
pub fn is_inside(&self, point: &Point2D) -> bool {
|
||
self.sdf(point) < 0.0
|
||
}
|
||
|
||
/// Checks if a point is on the boundary (|SDF| < tolerance)
|
||
#[must_use]
|
||
pub fn is_boundary(&self, point: &Point2D, tolerance: f64) -> bool {
|
||
self.sdf(point).abs() < tolerance
|
||
}
|
||
|
||
/// Samples random points inside the vessel
|
||
///
|
||
/// Uses rejection sampling within the bounding box.
|
||
#[must_use]
|
||
pub fn sample_interior(&self, n: usize, _seed: u64) -> Vec<Point2D> {
|
||
self.geometry.sample_interior(n)
|
||
}
|
||
|
||
/// Samples points on the vessel boundary
|
||
#[must_use]
|
||
pub fn sample_boundary(&self, n: usize) -> Vec<Point2D> {
|
||
self.geometry.sample_boundary(n)
|
||
}
|
||
|
||
/// Samples points at the inlet (x = 0)
|
||
#[must_use]
|
||
pub fn sample_inlet(&self, n: usize) -> Vec<Point2D> {
|
||
let mut points = Vec::with_capacity(n);
|
||
let r = self.geometry.base_radius();
|
||
|
||
for i in 0..n {
|
||
let y = r * (2.0 * (i as f64 + 0.5) / n as f64 - 1.0) * 0.99;
|
||
points.push(Point2D::new(0.0, y));
|
||
}
|
||
|
||
points
|
||
}
|
||
|
||
/// Samples points at the outlet (x = length)
|
||
#[must_use]
|
||
pub fn sample_outlet(&self, n: usize) -> Vec<Point2D> {
|
||
let mut points = Vec::with_capacity(n);
|
||
let length = self.geometry.length();
|
||
let r = self.geometry.local_radius(length);
|
||
|
||
for i in 0..n {
|
||
let y = r * (2.0 * (i as f64 + 0.5) / n as f64 - 1.0) * 0.99;
|
||
points.push(Point2D::new(length, y));
|
||
}
|
||
|
||
points
|
||
}
|
||
|
||
/// Returns the local vessel radius at x position
|
||
#[must_use]
|
||
pub fn local_radius(&self, x: f64) -> f64 {
|
||
self.geometry.local_radius(x)
|
||
}
|
||
|
||
/// Projects a point onto the vessel boundary
|
||
///
|
||
/// Uses gradient descent on the SDF.
|
||
#[must_use]
|
||
pub fn project_to_boundary(&self, point: &Point2D, max_iter: usize) -> Point2D {
|
||
let mut p = *point;
|
||
let eps = 1e-6;
|
||
|
||
for _ in 0..max_iter {
|
||
let sdf = self.sdf(&p);
|
||
if sdf.abs() < eps {
|
||
break;
|
||
}
|
||
|
||
let grad = self.sdf_gradient(&p, eps);
|
||
let grad_mag = grad.magnitude();
|
||
if grad_mag < eps {
|
||
break;
|
||
}
|
||
|
||
// Move along negative gradient direction
|
||
p = Point2D::new(p.x - sdf * grad.x / grad_mag, p.y - sdf * grad.y / grad_mag);
|
||
}
|
||
|
||
p
|
||
}
|
||
|
||
/// Modifies the vessel with a stenosis
|
||
#[must_use]
|
||
pub fn with_stenosis(self, stenosis: StenosisParams) -> Self {
|
||
let new_geom = self.geometry.with_stenosis(stenosis);
|
||
Self::new(new_geom)
|
||
}
|
||
|
||
/// Returns vessel statistics for validation
|
||
#[must_use]
|
||
pub fn statistics(&self) -> VesselStatistics {
|
||
// Sample points to compute statistics
|
||
let interior = self.sample_interior(1000, 42);
|
||
let boundary = self.sample_boundary(200);
|
||
|
||
let interior_sdfs: Vec<f64> = interior.iter().map(|p| self.sdf(p)).collect();
|
||
let boundary_sdfs: Vec<f64> = boundary.iter().map(|p| self.sdf(p)).collect();
|
||
|
||
let min_interior_sdf = interior_sdfs.iter().copied().fold(f64::INFINITY, f64::min);
|
||
let max_interior_sdf = interior_sdfs
|
||
.iter()
|
||
.copied()
|
||
.fold(f64::NEG_INFINITY, f64::max);
|
||
let max_boundary_sdf = boundary_sdfs
|
||
.iter()
|
||
.map(|s| s.abs())
|
||
.fold(0.0_f64, f64::max);
|
||
|
||
VesselStatistics {
|
||
length: self.geometry.length(),
|
||
base_radius: self.geometry.base_radius(),
|
||
min_interior_sdf,
|
||
max_interior_sdf,
|
||
max_boundary_sdf_error: max_boundary_sdf,
|
||
volume_estimate: self.estimate_volume(1000),
|
||
}
|
||
}
|
||
|
||
/// Estimates vessel volume using Monte Carlo integration
|
||
fn estimate_volume(&self, n_samples: usize) -> f64 {
|
||
let interior = self.sample_interior(n_samples, 42);
|
||
let n_inside = interior.iter().filter(|p| self.is_inside(p)).count();
|
||
|
||
// Volume = bbox_area * (n_inside / n_samples)
|
||
let bbox_width = self.bbox_max.x - self.bbox_min.x;
|
||
let bbox_height = self.bbox_max.y - self.bbox_min.y;
|
||
let bbox_area = bbox_width * bbox_height;
|
||
|
||
bbox_area * (n_inside as f64 / n_samples as f64)
|
||
}
|
||
}
|
||
|
||
/// Statistics about a vessel geometry
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct VesselStatistics {
|
||
/// Vessel length
|
||
pub length: f64,
|
||
/// Base radius
|
||
pub base_radius: f64,
|
||
/// Minimum SDF value in interior (should be negative)
|
||
pub min_interior_sdf: f64,
|
||
/// Maximum SDF value in interior (should be negative)
|
||
pub max_interior_sdf: f64,
|
||
/// Maximum |SDF| error at boundary points (should be ~0)
|
||
pub max_boundary_sdf_error: f64,
|
||
/// Estimated vessel volume
|
||
pub volume_estimate: f64,
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_vessel_sdf_straight() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
|
||
// Point on centerline should be inside
|
||
let center = Point2D::new(0.05, 0.0);
|
||
assert!(sdf.is_inside(¢er));
|
||
assert!(sdf.sdf(¢er) < 0.0);
|
||
|
||
// Point outside should have positive SDF
|
||
let outside = Point2D::new(0.05, 0.01);
|
||
assert!(!sdf.is_inside(&outside));
|
||
assert!(sdf.sdf(&outside) > 0.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_sdf_boundary() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
|
||
// Point on boundary should have SDF ~0
|
||
let boundary = Point2D::new(0.05, 0.005);
|
||
assert!(sdf.sdf(&boundary).abs() < 1e-10);
|
||
assert!(sdf.is_boundary(&boundary, 1e-6));
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_sdf_gradient() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
|
||
// At top boundary, normal should point up (positive y)
|
||
let top = Point2D::new(0.05, 0.005);
|
||
let normal = sdf.normal(&top);
|
||
assert!(normal.y > 0.9); // Should be nearly (0, 1)
|
||
|
||
// At bottom boundary, normal should point down
|
||
let bottom = Point2D::new(0.05, -0.005);
|
||
let normal = sdf.normal(&bottom);
|
||
assert!(normal.y < -0.9);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_sample_interior() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
let points = sdf.sample_interior(100, 42);
|
||
|
||
assert_eq!(points.len(), 100);
|
||
|
||
// All points should be inside
|
||
for p in &points {
|
||
assert!(
|
||
sdf.is_inside(p),
|
||
"Interior point should be inside: ({}, {}), SDF = {}",
|
||
p.x,
|
||
p.y,
|
||
sdf.sdf(p)
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_sample_boundary() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
let points = sdf.sample_boundary(50);
|
||
|
||
assert_eq!(points.len(), 50);
|
||
|
||
// All points should be on boundary
|
||
for p in &points {
|
||
assert!(
|
||
sdf.sdf(p).abs() < 1e-6,
|
||
"Boundary point should have SDF ~0: ({}, {}), SDF = {}",
|
||
p.x,
|
||
p.y,
|
||
sdf.sdf(p)
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_sample_inlet_outlet() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
|
||
let inlet = sdf.sample_inlet(20);
|
||
let outlet = sdf.sample_outlet(20);
|
||
|
||
// Inlet points should be at x=0
|
||
for p in &inlet {
|
||
assert!(p.x.abs() < f64::EPSILON);
|
||
assert!(sdf.is_inside(p));
|
||
}
|
||
|
||
// Outlet points should be at x=length
|
||
for p in &outlet {
|
||
assert!((p.x - 0.1).abs() < f64::EPSILON);
|
||
assert!(sdf.is_inside(p));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_project_to_boundary() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
|
||
// Project a point inside to boundary
|
||
let inside = Point2D::new(0.05, 0.002);
|
||
let projected = sdf.project_to_boundary(&inside, 100);
|
||
|
||
// Projected point should be on boundary
|
||
assert!(
|
||
sdf.sdf(&projected).abs() < 1e-5,
|
||
"Projected point should be on boundary, SDF = {}",
|
||
sdf.sdf(&projected)
|
||
);
|
||
|
||
// Project a point outside to boundary
|
||
let outside = Point2D::new(0.05, 0.01);
|
||
let projected = sdf.project_to_boundary(&outside, 100);
|
||
|
||
assert!(
|
||
sdf.sdf(&projected).abs() < 1e-5,
|
||
"Projected point should be on boundary, SDF = {}",
|
||
sdf.sdf(&projected)
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_with_stenosis() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
let stenosis = StenosisParams::new(0.5, 0.02, 0.05).unwrap();
|
||
let stenotic = sdf.with_stenosis(stenosis);
|
||
|
||
// At stenosis center, radius should be reduced
|
||
let r_center = stenotic.local_radius(0.05);
|
||
assert!((r_center - 0.0025).abs() < 0.001);
|
||
|
||
// Away from stenosis, radius should be normal
|
||
let r_inlet = stenotic.local_radius(0.0);
|
||
assert!((r_inlet - 0.005).abs() < f64::EPSILON);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_statistics() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
let stats = sdf.statistics();
|
||
|
||
assert!((stats.length - 0.1).abs() < f64::EPSILON);
|
||
assert!((stats.base_radius - 0.005).abs() < f64::EPSILON);
|
||
assert!(stats.min_interior_sdf < 0.0);
|
||
assert!(stats.max_interior_sdf < 0.0);
|
||
assert!(stats.max_boundary_sdf_error < 1e-5);
|
||
|
||
// For 2D simulation, "volume" is actually 2D area: L × 2R = 0.1 × 0.01 = 0.001
|
||
let expected_area = 0.1 * 2.0 * 0.005;
|
||
assert!(
|
||
(stats.volume_estimate - expected_area).abs() < expected_area * 0.2,
|
||
"Area estimate {} should be close to {}",
|
||
stats.volume_estimate,
|
||
expected_area
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_vessel_sdf_batch() {
|
||
let sdf = VesselSdf::straight(0.1, 0.005).unwrap();
|
||
let points = vec![
|
||
Point2D::new(0.05, 0.0), // Inside
|
||
Point2D::new(0.05, 0.005), // Boundary
|
||
Point2D::new(0.05, 0.01), // Outside
|
||
];
|
||
|
||
let sdfs = sdf.sdf_batch(&points);
|
||
|
||
assert!(sdfs[0] < 0.0); // Inside
|
||
assert!(sdfs[1].abs() < 1e-10); // Boundary
|
||
assert!(sdfs[2] > 0.0); // Outside
|
||
}
|
||
}
|