414 lines
12 KiB
Rust
414 lines
12 KiB
Rust
//! Sample data generation for NeuralOp Studio demos.
|
|
|
|
use neuralop_studio_shared::{
|
|
ActivationType, BCType, BoundaryCondition, Domain, InitialCondition, OperatorConfig,
|
|
OperatorType, PDEDefinition, PDEParameters, PDEType, SchedulerType, TrainingConfig,
|
|
};
|
|
|
|
/// Create a complete demo configuration.
|
|
#[must_use]
|
|
pub fn create_demo_config() -> (PDEDefinition, OperatorConfig, TrainingConfig) {
|
|
(
|
|
create_poisson_2d(),
|
|
create_fno_config(),
|
|
create_training_config(),
|
|
)
|
|
}
|
|
|
|
/// Create a 2D Poisson problem.
|
|
#[must_use]
|
|
pub fn create_poisson_2d() -> PDEDefinition {
|
|
PDEDefinition {
|
|
pde_type: PDEType::Poisson,
|
|
equation: r"-\nabla^2 u = \sin(\pi x) \sin(\pi y)".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 1D heat equation problem.
|
|
#[must_use]
|
|
pub fn create_heat_1d() -> PDEDefinition {
|
|
PDEDefinition {
|
|
pde_type: PDEType::Heat,
|
|
equation: r"\frac{\partial u}{\partial t} = 0.01 \nabla^2 u".to_string(),
|
|
domain: Domain {
|
|
dimensions: 1,
|
|
bounds: vec![(0.0, 1.0)],
|
|
resolution: vec![128],
|
|
time_bounds: Some((0.0, 0.5)),
|
|
time_resolution: Some(50),
|
|
},
|
|
boundary_conditions: vec![
|
|
BoundaryCondition {
|
|
boundary: "left".to_string(),
|
|
bc_type: BCType::Dirichlet,
|
|
value: "0".to_string(),
|
|
robin_coefficients: None,
|
|
},
|
|
BoundaryCondition {
|
|
boundary: "right".to_string(),
|
|
bc_type: BCType::Dirichlet,
|
|
value: "0".to_string(),
|
|
robin_coefficients: None,
|
|
},
|
|
],
|
|
initial_condition: Some(InitialCondition {
|
|
function: "sin(pi*x)".to_string(),
|
|
velocity: None,
|
|
}),
|
|
parameters: PDEParameters {
|
|
diffusion: Some(0.01),
|
|
..Default::default()
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create a 1D Burgers' equation problem.
|
|
#[must_use]
|
|
pub fn create_burgers_1d() -> PDEDefinition {
|
|
PDEDefinition {
|
|
pde_type: PDEType::Burgers,
|
|
equation:
|
|
r"\frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} = 0.01 \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: "all".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 2D Navier-Stokes problem.
|
|
#[must_use]
|
|
pub fn create_navier_stokes_2d() -> 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: "kolmogorov_flow".to_string(),
|
|
velocity: None,
|
|
}),
|
|
parameters: PDEParameters {
|
|
reynolds_number: Some(1000.0),
|
|
diffusion: Some(0.001),
|
|
..Default::default()
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create a 1D wave equation problem.
|
|
#[must_use]
|
|
pub fn create_wave_1d() -> PDEDefinition {
|
|
PDEDefinition {
|
|
pde_type: PDEType::Wave,
|
|
equation: r"\frac{\partial^2 u}{\partial t^2} = c^2 \nabla^2 u".to_string(),
|
|
domain: Domain {
|
|
dimensions: 1,
|
|
bounds: vec![(0.0, 1.0)],
|
|
resolution: vec![128],
|
|
time_bounds: Some((0.0, 2.0)),
|
|
time_resolution: Some(200),
|
|
},
|
|
boundary_conditions: vec![
|
|
BoundaryCondition {
|
|
boundary: "left".to_string(),
|
|
bc_type: BCType::Dirichlet,
|
|
value: "0".to_string(),
|
|
robin_coefficients: None,
|
|
},
|
|
BoundaryCondition {
|
|
boundary: "right".to_string(),
|
|
bc_type: BCType::Dirichlet,
|
|
value: "0".to_string(),
|
|
robin_coefficients: None,
|
|
},
|
|
],
|
|
initial_condition: Some(InitialCondition {
|
|
function: "gaussian".to_string(),
|
|
velocity: Some("0".to_string()),
|
|
}),
|
|
parameters: PDEParameters {
|
|
wave_speed: Some(1.0),
|
|
..Default::default()
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create FNO configuration.
|
|
#[must_use]
|
|
pub fn create_fno_config() -> OperatorConfig {
|
|
OperatorConfig {
|
|
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,
|
|
}
|
|
}
|
|
|
|
/// Create DeepONet configuration.
|
|
#[must_use]
|
|
pub fn create_deeponet_config() -> OperatorConfig {
|
|
OperatorConfig {
|
|
operator_type: OperatorType::DeepONet,
|
|
hidden_dim: 100,
|
|
num_layers: 6,
|
|
fourier_modes: None,
|
|
branch_width: Some(100),
|
|
trunk_width: Some(100),
|
|
physics_weight: None,
|
|
activation: ActivationType::Tanh,
|
|
residual: false,
|
|
dropout: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Create PINO configuration.
|
|
#[must_use]
|
|
pub fn create_pino_config() -> OperatorConfig {
|
|
OperatorConfig {
|
|
operator_type: OperatorType::PINO,
|
|
hidden_dim: 64,
|
|
num_layers: 4,
|
|
fourier_modes: Some(12),
|
|
branch_width: None,
|
|
trunk_width: None,
|
|
physics_weight: Some(0.1),
|
|
activation: ActivationType::GELU,
|
|
residual: true,
|
|
dropout: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Create training configuration.
|
|
#[must_use]
|
|
pub fn create_training_config() -> TrainingConfig {
|
|
TrainingConfig {
|
|
epochs: 100,
|
|
batch_size: 32,
|
|
learning_rate: 1e-3,
|
|
scheduler: SchedulerType::CosineAnnealing,
|
|
optimizer: neuralop_studio_shared::OptimizerType::AdamW,
|
|
weight_decay: 1e-4,
|
|
num_train_samples: 1000,
|
|
num_val_samples: 100,
|
|
num_test_samples: 100,
|
|
seed: Some(42),
|
|
}
|
|
}
|
|
|
|
/// Create a quick demo configuration (for testing).
|
|
#[must_use]
|
|
pub fn create_quick_demo_config() -> (PDEDefinition, OperatorConfig, TrainingConfig) {
|
|
let mut pde = create_poisson_2d();
|
|
pde.domain.resolution = vec![32, 32];
|
|
|
|
let config = create_fno_config();
|
|
|
|
let training = TrainingConfig {
|
|
epochs: 10,
|
|
batch_size: 16,
|
|
learning_rate: 1e-3,
|
|
scheduler: SchedulerType::Constant,
|
|
optimizer: neuralop_studio_shared::OptimizerType::Adam,
|
|
weight_decay: 0.0,
|
|
num_train_samples: 100,
|
|
num_val_samples: 20,
|
|
num_test_samples: 20,
|
|
seed: Some(42),
|
|
};
|
|
|
|
(pde, config, training)
|
|
}
|
|
|
|
/// Generate synthetic training data for a PDE.
|
|
#[must_use]
|
|
pub fn generate_training_data(
|
|
pde: &PDEDefinition,
|
|
num_samples: usize,
|
|
seed: u64,
|
|
) -> (Vec<Vec<f64>>, Vec<Vec<f64>>) {
|
|
let grid_size: usize = pde.domain.resolution.iter().product();
|
|
let mut rng_state = seed;
|
|
|
|
let random = |state: &mut u64| -> f64 {
|
|
*state = state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(*state >> 11) as f64 / (1u64 << 53) as f64
|
|
};
|
|
|
|
let random_normal = |state: &mut u64| -> f64 {
|
|
let u1 = random(state) + 1e-10;
|
|
let u2 = random(state);
|
|
(-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
|
|
};
|
|
|
|
let inputs: Vec<Vec<f64>> = (0..num_samples)
|
|
.map(|_| {
|
|
(0..grid_size)
|
|
.map(|j| {
|
|
let x = (j % pde.domain.resolution[0]) as f64 / pde.domain.resolution[0] as f64;
|
|
let y = if pde.domain.dimensions > 1 {
|
|
(j / pde.domain.resolution[0]) as f64
|
|
/ pde.domain.resolution.get(1).copied().unwrap_or(1) as f64
|
|
} else {
|
|
0.0
|
|
};
|
|
(std::f64::consts::PI * x).sin() * (std::f64::consts::PI * y).sin()
|
|
+ random_normal(&mut rng_state) * 0.1
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let outputs: Vec<Vec<f64>> = inputs
|
|
.iter()
|
|
.map(|input| {
|
|
// Simplified: solution is related to source
|
|
input.iter().map(|&v| v * 0.5).collect()
|
|
})
|
|
.collect();
|
|
|
|
(inputs, outputs)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_demo_config() {
|
|
let (pde, op, train) = create_demo_config();
|
|
assert_eq!(pde.pde_type, PDEType::Poisson);
|
|
assert_eq!(op.operator_type, OperatorType::FNO);
|
|
assert_eq!(train.epochs, 100);
|
|
}
|
|
|
|
#[test]
|
|
fn test_poisson_2d() {
|
|
let pde = create_poisson_2d();
|
|
assert_eq!(pde.domain.dimensions, 2);
|
|
assert_eq!(pde.domain.resolution, vec![64, 64]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_heat_1d() {
|
|
let pde = create_heat_1d();
|
|
assert_eq!(pde.domain.dimensions, 1);
|
|
assert!(pde.domain.time_bounds.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_burgers_1d() {
|
|
let pde = create_burgers_1d();
|
|
assert_eq!(pde.pde_type, PDEType::Burgers);
|
|
assert!(pde.parameters.diffusion.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_navier_stokes_2d() {
|
|
let pde = create_navier_stokes_2d();
|
|
assert_eq!(pde.pde_type, PDEType::NavierStokes);
|
|
assert!(pde.parameters.reynolds_number.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_wave_1d() {
|
|
let pde = create_wave_1d();
|
|
assert_eq!(pde.pde_type, PDEType::Wave);
|
|
assert!(pde.parameters.wave_speed.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_operator_configs() {
|
|
let fno = create_fno_config();
|
|
assert_eq!(fno.operator_type, OperatorType::FNO);
|
|
assert!(fno.fourier_modes.is_some());
|
|
|
|
let deeponet = create_deeponet_config();
|
|
assert_eq!(deeponet.operator_type, OperatorType::DeepONet);
|
|
assert!(deeponet.branch_width.is_some());
|
|
|
|
let pino = create_pino_config();
|
|
assert_eq!(pino.operator_type, OperatorType::PINO);
|
|
assert!(pino.physics_weight.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_training_config() {
|
|
let config = create_training_config();
|
|
assert_eq!(config.epochs, 100);
|
|
assert_eq!(config.batch_size, 32);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quick_demo_config() {
|
|
let (pde, _, train) = create_quick_demo_config();
|
|
assert_eq!(pde.domain.resolution, vec![32, 32]);
|
|
assert_eq!(train.epochs, 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_training_data() {
|
|
let pde = create_poisson_2d();
|
|
let (inputs, outputs) = generate_training_data(&pde, 10, 42);
|
|
|
|
assert_eq!(inputs.len(), 10);
|
|
assert_eq!(outputs.len(), 10);
|
|
assert_eq!(inputs[0].len(), 64 * 64);
|
|
}
|
|
}
|