rustytorch/demos/rtx-hemodynamics/src/training.rs

//! Training loop for hemodynamics PINN
//!
//! Implements the training procedure including:
//! - Physics loss (Navier-Stokes residuals)
//! - Data loss (velocity observations)
//! - Boundary condition loss

use crate::boundary::{BoundaryEnforcer, BoundaryLoss};
use crate::config::PinnConfig;
use crate::navier_stokes::NavierStokesResidual;
use crate::network::VesselPinn;
use crate::vessel::VesselSdf;

/// PINN Trainer for hemodynamics
#[derive(Debug)]
pub struct Trainer {
    /// Configuration
    config: PinnConfig,
    /// Vessel geometry
    vessel: VesselSdf,
    /// Navier-Stokes residual computer
    ns: NavierStokesResidual,
    /// Boundary enforcer
    boundary: BoundaryEnforcer,
    /// Training history
    history: TrainingHistory,
}

impl Trainer {
    /// Creates a new trainer
    pub fn new(config: PinnConfig, vessel: VesselSdf) -> Result<Self, String> {
        let ns = NavierStokesResidual::new(config.fluid(), true);

        let boundary = BoundaryEnforcer::poiseuille_inlet(
            0.1, // Inlet velocity
            0.0, // Outlet pressure
            vessel.clone(),
        )?;

        Ok(Self {
            config,
            vessel,
            ns,
            boundary,
            history: TrainingHistory::new(),
        })
    }

    /// Returns the configuration
    #[must_use]
    pub const fn config(&self) -> &PinnConfig {
        &self.config
    }

    /// Returns the training history
    #[must_use]
    pub const fn history(&self) -> &TrainingHistory {
        &self.history
    }

    /// Computes total loss for a model
    #[must_use]
    pub fn compute_loss(&self, model: &VesselPinn, time: f64) -> TotalLoss {
        let physics_loss = self.compute_physics_loss(model, time);
        let boundary_loss = self.compute_boundary_loss(model, time);

        let total = self.config.physics_weight * physics_loss.total
            + self.config.boundary_weight * boundary_loss.total;

        TotalLoss {
            total,
            physics: physics_loss,
            boundary: boundary_loss,
        }
    }

    /// Computes physics loss (Navier-Stokes residuals)
    fn compute_physics_loss(&self, model: &VesselPinn, time: f64) -> PhysicsLoss {
        let collocation = self.vessel.sample_interior(
            self.config.num_collocation_points,
            42, // Seed for reproducibility
        );

        let eps = 1e-5;
        let mut residuals = Vec::with_capacity(collocation.len());

        for point in &collocation {
            let (u, v, _p) = model.forward(point.x, point.y, time);
            let grads = model.gradients(point.x, point.y, time, eps);
            let grads2 = model.second_gradients(point.x, point.y, time, eps);

            let (rx, ry, rc) = self.ns.all_residuals(
                u,
                v,
                0.0,
                0.0, // Steady state
                grads.du_dx,
                grads.du_dy,
                grads.dv_dx,
                grads.dv_dy,
                grads2.d2u_dx2,
                grads2.d2u_dy2,
                grads2.d2v_dx2,
                grads2.d2v_dy2,
                grads.dp_dx,
                grads.dp_dy,
            );

            residuals.push((rx, ry, rc));
        }

        let (momentum_x, momentum_y, continuity) = self.ns.loss_components(&residuals);
        let total = self.ns.physics_loss(&residuals);

        PhysicsLoss {
            total,
            momentum_x,
            momentum_y,
            continuity,
        }
    }

    /// Computes boundary condition loss
    fn compute_boundary_loss(&self, model: &VesselPinn, time: f64) -> BoundaryLoss {
        let (inlet_points, outlet_points, wall_points) = self.boundary.sample_boundary_points(
            self.config.num_boundary_points / 4,
            self.config.num_boundary_points / 4,
            self.config.num_boundary_points / 2,
        );

        let inlet_preds = model.forward_batch(&inlet_points, time);
        let outlet_preds = model.forward_batch(&outlet_points, time);
        let wall_preds = model.forward_batch(&wall_points, time);

        self.boundary.boundary_loss(
            &inlet_points,
            &inlet_preds,
            &outlet_points,
            &outlet_preds,
            &wall_points,
            &wall_preds,
            time,
        )
    }

