rustytorch/demos/rtx-neural-operator-demo/src/inference.rs

//! Neural operator inference engine
//!
//! Provides the main demo functionality for FNO inference.

use std::path::Path;
use std::time::Instant;

use rtx_backend_cpu::{CpuBackend, CpuDevice};
use rtx_neural_operator::{FNO2d, load_fno2d_weights};
use rtx_neural_operator_shared::{
    config::{PDEConfig, PDEType},
    error::{NeuralOperatorError, Result},
    ipc::{ModelInfo, PerformanceMetrics, SolutionData},
};
use rtx_nn::GenericModule4D;
use rtx_tensor::GenericTensor;

/// Result from a solve operation
#[derive(Debug, Clone)]
pub struct SolveResult {
    /// Solution field (flattened [H, W])
    pub solution: Vec<f32>,
    /// Inference time in milliseconds
    pub inference_time_ms: f64,
}

/// Neural Operator Demo engine
///
/// Manages FNO model loading and inference for interactive PDE solving.
pub struct NeuralOperatorDemo {
    model: Option<FNO2d<CpuBackend>>,
    config: Option<PDEConfig>,
    metrics: PerformanceMetrics,
    last_solution: Option<Vec<f32>>,
}

impl NeuralOperatorDemo {
    /// Creates a new demo instance
    #[must_use]
    pub fn new() -> Self {
        Self {
            model: None,
            config: None,
            metrics: PerformanceMetrics::new(),
            last_solution: None,
        }
    }

    /// Initializes the demo with a PDE configuration
    ///
    /// # Arguments
    ///
    /// * `config` - PDE configuration specifying type and resolution
    ///
    /// # Errors
    ///
    /// Returns error if weight file not found or model loading fails.
    pub fn initialize(&mut self, config: PDEConfig) -> Result<()> {
        tracing::info!(
            "Initializing neural operator for {:?} at {}x{} resolution",
            config.pde_type,
            config.resolution,
            config.resolution
        );

        // For now, we'll create a model with default weights
        // In production, this would load from weight files
        self.config = Some(config);
        self.metrics = PerformanceMetrics::new();
        self.last_solution = None;

        Ok(())
    }

    /// Initializes with weights from a file path
    ///
    /// # Arguments
    ///
    /// * `config` - PDE configuration
    /// * `weights_path` - Path to `SafeTensors` weight file
    ///
    /// # Errors
    ///
    /// Returns error if weight loading fails.
    pub fn initialize_with_weights(
        &mut self,
        config: PDEConfig,
        weights_path: impl AsRef<Path>,
    ) -> Result<()> {
        tracing::info!("Loading FNO weights from {:?}", weights_path.as_ref());

        let weights = load_fno2d_weights(weights_path.as_ref())
            .map_err(|e| NeuralOperatorError::weight_load(e.to_string()))?;

        let device = CpuDevice::default();
        let model = FNO2d::from_weights(&weights, &device)
            .map_err(|e| NeuralOperatorError::weight_load(e.to_string()))?;

        self.model = Some(model);
        self.config = Some(config);
        self.metrics = PerformanceMetrics::new();
        self.last_solution = None;

        Ok(())
    }

    /// Solves the PDE with the given input field
    ///
    /// # Arguments
    ///
    /// * `input` - Input field (flattened [H, W] array)
    ///
    /// # Errors
    ///
    /// Returns error if model not initialized or inference fails.
    pub fn solve(&mut self, input: &[f32]) -> Result<SolveResult> {
        let config = self
            .config
            .as_ref()
            .ok_or(NeuralOperatorError::NotInitialized)?;

        let expected_size = (config.resolution * config.resolution) as usize;
        if input.len() != expected_size {
            return Err(NeuralOperatorError::invalid_dimensions(
                format!("[{expected_size}]"),
                format!("[{}]", input.len()),
            ));
        }

        let start = Instant::now();

        // If we have a loaded model, use it
        let solution = if let Some(ref model) = self.model {
            // Create input tensor [1, 1, H, W]
            let h = config.resolution as usize;
            let w = config.resolution as usize;
            let device = CpuDevice::default();

            let tensor: GenericTensor<CpuBackend, 4> =
                GenericTensor::from_slice(input, [1, 1, h, w], &device);

            // Run inference
            let output = model.forward_4d(&tensor);

            // Extract result
            output.to_vec()
        } else {
            // No model loaded - return a simple demo output
            // This is a placeholder that shows the demo UI works
            generate_demo_solution(input, config)
        };

        let inference_time_ms = start.elapsed().as_secs_f64() * 1000.0;

        // Record metrics
        self.metrics.record_inference(inference_time_ms);

        // Cache solution
        self.last_solution = Some(solution.clone());

        Ok(SolveResult {
            solution,
            inference_time_ms,
        })
    }

    /// Creates solution data from a solve result
    #[must_use]
    pub fn create_solution_data(&self, result: &SolveResult) -> Option<SolutionData> {
        let config = self.config.as_ref()?;
        Some(SolutionData::new(
            result.solution.clone(),
            config.resolution,
            config.resolution,
            result.inference_time_ms,
        ))
    }

    /// Returns the last computed solution
    #[must_use]
    pub fn last_solution(&self) -> Option<&Vec<f32>> {
        self.last_solution.as_ref()
    }

    /// Returns current performance metrics
    #[must_use]
    pub fn metrics(&self) -> &PerformanceMetrics {
        &self.metrics
    }

    /// Returns model information
    #[must_use]
    pub fn model_info(&self) -> Option<ModelInfo> {
        let config = self.config.as_ref()?;
        Some(ModelInfo::new(
            format!("FNO2d-{}", config.pde_type.name()),
            config.pde_type.name(),
            config.resolution,
            config.n_modes,
            config.model_width,
            config.n_layers,
        ))
    }

