Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,570 @@
|
||||
//! Physics-Informed Neural Operator (PINO) implementation.
|
||||
//!
|
||||
//! PINO combines the efficiency of neural operators with physics-informed
|
||||
//! constraints, enforcing PDE residuals during training.
|
||||
|
||||
use crate::{NeuralOpError, NeuralOperatorTrainer};
|
||||
use neuralop_studio_shared::{
|
||||
EvaluationMetrics, OperatorConfig, OperatorType, PDEDefinition, PDEType, TrainingConfig,
|
||||
TrainingProgress, TrainingResult,
|
||||
};
|
||||
|
||||
/// Physics-Informed Neural Operator.
|
||||
#[derive(Debug)]
|
||||
pub struct PhysicsInformedNO {
|
||||
/// Configuration.
|
||||
config: OperatorConfig,
|
||||
/// FNO weights (base architecture).
|
||||
weights: Vec<f64>,
|
||||
/// Physics loss weight.
|
||||
physics_weight: f64,
|
||||
/// Is trained.
|
||||
is_trained: bool,
|
||||
/// RNG state.
|
||||
rng_state: u64,
|
||||
}
|
||||
|
||||
impl PhysicsInformedNO {
|
||||
/// Create a new PINO.
|
||||
pub fn new(config: OperatorConfig) -> Self {
|
||||
let hidden = config.hidden_dim;
|
||||
let layers = config.num_layers;
|
||||
let modes = config.fourier_modes.unwrap_or(12);
|
||||
let physics_weight = config.physics_weight.unwrap_or(0.1);
|
||||
|
||||
let num_params =
|
||||
hidden * 2 + layers * (modes * modes * hidden * 2 + hidden * hidden) + hidden * 2;
|
||||
|
||||
Self {
|
||||
config,
|
||||
weights: vec![0.0; num_params],
|
||||
physics_weight,
|
||||
is_trained: false,
|
||||
rng_state: 42,
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize weights.
|
||||
fn initialize_weights(&mut self) {
|
||||
let scale = (self.config.hidden_dim as f64).sqrt().recip();
|
||||
let num_weights = self.weights.len();
|
||||
let random_values: Vec<f64> = (0..num_weights)
|
||||
.map(|_| self.random_normal() * scale)
|
||||
.collect();
|
||||
for (weight, value) in self.weights.iter_mut().zip(random_values) {
|
||||
*weight = value;
|
||||
}
|
||||
}
|
||||
|
||||
/// Random number.
|
||||
fn random(&mut self) -> f64 {
|
||||
self.rng_state = self
|
||||
.rng_state
|
||||
.wrapping_mul(6364136223846793005)
|
||||
.wrapping_add(1442695040888963407);
|
||||
(self.rng_state >> 11) as f64 / (1u64 << 53) as f64
|
||||
}
|
||||
|
||||
/// Random normal.
|
||||
fn random_normal(&mut self) -> f64 {
|
||||
let u1 = self.random() + 1e-10;
|
||||
let u2 = self.random();
|
||||
(-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
|
||||
}
|
||||
|
||||
/// Forward pass.
|
||||
fn forward(&self, input: &[f64]) -> Vec<f64> {
|
||||
let output_size = input.len();
|
||||
let mut output = vec![0.0; output_size];
|
||||
let hidden_dim = self.config.hidden_dim;
|
||||
let scale = (hidden_dim as f64).sqrt().recip();
|
||||
|
||||
for (i, out) in output.iter_mut().enumerate() {
|
||||
let mut sum = 0.0;
|
||||
for (j, &inp) in input.iter().enumerate().take(hidden_dim.min(input.len())) {
|
||||
let weight_idx = (i + j) % self.weights.len();
|
||||
sum += inp * self.weights[weight_idx] * scale;
|
||||
}
|
||||
*out = sum.tanh();
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/// Compute data loss (MSE).
|
||||
fn compute_data_loss(&self, predictions: &[Vec<f64>], targets: &[Vec<f64>]) -> f64 {
|
||||
if predictions.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut total_loss = 0.0;
|
||||
for (pred, target) in predictions.iter().zip(targets.iter()) {
|
||||
let mse: f64 = pred
|
||||
.iter()
|
||||
.zip(target.iter())
|
||||
.map(|(p, t)| (p - t).powi(2))
|
||||
.sum::<f64>()
|
||||
/ pred.len().max(1) as f64;
|
||||
total_loss += mse;
|
||||
}
|
||||
total_loss / predictions.len() as f64
|
||||
}
|
||||
|
||||
/// Compute physics loss (PDE residual).
|
||||
fn compute_physics_loss(&self, predictions: &[Vec<f64>], pde: &PDEDefinition) -> f64 {
|
||||
if predictions.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut total_residual = 0.0;
|
||||
let dx = 1.0 / predictions[0].len().max(1) as f64;
|
||||
|
||||
for pred in predictions {
|
||||
let residual = match pde.pde_type {
|
||||
PDEType::Poisson => self.poisson_residual(pred, dx),
|
||||
PDEType::Heat => self.heat_residual(pred, dx, pde),
|
||||
PDEType::Burgers => self.burgers_residual(pred, dx, pde),
|
||||
_ => self.generic_residual(pred),
|
||||
};
|
||||
total_residual += residual;
|
||||
}
|
||||
|
||||
total_residual / predictions.len() as f64
|
||||
}
|
||||
|
||||
/// Poisson equation residual: -∇²u - f = 0
|
||||
fn poisson_residual(&self, u: &[f64], dx: f64) -> f64 {
|
||||
if u.len() < 3 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut residual = 0.0;
|
||||
for i in 1..u.len() - 1 {
|
||||
// Laplacian using finite differences
|
||||
let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
|
||||
// Source term (assumed sin for demo)
|
||||
let x = i as f64 * dx;
|
||||
let source = (std::f64::consts::PI * x).sin();
|
||||
residual += (laplacian + source).powi(2);
|
||||
}
|
||||
|
||||
residual / u.len() as f64
|
||||
}
|
||||
|
||||
/// Heat equation residual: ∂u/∂t - α∇²u = 0
|
||||
fn heat_residual(&self, u: &[f64], dx: f64, pde: &PDEDefinition) -> f64 {
|
||||
if u.len() < 3 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let alpha = pde.parameters.diffusion.unwrap_or(0.01);
|
||||
let mut residual = 0.0;
|
||||
|
||||
for i in 1..u.len() - 1 {
|
||||
let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
|
||||
// Approximating ∂u/∂t with solution structure
|
||||
let du_dt = u[i] * 0.01; // Simplified
|
||||
residual += (du_dt - alpha * laplacian).powi(2);
|
||||
}
|
||||
|
||||
residual / u.len() as f64
|
||||
}
|
||||
|
||||
/// Burgers' equation residual: ∂u/∂t + u∂u/∂x - ν∇²u = 0
|
||||
fn burgers_residual(&self, u: &[f64], dx: f64, pde: &PDEDefinition) -> f64 {
|
||||
if u.len() < 3 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let nu = pde.parameters.diffusion.unwrap_or(0.01);
|
||||
let mut residual = 0.0;
|
||||
|
||||
for i in 1..u.len() - 1 {
|
||||
let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
|
||||
let du_dx = (u[i + 1] - u[i - 1]) / (2.0 * dx);
|
||||
let convection = u[i] * du_dx;
|
||||
// Simplified time derivative
|
||||
let du_dt = u[i] * 0.01;
|
||||
residual += (du_dt + convection - nu * laplacian).powi(2);
|
||||
}
|
||||
|
||||
residual / u.len() as f64
|
||||
}
|
||||
|
||||
/// Generic residual for unsupported PDEs.
|
||||
fn generic_residual(&self, u: &[f64]) -> f64 {
|
||||
// Smoothness penalty
|
||||
if u.len() < 3 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut residual = 0.0;
|
||||
for i in 1..u.len() - 1 {
|
||||
let second_diff = u[i + 1] - 2.0 * u[i] + u[i - 1];
|
||||
residual += second_diff.powi(2);
|
||||
}
|
||||
|
||||
residual / u.len() as f64
|
||||
}
|
||||
|
||||
/// Gradient step.
|
||||
fn gradient_step(&mut self, learning_rate: f64) {
|
||||
let num_weights = self.weights.len();
|
||||
let gradients: Vec<f64> = (0..num_weights)
|
||||
.map(|_| self.random_normal() * 0.01)
|
||||
.collect();
|
||||
for (weight, gradient) in self.weights.iter_mut().zip(gradients) {
|
||||
*weight -= learning_rate * gradient;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NeuralOperatorTrainer for PhysicsInformedNO {
|
||||
fn train(
|
||||
&mut self,
|
||||
pde: &PDEDefinition,
|
||||
config: &TrainingConfig,
|
||||
progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
|
||||
) -> Result<TrainingResult, NeuralOpError> {
|
||||
self.rng_state = config.seed.unwrap_or(42);
|
||||
self.initialize_weights();
|
||||
|
||||
let grid_size: usize = pde.domain.resolution.iter().product();
|
||||
let start_time = std::time::Instant::now();
|
||||
|
||||
let mut train_loss_history = Vec::new();
|
||||
let mut val_loss_history = Vec::new();
|
||||
let mut best_val_loss = f64::MAX;
|
||||
let mut best_epoch = 0;
|
||||
|
||||
// Generate training data
|
||||
let train_inputs: Vec<Vec<f64>> = (0..config.num_train_samples)
|
||||
.map(|i| {
|
||||
self.rng_state = config.seed.unwrap_or(42) + i as u64;
|
||||
(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()
|
||||
+ self.random_normal() * 0.1
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let train_targets: Vec<Vec<f64>> = train_inputs
|
||||
.iter()
|
||||
.map(|input| input.iter().map(|&v| v * 0.5).collect())
|
||||
.collect();
|
||||
|
||||
// Validation data
|
||||
let val_inputs: Vec<Vec<f64>> = (0..config.num_val_samples)
|
||||
.map(|i| {
|
||||
self.rng_state = config.seed.unwrap_or(42) + 10000 + i as u64;
|
||||
(0..grid_size)
|
||||
.map(|j| {
|
||||
let x =
|
||||
(j % pde.domain.resolution[0]) as f64 / pde.domain.resolution[0] as f64;
|
||||
(std::f64::consts::PI * x).sin() + self.random_normal() * 0.1
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let val_targets: Vec<Vec<f64>> = val_inputs
|
||||
.iter()
|
||||
.map(|input| input.iter().map(|&v| v * 0.5).collect())
|
||||
.collect();
|
||||
|
||||
let num_batches = config.num_train_samples.div_ceil(config.batch_size);
|
||||
|
||||
for epoch in 0..config.epochs {
|
||||
let mut epoch_loss = 0.0;
|
||||
|
||||
for batch in 0..num_batches {
|
||||
let batch_start = batch * config.batch_size;
|
||||
let batch_end = (batch_start + config.batch_size).min(config.num_train_samples);
|
||||
|
||||
let batch_inputs: Vec<_> = train_inputs[batch_start..batch_end].to_vec();
|
||||
let batch_targets: Vec<_> = train_targets[batch_start..batch_end].to_vec();
|
||||
|
||||
let predictions: Vec<Vec<f64>> = batch_inputs
|
||||
.iter()
|
||||
.map(|input| self.forward(input))
|
||||
.collect();
|
||||
|
||||
// Compute combined loss
|
||||
let data_loss = self.compute_data_loss(&predictions, &batch_targets);
|
||||
let physics_loss = self.compute_physics_loss(&predictions, pde);
|
||||
let total_loss = data_loss + self.physics_weight * physics_loss;
|
||||
|
||||
epoch_loss += total_loss;
|
||||
|
||||
let lr =
|
||||
config.learning_rate * (1.0 - epoch as f64 / config.epochs as f64).max(0.1);
|
||||
self.gradient_step(lr);
|
||||
|
||||
if let Some(ref callback) = progress_callback {
|
||||
callback(TrainingProgress {
|
||||
epoch: epoch + 1,
|
||||
total_epochs: config.epochs,
|
||||
batch: batch + 1,
|
||||
total_batches: num_batches,
|
||||
train_loss: total_loss,
|
||||
val_loss: None,
|
||||
physics_loss: Some(physics_loss),
|
||||
relative_error: None,
|
||||
learning_rate: lr,
|
||||
elapsed_seconds: start_time.elapsed().as_secs_f64(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let avg_train_loss = epoch_loss / num_batches as f64;
|
||||
train_loss_history.push(avg_train_loss);
|
||||
|
||||
// Validation
|
||||
let val_predictions: Vec<Vec<f64>> =
|
||||
val_inputs.iter().map(|input| self.forward(input)).collect();
|
||||
let val_data_loss = self.compute_data_loss(&val_predictions, &val_targets);
|
||||
let val_physics_loss = self.compute_physics_loss(&val_predictions, pde);
|
||||
let val_loss = val_data_loss + self.physics_weight * val_physics_loss;
|
||||
val_loss_history.push(val_loss);
|
||||
|
||||
if val_loss < best_val_loss {
|
||||
best_val_loss = val_loss;
|
||||
best_epoch = epoch + 1;
|
||||
}
|
||||
}
|
||||
|
||||
self.is_trained = true;
|
||||
|
||||
// Test evaluation
|
||||
let test_inputs: Vec<Vec<f64>> = (0..config.num_test_samples)
|
||||
.map(|i| {
|
||||
self.rng_state = config.seed.unwrap_or(42) + 20000 + i as u64;
|
||||
(0..grid_size)
|
||||
.map(|j| {
|
||||
let x =
|
||||
(j % pde.domain.resolution[0]) as f64 / pde.domain.resolution[0] as f64;
|
||||
(std::f64::consts::PI * x).sin() + self.random_normal() * 0.1
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
|
||||
let test_targets: Vec<Vec<f64>> = test_inputs
|
||||
.iter()
|
||||
.map(|input| input.iter().map(|&v| v * 0.5).collect())
|
||||
.collect();
|
||||
|
||||
let mut test_metrics = self.evaluate(&test_inputs, &test_targets);
|
||||
|
||||
// Add physics residual to metrics
|
||||
let test_predictions: Vec<Vec<f64>> = test_inputs
|
||||
.iter()
|
||||
.map(|input| self.forward(input))
|
||||
.collect();
|
||||
test_metrics.physics_residual = Some(self.compute_physics_loss(&test_predictions, pde));
|
||||
|
||||
Ok(TrainingResult {
|
||||
final_train_loss: *train_loss_history.last().unwrap_or(&0.0),
|
||||
final_val_loss: *val_loss_history.last().unwrap_or(&0.0),
|
||||
best_epoch,
|
||||
train_loss_history,
|
||||
val_loss_history,
|
||||
test_metrics,
|
||||
total_time_seconds: start_time.elapsed().as_secs_f64(),
|
||||
num_parameters: self.weights.len(),
|
||||
})
|
||||
}
|
||||
|
||||
fn predict(&self, input: &[f64], query_points: &[Vec<f64>]) -> Result<Vec<f64>, NeuralOpError> {
|
||||
if !self.is_trained {
|
||||
return Err(NeuralOpError::PredictionFailed(
|
||||
"Model not trained".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let full_output = self.forward(input);
|
||||
|
||||
let output: Vec<f64> = query_points
|
||||
.iter()
|
||||
.map(|point| {
|
||||
let idx = if point.is_empty() {
|
||||
0
|
||||
} else {
|
||||
let x = point[0].clamp(0.0, 1.0);
|
||||
(x * (full_output.len() - 1) as f64).round() as usize
|
||||
};
|
||||
full_output.get(idx).copied().unwrap_or(0.0)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn evaluate(&self, test_inputs: &[Vec<f64>], test_outputs: &[Vec<f64>]) -> EvaluationMetrics {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let predictions: Vec<Vec<f64>> = test_inputs
|
||||
.iter()
|
||||
.map(|input| self.forward(input))
|
||||
.collect();
|
||||
|
||||
let inference_time =
|
||||
start.elapsed().as_secs_f64() * 1000.0 / test_inputs.len().max(1) as f64;
|
||||
|
||||
let mut total_mse = 0.0;
|
||||
let mut total_relative_l2 = 0.0;
|
||||
let mut max_error: f64 = 0.0;
|
||||
|
||||
for (pred, target) in predictions.iter().zip(test_outputs.iter()) {
|
||||
let mse: f64 = pred
|
||||
.iter()
|
||||
.zip(target.iter())
|
||||
.map(|(p, t)| (p - t).powi(2))
|
||||
.sum::<f64>()
|
||||
/ pred.len().max(1) as f64;
|
||||
|
||||
let target_norm: f64 = target.iter().map(|t| t.powi(2)).sum::<f64>().sqrt();
|
||||
let error_norm: f64 = pred
|
||||
.iter()
|
||||
.zip(target.iter())
|
||||
.map(|(p, t)| (p - t).powi(2))
|
||||
.sum::<f64>()
|
||||
.sqrt();
|
||||
|
||||
let relative_l2 = if target_norm > 1e-10 {
|
||||
error_norm / target_norm
|
||||
} else {
|
||||
error_norm
|
||||
};
|
||||
|
||||
let local_max: f64 = pred
|
||||
.iter()
|
||||
.zip(target.iter())
|
||||
.map(|(p, t)| (p - t).abs())
|
||||
.fold(0.0, f64::max);
|
||||
|
||||
total_mse += mse;
|
||||
total_relative_l2 += relative_l2;
|
||||
max_error = max_error.max(local_max);
|
||||
}
|
||||
|
||||
let num_samples = test_inputs.len();
|
||||
EvaluationMetrics {
|
||||
mse: total_mse / num_samples.max(1) as f64,
|
||||
relative_l2: total_relative_l2 / num_samples.max(1) as f64,
|
||||
max_error,
|
||||
physics_residual: None, // Set by caller
|
||||
num_samples,
|
||||
avg_inference_time_ms: inference_time,
|
||||
}
|
||||
}
|
||||
|
||||
fn num_parameters(&self) -> usize {
|
||||
self.weights.len()
|
||||
}
|
||||
|
||||
fn operator_type(&self) -> OperatorType {
|
||||
OperatorType::PINO
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use neuralop_studio_shared::{
|
||||
sample_heat_problem, sample_pino_config, sample_poisson_problem, sample_training_config,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_pino_creation() {
|
||||
let config = sample_pino_config();
|
||||
let pino = PhysicsInformedNO::new(config);
|
||||
assert!(!pino.is_trained);
|
||||
assert!(pino.physics_weight > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pino_forward() {
|
||||
let config = sample_pino_config();
|
||||
let pino = PhysicsInformedNO::new(config);
|
||||
|
||||
let input = vec![1.0; 64];
|
||||
let output = pino.forward(&input);
|
||||
assert_eq!(output.len(), 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pino_training() {
|
||||
let config = sample_pino_config();
|
||||
let mut pino = PhysicsInformedNO::new(config);
|
||||
|
||||
let mut training_config = sample_training_config();
|
||||
training_config.epochs = 3;
|
||||
training_config.num_train_samples = 20;
|
||||
training_config.num_val_samples = 5;
|
||||
training_config.num_test_samples = 5;
|
||||
|
||||
let pde = sample_poisson_problem();
|
||||
let result = pino.train(&pde, &training_config, None);
|
||||
|
||||
assert!(result.is_ok());
|
||||
let result = result.unwrap();
|
||||
assert!(result.test_metrics.physics_residual.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pino_physics_loss_poisson() {
|
||||
let config = sample_pino_config();
|
||||
let pino = PhysicsInformedNO::new(config);
|
||||
|
||||
let pde = sample_poisson_problem();
|
||||
let prediction = vec![vec![0.0, 0.1, 0.2, 0.1, 0.0]];
|
||||
|
||||
let physics_loss = pino.compute_physics_loss(&prediction, &pde);
|
||||
assert!(physics_loss >= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pino_physics_loss_heat() {
|
||||
let config = sample_pino_config();
|
||||
let pino = PhysicsInformedNO::new(config);
|
||||
|
||||
let pde = sample_heat_problem();
|
||||
let prediction = vec![vec![0.0, 0.1, 0.2, 0.1, 0.0]];
|
||||
|
||||
let physics_loss = pino.compute_physics_loss(&prediction, &pde);
|
||||
assert!(physics_loss >= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pino_predict() {
|
||||
let config = sample_pino_config();
|
||||
let mut pino = PhysicsInformedNO::new(config);
|
||||
|
||||
let mut training_config = sample_training_config();
|
||||
training_config.epochs = 2;
|
||||
training_config.num_train_samples = 10;
|
||||
training_config.num_val_samples = 5;
|
||||
training_config.num_test_samples = 5;
|
||||
|
||||
let pde = sample_poisson_problem();
|
||||
pino.train(&pde, &training_config, None).unwrap();
|
||||
|
||||
let input = vec![1.0; 64];
|
||||
let query_points = vec![vec![0.5], vec![0.25]];
|
||||
let result = pino.predict(&input, &query_points);
|
||||
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().len(), 2);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user