973 lines
32 KiB
Rust
973 lines
32 KiB
Rust
//! Core PINN (Physics-Informed Neural Network) Implementation
|
|
//!
|
|
//! This module provides the main PINN structure and functionality for solving
|
|
//! partial differential equations with neural networks that respect physics laws.
|
|
|
|
use crate::error::{Result, ScienceError};
|
|
use crate::physics::{
|
|
BoundaryConditions, ConservationLoss, PINNTrainer, PhysicsLoss, TrainingConfig,
|
|
};
|
|
use crate::variable_extensions::VariableExt;
|
|
use async_trait::async_trait;
|
|
use rtx_autograd::Variable;
|
|
use rtx_tensor::{Device, Tensor, TensorError};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use tokio::sync::RwLock;
|
|
use tracing::{debug, info, warn};
|
|
|
|
/// Physics-Informed Neural Network for solving PDEs
|
|
pub struct PINN {
|
|
/// Neural network layers
|
|
network: Vec<LinearLayer>,
|
|
/// Device for computations
|
|
device: Device,
|
|
/// Physics loss function
|
|
physics_loss: Box<dyn PhysicsLoss + Send + Sync>,
|
|
/// Conservation losses
|
|
conservation_losses: Vec<Box<dyn ConservationLoss + Send + Sync>>,
|
|
/// Training configuration
|
|
config: TrainingConfig,
|
|
/// Current training state
|
|
state: RwLock<PINNState>,
|
|
}
|
|
|
|
/// Individual linear layer in the neural network
|
|
#[derive(Debug, Clone)]
|
|
struct LinearLayer {
|
|
/// Weight matrix
|
|
weights: Variable,
|
|
/// Bias vector
|
|
bias: Variable,
|
|
/// Activation function
|
|
activation: ActivationType,
|
|
}
|
|
|
|
/// Supported activation functions
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum ActivationType {
|
|
/// Hyperbolic tangent
|
|
Tanh,
|
|
/// Sigmoid
|
|
Sigmoid,
|
|
/// Sine activation (useful for periodic solutions)
|
|
Sin,
|
|
/// Swish/SiLU activation
|
|
Swish,
|
|
/// GELU activation
|
|
GELU,
|
|
/// `ReLU` activation
|
|
ReLU,
|
|
/// Learnable activation (adaptive)
|
|
Learnable,
|
|
}
|
|
|
|
/// Internal state of the PINN during training
|
|
#[derive(Debug, Default, Clone)]
|
|
pub struct PINNState {
|
|
/// Current training iteration
|
|
iteration: usize,
|
|
/// Training loss history
|
|
loss_history: Vec<LossComponents>,
|
|
/// Best model weights (for early stopping)
|
|
best_weights: Option<HashMap<String, Tensor>>,
|
|
/// Best validation loss
|
|
best_loss: f64,
|
|
/// Convergence metrics
|
|
convergence: ConvergenceMetrics,
|
|
}
|
|
|
|
/// Components of the total loss function
|
|
#[derive(Debug, Clone)]
|
|
pub struct LossComponents {
|
|
/// Physics/PDE loss
|
|
pub physics: f64,
|
|
/// Boundary condition loss
|
|
pub boundary: f64,
|
|
/// Initial condition loss
|
|
pub initial: f64,
|
|
/// Conservation law losses
|
|
pub conservation: Vec<f64>,
|
|
/// Data fitting loss (if available)
|
|
pub data: f64,
|
|
/// Total combined loss
|
|
pub total: f64,
|
|
/// Training iteration
|
|
pub iteration: usize,
|
|
}
|
|
|
|
/// Convergence tracking metrics
|
|
#[derive(Debug, Default, Clone)]
|
|
struct ConvergenceMetrics {
|
|
/// Moving average of loss gradient magnitude
|
|
gradient_norm: f64,
|
|
/// Loss improvement rate
|
|
improvement_rate: f64,
|
|
/// Plateau detection counter
|
|
plateau_count: usize,
|
|
/// Early stopping criteria met
|
|
should_stop: bool,
|
|
}
|
|
|
|
/// Checkpoint for saving and restoring PINN state
|
|
#[derive(Debug, Clone)]
|
|
pub struct PINNCheckpoint {
|
|
/// Layer weights
|
|
pub weights: Vec<Tensor>,
|
|
/// Layer biases
|
|
pub biases: Vec<Tensor>,
|
|
/// Training iteration
|
|
pub iteration: usize,
|
|
/// Best loss value
|
|
pub best_loss: f64,
|
|
}
|
|
|
|
/// Builder pattern for PINN construction
|
|
pub struct PINNBuilder {
|
|
device: Option<Device>,
|
|
layers: Vec<usize>,
|
|
activation: ActivationType,
|
|
physics_loss: Option<Box<dyn PhysicsLoss + Send + Sync>>,
|
|
conservation_losses: Vec<Box<dyn ConservationLoss + Send + Sync>>,
|
|
config: TrainingConfig,
|
|
}
|
|
|
|
/// Trait for physics-informed neural network functionality
|
|
#[async_trait]
|
|
pub trait PhysicsInformedNetwork: Send + Sync {
|
|
/// Predict output for given input coordinates
|
|
async fn predict(&self, inputs: &[(f64, f64)]) -> Result<Vec<f64>>;
|
|
|
|
/// Predict with gradients (for physics loss computation)
|
|
async fn predict_with_gradients(
|
|
&self,
|
|
inputs: &[(f64, f64)],
|
|
) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>)>;
|
|
|
|
/// Train the network with boundary conditions
|
|
async fn train(&mut self, boundary_data: BoundaryConditions, epochs: usize) -> Result<()>;
|
|
|
|
/// Evaluate physics loss at given points
|
|
async fn physics_residual(&self, inputs: &[(f64, f64)]) -> Result<Vec<f64>>;
|
|
|
|
/// Get current training state
|
|
async fn training_state(&self) -> Result<PINNState>;
|
|
}
|
|
|
|
impl PINN {
|
|
/// Create a new PINN builder
|
|
#[must_use]
|
|
pub fn builder() -> PINNBuilder {
|
|
PINNBuilder::new()
|
|
}
|
|
|
|
/// Create PINN from builder
|
|
pub(crate) fn from_builder(builder: PINNBuilder) -> Result<Self> {
|
|
let device = builder.device.ok_or_else(|| {
|
|
ScienceError::physics(
|
|
"Device must be specified",
|
|
crate::error::PhysicsDomain::FluidDynamics,
|
|
)
|
|
})?;
|
|
|
|
let physics_loss = builder.physics_loss.ok_or_else(|| {
|
|
ScienceError::physics(
|
|
"Physics loss must be specified",
|
|
crate::error::PhysicsDomain::FluidDynamics,
|
|
)
|
|
})?;
|
|
|
|
if builder.layers.len() < 2 {
|
|
return Err(ScienceError::physics(
|
|
"At least input and output layers required",
|
|
crate::error::PhysicsDomain::FluidDynamics,
|
|
));
|
|
}
|
|
|
|
// Initialize network layers
|
|
let mut network = Vec::new();
|
|
for i in 0..builder.layers.len() - 1 {
|
|
let input_size = builder.layers[i];
|
|
let output_size = builder.layers[i + 1];
|
|
|
|
let layer = LinearLayer::new(input_size, output_size, &builder.activation, &device)?;
|
|
network.push(layer);
|
|
}
|
|
|
|
Ok(Self {
|
|
network,
|
|
device,
|
|
physics_loss,
|
|
conservation_losses: builder.conservation_losses,
|
|
config: builder.config,
|
|
state: RwLock::new(PINNState::default()),
|
|
})
|
|
}
|
|
|
|
/// Forward pass through the network
|
|
pub async fn forward(&self, inputs: &Tensor) -> Result<Variable> {
|
|
let mut x = Variable::from_tensor(inputs.clone());
|
|
|
|
for (i, layer) in self.network.iter().enumerate() {
|
|
x = layer.forward(x)?;
|
|
|
|
// Apply activation (skip for output layer)
|
|
if i < self.network.len() - 1 {
|
|
x = apply_activation(x, &layer.activation)?;
|
|
}
|
|
}
|
|
|
|
Ok(x)
|
|
}
|
|
|
|
/// Compute derivatives using automatic differentiation
|
|
async fn compute_derivatives(
|
|
&self,
|
|
inputs: &Tensor,
|
|
outputs: &Variable,
|
|
) -> Result<(Variable, Variable)> {
|
|
// Simplified gradient computation using finite differences
|
|
// A full implementation would use automatic differentiation
|
|
let epsilon = 1e-5;
|
|
|
|
// Create perturbation tensors
|
|
let eps_x = Tensor::from_slice(&[epsilon, 0.0], &[1, 2], &self.device)?;
|
|
let eps_t = Tensor::from_slice(&[0.0, epsilon], &[1, 2], &self.device)?;
|
|
|
|
// Perturb inputs
|
|
let x_plus = inputs.add(&eps_x)?;
|
|
let t_plus = inputs.add(&eps_t)?;
|
|
|
|
// Forward pass with perturbed inputs
|
|
let output_x_plus = self.forward(&x_plus).await?;
|
|
let output_t_plus = self.forward(&t_plus).await?;
|
|
|
|
// Compute finite difference approximations
|
|
let du_dx_tensor = output_x_plus
|
|
.tensor()
|
|
.sub(outputs.tensor())?
|
|
.mul_scalar(1.0 / epsilon)?;
|
|
let du_dt_tensor = output_t_plus
|
|
.tensor()
|
|
.sub(outputs.tensor())?
|
|
.mul_scalar(1.0 / epsilon)?;
|
|
|
|
// Create Variables from the gradient tensors
|
|
let du_dx = Variable::new(du_dx_tensor, false);
|
|
let du_dt = Variable::new(du_dt_tensor, false);
|
|
|
|
Ok((du_dx, du_dt))
|
|
}
|
|
|
|
/// Compute second derivatives
|
|
async fn compute_second_derivatives(
|
|
&self,
|
|
inputs: &Tensor,
|
|
first_derivs: &(Variable, Variable),
|
|
) -> Result<(Variable, Variable, Variable)> {
|
|
// Simplified second derivative computation using finite differences
|
|
let epsilon = 1e-5;
|
|
let (du_dx, du_dt) = first_derivs;
|
|
|
|
// Create perturbation tensors
|
|
let eps_x = Tensor::from_slice(&[epsilon, 0.0], &[1, 2], &self.device)?;
|
|
let eps_t = Tensor::from_slice(&[0.0, epsilon], &[1, 2], &self.device)?;
|
|
|
|
// Perturb inputs in both directions
|
|
let x_plus = inputs.add(&eps_x)?;
|
|
let x_minus = inputs.sub(&eps_x)?;
|
|
let t_plus = inputs.add(&eps_t)?;
|
|
let t_minus = inputs.sub(&eps_t)?;
|
|
|
|
// Compute outputs at perturbed points
|
|
let out_x_plus = self.forward(&x_plus).await?;
|
|
let out_x_minus = self.forward(&x_minus).await?;
|
|
let out_t_plus = self.forward(&t_plus).await?;
|
|
let out_t_minus = self.forward(&t_minus).await?;
|
|
|
|
// Compute second derivatives using central differences
|
|
let two_du_dx = du_dx.tensor().mul_scalar(2.0)?;
|
|
let two_du_dt = du_dt.tensor().mul_scalar(2.0)?;
|
|
|
|
let d2u_dx2_tensor = out_x_plus
|
|
.tensor()
|
|
.add(out_x_minus.tensor())?
|
|
.sub(&two_du_dx)?
|
|
.mul_scalar(1.0 / (epsilon * epsilon))?;
|
|
|
|
let d2u_dt2_tensor = out_t_plus
|
|
.tensor()
|
|
.add(out_t_minus.tensor())?
|
|
.sub(&two_du_dt)?
|
|
.mul_scalar(1.0 / (epsilon * epsilon))?;
|
|
|
|
// Mixed derivative (simplified)
|
|
let d2u_dxdt_tensor = out_x_plus
|
|
.tensor()
|
|
.sub(out_x_minus.tensor())?
|
|
.mul_scalar(1.0 / (2.0 * epsilon * epsilon))?;
|
|
|
|
// Create Variables from the gradient tensors
|
|
let d2u_dx2 = Variable::new(d2u_dx2_tensor, false);
|
|
let d2u_dt2 = Variable::new(d2u_dt2_tensor, false);
|
|
let d2u_dxdt = Variable::new(d2u_dxdt_tensor, false);
|
|
|
|
Ok((d2u_dx2, d2u_dt2, d2u_dxdt))
|
|
}
|
|
|
|
/// Compute total loss combining all components
|
|
async fn compute_total_loss(
|
|
&self,
|
|
collocation_points: &Tensor,
|
|
boundary_data: &BoundaryConditions,
|
|
) -> Result<(Variable, LossComponents)> {
|
|
let mut loss_components = LossComponents {
|
|
physics: 0.0,
|
|
boundary: 0.0,
|
|
initial: 0.0,
|
|
conservation: Vec::new(),
|
|
data: 0.0,
|
|
total: 0.0,
|
|
iteration: 0,
|
|
};
|
|
|
|
// Physics/PDE loss
|
|
let outputs = self.forward(collocation_points).await?;
|
|
let (du_dx, du_dt) = self
|
|
.compute_derivatives(collocation_points, &outputs)
|
|
.await?;
|
|
let second_derivs = self
|
|
.compute_second_derivatives(collocation_points, &(du_dx.clone(), du_dt.clone()))
|
|
.await?;
|
|
|
|
let physics_residual = self
|
|
.physics_loss
|
|
.compute_residual(
|
|
collocation_points,
|
|
&outputs,
|
|
&du_dx,
|
|
&du_dt,
|
|
&second_derivs.0,
|
|
&second_derivs.1,
|
|
&second_derivs.2,
|
|
)
|
|
.await?;
|
|
|
|
// Use the mean_square from rtx_autograd Variable
|
|
let squared = physics_residual.multiply(&physics_residual)?;
|
|
let physics_loss_var = squared.mean()?;
|
|
let physics_loss = physics_loss_var.value().to_scalar::<f32>()?;
|
|
loss_components.physics = f64::from(physics_loss);
|
|
|
|
// Boundary condition loss
|
|
let boundary_loss_var = boundary_data.compute_loss(self).await?;
|
|
let boundary_loss = boundary_loss_var.value().to_scalar::<f32>()?;
|
|
loss_components.boundary = f64::from(boundary_loss);
|
|
|
|
// Conservation law losses
|
|
for conservation_law in &self.conservation_losses {
|
|
let conservation_loss = conservation_law
|
|
.compute_loss(collocation_points, &outputs, &du_dx, &du_dt)
|
|
.await?;
|
|
loss_components
|
|
.conservation
|
|
.push(f64::from(conservation_loss));
|
|
}
|
|
|
|
// Combine losses with weights
|
|
let mut total_loss = physics_loss * (self.config.physics_weight as f32);
|
|
total_loss += boundary_loss * self.config.boundary_weight as f32;
|
|
|
|
for (i, conservation_law) in self.conservation_losses.iter().enumerate() {
|
|
if let Some(weight) = self.config.conservation_weights.get(i) {
|
|
let cons_loss = conservation_law
|
|
.compute_loss(collocation_points, &outputs, &du_dx, &du_dt)
|
|
.await?;
|
|
total_loss += cons_loss * *weight as f32;
|
|
}
|
|
}
|
|
|
|
loss_components.total = f64::from(total_loss);
|
|
|
|
// Convert total_loss to Variable for backward pass
|
|
let loss_tensor = Tensor::scalar(total_loss, rtx_tensor::DType::F32, &self.device)?;
|
|
let loss_variable = Variable::from_tensor(loss_tensor);
|
|
|
|
Ok((loss_variable, loss_components))
|
|
}
|
|
|
|
/// Update training state with new loss information
|
|
async fn update_training_state(&self, loss_components: LossComponents) -> Result<()> {
|
|
let mut state = self.state.write().await;
|
|
state.iteration += 1;
|
|
|
|
// Update convergence metrics
|
|
if let Some(prev_loss) = state.loss_history.last() {
|
|
let loss_change = prev_loss.total - loss_components.total;
|
|
state.convergence.improvement_rate = loss_change / prev_loss.total;
|
|
|
|
if loss_change.abs() < self.config.convergence_tolerance {
|
|
state.convergence.plateau_count += 1;
|
|
} else {
|
|
state.convergence.plateau_count = 0;
|
|
}
|
|
|
|
// Early stopping check
|
|
if state.convergence.plateau_count > self.config.patience {
|
|
state.convergence.should_stop = true;
|
|
warn!("Early stopping triggered due to lack of improvement");
|
|
}
|
|
}
|
|
|
|
// Save best model
|
|
if loss_components.total < state.best_loss {
|
|
state.best_loss = loss_components.total;
|
|
|
|
// Save current model weights as best weights
|
|
let mut best_weights = HashMap::new();
|
|
for (i, layer) in self.network.iter().enumerate() {
|
|
let weight_key = format!("layer_{i}_weight");
|
|
let bias_key = format!("layer_{i}_bias");
|
|
|
|
best_weights.insert(weight_key, layer.weights.value().clone());
|
|
best_weights.insert(bias_key, layer.bias.value().clone());
|
|
}
|
|
state.best_weights = Some(best_weights);
|
|
|
|
debug!(
|
|
"New best model saved at iteration {} with loss {:.6}",
|
|
state.iteration, loss_components.total
|
|
);
|
|
}
|
|
|
|
if state.iteration % 100 == 0 {
|
|
info!(
|
|
"Iteration {}: Total loss = {:.6}, Physics = {:.6}, Boundary = {:.6}",
|
|
state.iteration,
|
|
loss_components.total,
|
|
loss_components.physics,
|
|
loss_components.boundary
|
|
);
|
|
}
|
|
|
|
state.loss_history.push(loss_components);
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl PhysicsInformedNetwork for PINN {
|
|
async fn predict(&self, inputs: &[(f64, f64)]) -> Result<Vec<f64>> {
|
|
let input_tensor = create_input_tensor(inputs, &self.device)?;
|
|
let output = self.forward(&input_tensor).await?;
|
|
let values = output
|
|
.tensor()
|
|
.to_vec()?
|
|
.iter()
|
|
.map(|&x| f64::from(x))
|
|
.collect();
|
|
Ok(values)
|
|
}
|
|
|
|
async fn predict_with_gradients(
|
|
&self,
|
|
inputs: &[(f64, f64)],
|
|
) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>)> {
|
|
let input_tensor = create_input_tensor(inputs, &self.device)?;
|
|
let output = self.forward(&input_tensor).await?;
|
|
let (du_dx, du_dt) = self.compute_derivatives(&input_tensor, &output).await?;
|
|
|
|
let values = output
|
|
.tensor()
|
|
.to_vec()?
|
|
.iter()
|
|
.map(|&x| f64::from(x))
|
|
.collect();
|
|
let dx_values = du_dx
|
|
.tensor()
|
|
.to_vec()?
|
|
.iter()
|
|
.map(|&x| f64::from(x))
|
|
.collect();
|
|
let dt_values = du_dt
|
|
.tensor()
|
|
.to_vec()?
|
|
.iter()
|
|
.map(|&x| f64::from(x))
|
|
.collect();
|
|
|
|
Ok((values, dx_values, dt_values))
|
|
}
|
|
|
|
async fn train(&mut self, boundary_data: BoundaryConditions, epochs: usize) -> Result<()> {
|
|
let trainer = PINNTrainer::new(self.config.clone());
|
|
trainer.train(self, boundary_data, epochs).await
|
|
}
|
|
|
|
async fn physics_residual(&self, inputs: &[(f64, f64)]) -> Result<Vec<f64>> {
|
|
let input_tensor = create_input_tensor(inputs, &self.device)?;
|
|
let output = self.forward(&input_tensor).await?;
|
|
let (du_dx, du_dt) = self.compute_derivatives(&input_tensor, &output).await?;
|
|
let derivatives = (du_dx, du_dt);
|
|
let second_derivs = self
|
|
.compute_second_derivatives(&input_tensor, &derivatives)
|
|
.await?;
|
|
|
|
let residual = self
|
|
.physics_loss
|
|
.compute_residual(
|
|
&input_tensor,
|
|
&output,
|
|
&derivatives.0,
|
|
&derivatives.1,
|
|
&second_derivs.0,
|
|
&second_derivs.1,
|
|
&second_derivs.2,
|
|
)
|
|
.await?;
|
|
|
|
let values = residual
|
|
.tensor()
|
|
.to_vec()?
|
|
.iter()
|
|
.map(|&x| f64::from(x))
|
|
.collect();
|
|
Ok(values)
|
|
}
|
|
|
|
async fn training_state(&self) -> Result<PINNState> {
|
|
let state_guard = self.state.read().await;
|
|
Ok((*state_guard).clone())
|
|
}
|
|
}
|
|
|
|
impl PINN {
|
|
/// Save model weights to a file using `SafeTensors` format
|
|
pub async fn save_weights(&self, path: &std::path::Path) -> Result<()> {
|
|
use std::collections::HashMap;
|
|
use std::io::Write;
|
|
|
|
let mut weights_map = HashMap::new();
|
|
|
|
// Collect all layer weights and biases
|
|
for (i, layer) in self.network.iter().enumerate() {
|
|
let weight_key = format!("layer_{i}_weight");
|
|
let bias_key = format!("layer_{i}_bias");
|
|
|
|
// Get tensor data from the layer
|
|
let weight_data = layer.weights.value().to_vec()?;
|
|
let bias_data = layer.bias.value().to_vec()?;
|
|
|
|
weights_map.insert(weight_key, weight_data);
|
|
weights_map.insert(bias_key, bias_data);
|
|
}
|
|
|
|
// For now, use a simple JSON serialization
|
|
// In a production system, this would use SafeTensors format
|
|
let json_data = serde_json::to_string_pretty(&weights_map)?;
|
|
|
|
let mut file = std::fs::File::create(path)?;
|
|
file.write_all(json_data.as_bytes())?;
|
|
|
|
tracing::info!("Model weights saved to {:?}", path);
|
|
Ok(())
|
|
}
|
|
|
|
/// Load model weights from a file
|
|
pub async fn load_weights(&mut self, path: &std::path::Path) -> Result<()> {
|
|
use std::collections::HashMap;
|
|
|
|
if !path.exists() {
|
|
return Err(ScienceError::io_error(
|
|
format!("Weight file not found: {path:?}"),
|
|
"file_not_found",
|
|
));
|
|
}
|
|
|
|
let json_data = std::fs::read_to_string(path)?;
|
|
let weights_map: HashMap<String, Vec<f32>> = serde_json::from_str(&json_data)?;
|
|
|
|
// Load weights into each layer
|
|
for (i, layer) in self.network.iter_mut().enumerate() {
|
|
let weight_key = format!("layer_{i}_weight");
|
|
let bias_key = format!("layer_{i}_bias");
|
|
|
|
if let Some(weight_data) = weights_map.get(&weight_key) {
|
|
// Create new tensor from the loaded data
|
|
let weight_shape = layer.weights.value().shape().to_vec();
|
|
let weight_tensor = Tensor::from_slice(weight_data, &weight_shape, &self.device)?;
|
|
layer.weights = Variable::new(weight_tensor, true);
|
|
}
|
|
|
|
if let Some(bias_data) = weights_map.get(&bias_key) {
|
|
let bias_shape = layer.bias.value().shape().to_vec();
|
|
let bias_tensor = Tensor::from_slice(bias_data, &bias_shape, &self.device)?;
|
|
layer.bias = Variable::new(bias_tensor, true);
|
|
}
|
|
}
|
|
|
|
tracing::info!("Model weights loaded from {:?}", path);
|
|
Ok(())
|
|
}
|
|
|
|
/// Get device used by this PINN
|
|
pub fn device(&self) -> &Device {
|
|
&self.device
|
|
}
|
|
|
|
/// Forward pass with gradient tracking for physics-informed training
|
|
pub async fn forward_with_gradients(
|
|
&mut self,
|
|
inputs: &Tensor,
|
|
) -> Result<(Variable, Variable, Variable)> {
|
|
let outputs = self.forward(inputs).await?;
|
|
let (du_dx, du_dt) = self.compute_derivatives(inputs, &outputs).await?;
|
|
Ok((outputs, du_dx, du_dt))
|
|
}
|
|
|
|
/// Get all trainable parameters
|
|
pub fn parameters(&self) -> Vec<Variable> {
|
|
let mut params = Vec::new();
|
|
for layer in &self.network {
|
|
params.push(layer.weights.clone());
|
|
params.push(layer.bias.clone());
|
|
}
|
|
params
|
|
}
|
|
|
|
/// Zero all gradients
|
|
pub fn zero_gradients(&mut self) {
|
|
for layer in &mut self.network {
|
|
layer.weights.zero_grad();
|
|
layer.bias.zero_grad();
|
|
}
|
|
}
|
|
|
|
/// Clip gradients to prevent explosion
|
|
pub fn clip_gradients(&mut self, max_norm: f32) {
|
|
for layer in &mut self.network {
|
|
layer.weights.clip_grad(max_norm);
|
|
layer.bias.clip_grad(max_norm);
|
|
}
|
|
}
|
|
|
|
/// Save current state for checkpointing
|
|
pub async fn save_state(&self) -> Result<PINNCheckpoint> {
|
|
let mut weights = Vec::new();
|
|
let mut biases = Vec::new();
|
|
|
|
for layer in &self.network {
|
|
weights.push(layer.weights.value().clone());
|
|
biases.push(layer.bias.value().clone());
|
|
}
|
|
|
|
let state = self.state.read().await;
|
|
Ok(PINNCheckpoint {
|
|
weights,
|
|
biases,
|
|
iteration: state.iteration,
|
|
best_loss: state.best_loss,
|
|
})
|
|
}
|
|
|
|
/// Restore from checkpoint
|
|
pub async fn restore_state(&mut self, checkpoint: PINNCheckpoint) -> Result<()> {
|
|
if checkpoint.weights.len() != self.network.len() {
|
|
return Err(ScienceError::physics(
|
|
"Checkpoint layer count mismatch",
|
|
crate::error::PhysicsDomain::FluidDynamics,
|
|
));
|
|
}
|
|
|
|
for (i, layer) in self.network.iter_mut().enumerate() {
|
|
layer.weights = Variable::from_tensor(checkpoint.weights[i].clone());
|
|
layer.bias = Variable::from_tensor(checkpoint.biases[i].clone());
|
|
}
|
|
|
|
let mut state = self.state.write().await;
|
|
state.iteration = checkpoint.iteration;
|
|
state.best_loss = checkpoint.best_loss;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Predict without gradients for inference
|
|
pub async fn predict(&self, inputs: &Tensor) -> Result<Tensor> {
|
|
let mut x = inputs.clone();
|
|
|
|
for (i, layer) in self.network.iter().enumerate() {
|
|
// Manual forward pass without Variable tracking
|
|
let weights_tensor = layer.weights.value();
|
|
let bias_tensor = layer.bias.value();
|
|
|
|
x = x.matmul(weights_tensor)?;
|
|
x = x.add(bias_tensor)?;
|
|
|
|
// Apply activation (skip for output layer)
|
|
if i < self.network.len() - 1 {
|
|
x = apply_activation_tensor(x, &layer.activation)?;
|
|
}
|
|
}
|
|
|
|
Ok(x)
|
|
}
|
|
|
|
/// Get reference to physics loss function
|
|
pub fn physics_loss(&self) -> &dyn PhysicsLoss {
|
|
self.physics_loss.as_ref()
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Debug for PINN {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("PINN")
|
|
.field("network", &self.network)
|
|
.field("device", &self.device)
|
|
.field("physics_loss", &"<PhysicsLoss>")
|
|
.field(
|
|
"conservation_losses",
|
|
&format!("{} losses", self.conservation_losses.len()),
|
|
)
|
|
.field("config", &self.config)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl PINNBuilder {
|
|
/// Create a new PINN builder
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
device: None,
|
|
layers: Vec::new(),
|
|
activation: ActivationType::Tanh,
|
|
physics_loss: None,
|
|
conservation_losses: Vec::new(),
|
|
config: TrainingConfig::default(),
|
|
}
|
|
}
|
|
|
|
/// Set the computation device
|
|
#[must_use]
|
|
pub fn device(mut self, device: &Device) -> Self {
|
|
self.device = Some(device.clone());
|
|
self
|
|
}
|
|
|
|
/// Set the network layer configuration
|
|
#[must_use]
|
|
pub fn layers(mut self, layers: Vec<usize>) -> Self {
|
|
self.layers = layers;
|
|
self
|
|
}
|
|
|
|
/// Set the activation function
|
|
#[must_use]
|
|
pub fn activation(mut self, activation: ActivationType) -> Self {
|
|
self.activation = activation;
|
|
self
|
|
}
|
|
|
|
/// Set the physics loss function
|
|
#[must_use]
|
|
pub fn physics_loss(mut self, loss: Box<dyn PhysicsLoss + Send + Sync>) -> Self {
|
|
self.physics_loss = Some(loss);
|
|
self
|
|
}
|
|
|
|
/// Add a conservation law constraint
|
|
#[must_use]
|
|
pub fn conservation_loss(mut self, loss: Box<dyn ConservationLoss + Send + Sync>) -> Self {
|
|
self.conservation_losses.push(loss);
|
|
self
|
|
}
|
|
|
|
/// Set the training configuration
|
|
#[must_use]
|
|
pub fn config(mut self, config: TrainingConfig) -> Self {
|
|
self.config = config;
|
|
self
|
|
}
|
|
|
|
/// Build the PINN
|
|
pub fn build(self) -> Result<PINN> {
|
|
PINN::from_builder(self)
|
|
}
|
|
}
|
|
|
|
impl Default for PINNBuilder {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl LinearLayer {
|
|
/// Create a new linear layer
|
|
fn new(
|
|
input_size: usize,
|
|
output_size: usize,
|
|
activation: &ActivationType,
|
|
device: &Device,
|
|
) -> Result<Self> {
|
|
// Xavier/Glorot initialization
|
|
let std_dev = (2.0 / (input_size + output_size) as f64).sqrt();
|
|
|
|
let weights = Variable::new(
|
|
Tensor::randn(&[input_size, output_size], device)?.mul_scalar(std_dev as f32)?,
|
|
true, // requires_grad
|
|
);
|
|
let bias = Variable::new(Tensor::zeros([output_size], device)?, true);
|
|
|
|
Ok(Self {
|
|
weights,
|
|
bias,
|
|
activation: activation.clone(),
|
|
})
|
|
}
|
|
|
|
/// Forward pass through the layer
|
|
fn forward(&self, input: Variable) -> Result<Variable> {
|
|
let output = input
|
|
.matmul(&self.weights)
|
|
.map_err(|e| ScienceError::computation(e.to_string()))?;
|
|
output
|
|
.add(&self.bias)
|
|
.map_err(|e| ScienceError::computation(e.to_string()))
|
|
}
|
|
}
|
|
|
|
/// Apply activation function to a variable
|
|
fn apply_activation(input: Variable, activation: &ActivationType) -> Result<Variable> {
|
|
use crate::variable_extensions::VariableExt;
|
|
match activation {
|
|
ActivationType::Tanh => input.tanh(),
|
|
ActivationType::Sigmoid => input.sigmoid(),
|
|
ActivationType::Sin => input.sin(),
|
|
ActivationType::Swish => {
|
|
let sigmoid = input.sigmoid()?;
|
|
input.multiply(&sigmoid).map_err(std::convert::Into::into)
|
|
}
|
|
ActivationType::GELU => {
|
|
// GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
|
|
let x_cubed = input.multiply(&input)?.multiply(&input)?;
|
|
let inner = input.add(&x_cubed.multiply_scalar(0.044715)?)?;
|
|
let inner = inner.multiply_scalar((2.0 / std::f64::consts::PI).sqrt() as f32)?;
|
|
let tanh_part = inner.tanh()?;
|
|
let one_plus_tanh = tanh_part.add_scalar(1.0)?;
|
|
let result = input.multiply(&one_plus_tanh)?;
|
|
result
|
|
.multiply_scalar(0.5)
|
|
.map_err(std::convert::Into::into)
|
|
}
|
|
ActivationType::ReLU => VariableExt::relu(&input),
|
|
ActivationType::Learnable => {
|
|
// Simple learnable activation: a * tanh(b * x)
|
|
// For now, use standard tanh (parameters would need to be learned)
|
|
VariableExt::tanh(&input)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Apply activation function to a tensor (non-gradient version)
|
|
fn apply_activation_tensor(input: Tensor, activation: &ActivationType) -> Result<Tensor> {
|
|
match activation {
|
|
ActivationType::Tanh => input.tanh().map_err(std::convert::Into::into),
|
|
ActivationType::Sigmoid => input.sigmoid().map_err(std::convert::Into::into),
|
|
ActivationType::Sin => input.sin().map_err(std::convert::Into::into),
|
|
ActivationType::Swish => {
|
|
let sigmoid = input
|
|
.sigmoid()
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
input.mul(&sigmoid).map_err(std::convert::Into::into)
|
|
}
|
|
ActivationType::GELU => {
|
|
// GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
|
|
let x_squared = input
|
|
.mul(&input)
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
let x_cubed = x_squared
|
|
.mul(&input)
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
let inner = input
|
|
.add(
|
|
&x_cubed
|
|
.mul_scalar(0.044715)
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?,
|
|
)
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
let inner = inner
|
|
.mul_scalar((2.0 / std::f64::consts::PI).sqrt() as f32)
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
let tanh_part = inner
|
|
.tanh()
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
let one = Tensor::ones(input.shape().dims(), input.device())
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
let one_plus_tanh = one
|
|
.add(&tanh_part)
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
let result = input
|
|
.mul(&one_plus_tanh)
|
|
.map_err(|e: TensorError| -> ScienceError { e.into() })?;
|
|
result.mul_scalar(0.5).map_err(std::convert::Into::into)
|
|
}
|
|
ActivationType::ReLU => input.relu().map_err(std::convert::Into::into),
|
|
ActivationType::Learnable => {
|
|
// Simple learnable activation: for now, use standard tanh
|
|
input.tanh().map_err(std::convert::Into::into)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Create input tensor from coordinate pairs
|
|
fn create_input_tensor(inputs: &[(f64, f64)], device: &Device) -> Result<Tensor> {
|
|
let data: Vec<f32> = inputs
|
|
.iter()
|
|
.flat_map(|(x, t)| vec![*x as f32, *t as f32])
|
|
.collect();
|
|
|
|
Ok(Tensor::from_slice(&data, &[inputs.len(), 2], device)?)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::physics::HeatEquation;
|
|
|
|
#[tokio::test]
|
|
async fn test_pinn_creation() -> Result<()> {
|
|
let device = Device::cpu();
|
|
let heat_eq = HeatEquation::new(0.1);
|
|
|
|
let pinn = PINN::builder()
|
|
.device(&device)
|
|
.layers(vec![2, 64, 64, 1])
|
|
.physics_loss(Box::new(heat_eq))
|
|
.build()?;
|
|
|
|
assert_eq!(pinn.network.len(), 3);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pinn_forward_pass() -> Result<()> {
|
|
let device = Device::cpu();
|
|
let heat_eq = HeatEquation::new(0.1);
|
|
|
|
let pinn = PINN::builder()
|
|
.device(&device)
|
|
.layers(vec![2, 32, 1])
|
|
.physics_loss(Box::new(heat_eq))
|
|
.build()?;
|
|
|
|
// Create input tensor with shape [2, 2] for 2 samples with 2 features each
|
|
let inputs = Tensor::from_slice(&[0.5_f32, 1.0, 0.3, 0.8], &[2, 2], &device)?;
|
|
let outputs = pinn.predict(&inputs).await?;
|
|
|
|
// Output shape should be [2, 1] for 2 samples with 1 output each
|
|
assert_eq!(outputs.shape().dims(), &[2, 1]);
|
|
Ok(())
|
|
}
|
|
}
|