    /// Records a training step in history
    pub fn record_step(&mut self, epoch: usize, loss: &TotalLoss, lr: f64) {
        self.history.add_entry(TrainingEntry {
            epoch,
            total_loss: loss.total,
            physics_loss: loss.physics.total,
            boundary_loss: loss.boundary.total,
            momentum_x_loss: loss.physics.momentum_x,
            momentum_y_loss: loss.physics.momentum_y,
            continuity_loss: loss.physics.continuity,
            learning_rate: lr,
        });
    }

    /// Checks convergence based on loss history
    #[must_use]
    pub fn is_converged(&self, tolerance: f64, window: usize) -> bool {
        self.history.is_converged(tolerance, window)
    }
}

/// Total loss combining physics and boundary losses
#[derive(Debug, Clone, Copy)]
pub struct TotalLoss {
    /// Total weighted loss
    pub total: f64,
    /// Physics loss components
    pub physics: PhysicsLoss,
    /// Boundary loss components
    pub boundary: BoundaryLoss,
}

/// Physics loss components
#[derive(Debug, Clone, Copy)]
pub struct PhysicsLoss {
    /// Total physics loss
    pub total: f64,
    /// X-momentum residual loss
    pub momentum_x: f64,
    /// Y-momentum residual loss
    pub momentum_y: f64,
    /// Continuity residual loss
    pub continuity: f64,
}

/// Training history for monitoring
#[derive(Debug, Clone, Default)]
pub struct TrainingHistory {
    /// Training entries
    entries: Vec<TrainingEntry>,
}

impl TrainingHistory {
    /// Creates empty history
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds an entry to history
    pub fn add_entry(&mut self, entry: TrainingEntry) {
        self.entries.push(entry);
    }

    /// Returns all entries
    #[must_use]
    pub fn entries(&self) -> &[TrainingEntry] {
        &self.entries
    }

    /// Returns the latest entry
    #[must_use]
    pub fn latest(&self) -> Option<&TrainingEntry> {
        self.entries.last()
    }

    /// Returns the best (minimum) total loss
    #[must_use]
    pub fn best_loss(&self) -> f64 {
        self.entries
            .iter()
            .map(|e| e.total_loss)
            .fold(f64::INFINITY, f64::min)
    }

    /// Checks if loss has converged
    #[must_use]
    pub fn is_converged(&self, tolerance: f64, window: usize) -> bool {
        if self.entries.len() < window {
            return false;
        }

        let recent: Vec<f64> = self
            .entries
            .iter()
            .rev()
            .take(window)
            .map(|e| e.total_loss)
            .collect();

        let mean = recent.iter().sum::<f64>() / window as f64;
        let variance: f64 = recent.iter().map(|&l| (l - mean).powi(2)).sum::<f64>() / window as f64;
        let std_dev = variance.sqrt();

        std_dev / mean < tolerance
    }

    /// Returns loss values for plotting
    #[must_use]
    pub fn loss_curve(&self) -> Vec<(usize, f64)> {
        self.entries
            .iter()
            .map(|e| (e.epoch, e.total_loss))
            .collect()
    }
}

/// Single training entry
#[derive(Debug, Clone, Copy)]
pub struct TrainingEntry {
    /// Epoch number
    pub epoch: usize,
    /// Total loss
    pub total_loss: f64,
    /// Physics loss
    pub physics_loss: f64,
    /// Boundary loss
    pub boundary_loss: f64,
    /// X-momentum loss
    pub momentum_x_loss: f64,
    /// Y-momentum loss
    pub momentum_y_loss: f64,
    /// Continuity loss
    pub continuity_loss: f64,
    /// Learning rate
    pub learning_rate: f64,
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_test_trainer() -> Trainer {
        let config = PinnConfig::default()
            .with_layers(2)
            .with_hidden_dim(16)
            .with_collocation_points(100)
            .with_boundary_points(50);

        let vessel = VesselSdf::straight(0.1, 0.005).unwrap();
        Trainer::new(config, vessel).unwrap()
    }