    /// Returns the current configuration
    #[must_use]
    pub fn config(&self) -> Option<&PDEConfig> {
        self.config.as_ref()
    }

    /// Returns whether the model is initialized
    #[must_use]
    pub fn is_initialized(&self) -> bool {
        self.config.is_some()
    }

    /// Returns whether a trained model is loaded
    #[must_use]
    pub fn has_model(&self) -> bool {
        self.model.is_some()
    }

    /// Resets the demo to initial state
    pub fn reset(&mut self) {
        self.model = None;
        self.config = None;
        self.metrics = PerformanceMetrics::new();
        self.last_solution = None;
    }
}

impl Default for NeuralOperatorDemo {
    fn default() -> Self {
        Self::new()
    }
}

/// Generates a demo solution when no model is loaded
///
/// This creates a visually interesting output that demonstrates the UI
/// without requiring a trained model.
fn generate_demo_solution(input: &[f32], config: &PDEConfig) -> Vec<f32> {
    let n = config.resolution as usize;
    let mut solution = vec![0.0; n * n];

    // Create a simple diffusion-like response to the input
    // This mimics what a PDE solver would produce
    for i in 0..n {
        for j in 0..n {
            let idx = i * n + j;
            let x = j as f32 / n as f32;
            let y = i as f32 / n as f32;

            // Combine input with a smooth basis function
            let input_val = input[idx];
            let smooth = (std::f32::consts::PI * x).sin() * (std::f32::consts::PI * y).sin();

            // Different responses for different PDE types
            let response = match config.pde_type {
                PDEType::DarcyFlow => {
                    // Pressure-like response
                    input_val * smooth * 0.5 + (1.0 - x) * 0.3
                }
                PDEType::HeatEquation => {
                    // Temperature diffusion
                    input_val * smooth.powi(2) * 0.8
                }
                PDEType::Poisson => {
                    // Potential field
                    input_val * smooth * 0.6 + smooth * 0.2
                }
                PDEType::NavierStokes => {
                    // Velocity-like field
                    let vortex = ((x - 0.5).powi(2) + (y - 0.5).powi(2)).sqrt();
                    input_val * (1.0 - vortex).max(0.0) * 0.7
                }
            };

            solution[idx] = response;
        }
    }

    solution
}

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

    #[test]
    fn test_demo_creation() {
        let demo = NeuralOperatorDemo::new();
        assert!(!demo.is_initialized());
        assert!(!demo.has_model());
    }

    #[test]
    fn test_initialize() {
        let mut demo = NeuralOperatorDemo::new();
        let config = PDEConfig::darcy(64);
        demo.initialize(config).unwrap();

        assert!(demo.is_initialized());
        assert!(!demo.has_model()); // No weights loaded
    }

    #[test]
    fn test_solve_without_model() {
        let mut demo = NeuralOperatorDemo::new();
        let config = PDEConfig::darcy(32);
        demo.initialize(config).unwrap();

        let input = vec![1.0; 32 * 32];
        let result = demo.solve(&input).unwrap();

        assert_eq!(result.solution.len(), 32 * 32);
        assert!(result.inference_time_ms > 0.0);
    }

    #[test]
    fn test_solve_dimension_mismatch() {
        let mut demo = NeuralOperatorDemo::new();
        let config = PDEConfig::darcy(64);
        demo.initialize(config).unwrap();

        let input = vec![1.0; 32 * 32]; // Wrong size
        let result = demo.solve(&input);

        assert!(result.is_err());
    }

    #[test]
    fn test_solve_not_initialized() {
        let mut demo = NeuralOperatorDemo::new();
        let input = vec![1.0; 64 * 64];
        let result = demo.solve(&input);

        assert!(result.is_err());
    }

    #[test]
    fn test_metrics_recording() {
        let mut demo = NeuralOperatorDemo::new();
        let config = PDEConfig::darcy(16);
        demo.initialize(config).unwrap();

        let input = vec![1.0; 16 * 16];
        demo.solve(&input).unwrap();
        demo.solve(&input).unwrap();

        let metrics = demo.metrics();
        assert_eq!(metrics.inference_count, 2);
        assert!(metrics.avg_inference_time_ms > 0.0);
    }

    #[test]
    fn test_model_info() {
        let mut demo = NeuralOperatorDemo::new();
        let config = PDEConfig::darcy(64)
            .with_modes(12, 12)
            .with_width(32)
            .with_layers(4);
        demo.initialize(config).unwrap();

        let info = demo.model_info().unwrap();
        assert_eq!(info.resolution, 64);
        assert_eq!(info.n_modes, (12, 12));
        assert_eq!(info.model_width, 32);
        assert_eq!(info.n_layers, 4);
    }

    #[test]
    fn test_reset() {
        let mut demo = NeuralOperatorDemo::new();
        let config = PDEConfig::darcy(64);
        demo.initialize(config).unwrap();

        let input = vec![1.0; 64 * 64];
        demo.solve(&input).unwrap();

        demo.reset();

        assert!(!demo.is_initialized());
        assert!(demo.last_solution().is_none());
    }

    #[test]
    fn test_different_pde_types() {
        let mut demo = NeuralOperatorDemo::new();
        let input = vec![1.0; 16 * 16];

        for pde_type in [
            PDEType::DarcyFlow,
            PDEType::HeatEquation,
            PDEType::Poisson,
            PDEType::NavierStokes,
        ] {
            let config = PDEConfig::new(pde_type, 16);
            demo.initialize(config).unwrap();

            let result = demo.solve(&input).unwrap();
            assert_eq!(result.solution.len(), 16 * 16);
        }
    }
}