rustytorch/demos/rtx-neuralop-studio-demo/src/lib.rs

//! NeuralOp Studio - Neural Operator Workbench.
//!
//! This demo showcases neural operators for solving PDEs, including
//! FNO, DeepONet, and PINO implementations.

pub mod deeponet;
pub mod fno;
pub mod pde_solver;
pub mod pino;
pub mod sample_data;

use neuralop_studio_shared::{
    EvaluationMetrics, OperatorConfig, OperatorType, PDEDefinition, TrainingConfig,
    TrainingProgress, TrainingResult,
};
use thiserror::Error;

/// Errors that can occur in NeuralOp Studio.
#[derive(Debug, Error)]
pub enum NeuralOpError {
    /// Invalid PDE configuration.
    #[error("Invalid PDE configuration: {0}")]
    InvalidPDE(String),

    /// Invalid operator configuration.
    #[error("Invalid operator configuration: {0}")]
    InvalidOperator(String),

    /// Training failed.
    #[error("Training failed: {0}")]
    TrainingFailed(String),

    /// Prediction failed.
    #[error("Prediction failed: {0}")]
    PredictionFailed(String),

    /// Unsupported feature.
    #[error("Unsupported feature: {0}")]
    Unsupported(String),
}

/// Neural operator trainer.
pub trait NeuralOperatorTrainer: std::fmt::Debug {
    /// Train the neural operator.
    fn train(
        &mut self,
        pde: &PDEDefinition,
        config: &TrainingConfig,
        progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
    ) -> Result<TrainingResult, NeuralOpError>;

    /// Predict solution.
    fn predict(&self, input: &[f64], query_points: &[Vec<f64>]) -> Result<Vec<f64>, NeuralOpError>;

    /// Evaluate on test data.
    fn evaluate(&self, test_inputs: &[Vec<f64>], test_outputs: &[Vec<f64>]) -> EvaluationMetrics;

    /// Get number of parameters.
    fn num_parameters(&self) -> usize;

    /// Get operator type.
    fn operator_type(&self) -> OperatorType;
}

/// Main NeuralOp Studio system.
#[derive(Debug)]
pub struct NeuralOpStudio {
    /// Current operator.
    operator: Option<Box<dyn NeuralOperatorTrainer>>,
    /// Current PDE definition.
    pde: Option<PDEDefinition>,
    /// Training history.
    training_history: Vec<TrainingResult>,
}

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

impl NeuralOpStudio {
    /// Create a new NeuralOp Studio.
    #[must_use]
    pub fn new() -> Self {
        Self {
            operator: None,
            pde: None,
            training_history: vec![],
        }
    }

    /// Set the PDE to solve.
    pub fn set_pde(&mut self, pde: PDEDefinition) -> Result<(), NeuralOpError> {
        // Validate PDE
        if pde.domain.dimensions == 0 {
            return Err(NeuralOpError::InvalidPDE(
                "Domain must have at least 1 dimension".to_string(),
            ));
        }
        if pde.domain.resolution.len() != pde.domain.dimensions {
            return Err(NeuralOpError::InvalidPDE(
                "Resolution must match dimensions".to_string(),
            ));
        }
        self.pde = Some(pde);
        Ok(())
    }

    /// Create an operator with the given configuration.
    pub fn create_operator(&mut self, config: OperatorConfig) -> Result<(), NeuralOpError> {
        let operator: Box<dyn NeuralOperatorTrainer> = match config.operator_type {
            OperatorType::FNO => Box::new(fno::FourierNeuralOperator::new(config)),
            OperatorType::DeepONet => Box::new(deeponet::DeepONet::new(config)),
            OperatorType::PINO => Box::new(pino::PhysicsInformedNO::new(config)),
            OperatorType::SINO => {
                return Err(NeuralOpError::Unsupported(
                    "SINO not yet implemented".to_string(),
                ));
            }
            OperatorType::GalerkinTransformer => {
                return Err(NeuralOpError::Unsupported(
                    "Galerkin Transformer not yet implemented".to_string(),
                ));
            }
            OperatorType::MPNO => {
                return Err(NeuralOpError::Unsupported(
                    "MPNO not yet implemented".to_string(),
                ));
            }
        };

        self.operator = Some(operator);
        Ok(())
    }

    /// Train the operator.
    pub fn train(
        &mut self,
        config: &TrainingConfig,
        progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
    ) -> Result<TrainingResult, NeuralOpError> {
        let pde = self
            .pde
            .as_ref()
            .ok_or_else(|| NeuralOpError::InvalidPDE("No PDE set".to_string()))?;

        let operator = self
            .operator
            .as_mut()
            .ok_or_else(|| NeuralOpError::InvalidOperator("No operator created".to_string()))?;

        let result = operator.train(pde, config, progress_callback)?;
        self.training_history.push(result.clone());
        Ok(result)
    }