    #[test]
    fn test_trainer_creation() {
        let trainer = create_test_trainer();
        assert_eq!(trainer.config().num_collocation_points, 100);
    }

    #[test]
    fn test_compute_loss() {
        let trainer = create_test_trainer();
        let model = VesselPinn::new(trainer.config().clone());

        let loss = trainer.compute_loss(&model, 0.0);

        assert!(loss.total.is_finite());
        assert!(loss.total >= 0.0);
        assert!(loss.physics.total >= 0.0);
        assert!(loss.boundary.total >= 0.0);
    }

    #[test]
    fn test_training_history() {
        let mut history = TrainingHistory::new();

        history.add_entry(TrainingEntry {
            epoch: 0,
            total_loss: 1.0,
            physics_loss: 0.5,
            boundary_loss: 0.5,
            momentum_x_loss: 0.2,
            momentum_y_loss: 0.2,
            continuity_loss: 0.1,
            learning_rate: 0.001,
        });

        history.add_entry(TrainingEntry {
            epoch: 1,
            total_loss: 0.5,
            physics_loss: 0.25,
            boundary_loss: 0.25,
            momentum_x_loss: 0.1,
            momentum_y_loss: 0.1,
            continuity_loss: 0.05,
            learning_rate: 0.001,
        });

        assert_eq!(history.entries().len(), 2);
        assert!((history.best_loss() - 0.5).abs() < f64::EPSILON);
        assert_eq!(history.latest().unwrap().epoch, 1);
    }

    #[test]
    fn test_convergence_check() {
        let mut history = TrainingHistory::new();

        // Add entries with decreasing loss
        for i in 0..10 {
            history.add_entry(TrainingEntry {
                epoch: i,
                total_loss: 1.0 / (i + 1) as f64,
                physics_loss: 0.0,
                boundary_loss: 0.0,
                momentum_x_loss: 0.0,
                momentum_y_loss: 0.0,
                continuity_loss: 0.0,
                learning_rate: 0.001,
            });
        }

        // Should not be converged with high variance
        assert!(!history.is_converged(0.01, 5));

        // Add stable entries
        for i in 10..20 {
            history.add_entry(TrainingEntry {
                epoch: i,
                total_loss: 0.1,
                physics_loss: 0.0,
                boundary_loss: 0.0,
                momentum_x_loss: 0.0,
                momentum_y_loss: 0.0,
                continuity_loss: 0.0,
                learning_rate: 0.001,
            });
        }

        // Should be converged now
        assert!(history.is_converged(0.01, 5));
    }

    #[test]
    fn test_loss_curve() {
        let mut history = TrainingHistory::new();

        for i in 0..5 {
            history.add_entry(TrainingEntry {
                epoch: i,
                total_loss: (5 - i) as f64,
                physics_loss: 0.0,
                boundary_loss: 0.0,
                momentum_x_loss: 0.0,
                momentum_y_loss: 0.0,
                continuity_loss: 0.0,
                learning_rate: 0.001,
            });
        }

        let curve = history.loss_curve();
        assert_eq!(curve.len(), 5);
        assert_eq!(curve[0], (0, 5.0));
        assert_eq!(curve[4], (4, 1.0));
    }

    #[test]
    fn test_record_step() {
        let mut trainer = create_test_trainer();
        let model = VesselPinn::new(trainer.config().clone());
        let loss = trainer.compute_loss(&model, 0.0);

        trainer.record_step(0, &loss, 0.001);

        assert_eq!(trainer.history().entries().len(), 1);
        assert_eq!(trainer.history().latest().unwrap().epoch, 0);
    }
}