617 lines
19 KiB
Rust
617 lines
19 KiB
Rust
//! Physics-Informed Neural Network for Structural Mechanics
|
|
//!
|
|
//! This module implements the neural network architecture and physics losses
|
|
//! for solving structural mechanics problems using PINNs.
|
|
|
|
use structural_shared::{BoundaryCondition, Bounds2D, DisplacementField, Material, StressField};
|
|
|
|
// ============================================================================
|
|
// Activation Functions
|
|
// ============================================================================
|
|
|
|
/// Activation function type.
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub enum Activation {
|
|
Tanh,
|
|
Relu,
|
|
Sigmoid,
|
|
Swish,
|
|
Gelu,
|
|
Sin, // For Fourier-like feature embedding
|
|
}
|
|
|
|
impl Activation {
|
|
/// Parse activation from string.
|
|
pub fn from_str(s: &str) -> Self {
|
|
match s.to_lowercase().as_str() {
|
|
"tanh" => Self::Tanh,
|
|
"relu" => Self::Relu,
|
|
"sigmoid" => Self::Sigmoid,
|
|
"swish" => Self::Swish,
|
|
"gelu" => Self::Gelu,
|
|
"sin" => Self::Sin,
|
|
_ => Self::Tanh,
|
|
}
|
|
}
|
|
|
|
/// Apply activation function.
|
|
pub fn apply(&self, x: f64) -> f64 {
|
|
match self {
|
|
Self::Tanh => x.tanh(),
|
|
Self::Relu => x.max(0.0),
|
|
Self::Sigmoid => 1.0 / (1.0 + (-x).exp()),
|
|
Self::Swish => x * (1.0 / (1.0 + (-x).exp())),
|
|
Self::Gelu => {
|
|
0.5 * x
|
|
* (1.0
|
|
+ (std::f64::consts::FRAC_2_SQRT_PI * (x + 0.044715 * x.powi(3))
|
|
/ std::f64::consts::SQRT_2)
|
|
.tanh())
|
|
}
|
|
Self::Sin => x.sin(),
|
|
}
|
|
}
|
|
|
|
/// Apply derivative of activation function.
|
|
pub fn derivative(&self, x: f64) -> f64 {
|
|
match self {
|
|
Self::Tanh => 1.0 - x.tanh().powi(2),
|
|
Self::Relu => {
|
|
if x > 0.0 {
|
|
1.0
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
Self::Sigmoid => {
|
|
let s = 1.0 / (1.0 + (-x).exp());
|
|
s * (1.0 - s)
|
|
}
|
|
Self::Swish => {
|
|
let s = 1.0 / (1.0 + (-x).exp());
|
|
s + x * s * (1.0 - s)
|
|
}
|
|
Self::Gelu => {
|
|
// Approximate derivative
|
|
let cdf = 0.5 * (1.0 + (x / std::f64::consts::SQRT_2).tanh());
|
|
let pdf = (-0.5 * x * x).exp() / (2.0 * std::f64::consts::PI).sqrt();
|
|
cdf + x * pdf
|
|
}
|
|
Self::Sin => x.cos(),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Dense Layer
|
|
// ============================================================================
|
|
|
|
/// Dense (fully connected) layer.
|
|
#[derive(Debug, Clone)]
|
|
pub struct DenseLayer {
|
|
/// Weight matrix (out_dim x in_dim).
|
|
weights: Vec<Vec<f64>>,
|
|
/// Bias vector.
|
|
biases: Vec<f64>,
|
|
/// Input dimension.
|
|
in_dim: usize,
|
|
/// Output dimension.
|
|
out_dim: usize,
|
|
}
|
|
|
|
impl DenseLayer {
|
|
/// Create a new dense layer with Xavier initialization.
|
|
pub fn new(in_dim: usize, out_dim: usize, seed: u64) -> Self {
|
|
let mut rng_state = seed;
|
|
let scale = (2.0 / (in_dim + out_dim) as f64).sqrt();
|
|
|
|
let mut weights = vec![vec![0.0; in_dim]; out_dim];
|
|
let mut biases = vec![0.0; out_dim];
|
|
|
|
// Xavier initialization
|
|
for row in &mut weights {
|
|
for w in row {
|
|
rng_state = rng_state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
let u = (rng_state >> 11) as f64 / (1u64 << 53) as f64;
|
|
*w = (u * 2.0 - 1.0) * scale;
|
|
}
|
|
}
|
|
|
|
for b in &mut biases {
|
|
*b = 0.0; // Initialize biases to zero
|
|
}
|
|
|
|
Self {
|
|
weights,
|
|
biases,
|
|
in_dim,
|
|
out_dim,
|
|
}
|
|
}
|
|
|
|
/// Forward pass.
|
|
pub fn forward(&self, input: &[f64]) -> Vec<f64> {
|
|
assert_eq!(input.len(), self.in_dim);
|
|
|
|
let mut output = vec![0.0; self.out_dim];
|
|
for i in 0..self.out_dim {
|
|
let mut sum = self.biases[i];
|
|
for j in 0..self.in_dim {
|
|
sum += self.weights[i][j] * input[j];
|
|
}
|
|
output[i] = sum;
|
|
}
|
|
output
|
|
}
|
|
|
|
/// Update weights with gradient descent (simulated).
|
|
pub fn update(&mut self, learning_rate: f64) {
|
|
// Simulated weight update with small perturbation
|
|
for row in &mut self.weights {
|
|
for w in row {
|
|
*w *= 1.0 - learning_rate * 0.001;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Structural Network (MLP)
|
|
// ============================================================================
|
|
|
|
/// Multi-layer perceptron for structural mechanics.
|
|
#[derive(Debug)]
|
|
pub struct StructuralNetwork {
|
|
/// Hidden layers.
|
|
layers: Vec<DenseLayer>,
|
|
/// Activation function.
|
|
activation: Activation,
|
|
/// Number of hidden layers.
|
|
num_layers: usize,
|
|
/// Hidden dimension.
|
|
hidden_dim: usize,
|
|
/// Fourier feature frequencies for positional encoding.
|
|
fourier_frequencies: Vec<f64>,
|
|
}
|
|
|
|
impl StructuralNetwork {
|
|
/// Create a new structural network.
|
|
pub fn new(num_layers: usize, hidden_dim: usize, activation_str: &str) -> Self {
|
|
let activation = Activation::from_str(activation_str);
|
|
|
|
// Input: (x, y) + Fourier features
|
|
let num_frequencies = 10;
|
|
let fourier_dim = 2 * 2 * num_frequencies; // sin and cos for x and y
|
|
let input_dim = 2 + fourier_dim;
|
|
|
|
// Output: (u_x, u_y)
|
|
let output_dim = 2;
|
|
|
|
let mut layers = Vec::with_capacity(num_layers + 1);
|
|
let mut seed = 42u64;
|
|
|
|
// Input layer
|
|
layers.push(DenseLayer::new(input_dim, hidden_dim, seed));
|
|
seed += 1;
|
|
|
|
// Hidden layers
|
|
for _ in 1..num_layers {
|
|
layers.push(DenseLayer::new(hidden_dim, hidden_dim, seed));
|
|
seed += 1;
|
|
}
|
|
|
|
// Output layer
|
|
layers.push(DenseLayer::new(hidden_dim, output_dim, seed));
|
|
|
|
// Fourier frequencies for positional encoding
|
|
let fourier_frequencies: Vec<f64> = (0..num_frequencies)
|
|
.map(|i| 2.0f64.powi(i as i32) * std::f64::consts::PI)
|
|
.collect();
|
|
|
|
Self {
|
|
layers,
|
|
activation,
|
|
num_layers,
|
|
hidden_dim,
|
|
fourier_frequencies,
|
|
}
|
|
}
|
|
|
|
/// Apply Fourier positional encoding.
|
|
fn fourier_features(&self, x: f64, y: f64) -> Vec<f64> {
|
|
let mut features = vec![x, y];
|
|
|
|
for &freq in &self.fourier_frequencies {
|
|
features.push((freq * x).sin());
|
|
features.push((freq * x).cos());
|
|
features.push((freq * y).sin());
|
|
features.push((freq * y).cos());
|
|
}
|
|
|
|
features
|
|
}
|
|
|
|
/// Forward pass for 2D input.
|
|
pub fn forward_2d(&self, x: f64, y: f64) -> (f64, f64) {
|
|
// Apply Fourier features
|
|
let mut current = self.fourier_features(x, y);
|
|
|
|
// Forward through hidden layers
|
|
for (i, layer) in self.layers.iter().enumerate() {
|
|
current = layer.forward(¤t);
|
|
|
|
// Apply activation to all but last layer
|
|
if i < self.layers.len() - 1 {
|
|
for val in &mut current {
|
|
*val = self.activation.apply(*val);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Output: (u_x, u_y)
|
|
(current[0], current[1])
|
|
}
|
|
|
|
/// Update weights.
|
|
pub fn update_weights(&mut self, learning_rate: f64) {
|
|
for layer in &mut self.layers {
|
|
layer.update(learning_rate);
|
|
}
|
|
}
|
|
|
|
/// Get number of parameters.
|
|
pub fn num_parameters(&self) -> usize {
|
|
let mut count = 0;
|
|
for layer in &self.layers {
|
|
count += layer.weights.len() * layer.weights[0].len();
|
|
count += layer.biases.len();
|
|
}
|
|
count
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Elasticity Loss (Navier-Cauchy Equations)
|
|
// ============================================================================
|
|
|
|
/// Elasticity loss implementing Navier-Cauchy equations.
|
|
///
|
|
/// The Navier-Cauchy equations for linear elasticity are:
|
|
/// (lambda + mu) * grad(div(u)) + mu * laplacian(u) + f = 0
|
|
///
|
|
/// In component form (2D, no body forces):
|
|
/// (lambda + 2*mu) * d2u/dx2 + (lambda + mu) * d2v/dxdy + mu * d2u/dy2 = 0
|
|
/// (lambda + mu) * d2u/dxdy + mu * d2v/dx2 + (lambda + 2*mu) * d2v/dy2 = 0
|
|
#[derive(Debug)]
|
|
pub struct ElasticityLoss {
|
|
/// Lame's first parameter.
|
|
lambda: f64,
|
|
/// Lame's second parameter (shear modulus).
|
|
mu: f64,
|
|
}
|
|
|
|
impl ElasticityLoss {
|
|
/// Create a new elasticity loss.
|
|
pub fn new(material: Material) -> Self {
|
|
Self {
|
|
lambda: material.lame_lambda(),
|
|
mu: material.lame_mu(),
|
|
}
|
|
}
|
|
|
|
/// Evaluate the physics loss for 2D linear elasticity.
|
|
pub fn evaluate_2d(
|
|
&self,
|
|
displacement: &DisplacementField,
|
|
_stress: &StressField,
|
|
bounds: &Bounds2D,
|
|
) -> f64 {
|
|
let n = displacement.u_x.len();
|
|
if n < 3 {
|
|
return 0.0;
|
|
}
|
|
|
|
let dx = bounds.width() / (n as f64).sqrt();
|
|
let dy = bounds.height() / (n as f64).sqrt();
|
|
|
|
let mut residual = 0.0;
|
|
let mut count = 0;
|
|
|
|
// Evaluate residual at interior points
|
|
for i in 1..n - 1 {
|
|
// Approximate second derivatives using central differences
|
|
let u = &displacement.u_x;
|
|
let v = &displacement.u_y;
|
|
|
|
// d2u/dx2 approximation (using neighboring indices)
|
|
let d2u_dx2 = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
|
|
// d2u/dy2 approximation
|
|
let d2u_dy2 = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dy * dy);
|
|
// d2v/dx2 approximation
|
|
let d2v_dx2 = (v[i + 1] - 2.0 * v[i] + v[i - 1]) / (dx * dx);
|
|
// d2v/dy2 approximation
|
|
let d2v_dy2 = (v[i + 1] - 2.0 * v[i] + v[i - 1]) / (dy * dy);
|
|
|
|
// Mixed derivative approximation (simplified)
|
|
let d2u_dxdy = (u[i + 1] - u[i - 1]) / (4.0 * dx * dy);
|
|
let d2v_dxdy = (v[i + 1] - v[i - 1]) / (4.0 * dx * dy);
|
|
|
|
// Navier-Cauchy residuals
|
|
let res_x = (self.lambda + 2.0 * self.mu) * d2u_dx2
|
|
+ (self.lambda + self.mu) * d2v_dxdy
|
|
+ self.mu * d2u_dy2;
|
|
|
|
let res_y = (self.lambda + self.mu) * d2u_dxdy
|
|
+ self.mu * d2v_dx2
|
|
+ (self.lambda + 2.0 * self.mu) * d2v_dy2;
|
|
|
|
// Normalized residual
|
|
let norm = (self.lambda + 2.0 * self.mu).max(1.0);
|
|
residual += (res_x / norm).powi(2) + (res_y / norm).powi(2);
|
|
count += 1;
|
|
}
|
|
|
|
if count > 0 {
|
|
(residual / count as f64).sqrt()
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
|
|
/// Evaluate equilibrium equations (stress divergence form).
|
|
pub fn evaluate_equilibrium(&self, stress: &StressField, bounds: &Bounds2D) -> f64 {
|
|
let n = stress.sigma_xx.len();
|
|
if n < 3 {
|
|
return 0.0;
|
|
}
|
|
|
|
let dx = bounds.width() / (n as f64).sqrt();
|
|
let dy = bounds.height() / (n as f64).sqrt();
|
|
|
|
let mut residual = 0.0;
|
|
let mut count = 0;
|
|
|
|
// Evaluate equilibrium: div(sigma) = 0
|
|
// d(sigma_xx)/dx + d(sigma_xy)/dy = 0
|
|
// d(sigma_xy)/dx + d(sigma_yy)/dy = 0
|
|
for i in 1..n - 1 {
|
|
let dsxx_dx = (stress.sigma_xx[i + 1] - stress.sigma_xx[i - 1]) / (2.0 * dx);
|
|
let dsxy_dy = (stress.sigma_xy[i + 1] - stress.sigma_xy[i - 1]) / (2.0 * dy);
|
|
let dsxy_dx = (stress.sigma_xy[i + 1] - stress.sigma_xy[i - 1]) / (2.0 * dx);
|
|
let dsyy_dy = (stress.sigma_yy[i + 1] - stress.sigma_yy[i - 1]) / (2.0 * dy);
|
|
|
|
let res_x = dsxx_dx + dsxy_dy;
|
|
let res_y = dsxy_dx + dsyy_dy;
|
|
|
|
// Normalize by characteristic stress
|
|
let char_stress = stress.sigma_xx[i]
|
|
.abs()
|
|
.max(stress.sigma_yy[i].abs())
|
|
.max(1.0);
|
|
residual += (res_x / char_stress).powi(2) + (res_y / char_stress).powi(2);
|
|
count += 1;
|
|
}
|
|
|
|
if count > 0 {
|
|
(residual / count as f64).sqrt()
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Boundary Loss
|
|
// ============================================================================
|
|
|
|
/// Loss for enforcing boundary conditions.
|
|
#[derive(Debug)]
|
|
pub struct BoundaryLoss {
|
|
/// Penalty factor for Dirichlet BCs.
|
|
dirichlet_penalty: f64,
|
|
/// Penalty factor for Neumann BCs.
|
|
neumann_penalty: f64,
|
|
}
|
|
|
|
impl BoundaryLoss {
|
|
/// Create a new boundary loss.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
dirichlet_penalty: 1.0,
|
|
neumann_penalty: 1.0,
|
|
}
|
|
}
|
|
|
|
/// Evaluate boundary condition loss.
|
|
pub fn evaluate(
|
|
&self,
|
|
displacement: &DisplacementField,
|
|
boundary_conditions: &[BoundaryCondition],
|
|
) -> f64 {
|
|
let mut total_loss = 0.0;
|
|
let mut count = 0;
|
|
|
|
for bc in boundary_conditions {
|
|
match bc {
|
|
BoundaryCondition::Dirichlet {
|
|
nodes,
|
|
displacement: prescribed,
|
|
constrained,
|
|
} => {
|
|
for &node in nodes {
|
|
if node < displacement.u_x.len() {
|
|
// X-displacement
|
|
if !constrained.is_empty() && constrained[0] {
|
|
let target = if !prescribed.is_empty() {
|
|
prescribed[0]
|
|
} else {
|
|
0.0
|
|
};
|
|
let error = displacement.u_x[node] - target;
|
|
total_loss += self.dirichlet_penalty * error.powi(2);
|
|
count += 1;
|
|
}
|
|
|
|
// Y-displacement
|
|
if constrained.len() > 1 && constrained[1] {
|
|
let target = if prescribed.len() > 1 {
|
|
prescribed[1]
|
|
} else {
|
|
0.0
|
|
};
|
|
let error = displacement.u_y[node] - target;
|
|
total_loss += self.dirichlet_penalty * error.powi(2);
|
|
count += 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
BoundaryCondition::Neumann { nodes, traction } => {
|
|
// For Neumann BCs, we would need stress at boundary
|
|
// Simplified: just count the contribution
|
|
let _traction_mag: f64 = traction.iter().map(|t| t.abs()).sum();
|
|
for &_node in nodes {
|
|
count += 1;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
if count > 0 {
|
|
(total_loss / count as f64).sqrt()
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
|
|
/// Set Dirichlet penalty factor.
|
|
pub fn set_dirichlet_penalty(&mut self, penalty: f64) {
|
|
self.dirichlet_penalty = penalty;
|
|
}
|
|
|
|
/// Set Neumann penalty factor.
|
|
pub fn set_neumann_penalty(&mut self, penalty: f64) {
|
|
self.neumann_penalty = penalty;
|
|
}
|
|
}
|
|
|
|
impl Default for BoundaryLoss {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_activation_tanh() {
|
|
let act = Activation::Tanh;
|
|
assert!((act.apply(0.0)).abs() < 1e-10);
|
|
assert!((act.apply(1.0) - 1.0f64.tanh()).abs() < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_activation_relu() {
|
|
let act = Activation::Relu;
|
|
assert!((act.apply(-1.0)).abs() < 1e-10);
|
|
assert!((act.apply(1.0) - 1.0).abs() < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_activation_gelu() {
|
|
let act = Activation::Gelu;
|
|
let val = act.apply(1.0);
|
|
assert!(val > 0.0 && val < 1.5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dense_layer() {
|
|
let layer = DenseLayer::new(3, 2, 42);
|
|
let input = vec![1.0, 0.5, -0.5];
|
|
let output = layer.forward(&input);
|
|
assert_eq!(output.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_structural_network() {
|
|
let network = StructuralNetwork::new(3, 32, "tanh");
|
|
let (u, v) = network.forward_2d(0.5, 0.5);
|
|
assert!(u.is_finite());
|
|
assert!(v.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_network_parameters() {
|
|
let network = StructuralNetwork::new(4, 64, "tanh");
|
|
let params = network.num_parameters();
|
|
assert!(params > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_elasticity_loss() {
|
|
let material = Material::steel();
|
|
let loss = ElasticityLoss::new(material);
|
|
|
|
let mut displacement = DisplacementField::with_size(100);
|
|
// Linear displacement (zero stress state)
|
|
for i in 0..100 {
|
|
let x = i as f64 / 99.0;
|
|
displacement.u_x[i] = 0.001 * x;
|
|
displacement.u_y[i] = 0.0;
|
|
}
|
|
|
|
let stress = StressField::with_size(100);
|
|
let bounds = Bounds2D::default();
|
|
let residual = loss.evaluate_2d(&displacement, &stress, &bounds);
|
|
assert!(residual.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_boundary_loss() {
|
|
let loss = BoundaryLoss::new();
|
|
|
|
let mut displacement = DisplacementField::with_size(10);
|
|
displacement.u_x = vec![0.0; 10];
|
|
displacement.u_y = vec![0.0; 10];
|
|
|
|
let bc = BoundaryCondition::fixed(vec![0, 1, 2], 2);
|
|
let bcs = vec![bc];
|
|
|
|
let residual = loss.evaluate(&displacement, &bcs);
|
|
assert!(residual < 1e-10); // Perfect enforcement
|
|
}
|
|
|
|
#[test]
|
|
fn test_fourier_features() {
|
|
let network = StructuralNetwork::new(2, 32, "tanh");
|
|
let features = network.fourier_features(0.5, 0.5);
|
|
// 2 base + 4 * num_frequencies
|
|
assert!(features.len() > 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_activation_derivatives() {
|
|
let activations = vec![
|
|
Activation::Tanh,
|
|
Activation::Relu,
|
|
Activation::Sigmoid,
|
|
Activation::Swish,
|
|
Activation::Sin,
|
|
];
|
|
|
|
for act in activations {
|
|
let deriv = act.derivative(0.5);
|
|
assert!(deriv.is_finite());
|
|
}
|
|
}
|
|
}
|