Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,413 @@
+//! 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);
+    }
+}