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

//! MRE solver configuration with non-dimensionalization
//!
//! Non-dimensionalization is critical for numerical stability:
//! - Coordinates scaled by domain size
//! - Stiffness scaled by reference value (~3 kPa)
//! - Wave amplitude scaled by typical MRE displacement (~1 um)

use mre_shared::TissueProperties;
use serde::{Deserialize, Serialize};

/// Non-dimensionalization scales for numerical stability
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct NonDimScales {
    /// Reference length (domain size) in meters
    pub length_ref: f32,
    /// Reference stiffness in Pascals (not kPa!)
    pub stiffness_ref: f32,
    /// Reference wave amplitude in meters
    pub wave_ref: f32,
}

impl NonDimScales {
    /// Create scales from domain and typical stiffness
    #[must_use]
    pub fn from_domain(domain_size: f32, typical_stiffness_kpa: f32) -> Self {
        Self {
            length_ref: domain_size,
            stiffness_ref: typical_stiffness_kpa * 1000.0, // Convert to Pa
            wave_ref: 1e-6,                                // 1 micrometer typical MRE displacement
        }
    }

    /// Non-dimensionalize coordinates: x' = x / L_ref
    #[must_use]
    pub fn nondim_length(&self, x: f32) -> f32 {
        x / self.length_ref
    }

    /// Dimensionalize coordinates: x = x' * L_ref
    #[must_use]
    pub fn dim_length(&self, x_nondim: f32) -> f32 {
        x_nondim * self.length_ref
    }

    /// Non-dimensionalize stiffness: mu' = mu / mu_ref
    #[must_use]
    pub fn nondim_stiffness(&self, mu_pa: f32) -> f32 {
        mu_pa / self.stiffness_ref
    }

    /// Dimensionalize stiffness: mu = mu' * mu_ref (returns Pa)
    #[must_use]
    pub fn dim_stiffness_pa(&self, mu_nondim: f32) -> f32 {
        mu_nondim * self.stiffness_ref
    }

    /// Dimensionalize stiffness to kPa
    #[must_use]
    pub fn dim_stiffness_kpa(&self, mu_nondim: f32) -> f32 {
        self.dim_stiffness_pa(mu_nondim) / 1000.0
    }

    /// Non-dimensionalize stiffness from kPa
    #[must_use]
    pub fn nondim_stiffness_kpa(&self, mu_kpa: f32) -> f32 {
        self.nondim_stiffness(mu_kpa * 1000.0)
    }

    /// Non-dimensionalize wave amplitude
    #[must_use]
    pub fn nondim_wave(&self, u: f32) -> f32 {
        u / self.wave_ref
    }

    /// Dimensionalize wave amplitude
    #[must_use]
    pub fn dim_wave(&self, u_nondim: f32) -> f32 {
        u_nondim * self.wave_ref
    }

    /// Get the non-dimensional rho*omega^2 coefficient
    /// In non-dim form: rho' * omega'^2 = (rho * L^2 * omega^2) / mu_ref
    #[must_use]
    pub fn nondim_rho_omega_sq(&self, tissue: &TissueProperties) -> f32 {
        let rho = tissue.density;
        let omega_sq = tissue.omega_squared();
        (rho * self.length_ref * self.length_ref * omega_sq) / self.stiffness_ref
    }
}

impl Default for NonDimScales {
    fn default() -> Self {
        Self::from_domain(0.1, 3.0) // 10cm domain, 3 kPa reference
    }
}

/// Complete MRE solver configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MreConfig {
    // --- Wave Net Architecture ---
    /// Number of hidden layers
    pub wave_net_layers: usize,
    /// Hidden dimension per layer
    pub wave_net_hidden: usize,
    /// Number of Fourier feature frequencies
    pub fourier_features: usize,
    /// Scale for Fourier features (controls frequency range)
    pub fourier_scale: f32,

    // --- Stiffness Texture ---
    /// Grid resolution for stiffness map (width)
    pub stiffness_nx: usize,
    /// Grid resolution for stiffness map (height)
    pub stiffness_ny: usize,
    /// Minimum stiffness (non-dimensionalized) for stability
    pub min_stiffness: f32,

    // --- Physics ---
    /// Tissue properties (density, frequency)
    pub tissue: TissueProperties,

    // --- Non-dimensionalization ---
    /// Scaling factors for numerical stability
    pub nondim: NonDimScales,

    // --- Training ---
    /// Learning rate for Wave Net parameters
    pub learning_rate_wave: f32,
    /// Learning rate for stiffness texture (typically higher)
    pub learning_rate_stiffness: f32,
    /// Weight for physics loss (PDE residual)
    pub physics_weight: f32,
    /// Weight for data loss (wave field matching)
    pub data_weight: f32,
    /// Weight for Total Variation regularization on stiffness
    pub tv_weight: f32,

    // --- Collocation ---
    /// Number of collocation points for physics loss
    pub num_collocation: usize,
    /// Number of data points to sample per step
    pub num_data_points: usize,

    // --- Domain ---
    /// Physical domain width in meters
    pub domain_width: f32,
    /// Physical domain height in meters
    pub domain_height: f32,
}

