693 lines
20 KiB
Rust
693 lines
20 KiB
Rust
//! 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);
|
||
}
|
||
}
|