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