impl Default for MreConfig {
    fn default() -> Self {
        let domain = 0.1; // 10 cm
        Self {
            // Wave Net
            wave_net_layers: 4,
            wave_net_hidden: 128,
            fourier_features: 32,
            fourier_scale: 10.0,

            // Stiffness
            stiffness_nx: 64,
            stiffness_ny: 64,
            min_stiffness: 0.1, // 10% of reference

            // Physics
            tissue: TissueProperties::soft_tissue(),

            // Non-dim
            nondim: NonDimScales::from_domain(domain, 3.0),

            // Training
            learning_rate_wave: 1e-4,
            learning_rate_stiffness: 1e-3,
            physics_weight: 1.0,
            data_weight: 10.0,
            tv_weight: 0.01,

            // Collocation
            num_collocation: 2000,
            num_data_points: 1000,

            // Domain
            domain_width: domain,
            domain_height: domain,
        }
    }
}

impl MreConfig {
    /// Create a new config with default values
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Builder: set wave net depth
    #[must_use]
    pub const fn with_wave_net_layers(mut self, layers: usize) -> Self {
        self.wave_net_layers = layers;
        self
    }

    /// Builder: set wave net width
    #[must_use]
    pub const fn with_wave_net_hidden(mut self, hidden: usize) -> Self {
        self.wave_net_hidden = hidden;
        self
    }

    /// Builder: set Fourier features
    #[must_use]
    pub const fn with_fourier_features(mut self, features: usize) -> Self {
        self.fourier_features = features;
        self
    }

    /// Builder: set stiffness resolution
    #[must_use]
    pub const fn with_stiffness_resolution(mut self, nx: usize, ny: usize) -> Self {
        self.stiffness_nx = nx;
        self.stiffness_ny = ny;
        self
    }

    /// Builder: set tissue properties
    #[must_use]
    pub const fn with_tissue(mut self, tissue: TissueProperties) -> Self {
        self.tissue = tissue;
        self
    }

    /// Builder: set learning rates
    #[must_use]
    pub const fn with_learning_rates(mut self, wave_lr: f32, stiffness_lr: f32) -> Self {
        self.learning_rate_wave = wave_lr;
        self.learning_rate_stiffness = stiffness_lr;
        self
    }

    /// Builder: set loss weights
    #[must_use]
    pub const fn with_loss_weights(mut self, physics: f32, data: f32) -> Self {
        self.physics_weight = physics;
        self.data_weight = data;
        self
    }

    /// Builder: set domain size
    #[must_use]
    pub fn with_domain(mut self, width: f32, height: f32) -> Self {
        self.domain_width = width;
        self.domain_height = height;
        self.nondim = NonDimScales::from_domain(width.max(height), 3.0);
        self
    }

    /// Input dimension for Wave Net (2: x, y)
    #[must_use]
    pub const fn input_dim(&self) -> usize {
        2
    }

    /// Output dimension for Wave Net (2: real, imag)
    #[must_use]
    pub const fn output_dim(&self) -> usize {
        2
    }

    /// Fourier encoding output dimension (2 * fourier_features for sin + cos)
    #[must_use]
    pub const fn fourier_output_dim(&self) -> usize {
        2 * self.fourier_features
    }

    /// Get the non-dimensional rho*omega^2 coefficient
    #[must_use]
    pub fn nondim_rho_omega_sq(&self) -> f32 {
        self.nondim.nondim_rho_omega_sq(&self.tissue)
    }

    /// Convert physical coordinates to non-dimensional
    #[must_use]
    pub fn to_nondim_coords(&self, x: f32, y: f32) -> (f32, f32) {
        (self.nondim.nondim_length(x), self.nondim.nondim_length(y))
    }

    /// Configuration for high accuracy (slower)
    #[must_use]
    pub fn high_accuracy() -> Self {
        Self::default()
            .with_wave_net_layers(6)
            .with_wave_net_hidden(256)
            .with_fourier_features(64)
            .with_stiffness_resolution(128, 128)
    }

    /// Configuration for fast inference (demo)
    #[must_use]
    pub fn fast() -> Self {
        Self::default()
            .with_wave_net_layers(3)
            .with_wave_net_hidden(64)
            .with_fourier_features(16)
            .with_stiffness_resolution(32, 32)
    }
}

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

    #[test]
    fn test_nondim_roundtrip() {
        let scales = NonDimScales::from_domain(0.1, 3.0);

        // Length roundtrip
        let x = 0.05;
        let x_nd = scales.nondim_length(x);
        let x_back = scales.dim_length(x_nd);
        assert!((x - x_back).abs() < 1e-6);

        // Stiffness roundtrip
        let mu = 5.0; // kPa
        let mu_nd = scales.nondim_stiffness_kpa(mu);
        let mu_back = scales.dim_stiffness_kpa(mu_nd);
        assert!((mu - mu_back).abs() < 1e-6);
    }

    #[test]
    fn test_config_builder() {
        let config = MreConfig::new()
            .with_wave_net_layers(6)
            .with_wave_net_hidden(512)
            .with_stiffness_resolution(128, 128);

        assert_eq!(config.wave_net_layers, 6);
        assert_eq!(config.wave_net_hidden, 512);
        assert_eq!(config.stiffness_nx, 128);
    }

    #[test]
    fn test_rho_omega_sq() {
        let config = MreConfig::default();
        let coeff = config.nondim_rho_omega_sq();
        // Should be positive and reasonable magnitude
        assert!(coeff > 0.0);
        assert!(coeff < 1000.0);
    }
}