Files
rustytorch/demos/neuralop-studio-shared/src/lib.rs
T
2026-03-04 00:08:42 +00:00

693 lines
20 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Shared types for NeuralOp Studio - Neural Operator Workbench.
//!
//! This crate defines the IPC types for training neural operators on PDEs.
use serde::{Deserialize, Serialize};
// ============================================================================
// PDE Types
// ============================================================================
/// Type of PDE to solve.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum PDEType {
/// Poisson equation: -∇²u = f
Poisson,
/// Heat equation: ∂u/∂t = α∇²u
Heat,
/// Wave equation: ∂²u/∂t² = c²∇²u
Wave,
/// Burgers' equation: ∂u/∂t + u∂u/∂x = ν∇²u
Burgers,
/// Navier-Stokes equations
NavierStokes,
/// Advection equation: ∂u/∂t + c·∇u = 0
Advection,
/// Diffusion-reaction: ∂u/∂t = D∇²u + R(u)
DiffusionReaction,
/// Custom PDE defined by user
Custom,
}
/// PDE definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PDEDefinition {
/// Type of PDE.
pub pde_type: PDEType,
/// Equation string (LaTeX format).
pub equation: String,
/// Domain description.
pub domain: Domain,
/// Boundary conditions.
pub boundary_conditions: Vec<BoundaryCondition>,
/// Initial condition (for time-dependent PDEs).
pub initial_condition: Option<InitialCondition>,
/// Physical parameters.
pub parameters: PDEParameters,
}
/// Domain definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Domain {
/// Spatial dimensions (1D, 2D, 3D).
pub dimensions: usize,
/// Domain bounds: [(x_min, x_max), (y_min, y_max), ...]
pub bounds: Vec<(f64, f64)>,
/// Grid resolution per dimension.
pub resolution: Vec<usize>,
/// Time bounds (for time-dependent PDEs).
pub time_bounds: Option<(f64, f64)>,
/// Time resolution (for time-dependent PDEs).
pub time_resolution: Option<usize>,
}
/// Boundary condition type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum BCType {
/// Dirichlet: u = g on boundary
Dirichlet,
/// Neumann: ∂u/∂n = g on boundary
Neumann,
/// Robin: au + b∂u/∂n = g on boundary
Robin,
/// Periodic: u(x) = u(x + L)
Periodic,
}
/// Boundary condition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BoundaryCondition {
/// Boundary identifier (e.g., "left", "right", "top").
pub boundary: String,
/// Type of boundary condition.
pub bc_type: BCType,
/// Value or function description.
pub value: String,
/// Coefficients for Robin BC: (a, b).
pub robin_coefficients: Option<(f64, f64)>,
}
/// Initial condition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InitialCondition {
/// Function description.
pub function: String,
/// Initial velocity (for wave equation).
pub velocity: Option<String>,
}
/// Physical parameters for PDEs.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PDEParameters {
/// Diffusion coefficient (α, D, ν).
pub diffusion: Option<f64>,
/// Wave speed (c).
pub wave_speed: Option<f64>,
/// Reaction rate.
pub reaction_rate: Option<f64>,
/// Reynolds number (for Navier-Stokes).
pub reynolds_number: Option<f64>,
/// Source term.
pub source_term: Option<String>,
/// Custom parameters.
pub custom: Vec<(String, f64)>,
}
impl Default for PDEParameters {
fn default() -> Self {
Self {
diffusion: Some(1.0),
wave_speed: Some(1.0),
reaction_rate: None,
reynolds_number: None,
source_term: None,
custom: vec![],
}
}
}
// ============================================================================
// Neural Operator Types
// ============================================================================
/// Neural operator architecture type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum OperatorType {
/// Fourier Neural Operator.
FNO,
/// DeepONet (Deep Operator Network).
DeepONet,
/// Physics-Informed Neural Operator (PINO).
PINO,
/// Spectral-Inspired Neural Operator (SINO).
SINO,
/// Galerkin Transformer.
GalerkinTransformer,
/// Message Passing Neural Operator.
MPNO,
}
/// Neural operator configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OperatorConfig {
/// Operator type.
pub operator_type: OperatorType,
/// Hidden dimension.
pub hidden_dim: usize,
/// Number of layers.
pub num_layers: usize,
/// Number of Fourier modes (for FNO).
pub fourier_modes: Option<usize>,
/// Branch network width (for DeepONet).
pub branch_width: Option<usize>,
/// Trunk network width (for DeepONet).
pub trunk_width: Option<usize>,
/// Physics loss weight (for PINO).
pub physics_weight: Option<f64>,
/// Activation function.
pub activation: ActivationType,
/// Use residual connections.
pub residual: bool,
/// Dropout rate.
pub dropout: f64,
}
impl Default for OperatorConfig {
fn default() -> Self {
Self {
operator_type: OperatorType::FNO,
hidden_dim: 64,
num_layers: 4,
fourier_modes: Some(12),
branch_width: None,
trunk_width: None,
physics_weight: None,
activation: ActivationType::GELU,
residual: true,
dropout: 0.0,
}
}
}
/// Activation function type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum ActivationType {
/// ReLU.
ReLU,
/// GELU.
GELU,
/// Tanh.
Tanh,
/// Swish / SiLU.
Swish,
/// LeakyReLU.
LeakyReLU,
}
// ============================================================================
// Training Types
// ============================================================================
/// Training configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
/// Number of training epochs.
pub epochs: usize,
/// Batch size.
pub batch_size: usize,
/// Learning rate.
pub learning_rate: f64,
/// Learning rate scheduler.
pub scheduler: SchedulerType,
/// Optimizer type.
pub optimizer: OptimizerType,
/// Weight decay.
pub weight_decay: f64,
/// Number of training samples.
pub num_train_samples: usize,
/// Number of validation samples.
pub num_val_samples: usize,
/// Number of test samples.
pub num_test_samples: usize,
/// Random seed.
pub seed: Option<u64>,
}
impl Default for TrainingConfig {
fn default() -> Self {
Self {
epochs: 100,
batch_size: 32,
learning_rate: 1e-3,
scheduler: SchedulerType::CosineAnnealing,
optimizer: OptimizerType::AdamW,
weight_decay: 1e-4,
num_train_samples: 1000,
num_val_samples: 100,
num_test_samples: 100,
seed: Some(42),
}
}
}
/// Learning rate scheduler type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum SchedulerType {
/// Constant learning rate.
Constant,
/// Step decay.
StepDecay,
/// Exponential decay.
ExponentialDecay,
/// Cosine annealing.
CosineAnnealing,
/// Warmup with linear decay.
WarmupLinearDecay,
}
/// Optimizer type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum OptimizerType {
/// SGD.
SGD,
/// Adam.
Adam,
/// AdamW.
AdamW,
/// RMSprop.
RMSprop,
}
/// Training progress update.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingProgress {
/// Current epoch.
pub epoch: usize,
/// Total epochs.
pub total_epochs: usize,
/// Current batch.
pub batch: usize,
/// Total batches.
pub total_batches: usize,
/// Training loss.
pub train_loss: f64,
/// Validation loss.
pub val_loss: Option<f64>,
/// Physics loss (for PINO).
pub physics_loss: Option<f64>,
/// Relative L2 error.
pub relative_error: Option<f64>,
/// Current learning rate.
pub learning_rate: f64,
/// Elapsed time (seconds).
pub elapsed_seconds: f64,
}
// ============================================================================
// Prediction and Results Types
// ============================================================================
/// Prediction request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionRequest {
/// Input function samples.
pub input_function: Vec<f64>,
/// Query points.
pub query_points: Vec<Vec<f64>>,
/// Time point (for time-dependent PDEs).
pub time: Option<f64>,
}
/// Prediction result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionResult {
/// Predicted solution values.
pub solution: Vec<f64>,
/// Gradient (if computed).
pub gradient: Option<Vec<Vec<f64>>>,
/// Uncertainty estimate (if available).
pub uncertainty: Option<Vec<f64>>,
/// Inference time (milliseconds).
pub inference_time_ms: f64,
}
/// Model evaluation metrics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EvaluationMetrics {
/// Mean squared error.
pub mse: f64,
/// Relative L2 error.
pub relative_l2: f64,
/// Maximum absolute error.
pub max_error: f64,
/// Physics residual (for PINO).
pub physics_residual: Option<f64>,
/// Number of test samples.
pub num_samples: usize,
/// Average inference time (ms).
pub avg_inference_time_ms: f64,
}
/// Training result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingResult {
/// Final training loss.
pub final_train_loss: f64,
/// Final validation loss.
pub final_val_loss: f64,
/// Best epoch.
pub best_epoch: usize,
/// Training loss history.
pub train_loss_history: Vec<f64>,
/// Validation loss history.
pub val_loss_history: Vec<f64>,
/// Evaluation metrics on test set.
pub test_metrics: EvaluationMetrics,
/// Total training time (seconds).
pub total_time_seconds: f64,
/// Number of parameters.
pub num_parameters: usize,
}
// ============================================================================
// Visualization Types
// ============================================================================
/// Solution field for visualization.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SolutionField {
/// Grid coordinates (flattened).
pub coordinates: Vec<Vec<f64>>,
/// Solution values.
pub values: Vec<f64>,
/// Field name.
pub name: String,
/// Time point (for time-dependent solutions).
pub time: Option<f64>,
}
/// Comparison between ground truth and prediction.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SolutionComparison {
/// Ground truth solution.
pub ground_truth: SolutionField,
/// Predicted solution.
pub prediction: SolutionField,
/// Point-wise error.
pub error: SolutionField,
/// Relative error statistics.
pub error_stats: ErrorStatistics,
}
/// Error statistics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErrorStatistics {
/// Mean error.
pub mean: f64,
/// Standard deviation.
pub std: f64,
/// Maximum error.
pub max: f64,
/// Minimum error.
pub min: f64,
/// 95th percentile.
pub p95: f64,
}
// ============================================================================
// Sample Data Functions
// ============================================================================
/// Create a sample Poisson problem.
#[must_use]
pub fn sample_poisson_problem() -> PDEDefinition {
PDEDefinition {
pde_type: PDEType::Poisson,
equation: r"-\nabla^2 u = f".to_string(),
domain: Domain {
dimensions: 2,
bounds: vec![(0.0, 1.0), (0.0, 1.0)],
resolution: vec![64, 64],
time_bounds: None,
time_resolution: None,
},
boundary_conditions: vec![BoundaryCondition {
boundary: "all".to_string(),
bc_type: BCType::Dirichlet,
value: "0".to_string(),
robin_coefficients: None,
}],
initial_condition: None,
parameters: PDEParameters {
source_term: Some("sin(pi*x)*sin(pi*y)".to_string()),
..Default::default()
},
}
}
/// Create a sample heat equation problem.
#[must_use]
pub fn sample_heat_problem() -> PDEDefinition {
PDEDefinition {
pde_type: PDEType::Heat,
equation: r"\frac{\partial u}{\partial t} = \alpha \nabla^2 u".to_string(),
domain: Domain {
dimensions: 2,
bounds: vec![(0.0, 1.0), (0.0, 1.0)],
resolution: vec![64, 64],
time_bounds: Some((0.0, 1.0)),
time_resolution: Some(100),
},
boundary_conditions: vec![BoundaryCondition {
boundary: "all".to_string(),
bc_type: BCType::Dirichlet,
value: "0".to_string(),
robin_coefficients: None,
}],
initial_condition: Some(InitialCondition {
function: "sin(pi*x)*sin(pi*y)".to_string(),
velocity: None,
}),
parameters: PDEParameters {
diffusion: Some(0.01),
..Default::default()
},
}
}
/// Create a sample Burgers' equation problem.
#[must_use]
pub fn sample_burgers_problem() -> PDEDefinition {
PDEDefinition {
pde_type: PDEType::Burgers,
equation:
r"\frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} = \nu \nabla^2 u"
.to_string(),
domain: Domain {
dimensions: 1,
bounds: vec![(0.0, 2.0 * std::f64::consts::PI)],
resolution: vec![256],
time_bounds: Some((0.0, 1.0)),
time_resolution: Some(100),
},
boundary_conditions: vec![BoundaryCondition {
boundary: "left".to_string(),
bc_type: BCType::Periodic,
value: "periodic".to_string(),
robin_coefficients: None,
}],
initial_condition: Some(InitialCondition {
function: "sin(x)".to_string(),
velocity: None,
}),
parameters: PDEParameters {
diffusion: Some(0.01),
..Default::default()
},
}
}
/// Create a sample Navier-Stokes problem.
#[must_use]
pub fn sample_navier_stokes_problem() -> PDEDefinition {
PDEDefinition {
pde_type: PDEType::NavierStokes,
equation: r"\frac{\partial \mathbf{u}}{\partial t} + (\mathbf{u} \cdot \nabla)\mathbf{u} = -\nabla p + \nu \nabla^2 \mathbf{u}".to_string(),
domain: Domain {
dimensions: 2,
bounds: vec![(0.0, 2.0 * std::f64::consts::PI), (0.0, 2.0 * std::f64::consts::PI)],
resolution: vec![64, 64],
time_bounds: Some((0.0, 10.0)),
time_resolution: Some(100),
},
boundary_conditions: vec![
BoundaryCondition {
boundary: "all".to_string(),
bc_type: BCType::Periodic,
value: "periodic".to_string(),
robin_coefficients: None,
},
],
initial_condition: Some(InitialCondition {
function: "vortex_pair".to_string(),
velocity: None,
}),
parameters: PDEParameters {
reynolds_number: Some(1000.0),
diffusion: Some(0.001),
..Default::default()
},
}
}
/// Create a sample FNO configuration.
#[must_use]
pub fn sample_fno_config() -> OperatorConfig {
OperatorConfig {
operator_type: OperatorType::FNO,
hidden_dim: 64,
num_layers: 4,
fourier_modes: Some(12),
activation: ActivationType::GELU,
residual: true,
dropout: 0.0,
..Default::default()
}
}
/// Create a sample DeepONet configuration.
#[must_use]
pub fn sample_deeponet_config() -> OperatorConfig {
OperatorConfig {
operator_type: OperatorType::DeepONet,
hidden_dim: 100,
num_layers: 6,
branch_width: Some(100),
trunk_width: Some(100),
activation: ActivationType::Tanh,
residual: false,
dropout: 0.0,
..Default::default()
}
}
/// Create a sample PINO configuration.
#[must_use]
pub fn sample_pino_config() -> OperatorConfig {
OperatorConfig {
operator_type: OperatorType::PINO,
hidden_dim: 64,
num_layers: 4,
fourier_modes: Some(12),
physics_weight: Some(0.1),
activation: ActivationType::GELU,
residual: true,
dropout: 0.0,
..Default::default()
}
}
/// Create a sample training configuration.
#[must_use]
pub fn sample_training_config() -> TrainingConfig {
TrainingConfig {
epochs: 100,
batch_size: 32,
learning_rate: 1e-3,
scheduler: SchedulerType::CosineAnnealing,
optimizer: OptimizerType::AdamW,
weight_decay: 1e-4,
num_train_samples: 1000,
num_val_samples: 100,
num_test_samples: 100,
seed: Some(42),
}
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pde_types() {
assert_eq!(PDEType::Poisson, PDEType::Poisson);
assert_ne!(PDEType::Heat, PDEType::Wave);
}
#[test]
fn test_sample_poisson() {
let problem = sample_poisson_problem();
assert_eq!(problem.pde_type, PDEType::Poisson);
assert_eq!(problem.domain.dimensions, 2);
assert!(problem.initial_condition.is_none());
}
#[test]
fn test_sample_heat() {
let problem = sample_heat_problem();
assert_eq!(problem.pde_type, PDEType::Heat);
assert!(problem.initial_condition.is_some());
assert!(problem.domain.time_bounds.is_some());
}
#[test]
fn test_sample_burgers() {
let problem = sample_burgers_problem();
assert_eq!(problem.pde_type, PDEType::Burgers);
assert_eq!(problem.domain.dimensions, 1);
}
#[test]
fn test_sample_navier_stokes() {
let problem = sample_navier_stokes_problem();
assert_eq!(problem.pde_type, PDEType::NavierStokes);
assert!(problem.parameters.reynolds_number.is_some());
}
#[test]
fn test_operator_configs() {
let fno = sample_fno_config();
assert_eq!(fno.operator_type, OperatorType::FNO);
assert!(fno.fourier_modes.is_some());
let deeponet = sample_deeponet_config();
assert_eq!(deeponet.operator_type, OperatorType::DeepONet);
assert!(deeponet.branch_width.is_some());
let pino = sample_pino_config();
assert_eq!(pino.operator_type, OperatorType::PINO);
assert!(pino.physics_weight.is_some());
}
#[test]
fn test_training_config() {
let config = sample_training_config();
assert_eq!(config.epochs, 100);
assert_eq!(config.optimizer, OptimizerType::AdamW);
}
#[test]
fn test_serialization() {
let problem = sample_poisson_problem();
let json = serde_json::to_string(&problem).unwrap();
assert!(json.contains("Poisson"));
let parsed: PDEDefinition = serde_json::from_str(&json).unwrap();
assert_eq!(parsed.pde_type, problem.pde_type);
}
#[test]
fn test_operator_config_serialization() {
let config = sample_fno_config();
let json = serde_json::to_string(&config).unwrap();
assert!(json.contains("FNO"));
let parsed: OperatorConfig = serde_json::from_str(&json).unwrap();
assert_eq!(parsed.operator_type, config.operator_type);
}
}