rustytorch/demos/piddm-shared/src/ipc.rs

//! IPC message types for PIDDM demo.

use serde::{Deserialize, Serialize};

use crate::config::{PdeType, PiddmSamplingConfig, PiddmTrainingConfig, SchedulerType};

/// Training progress update.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingProgress {
    /// Current epoch
    pub epoch: u32,
    /// Total epochs
    pub total_epochs: u32,
    /// Current batch within epoch
    pub batch: u32,
    /// Total batches per epoch
    pub total_batches: u32,
    /// Current diffusion loss
    pub diffusion_loss: f64,
    /// Current physics loss
    pub physics_loss: f64,
    /// Combined loss
    pub total_loss: f64,
    /// Current learning rate
    pub learning_rate: f64,
    /// Training samples per second
    pub samples_per_second: f64,
    /// Estimated time remaining in seconds
    pub eta_seconds: f64,
    /// Current device being used
    pub device: String,
}

impl Default for TrainingProgress {
    fn default() -> Self {
        Self {
            epoch: 0,
            total_epochs: 0,
            batch: 0,
            total_batches: 0,
            diffusion_loss: 0.0,
            physics_loss: 0.0,
            total_loss: 0.0,
            learning_rate: 0.0,
            samples_per_second: 0.0,
            eta_seconds: 0.0,
            device: "CPU".to_string(),
        }
    }
}

/// Sampling progress update.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SamplingProgress {
    /// Current denoising step
    pub step: u32,
    /// Total denoising steps
    pub total_steps: u32,
    /// Current sample index
    pub sample: u32,
    /// Total samples
    pub total_samples: u32,
    /// Current noise level (sigma)
    pub noise_level: f64,
    /// Physics residual of current sample
    pub physics_residual: f64,
}

impl Default for SamplingProgress {
    fn default() -> Self {
        Self {
            step: 0,
            total_steps: 0,
            sample: 0,
            total_samples: 0,
            noise_level: 0.0,
            physics_residual: 0.0,
        }
    }
}

/// Training result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingResult {
    /// Final diffusion loss
    pub final_diffusion_loss: f64,
    /// Final physics loss
    pub final_physics_loss: f64,
    /// Total training time in seconds
    pub training_time_seconds: f64,
    /// Path to saved model weights
    pub weights_path: Option<String>,
    /// Training history (loss per epoch)
    pub loss_history: Vec<LossRecord>,
}

/// Loss record for a single epoch.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LossRecord {
    /// Epoch number
    pub epoch: u32,
    /// Diffusion loss
    pub diffusion_loss: f64,
    /// Physics loss
    pub physics_loss: f64,
    /// Total loss
    pub total_loss: f64,
}

/// Generated sample from diffusion model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneratedSample {
    /// Sample index
    pub index: u32,
    /// Field values as flattened array [H, W]
    pub field: Vec<f32>,
    /// Grid resolution
    pub resolution: u32,
    /// Physics residual (L2 norm of PDE residual)
    pub physics_residual: f64,
    /// Maximum value in field
    pub max_value: f64,
    /// Minimum value in field
    pub min_value: f64,
}

/// Sampling result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SamplingResult {
    /// Generated samples
    pub samples: Vec<GeneratedSample>,
    /// Total sampling time in seconds
    pub sampling_time_seconds: f64,
    /// Average physics residual across samples
    pub avg_physics_residual: f64,
    /// PDE type used
    pub pde_type: PdeType,
}

/// Request to start training.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StartTrainingRequest {
    /// Training configuration
    pub config: PiddmTrainingConfig,
}

/// Request to start sampling.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StartSamplingRequest {
    /// Sampling configuration
    pub config: PiddmSamplingConfig,
    /// PDE type (must match trained model)
    pub pde_type: PdeType,
    /// Optional path to model weights (uses latest if not specified)
    pub weights_path: Option<String>,
}

/// Demo state for UI synchronization.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PiddmDemoState {
    /// Whether model is loaded
    pub model_loaded: bool,
    /// Current PDE type
    pub pde_type: PdeType,
    /// Current scheduler type
    pub scheduler: SchedulerType,
    /// Grid resolution
    pub resolution: u32,
    /// Whether training is in progress
    pub is_training: bool,
    /// Whether sampling is in progress
    pub is_sampling: bool,
    /// Last training result
    pub last_training_result: Option<TrainingResult>,
    /// Last generated samples
    pub last_samples: Vec<GeneratedSample>,
}

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

    #[test]
    fn test_training_progress_default() {
        let progress = TrainingProgress::default();
        assert_eq!(progress.epoch, 0);
        assert_eq!(progress.device, "CPU");
    }

    #[test]
    fn test_sampling_progress_default() {
        let progress = SamplingProgress::default();
        assert_eq!(progress.step, 0);
        assert_eq!(progress.noise_level, 0.0);
    }

    #[test]
    fn test_demo_state_default() {
        let state = PiddmDemoState::default();
        assert!(!state.model_loaded);
        assert!(!state.is_training);
    }

    #[test]
    fn test_serialization() {
        let sample = GeneratedSample {
            index: 0,
            field: vec![0.0, 1.0, 2.0, 3.0],
            resolution: 2,
            physics_residual: 0.01,
            max_value: 3.0,
            min_value: 0.0,
        };

        let json = serde_json::to_string(&sample).unwrap();
        let deserialized: GeneratedSample = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.resolution, 2);
    }
}