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

//! Configuration for the bioheat solver

use bioheat_shared::{BioheatParams, BoundingBox3D, Point3D, ProbeGeometry, TissueType};
use serde::{Deserialize, Serialize};

/// Configuration for the bioheat PINN solver
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BioheatConfig {
    /// Bioheat equation parameters (tissue properties, blood properties)
    pub physics: BioheatParams,

    /// Probe configuration
    pub probe: ProbeGeometry,

    /// Probe power in Watts
    pub probe_power: f32,

    /// Domain bounds in meters
    pub domain: BoundingBox3D,

    /// Network configuration
    pub network: NetworkConfig,

    /// Training configuration
    pub training: TrainingConfig,

    /// Simulation time parameters
    pub time: TimeConfig,
}

/// Neural network architecture configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkConfig {
    /// Number of Fourier features for positional encoding
    pub fourier_features: usize,

    /// Scale for Fourier feature frequencies
    pub fourier_scale: f32,

    /// Hidden layer sizes
    pub hidden_layers: Vec<usize>,

    /// Activation function
    pub activation: Activation,
}

/// Available activation functions
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
pub enum Activation {
    #[default]
    Tanh,
    Swish,
    Gelu,
    Sin,
}

/// Training hyperparameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
    /// Learning rate
    pub learning_rate: f32,

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

    /// Number of boundary points
    pub num_boundary: usize,

    /// Number of initial condition points
    pub num_initial: usize,

    /// Loss weights
    pub weights: LossWeights,
}

/// Weights for different loss components
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LossWeights {
    /// Weight for physics (PDE residual) loss
    pub physics: f32,

    /// Weight for boundary condition loss
    pub boundary: f32,

    /// Weight for initial condition loss
    pub initial: f32,

    /// Weight for probe heat source condition
    pub probe: f32,
}

/// Time domain configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeConfig {
    /// Start time (usually 0)
    pub t_start: f32,

    /// End time in seconds
    pub t_end: f32,

    /// Time step for visualization
    pub dt_vis: f32,
}

impl Default for NetworkConfig {
    fn default() -> Self {
        Self {
            fourier_features: 64,
            fourier_scale: 4.0,
            hidden_layers: vec![128, 128, 128, 64],
            activation: Activation::Tanh,
        }
    }
}

impl Default for LossWeights {
    fn default() -> Self {
        Self {
            physics: 1.0,
            boundary: 10.0, // Higher weight for boundary conditions
            initial: 10.0,  // Higher weight for initial condition
            probe: 5.0,
        }
    }
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            learning_rate: 1e-3,
            num_collocation: 4096,
            num_boundary: 512,
            num_initial: 512,
            weights: LossWeights::default(),
        }
    }
}

impl Default for TimeConfig {
    fn default() -> Self {
        Self {
            t_start: 0.0,
            t_end: 300.0, // 5 minutes treatment
            dt_vis: 10.0, // Visualize every 10 seconds
        }
    }
}

impl BioheatConfig {
    /// Create default configuration for liver ablation
    #[must_use]
    pub fn liver_default() -> Self {
        Self {
            physics: BioheatParams::liver_ablation(),
            probe: ProbeGeometry::rf_needle(Point3D::origin()),
            probe_power: 15.0,
            domain: BoundingBox3D::from_dimensions(0.1, 0.1, 0.1),
            network: NetworkConfig::default(),
            training: TrainingConfig::default(),
            time: TimeConfig::default(),
        }
    }

    /// Create configuration for tumor ablation
    #[must_use]
    pub fn tumor_default() -> Self {
        Self {
            physics: BioheatParams::liver_tumor(),
            probe: ProbeGeometry::rf_needle(Point3D::origin()),
            probe_power: 20.0,
            domain: BoundingBox3D::from_dimensions(0.1, 0.1, 0.1),
            network: NetworkConfig::default(),
            training: TrainingConfig::default(),
            time: TimeConfig {
                t_end: 600.0, // 10 minutes for tumor
                ..Default::default()
            },
        }
    }

    /// Create configuration from tissue type
    #[must_use]
    pub fn from_tissue_type(tissue_type: TissueType) -> Self {
        Self {
            physics: BioheatParams::from_tissue_type(tissue_type),
            ..Self::liver_default()
        }
    }

    /// Create a minimal configuration for testing/benchmarking
    #[must_use]
    pub fn benchmark() -> Self {
        Self {
            physics: BioheatParams::liver_ablation(),
            probe: ProbeGeometry::rf_needle(Point3D::origin()),
            probe_power: 15.0,
            domain: BoundingBox3D::from_dimensions(0.05, 0.05, 0.05),
            network: NetworkConfig {
                fourier_features: 32,
                hidden_layers: vec![64, 64],
                ..Default::default()
            },
            training: TrainingConfig {
                num_collocation: 1024,
                num_boundary: 128,
                num_initial: 128,
                ..Default::default()
            },
            time: TimeConfig {
                t_end: 60.0,
                ..Default::default()
            },
        }
    }
}

impl Default for BioheatConfig {
    fn default() -> Self {
        Self::liver_default()
    }
}

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

    #[test]
    fn test_default_config() {
        let config = BioheatConfig::default();
        assert!(config.probe_power > 0.0);
        assert!(config.time.t_end > 0.0);
    }

    #[test]
    fn test_benchmark_config() {
        let config = BioheatConfig::benchmark();
        assert!(config.network.hidden_layers.len() < 4);
        assert!(config.training.num_collocation < 2048);
    }

    #[test]
    fn test_serialize_config() {
        let config = BioheatConfig::default();
        let json = serde_json::to_string(&config).unwrap();
        let _: BioheatConfig = serde_json::from_str(&json).unwrap();
    }
}