328 lines
11 KiB
Rust
328 lines
11 KiB
Rust
//! 4D PINN network for temperature prediction T(x, y, z, t)
|
|
//!
|
|
//! Uses Learnable Fourier Feature Network (LFFN) architecture with
|
|
//! 4-dimensional input for spatial coordinates and time.
|
|
|
|
use crate::config::{Activation, NetworkConfig};
|
|
use bioheat_shared::Point3D;
|
|
use rand::Rng;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// 4D coordinate input to the network
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct Coordinate4D {
|
|
pub x: f64,
|
|
pub y: f64,
|
|
pub z: f64,
|
|
pub t: f64,
|
|
}
|
|
|
|
impl Coordinate4D {
|
|
/// Create new 4D coordinate
|
|
#[must_use]
|
|
pub const fn new(x: f64, y: f64, z: f64, t: f64) -> Self {
|
|
Self { x, y, z, t }
|
|
}
|
|
|
|
/// Create from Point3D and time
|
|
#[must_use]
|
|
pub fn from_point_and_time(p: Point3D, t: f32) -> Self {
|
|
Self::new(p.x as f64, p.y as f64, p.z as f64, t as f64)
|
|
}
|
|
|
|
/// Convert to array
|
|
#[must_use]
|
|
pub fn to_array(&self) -> [f64; 4] {
|
|
[self.x, self.y, self.z, self.t]
|
|
}
|
|
}
|
|
|
|
/// Thermal PINN network for predicting temperature T(x,y,z,t)
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ThermalPinn {
|
|
/// Configuration
|
|
config: NetworkConfig,
|
|
/// Fourier feature matrix B [4 x fourier_features]
|
|
/// Each row corresponds to frequencies for (x, y, z, t)
|
|
fourier_b: Vec<Vec<f64>>,
|
|
/// MLP weights for each layer
|
|
weights: Vec<Vec<Vec<f64>>>,
|
|
/// MLP biases for each layer
|
|
biases: Vec<Vec<f64>>,
|
|
}
|
|
|
|
impl ThermalPinn {
|
|
/// Create a new ThermalPinn with given configuration
|
|
#[must_use]
|
|
pub fn new(config: NetworkConfig) -> Self {
|
|
let mut rng = rand::thread_rng();
|
|
|
|
// Initialize Fourier feature matrix B
|
|
// Shape: [4, fourier_features]
|
|
let fourier_b: Vec<Vec<f64>> = (0..4)
|
|
.map(|_| {
|
|
(0..config.fourier_features)
|
|
.map(|_| rng.r#gen::<f64>() * config.fourier_scale as f64)
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
// Build MLP layer sizes
|
|
// Input: 2 * fourier_features (cos + sin features)
|
|
let input_dim = 2 * config.fourier_features;
|
|
let mut layer_sizes = vec![input_dim];
|
|
layer_sizes.extend(config.hidden_layers.iter().copied());
|
|
layer_sizes.push(1); // Output: temperature
|
|
|
|
// Initialize weights and biases with Xavier initialization
|
|
let mut weights = Vec::new();
|
|
let mut biases = Vec::new();
|
|
|
|
for i in 0..layer_sizes.len() - 1 {
|
|
let fan_in = layer_sizes[i];
|
|
let fan_out = layer_sizes[i + 1];
|
|
let scale = (6.0 / (fan_in + fan_out) as f64).sqrt();
|
|
|
|
// Weight matrix [fan_out x fan_in]
|
|
let w: Vec<Vec<f64>> = (0..fan_out)
|
|
.map(|_| (0..fan_in).map(|_| rng.gen_range(-scale..scale)).collect())
|
|
.collect();
|
|
weights.push(w);
|
|
|
|
// Bias vector [fan_out]
|
|
let b: Vec<f64> = (0..fan_out).map(|_| 0.0).collect();
|
|
biases.push(b);
|
|
}
|
|
|
|
Self {
|
|
config,
|
|
fourier_b,
|
|
weights,
|
|
biases,
|
|
}
|
|
}
|
|
|
|
/// Create with default configuration
|
|
#[must_use]
|
|
pub fn default_network() -> Self {
|
|
Self::new(NetworkConfig::default())
|
|
}
|
|
|
|
/// Compute Fourier features for input coordinate
|
|
fn fourier_features(&self, coord: &Coordinate4D) -> Vec<f64> {
|
|
let input = coord.to_array();
|
|
let mut features = Vec::with_capacity(2 * self.config.fourier_features);
|
|
|
|
for j in 0..self.config.fourier_features {
|
|
// Compute dot product: B[:, j] · input
|
|
let dot: f64 = (0..4).map(|i| self.fourier_b[i][j] * input[i]).sum();
|
|
let angle = 2.0 * std::f64::consts::PI * dot;
|
|
|
|
// Append cos and sin features
|
|
features.push(angle.cos());
|
|
features.push(angle.sin());
|
|
}
|
|
|
|
features
|
|
}
|
|
|
|
/// Apply activation function
|
|
fn activate(&self, x: f64) -> f64 {
|
|
match self.config.activation {
|
|
Activation::Tanh => x.tanh(),
|
|
Activation::Swish => x * (1.0 / (1.0 + (-x).exp())), // x * sigmoid(x)
|
|
Activation::Gelu => {
|
|
// Approximate GELU
|
|
0.5 * x * (1.0 + (0.7978845608 * (x + 0.044715 * x.powi(3))).tanh())
|
|
}
|
|
Activation::Sin => x.sin(),
|
|
}
|
|
}
|
|
|
|
/// Forward pass through MLP
|
|
fn mlp_forward(&self, features: &[f64]) -> f64 {
|
|
let mut current = features.to_vec();
|
|
|
|
for (layer_idx, (w, b)) in self.weights.iter().zip(&self.biases).enumerate() {
|
|
let is_last = layer_idx == self.weights.len() - 1;
|
|
|
|
let mut next = Vec::with_capacity(w.len());
|
|
for (row, bias) in w.iter().zip(b) {
|
|
let sum: f64 = row.iter().zip(¤t).map(|(wi, xi)| wi * xi).sum();
|
|
let activated = if is_last {
|
|
sum + bias // No activation on output layer
|
|
} else {
|
|
self.activate(sum + bias)
|
|
};
|
|
next.push(activated);
|
|
}
|
|
current = next;
|
|
}
|
|
|
|
// Output is a scalar (temperature)
|
|
current[0]
|
|
}
|
|
|
|
/// Predict temperature at a 4D coordinate
|
|
#[must_use]
|
|
pub fn forward(&self, coord: &Coordinate4D) -> f64 {
|
|
let features = self.fourier_features(coord);
|
|
self.mlp_forward(&features)
|
|
}
|
|
|
|
/// Predict temperature from components
|
|
#[must_use]
|
|
pub fn predict(&self, x: f64, y: f64, z: f64, t: f64) -> f64 {
|
|
self.forward(&Coordinate4D::new(x, y, z, t))
|
|
}
|
|
|
|
/// Compute gradients using finite differences
|
|
/// Returns (dT/dx, dT/dy, dT/dz, dT/dt)
|
|
#[must_use]
|
|
pub fn gradients(&self, coord: &Coordinate4D, eps: f64) -> (f64, f64, f64, f64) {
|
|
let _t0 = self.forward(coord);
|
|
|
|
// dT/dx
|
|
let tx_plus = self.forward(&Coordinate4D::new(coord.x + eps, coord.y, coord.z, coord.t));
|
|
let tx_minus = self.forward(&Coordinate4D::new(coord.x - eps, coord.y, coord.z, coord.t));
|
|
let dt_dx = (tx_plus - tx_minus) / (2.0 * eps);
|
|
|
|
// dT/dy
|
|
let ty_plus = self.forward(&Coordinate4D::new(coord.x, coord.y + eps, coord.z, coord.t));
|
|
let ty_minus = self.forward(&Coordinate4D::new(coord.x, coord.y - eps, coord.z, coord.t));
|
|
let dt_dy = (ty_plus - ty_minus) / (2.0 * eps);
|
|
|
|
// dT/dz
|
|
let tz_plus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z + eps, coord.t));
|
|
let tz_minus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z - eps, coord.t));
|
|
let dt_dz = (tz_plus - tz_minus) / (2.0 * eps);
|
|
|
|
// dT/dt
|
|
let tt_plus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z, coord.t + eps));
|
|
let tt_minus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z, coord.t - eps));
|
|
let dt_dt = (tt_plus - tt_minus) / (2.0 * eps);
|
|
|
|
(dt_dx, dt_dy, dt_dz, dt_dt)
|
|
}
|
|
|
|
/// Compute second derivatives (for Laplacian)
|
|
/// Returns (d²T/dx², d²T/dy², d²T/dz²)
|
|
#[must_use]
|
|
pub fn second_derivatives(&self, coord: &Coordinate4D, eps: f64) -> (f64, f64, f64) {
|
|
let t0 = self.forward(coord);
|
|
|
|
// d²T/dx²
|
|
let tx_plus = self.forward(&Coordinate4D::new(coord.x + eps, coord.y, coord.z, coord.t));
|
|
let tx_minus = self.forward(&Coordinate4D::new(coord.x - eps, coord.y, coord.z, coord.t));
|
|
let d2t_dx2 = (tx_plus - 2.0 * t0 + tx_minus) / (eps * eps);
|
|
|
|
// d²T/dy²
|
|
let ty_plus = self.forward(&Coordinate4D::new(coord.x, coord.y + eps, coord.z, coord.t));
|
|
let ty_minus = self.forward(&Coordinate4D::new(coord.x, coord.y - eps, coord.z, coord.t));
|
|
let d2t_dy2 = (ty_plus - 2.0 * t0 + ty_minus) / (eps * eps);
|
|
|
|
// d²T/dz²
|
|
let tz_plus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z + eps, coord.t));
|
|
let tz_minus = self.forward(&Coordinate4D::new(coord.x, coord.y, coord.z - eps, coord.t));
|
|
let d2t_dz2 = (tz_plus - 2.0 * t0 + tz_minus) / (eps * eps);
|
|
|
|
(d2t_dx2, d2t_dy2, d2t_dz2)
|
|
}
|
|
|
|
/// Compute Laplacian ∇²T = d²T/dx² + d²T/dy² + d²T/dz²
|
|
#[must_use]
|
|
pub fn laplacian(&self, coord: &Coordinate4D, eps: f64) -> f64 {
|
|
let (d2x, d2y, d2z) = self.second_derivatives(coord, eps);
|
|
d2x + d2y + d2z
|
|
}
|
|
|
|
/// Get number of trainable parameters
|
|
#[must_use]
|
|
pub fn num_parameters(&self) -> usize {
|
|
let fourier_params = 4 * self.config.fourier_features;
|
|
let mlp_params: usize = self
|
|
.weights
|
|
.iter()
|
|
.zip(&self.biases)
|
|
.map(|(w, b)| w.len() * w[0].len() + b.len())
|
|
.sum();
|
|
fourier_params + mlp_params
|
|
}
|
|
|
|
/// Get mutable access to weights for training
|
|
pub fn weights_mut(&mut self) -> &mut Vec<Vec<Vec<f64>>> {
|
|
&mut self.weights
|
|
}
|
|
|
|
/// Get mutable access to biases for training
|
|
pub fn biases_mut(&mut self) -> &mut Vec<Vec<f64>> {
|
|
&mut self.biases
|
|
}
|
|
|
|
/// Get mutable access to Fourier features for training
|
|
pub fn fourier_b_mut(&mut self) -> &mut Vec<Vec<f64>> {
|
|
&mut self.fourier_b
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn test_config() -> NetworkConfig {
|
|
NetworkConfig {
|
|
fourier_features: 16,
|
|
fourier_scale: 2.0,
|
|
hidden_layers: vec![32, 32],
|
|
activation: Activation::Tanh,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_network_creation() {
|
|
let net = ThermalPinn::new(test_config());
|
|
assert!(net.num_parameters() > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_forward_pass() {
|
|
let net = ThermalPinn::new(test_config());
|
|
let coord = Coordinate4D::new(0.0, 0.0, 0.0, 0.0);
|
|
let t = net.forward(&coord);
|
|
// Output should be a finite number
|
|
assert!(t.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_gradients() {
|
|
let net = ThermalPinn::new(test_config());
|
|
let coord = Coordinate4D::new(0.05, 0.05, 0.05, 100.0);
|
|
let (dx, dy, dz, dt) = net.gradients(&coord, 1e-5);
|
|
|
|
// Gradients should be finite
|
|
assert!(dx.is_finite());
|
|
assert!(dy.is_finite());
|
|
assert!(dz.is_finite());
|
|
assert!(dt.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_laplacian() {
|
|
let net = ThermalPinn::new(test_config());
|
|
let coord = Coordinate4D::new(0.05, 0.05, 0.05, 100.0);
|
|
let lap = net.laplacian(&coord, 1e-4);
|
|
assert!(lap.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_batch_consistency() {
|
|
let net = ThermalPinn::new(test_config());
|
|
let coord = Coordinate4D::new(0.03, 0.04, 0.02, 50.0);
|
|
|
|
// Forward pass should be deterministic
|
|
let t1 = net.forward(&coord);
|
|
let t2 = net.forward(&coord);
|
|
assert!((t1 - t2).abs() < 1e-10);
|
|
}
|
|
}
|