    /// Predict solution for given input.
    pub fn predict(
        &self,
        input: &[f64],
        query_points: &[Vec<f64>],
    ) -> Result<Vec<f64>, NeuralOpError> {
        let operator = self
            .operator
            .as_ref()
            .ok_or_else(|| NeuralOpError::InvalidOperator("No operator created".to_string()))?;

        operator.predict(input, query_points)
    }

    /// Get the current PDE.
    #[must_use]
    pub fn pde(&self) -> Option<&PDEDefinition> {
        self.pde.as_ref()
    }

    /// Get the training history.
    #[must_use]
    pub fn training_history(&self) -> &[TrainingResult] {
        &self.training_history
    }

    /// Get the current operator type.
    #[must_use]
    pub fn current_operator_type(&self) -> Option<OperatorType> {
        self.operator.as_ref().map(|op| op.operator_type())
    }
}

/// Run the full demo.
pub fn run_demo() -> Result<TrainingResult, NeuralOpError> {
    use neuralop_studio_shared::{
        sample_fno_config, sample_poisson_problem, sample_training_config,
    };

    let mut studio = NeuralOpStudio::new();

    // Set up problem
    studio.set_pde(sample_poisson_problem())?;
    studio.create_operator(sample_fno_config())?;

    // Train with a simple configuration
    let mut config = sample_training_config();
    config.epochs = 10; // Quick demo
    config.num_train_samples = 100;
    config.num_val_samples = 20;
    config.num_test_samples = 20;

    // Train
    let result = studio.train(&config, None)?;

    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use neuralop_studio_shared::{
        sample_fno_config, sample_poisson_problem, sample_training_config,
    };

    #[test]
    fn test_studio_creation() {
        let studio = NeuralOpStudio::new();
        assert!(studio.pde().is_none());
        assert!(studio.current_operator_type().is_none());
    }

    #[test]
    fn test_set_pde() {
        let mut studio = NeuralOpStudio::new();
        let pde = sample_poisson_problem();
        assert!(studio.set_pde(pde).is_ok());
        assert!(studio.pde().is_some());
    }

    #[test]
    fn test_create_fno() {
        let mut studio = NeuralOpStudio::new();
        let config = sample_fno_config();
        assert!(studio.create_operator(config).is_ok());
        assert_eq!(studio.current_operator_type(), Some(OperatorType::FNO));
    }

    #[test]
    fn test_create_deeponet() {
        let mut studio = NeuralOpStudio::new();
        let config = neuralop_studio_shared::sample_deeponet_config();
        assert!(studio.create_operator(config).is_ok());
        assert_eq!(studio.current_operator_type(), Some(OperatorType::DeepONet));
    }

    #[test]
    fn test_create_pino() {
        let mut studio = NeuralOpStudio::new();
        let config = neuralop_studio_shared::sample_pino_config();
        assert!(studio.create_operator(config).is_ok());
        assert_eq!(studio.current_operator_type(), Some(OperatorType::PINO));
    }

    #[test]
    fn test_train() {
        let mut studio = NeuralOpStudio::new();
        studio.set_pde(sample_poisson_problem()).unwrap();
        studio.create_operator(sample_fno_config()).unwrap();

        let mut config = sample_training_config();
        config.epochs = 5;
        config.num_train_samples = 50;
        config.num_val_samples = 10;
        config.num_test_samples = 10;

        let result = studio.train(&config, None);
        assert!(result.is_ok());

        let result = result.unwrap();
        assert!(result.final_train_loss > 0.0);
        assert_eq!(result.train_loss_history.len(), 5);
    }

    #[test]
    fn test_predict() {
        let mut studio = NeuralOpStudio::new();
        studio.set_pde(sample_poisson_problem()).unwrap();
        studio.create_operator(sample_fno_config()).unwrap();

        let mut config = sample_training_config();
        config.epochs = 2;
        config.num_train_samples = 20;
        config.num_val_samples = 5;
        config.num_test_samples = 5;

        studio.train(&config, None).unwrap();

        let input = vec![1.0; 64 * 64]; // Dummy input
        let query_points: Vec<Vec<f64>> = (0..10)
            .map(|i| vec![i as f64 / 10.0, i as f64 / 10.0])
            .collect();

        let result = studio.predict(&input, &query_points);
        assert!(result.is_ok());
        assert_eq!(result.unwrap().len(), 10);
    }

    #[test]
    fn test_run_demo() {
        let result = run_demo();
        assert!(result.is_ok());
    }

    #[test]
    fn test_unsupported_operators() {
        let mut studio = NeuralOpStudio::new();

        let mut config = sample_fno_config();
        config.operator_type = OperatorType::SINO;
        assert!(matches!(
            studio.create_operator(config),
            Err(NeuralOpError::Unsupported(_))
        ));
    }
}