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

//! PINN configuration parameters
//!
//! This module defines configuration for the hemodynamics PINN model
//! including network architecture, training parameters, and physics settings.

use rtx_hemodynamics_shared::physics::FluidProperties;
use serde::{Deserialize, Serialize};

/// PINN model configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PinnConfig {
    /// Number of hidden layers in the MLP
    pub num_layers: usize,
    /// Number of neurons per hidden layer
    pub hidden_dim: usize,
    /// Dimension of Fourier feature encoding
    pub fourier_dim: usize,
    /// Scale for Fourier features
    pub fourier_scale: f64,
    /// Learning rate
    pub learning_rate: f64,
    /// Weight for physics loss (momentum + continuity)
    pub physics_weight: f64,
    /// Weight for data loss
    pub data_weight: f64,
    /// Weight for boundary condition loss
    pub boundary_weight: f64,
    /// Fluid properties
    pub fluid: FluidProperties,
    /// Number of collocation points for physics loss
    pub num_collocation_points: usize,
    /// Number of boundary points
    pub num_boundary_points: usize,
}

impl Default for PinnConfig {
    fn default() -> Self {
        Self {
            num_layers: 4,
            hidden_dim: 256,
            fourier_dim: 64,
            fourier_scale: 10.0,
            learning_rate: 1e-3,
            physics_weight: 1.0,
            data_weight: 1.0,
            boundary_weight: 10.0,
            fluid: FluidProperties::blood(),
            num_collocation_points: 2000,
            num_boundary_points: 500,
        }
    }
}

impl PinnConfig {
    /// Creates a new configuration with default values
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the number of hidden layers
    #[must_use]
    pub const fn with_layers(mut self, num_layers: usize) -> Self {
        self.num_layers = num_layers;
        self
    }

    /// Sets the hidden dimension
    #[must_use]
    pub const fn with_hidden_dim(mut self, dim: usize) -> Self {
        self.hidden_dim = dim;
        self
    }

    /// Sets the Fourier feature dimension
    #[must_use]
    pub const fn with_fourier_dim(mut self, dim: usize) -> Self {
        self.fourier_dim = dim;
        self
    }

    /// Sets the learning rate
    #[must_use]
    pub const fn with_learning_rate(mut self, lr: f64) -> Self {
        self.learning_rate = lr;
        self
    }

    /// Sets the physics loss weight
    #[must_use]
    pub const fn with_physics_weight(mut self, weight: f64) -> Self {
        self.physics_weight = weight;
        self
    }

    /// Sets the data loss weight
    #[must_use]
    pub const fn with_data_weight(mut self, weight: f64) -> Self {
        self.data_weight = weight;
        self
    }

    /// Sets the boundary loss weight
    #[must_use]
    pub const fn with_boundary_weight(mut self, weight: f64) -> Self {
        self.boundary_weight = weight;
        self
    }

    /// Sets the fluid properties
    #[must_use]
    pub const fn with_fluid(mut self, fluid: FluidProperties) -> Self {
        self.fluid = fluid;
        self
    }

    /// Sets the number of collocation points
    #[must_use]
    pub const fn with_collocation_points(mut self, n: usize) -> Self {
        self.num_collocation_points = n;
        self
    }

    /// Sets the number of boundary points
    #[must_use]
    pub const fn with_boundary_points(mut self, n: usize) -> Self {
        self.num_boundary_points = n;
        self
    }

    /// Returns the total input dimension (x, y, t)
    #[must_use]
    pub const fn input_dim(&self) -> usize {
        3 // x, y, t
    }

    /// Returns the output dimension (u, v, p)
    #[must_use]
    pub const fn output_dim(&self) -> usize {
        3 // u, v, p (velocity components + pressure)
    }

    /// Returns the fluid properties
    #[must_use]
    pub const fn fluid(&self) -> &FluidProperties {
        &self.fluid
    }
}

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

    #[test]
    fn test_config_default() {
        let config = PinnConfig::default();
        assert_eq!(config.num_layers, 4);
        assert_eq!(config.hidden_dim, 256);
        assert_eq!(config.input_dim(), 3);
        assert_eq!(config.output_dim(), 3);
    }

    #[test]
    fn test_config_builder() {
        let config = PinnConfig::new()
            .with_layers(6)
            .with_hidden_dim(512)
            .with_learning_rate(1e-4);

        assert_eq!(config.num_layers, 6);
        assert_eq!(config.hidden_dim, 512);
        assert!((config.learning_rate - 1e-4).abs() < f64::EPSILON);
    }
}