Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,616 @@
|
||||
//! 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());
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user