rustytorch/demos/pinn-benchmark-shared/src/config.rs

//! Configuration types for PINN benchmark

use serde::{Deserialize, Serialize};

/// Type of PDE problem to solve
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ProblemType {
    /// 1D Burgers equation: `u_t` + u*`u_x` = nu*`u_xx`
    Burgers1D,

    /// 1D Heat equation: `u_t` = alpha*`u_xx`
    Heat1D,

    /// 2D Heat equation: `u_t` = alpha*(`u_xx` + `u_yy`)
    Heat2D,

    /// 2D Poisson equation: `u_xx` + `u_yy` = f
    Poisson2D,

    /// 2D Navier-Stokes equations (advanced)
    NavierStokes2D,
}

impl ProblemType {
    /// Returns the number of input dimensions for this problem
    #[must_use]
    pub const fn input_dim(&self) -> usize {
        match self {
            Self::Burgers1D | Self::Heat1D => 2, // (x, t)
            Self::Heat2D | Self::Poisson2D | Self::NavierStokes2D => 3, // (x, y, t) or (x, y) for Poisson
        }
    }

    /// Returns the number of output dimensions for this problem
    #[must_use]
    pub const fn output_dim(&self) -> usize {
        match self {
            Self::Burgers1D | Self::Heat1D | Self::Heat2D | Self::Poisson2D => 1, // scalar field
            Self::NavierStokes2D => 3, // (u, v, p) - velocity components and pressure
        }
    }

    /// Returns a human-readable name for this problem
    #[must_use]
    pub const fn name(&self) -> &'static str {
        match self {
            Self::Burgers1D => "1D Burgers Equation",
            Self::Heat1D => "1D Heat Equation",
            Self::Heat2D => "2D Heat Equation",
            Self::Poisson2D => "2D Poisson Equation",
            Self::NavierStokes2D => "2D Navier-Stokes",
        }
    }

    /// Returns a brief description of the PDE
    #[must_use]
    pub const fn description(&self) -> &'static str {
        match self {
            Self::Burgers1D => "Nonlinear advection-diffusion equation with shock formation",
            Self::Heat1D => "Linear parabolic PDE modeling heat diffusion",
            Self::Heat2D => "2D linear parabolic PDE modeling heat diffusion",
            Self::Poisson2D => "Elliptic PDE for steady-state problems",
            Self::NavierStokes2D => "Incompressible fluid dynamics equations",
        }
    }
}

/// Configuration for PINN benchmark
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkConfig {
    /// Type of PDE problem to solve
    pub problem: ProblemType,

    /// Hidden layer sizes (e.g., [64, 64, 64])
    pub hidden_layers: Vec<usize>,

    /// Learning rate for optimizer
    pub learning_rate: f64,

    /// Number of training epochs
    pub num_epochs: usize,

    /// Number of collocation points for physics loss
    pub num_collocation_points: usize,

    /// Number of boundary/initial condition points
    pub num_boundary_points: usize,

    /// Device to use: "cpu", "cuda", "metal"
    pub device: String,
}

impl BenchmarkConfig {
    /// Creates a default configuration for the given problem type
    #[must_use]
    pub fn default_for_problem(problem: ProblemType) -> Self {
        Self {
            problem,
            hidden_layers: vec![64, 64, 64],
            learning_rate: 0.001,
            num_epochs: 1000,
            num_collocation_points: 10_000,
            num_boundary_points: 100,
            device: "cpu".to_string(),
        }
    }

    /// Validates the configuration
    ///
    /// # Errors
    ///
    /// Returns an error if the configuration is invalid
    pub fn validate(&self) -> Result<(), String> {
        if self.hidden_layers.is_empty() {
            return Err("Hidden layers cannot be empty".to_string());
        }

        if self.hidden_layers.contains(&0) {
            return Err("Hidden layer sizes must be positive".to_string());
        }

        if self.learning_rate <= 0.0 || self.learning_rate > 1.0 {
            return Err("Learning rate must be in (0, 1]".to_string());
        }

        if self.num_epochs == 0 {
            return Err("Number of epochs must be positive".to_string());
        }

        if self.num_collocation_points == 0 {
            return Err("Number of collocation points must be positive".to_string());
        }

        if self.num_boundary_points == 0 {
            return Err("Number of boundary points must be positive".to_string());
        }

        Ok(())
    }
}

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

    #[test]
    fn test_problem_type_dimensions() {
        assert_eq!(ProblemType::Burgers1D.input_dim(), 2);
        assert_eq!(ProblemType::Burgers1D.output_dim(), 1);

        assert_eq!(ProblemType::Heat2D.input_dim(), 3);
        assert_eq!(ProblemType::Heat2D.output_dim(), 1);

        assert_eq!(ProblemType::NavierStokes2D.input_dim(), 3);
        assert_eq!(ProblemType::NavierStokes2D.output_dim(), 3);
    }

    #[test]
    fn test_problem_type_metadata() {
        let problem = ProblemType::Burgers1D;
        assert_eq!(problem.name(), "1D Burgers Equation");
        assert!(!problem.description().is_empty());
    }

    #[test]
    fn test_benchmark_config_default() {
        let config = BenchmarkConfig::default_for_problem(ProblemType::Heat1D);
        assert_eq!(config.problem, ProblemType::Heat1D);
        assert!(!config.hidden_layers.is_empty());
        assert!(config.learning_rate > 0.0);
    }

    #[test]
    fn test_benchmark_config_validation_success() {
        let config = BenchmarkConfig::default_for_problem(ProblemType::Burgers1D);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_benchmark_config_validation_empty_layers() {
        let mut config = BenchmarkConfig::default_for_problem(ProblemType::Heat1D);
        config.hidden_layers = vec![];
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_benchmark_config_validation_zero_layer_size() {
        let mut config = BenchmarkConfig::default_for_problem(ProblemType::Heat1D);
        config.hidden_layers = vec![64, 0, 32];
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_benchmark_config_validation_invalid_lr() {
        let mut config = BenchmarkConfig::default_for_problem(ProblemType::Heat1D);
        config.learning_rate = 0.0;
        assert!(config.validate().is_err());

        config.learning_rate = 1.5;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_benchmark_config_validation_zero_epochs() {
        let mut config = BenchmarkConfig::default_for_problem(ProblemType::Heat1D);
        config.num_epochs = 0;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_benchmark_config_validation_zero_points() {
        let mut config = BenchmarkConfig::default_for_problem(ProblemType::Heat1D);
        config.num_collocation_points = 0;
        assert!(config.validate().is_err());

        config.num_collocation_points = 100;
        config.num_boundary_points = 0;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_serde_problem_type() {
        let problem = ProblemType::Poisson2D;
        let json = serde_json::to_string(&problem).expect("Failed to serialize");
        let deserialized: ProblemType = serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(problem, deserialized);
    }

    #[test]
    fn test_serde_benchmark_config() {
        let config = BenchmarkConfig {
            problem: ProblemType::Heat2D,
            hidden_layers: vec![32, 32, 32],
            learning_rate: 0.0005,
            num_epochs: 500,
            num_collocation_points: 5000,
            num_boundary_points: 200,
            device: "cuda".to_string(),
        };

        let json = serde_json::to_string(&config).expect("Failed to serialize");
        let deserialized: BenchmarkConfig =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(config.problem, deserialized.problem);
        assert_eq!(config.hidden_layers, deserialized.hidden_layers);
        assert_eq!(config.device, deserialized.device);
    }
}