Initial commit

demos/rtx-hemodynamics/src/boundary.rs

@@ -0,0 +1,375 @@
+//! Boundary condition enforcement for hemodynamics PINN
+//!
+//! This module handles the enforcement of boundary conditions:
+//! - Inlet: prescribed velocity profile
+//! - Outlet: prescribed pressure
+//! - Wall: no-slip condition (velocity = 0)
+
+use crate::vessel::VesselSdf;
+use rtx_hemodynamics_shared::geometry::Point2D;
+use rtx_hemodynamics_shared::physics::{BoundaryCondition, BoundaryType};
+
+/// Boundary condition enforcer for PINN training
+#[derive(Debug)]
+pub struct BoundaryEnforcer {
+    /// Inlet boundary condition
+    inlet: BoundaryCondition,
+    /// Outlet boundary condition
+    outlet: BoundaryCondition,
+    /// Vessel SDF for wall detection
+    vessel: VesselSdf,
+    /// Tolerance for boundary detection
+    tolerance: f64,
+}
+
+impl BoundaryEnforcer {
+    /// Creates a new boundary enforcer
+    #[must_use]
+    pub fn new(inlet: BoundaryCondition, outlet: BoundaryCondition, vessel: VesselSdf) -> Self {
+        Self {
+            inlet,
+            outlet,
+            vessel,
+            tolerance: 1e-6,
+        }
+    }
+
+    /// Creates a boundary enforcer with Poiseuille inlet profile
+    pub fn poiseuille_inlet(
+        inlet_velocity: f64,
+        outlet_pressure: f64,
+        vessel: VesselSdf,
+    ) -> Result<Self, String> {
+        let inlet = BoundaryCondition::inlet_velocity(inlet_velocity).map_err(|e| e.to_string())?;
+        let outlet =
+            BoundaryCondition::outlet_pressure(outlet_pressure).map_err(|e| e.to_string())?;
+
+        Ok(Self::new(inlet, outlet, vessel))
+    }
+
+    /// Returns the inlet boundary condition
+    #[must_use]
+    pub const fn inlet(&self) -> &BoundaryCondition {
+        &self.inlet
+    }
+
+    /// Returns the outlet boundary condition
+    #[must_use]
+    pub const fn outlet(&self) -> &BoundaryCondition {
+        &self.outlet
+    }
+
+    /// Returns the vessel SDF
+    #[must_use]
+    pub const fn vessel(&self) -> &VesselSdf {
+        &self.vessel
+    }
+
+    /// Classifies a point as inlet, outlet, wall, or interior
+    #[must_use]
+    pub fn classify_point(&self, point: &Point2D) -> BoundaryLocation {
+        let length = self.vessel.geometry().length();
+
+        // Check if at inlet (x ≈ 0)
+        if point.x.abs() < self.tolerance {
+            return BoundaryLocation::Inlet;
+        }
+
+        // Check if at outlet (x ≈ length)
+        if (point.x - length).abs() < self.tolerance {
+            return BoundaryLocation::Outlet;
+        }
+
+        // Check if on wall
+        if self.vessel.is_boundary(point, self.tolerance) {
+            return BoundaryLocation::Wall;
+        }
+
+        // Check if outside
+        if !self.vessel.is_inside(point) {
+            return BoundaryLocation::Outside;
+        }
+
+        BoundaryLocation::Interior
+    }
+
+    /// Computes the target velocity at inlet using Poiseuille profile
+    ///
+    /// u(y) = `u_max` * (1 - (y/R)²)
+    #[must_use]
+    pub fn inlet_velocity_target(&self, y: f64, time: f64) -> (f64, f64) {
+        let r = self.vessel.geometry().base_radius();
+
+        // Get inlet velocity from boundary condition
+        let u_max = match self.inlet.boundary_type() {
+            BoundaryType::InletVelocity(v) => *v,
+            BoundaryType::PulsatileInlet { .. } => self.inlet.evaluate(time),
+            _ => 0.0,
+        };
+
+        // Parabolic profile: u(y) = u_max * (1 - (y/R)²)
+        let normalized_y = y / r;
+        let u = u_max * (1.0 - normalized_y.powi(2));
+
+        (u, 0.0) // v = 0 at inlet for fully-developed flow
+    }
+
+    /// Computes the target pressure at outlet
+    #[must_use]
+    pub fn outlet_pressure_target(&self, _time: f64) -> f64 {
+        match self.outlet.boundary_type() {
+            BoundaryType::OutletPressure(p) => *p,
+            _ => 0.0,
+        }
+    }
+
+    /// Computes the boundary loss for a set of predicted values
+    ///
+    /// `L_bc` = `L_inlet` + `L_outlet` + `L_wall`
+    #[must_use]
+    pub fn boundary_loss(
+        &self,
+        inlet_points: &[Point2D],
+        inlet_predictions: &[(f64, f64, f64)],
+        _outlet_points: &[Point2D],
+        outlet_predictions: &[(f64, f64, f64)],
+        _wall_points: &[Point2D],
+        wall_predictions: &[(f64, f64, f64)],
+        time: f64,
+    ) -> BoundaryLoss {
+        let inlet_loss = self.inlet_loss(inlet_points, inlet_predictions, time);
+        let outlet_loss = self.outlet_loss(outlet_predictions, time);
+        let wall_loss = self.wall_loss(wall_predictions);
+
+        BoundaryLoss {
+            inlet: inlet_loss,
+            outlet: outlet_loss,
+            wall: wall_loss,
+            total: inlet_loss + outlet_loss + wall_loss,
+        }
+    }
+
+    /// Computes inlet loss (velocity MSE)
+    fn inlet_loss(&self, points: &[Point2D], predictions: &[(f64, f64, f64)], time: f64) -> f64 {
+        if points.is_empty() {
+            return 0.0;
+        }
+
+        let mut loss = 0.0;
+        for (point, (u_pred, v_pred, _)) in points.iter().zip(predictions.iter()) {
+            let (u_target, v_target) = self.inlet_velocity_target(point.y, time);
+            loss += (u_pred - u_target).powi(2) + (v_pred - v_target).powi(2);
+        }
+
+        loss / points.len() as f64
+    }
+
+    /// Computes outlet loss (pressure MSE)
+    fn outlet_loss(&self, predictions: &[(f64, f64, f64)], time: f64) -> f64 {
+        if predictions.is_empty() {
+            return 0.0;
+        }
+
+        let p_target = self.outlet_pressure_target(time);
+        let mut loss = 0.0;
+
+        for (_, _, p_pred) in predictions {
+            loss += (p_pred - p_target).powi(2);
+        }
+
+        loss / predictions.len() as f64
+    }
+
+    /// Computes wall loss (no-slip: u = v = 0)
+    fn wall_loss(&self, predictions: &[(f64, f64, f64)]) -> f64 {
+        if predictions.is_empty() {
+            return 0.0;
+        }
+
+        let mut loss = 0.0;
+        for (u, v, _) in predictions {
+            loss += u.powi(2) + v.powi(2);
+        }
+
+        loss / predictions.len() as f64
+    }
+
+    /// Samples boundary points for training
+    #[must_use]
+    pub fn sample_boundary_points(
+        &self,
+        n_inlet: usize,
+        n_outlet: usize,
+        n_wall: usize,
+    ) -> (Vec<Point2D>, Vec<Point2D>, Vec<Point2D>) {
+        let inlet = self.vessel.sample_inlet(n_inlet);
+        let outlet = self.vessel.sample_outlet(n_outlet);
+        let wall = self.vessel.sample_boundary(n_wall);
+
+        (inlet, outlet, wall)
+    }
+}
+
+/// Classification of point location relative to boundaries
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BoundaryLocation {
+    /// At inlet (x = 0)
+    Inlet,
+    /// At outlet (x = length)
+    Outlet,
+    /// On vessel wall
+    Wall,
+    /// Inside vessel (not on boundary)
+    Interior,
+    /// Outside vessel domain
+    Outside,
+}
+
+/// Boundary loss components
+#[derive(Debug, Clone, Copy)]
+pub struct BoundaryLoss {
+    /// Inlet velocity loss
+    pub inlet: f64,
+    /// Outlet pressure loss
+    pub outlet: f64,
+    /// Wall no-slip loss
+    pub wall: f64,
+    /// Total boundary loss
+    pub total: f64,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn create_test_enforcer() -> BoundaryEnforcer {
+        let vessel = VesselSdf::straight(0.1, 0.005).unwrap();
+        BoundaryEnforcer::poiseuille_inlet(0.1, 0.0, vessel).unwrap()
+    }
+
+    #[test]
+    fn test_boundary_enforcer_creation() {
+        let enforcer = create_test_enforcer();
+        assert!(enforcer.inlet().is_inlet());
+        assert!(enforcer.outlet().is_outlet());
+    }
+
+    #[test]
+    fn test_classify_point_inlet() {
+        let enforcer = create_test_enforcer();
+        let inlet_point = Point2D::new(0.0, 0.002);
+
+        assert_eq!(
+            enforcer.classify_point(&inlet_point),
+            BoundaryLocation::Inlet
+        );
+    }
+
+    #[test]
+    fn test_classify_point_outlet() {
+        let enforcer = create_test_enforcer();
+        let outlet_point = Point2D::new(0.1, 0.002);
+
+        assert_eq!(
+            enforcer.classify_point(&outlet_point),
+            BoundaryLocation::Outlet
+        );
+    }
+
+    #[test]
+    fn test_classify_point_interior() {
+        let enforcer = create_test_enforcer();
+        let interior_point = Point2D::new(0.05, 0.002);
+
+        assert_eq!(
+            enforcer.classify_point(&interior_point),
+            BoundaryLocation::Interior
+        );
+    }
+
+    #[test]
+    fn test_inlet_velocity_profile() {
+        let enforcer = create_test_enforcer();
+
+        // Maximum at centerline
+        let (u_center, v_center) = enforcer.inlet_velocity_target(0.0, 0.0);
+        assert!((u_center - 0.1).abs() < f64::EPSILON);
+        assert!(v_center.abs() < f64::EPSILON);
+
+        // Zero at wall
+        let (u_wall, _) = enforcer.inlet_velocity_target(0.005, 0.0);
+        assert!(u_wall.abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn test_wall_loss() {
+        let enforcer = create_test_enforcer();
+
+        // Perfect no-slip should have zero loss
+        let perfect = vec![(0.0, 0.0, 100.0)];
+        let loss = enforcer.wall_loss(&perfect);
+        assert!(loss.abs() < f64::EPSILON);
+
+        // Non-zero velocity should have positive loss
+        let imperfect = vec![(0.01, 0.0, 100.0)];
+        let loss = enforcer.wall_loss(&imperfect);
+        assert!((loss - 0.0001).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn test_sample_boundary_points() {
+        let enforcer = create_test_enforcer();
+
+        let (inlet, outlet, wall) = enforcer.sample_boundary_points(10, 10, 20);
+
+        assert_eq!(inlet.len(), 10);
+        assert_eq!(outlet.len(), 10);
+        assert_eq!(wall.len(), 20);
+
+        // Inlet points should be at x=0
+        for p in &inlet {
+            assert!(p.x.abs() < f64::EPSILON);
+        }
+
+        // Outlet points should be at x=length
+        for p in &outlet {
+            assert!((p.x - 0.1).abs() < f64::EPSILON);
+        }
+    }
+
+    #[test]
+    fn test_boundary_loss_components() {
+        let enforcer = create_test_enforcer();
+
+        let inlet_points = vec![Point2D::new(0.0, 0.0)];
+        let inlet_preds = vec![(0.1, 0.0, 50.0)]; // Perfect velocity
+
+        let outlet_points = vec![Point2D::new(0.1, 0.0)];
+        let outlet_preds = vec![(0.05, 0.0, 0.0)]; // Perfect pressure
+
+        let wall_points = vec![Point2D::new(0.05, 0.005)];
+        let wall_preds = vec![(0.0, 0.0, 50.0)]; // Perfect no-slip
+
+        let loss = enforcer.boundary_loss(
+            &inlet_points,
+            &inlet_preds,
+            &outlet_points,
+            &outlet_preds,
+            &wall_points,
+            &wall_preds,
+            0.0,
+        );
+
+        assert!(
+            loss.inlet.abs() < f64::EPSILON,
+            "inlet loss: {}",
+            loss.inlet
+        );
+        assert!(
+            loss.outlet.abs() < f64::EPSILON,
+            "outlet loss: {}",
+            loss.outlet
+        );
+        assert!(loss.wall.abs() < f64::EPSILON, "wall loss: {}", loss.wall);
+    }
+}