618 lines
18 KiB
Rust
618 lines
18 KiB
Rust
//! Neural network architecture for source localization PINN.
|
||
//!
|
||
//! Implements a physics-informed network with Fourier feature encoding
|
||
//! for learning neural source distributions from MEG/EEG data.
|
||
|
||
use crate::error::{PinnError, PinnResult};
|
||
use ndarray::{Array1, Array2};
|
||
use serde::{Deserialize, Serialize};
|
||
|
||
/// Fourier feature encoding for positional inputs
|
||
///
|
||
/// Maps 3D coordinates to higher-dimensional space using:
|
||
/// γ(x) = [sin(2πBx), cos(2πBx)]
|
||
///
|
||
/// where B is a matrix of learnable or random frequencies.
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct FourierFeatures {
|
||
/// Frequency matrix [n_features, 3]
|
||
frequencies: Array2<f64>,
|
||
/// Scale factor for frequencies
|
||
scale: f64,
|
||
}
|
||
|
||
impl FourierFeatures {
|
||
/// Create new Fourier features with random frequencies
|
||
pub fn new(n_features: usize, scale: f64) -> Self {
|
||
// Initialize with random frequencies (normally distributed)
|
||
let mut frequencies = Array2::zeros((n_features, 3));
|
||
for i in 0..n_features {
|
||
for j in 0..3 {
|
||
// Simple pseudo-random initialization
|
||
let seed = (i * 3 + j) as f64;
|
||
frequencies[[i, j]] = (seed * 0.618033988749895).fract() * 2.0 - 1.0;
|
||
}
|
||
}
|
||
|
||
Self {
|
||
frequencies: frequencies * scale,
|
||
scale,
|
||
}
|
||
}
|
||
|
||
/// Create with specific frequencies
|
||
pub fn with_frequencies(frequencies: Array2<f64>, scale: f64) -> Self {
|
||
Self { frequencies, scale }
|
||
}
|
||
|
||
/// Number of output features (2x input for sin and cos)
|
||
pub fn output_dim(&self) -> usize {
|
||
self.frequencies.nrows() * 2
|
||
}
|
||
|
||
/// Encode 3D coordinates
|
||
pub fn encode(&self, coords: &Array2<f64>) -> Array2<f64> {
|
||
let n_points = coords.nrows();
|
||
let n_freqs = self.frequencies.nrows();
|
||
let mut output = Array2::zeros((n_points, n_freqs * 2));
|
||
|
||
for p in 0..n_points {
|
||
for f in 0..n_freqs {
|
||
// Compute dot product: B[f] · x[p]
|
||
let mut dot = 0.0;
|
||
for d in 0..3 {
|
||
dot += self.frequencies[[f, d]] * coords[[p, d]];
|
||
}
|
||
let angle = 2.0 * std::f64::consts::PI * dot;
|
||
output[[p, f]] = angle.sin();
|
||
output[[p, f + n_freqs]] = angle.cos();
|
||
}
|
||
}
|
||
|
||
output
|
||
}
|
||
|
||
/// Encode a single point
|
||
pub fn encode_point(&self, point: &[f64; 3]) -> Array1<f64> {
|
||
let n_freqs = self.frequencies.nrows();
|
||
let mut output = Array1::zeros(n_freqs * 2);
|
||
|
||
for f in 0..n_freqs {
|
||
let mut dot = 0.0;
|
||
for d in 0..3 {
|
||
dot += self.frequencies[[f, d]] * point[d];
|
||
}
|
||
let angle = 2.0 * std::f64::consts::PI * dot;
|
||
output[f] = angle.sin();
|
||
output[f + n_freqs] = angle.cos();
|
||
}
|
||
|
||
output
|
||
}
|
||
}
|
||
|
||
/// Activation function types
|
||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
|
||
pub enum Activation {
|
||
/// Hyperbolic tangent
|
||
Tanh,
|
||
/// Rectified Linear Unit
|
||
ReLU,
|
||
/// Sigmoid
|
||
Sigmoid,
|
||
/// Swish (x * sigmoid(x))
|
||
Swish,
|
||
/// GELU (Gaussian Error Linear Unit)
|
||
GELU,
|
||
}
|
||
|
||
impl Activation {
|
||
/// Apply activation function
|
||
pub fn apply(&self, x: f64) -> f64 {
|
||
match self {
|
||
Activation::Tanh => x.tanh(),
|
||
Activation::ReLU => x.max(0.0),
|
||
Activation::Sigmoid => 1.0 / (1.0 + (-x).exp()),
|
||
Activation::Swish => x / (1.0 + (-x).exp()),
|
||
Activation::GELU => {
|
||
let sqrt2 = std::f64::consts::SQRT_2;
|
||
0.5 * x * (1.0 + libm::erf(x / sqrt2))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Apply derivative for backpropagation
|
||
pub fn derivative(&self, x: f64) -> f64 {
|
||
match self {
|
||
Activation::Tanh => {
|
||
let t = x.tanh();
|
||
1.0 - t * t
|
||
}
|
||
Activation::ReLU => {
|
||
if x > 0.0 {
|
||
1.0
|
||
} else {
|
||
0.0
|
||
}
|
||
}
|
||
Activation::Sigmoid => {
|
||
let s = 1.0 / (1.0 + (-x).exp());
|
||
s * (1.0 - s)
|
||
}
|
||
Activation::Swish => {
|
||
let s = 1.0 / (1.0 + (-x).exp());
|
||
s + x * s * (1.0 - s)
|
||
}
|
||
Activation::GELU => {
|
||
let sqrt2 = std::f64::consts::SQRT_2;
|
||
let sqrt2pi = (2.0 * std::f64::consts::PI).sqrt();
|
||
let erf_term = libm::erf(x / sqrt2);
|
||
let exp_term = (-0.5 * x * x).exp();
|
||
0.5 * (1.0 + erf_term) + x * exp_term / sqrt2pi
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Dense (fully connected) layer
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct DenseLayer {
|
||
/// Weight matrix [in_features, out_features]
|
||
weights: Array2<f64>,
|
||
/// Bias vector [out_features]
|
||
biases: Array1<f64>,
|
||
/// Activation function
|
||
activation: Option<Activation>,
|
||
}
|
||
|
||
impl DenseLayer {
|
||
/// Create a new dense layer with Xavier initialization
|
||
pub fn new(in_features: usize, out_features: usize, activation: Option<Activation>) -> Self {
|
||
// Xavier/Glorot initialization
|
||
let scale = (2.0 / (in_features + out_features) as f64).sqrt();
|
||
let mut weights = Array2::zeros((in_features, out_features));
|
||
|
||
// Simple deterministic initialization for reproducibility
|
||
for i in 0..in_features {
|
||
for j in 0..out_features {
|
||
let seed = (i * out_features + j) as f64;
|
||
weights[[i, j]] = ((seed * 0.618033988749895).fract() * 2.0 - 1.0) * scale;
|
||
}
|
||
}
|
||
|
||
let biases = Array1::zeros(out_features);
|
||
|
||
Self {
|
||
weights,
|
||
biases,
|
||
activation,
|
||
}
|
||
}
|
||
|
||
/// Forward pass
|
||
pub fn forward(&self, input: &Array2<f64>) -> Array2<f64> {
|
||
let linear = input.dot(&self.weights) + &self.biases;
|
||
|
||
if let Some(act) = self.activation {
|
||
linear.mapv(|x| act.apply(x))
|
||
} else {
|
||
linear
|
||
}
|
||
}
|
||
|
||
/// Forward pass for single input
|
||
pub fn forward_single(&self, input: &Array1<f64>) -> Array1<f64> {
|
||
let linear = input.dot(&self.weights) + &self.biases;
|
||
|
||
if let Some(act) = self.activation {
|
||
linear.mapv(|x| act.apply(x))
|
||
} else {
|
||
linear
|
||
}
|
||
}
|
||
|
||
/// Get weights (for updates)
|
||
pub fn weights(&self) -> &Array2<f64> {
|
||
&self.weights
|
||
}
|
||
|
||
/// Get mutable weights
|
||
pub fn weights_mut(&mut self) -> &mut Array2<f64> {
|
||
&mut self.weights
|
||
}
|
||
|
||
/// Get biases
|
||
pub fn biases(&self) -> &Array1<f64> {
|
||
&self.biases
|
||
}
|
||
|
||
/// Get mutable biases
|
||
pub fn biases_mut(&mut self) -> &mut Array1<f64> {
|
||
&mut self.biases
|
||
}
|
||
}
|
||
|
||
/// Source network configuration
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SourceNetworkConfig {
|
||
/// Hidden layer sizes
|
||
pub hidden_layers: Vec<usize>,
|
||
/// Number of Fourier features
|
||
pub n_fourier_features: usize,
|
||
/// Fourier feature scale
|
||
pub fourier_scale: f64,
|
||
/// Activation function
|
||
pub activation: Activation,
|
||
/// Output current density components (typically 3 for Jx, Jy, Jz)
|
||
pub output_dim: usize,
|
||
/// Whether to also output conductivity
|
||
pub output_conductivity: bool,
|
||
}
|
||
|
||
impl Default for SourceNetworkConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
hidden_layers: vec![64, 128, 128, 64],
|
||
n_fourier_features: 32,
|
||
fourier_scale: 10.0,
|
||
activation: Activation::Tanh,
|
||
output_dim: 3,
|
||
output_conductivity: false,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Neural network for source localization
|
||
///
|
||
/// Takes 3D coordinates as input and outputs:
|
||
/// - Current density components (Jx, Jy, Jz)
|
||
/// - Optionally, tissue conductivity σ(x)
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SourceNetwork {
|
||
/// Configuration
|
||
config: SourceNetworkConfig,
|
||
/// Fourier feature encoder
|
||
fourier: FourierFeatures,
|
||
/// Hidden layers
|
||
layers: Vec<DenseLayer>,
|
||
/// Current density output layer
|
||
current_output: DenseLayer,
|
||
/// Conductivity output layer (optional)
|
||
conductivity_output: Option<DenseLayer>,
|
||
}
|
||
|
||
impl SourceNetwork {
|
||
/// Create a new source network
|
||
pub fn new(config: SourceNetworkConfig) -> Self {
|
||
let fourier = FourierFeatures::new(config.n_fourier_features, config.fourier_scale);
|
||
|
||
// Build hidden layers
|
||
let input_dim = fourier.output_dim();
|
||
let mut layers = Vec::new();
|
||
let mut prev_dim = input_dim;
|
||
|
||
for &hidden_dim in &config.hidden_layers {
|
||
layers.push(DenseLayer::new(
|
||
prev_dim,
|
||
hidden_dim,
|
||
Some(config.activation),
|
||
));
|
||
prev_dim = hidden_dim;
|
||
}
|
||
|
||
// Output layer for current density (no activation - unbounded output)
|
||
let current_output = DenseLayer::new(prev_dim, config.output_dim, None);
|
||
|
||
// Optional conductivity output (sigmoid for bounded positive output)
|
||
let conductivity_output = if config.output_conductivity {
|
||
Some(DenseLayer::new(prev_dim, 1, Some(Activation::Sigmoid)))
|
||
} else {
|
||
None
|
||
};
|
||
|
||
Self {
|
||
config,
|
||
fourier,
|
||
layers,
|
||
current_output,
|
||
conductivity_output,
|
||
}
|
||
}
|
||
|
||
/// Forward pass for multiple points
|
||
///
|
||
/// # Arguments
|
||
/// * `coords` - Input coordinates [n_points, 3]
|
||
///
|
||
/// # Returns
|
||
/// * Current density [n_points, 3] and optionally conductivity [n_points]
|
||
pub fn forward(&self, coords: &Array2<f64>) -> (Array2<f64>, Option<Array1<f64>>) {
|
||
// Fourier encoding
|
||
let mut x = self.fourier.encode(coords);
|
||
|
||
// Hidden layers
|
||
for layer in &self.layers {
|
||
x = layer.forward(&x);
|
||
}
|
||
|
||
// Current density output
|
||
let current = self.current_output.forward(&x);
|
||
|
||
// Conductivity output (if enabled)
|
||
let conductivity = self.conductivity_output.as_ref().map(|layer| {
|
||
let sigma = layer.forward(&x);
|
||
// Convert to 1D array and scale to reasonable conductivity range (0.01 - 2.0 S/m)
|
||
sigma.column(0).mapv(|s| 0.01 + s * 1.99).to_owned()
|
||
});
|
||
|
||
(current, conductivity)
|
||
}
|
||
|
||
/// Forward pass for single point
|
||
pub fn forward_single(&self, point: &[f64; 3]) -> (Array1<f64>, Option<f64>) {
|
||
// Fourier encoding
|
||
let mut x = self.fourier.encode_point(point);
|
||
|
||
// Hidden layers
|
||
for layer in &self.layers {
|
||
x = layer.forward_single(&x);
|
||
}
|
||
|
||
// Current density output
|
||
let current = self.current_output.forward_single(&x);
|
||
|
||
// Conductivity output
|
||
let conductivity = self.conductivity_output.as_ref().map(|layer| {
|
||
let sigma = layer.forward_single(&x);
|
||
0.01 + sigma[0] * 1.99
|
||
});
|
||
|
||
(current, conductivity)
|
||
}
|
||
|
||
/// Get network configuration
|
||
pub fn config(&self) -> &SourceNetworkConfig {
|
||
&self.config
|
||
}
|
||
|
||
/// Total number of parameters
|
||
pub fn n_parameters(&self) -> usize {
|
||
let mut count = 0;
|
||
|
||
for layer in &self.layers {
|
||
count += layer.weights().len();
|
||
count += layer.biases().len();
|
||
}
|
||
|
||
count += self.current_output.weights().len();
|
||
count += self.current_output.biases().len();
|
||
|
||
if let Some(ref layer) = self.conductivity_output {
|
||
count += layer.weights().len();
|
||
count += layer.biases().len();
|
||
}
|
||
|
||
count
|
||
}
|
||
|
||
/// Get all layer weights as flat vector (for optimization)
|
||
pub fn get_parameters(&self) -> Vec<f64> {
|
||
let mut params = Vec::with_capacity(self.n_parameters());
|
||
|
||
for layer in &self.layers {
|
||
params.extend(layer.weights().iter());
|
||
params.extend(layer.biases().iter());
|
||
}
|
||
|
||
params.extend(self.current_output.weights().iter());
|
||
params.extend(self.current_output.biases().iter());
|
||
|
||
if let Some(ref layer) = self.conductivity_output {
|
||
params.extend(layer.weights().iter());
|
||
params.extend(layer.biases().iter());
|
||
}
|
||
|
||
params
|
||
}
|
||
|
||
/// Set all layer weights from flat vector
|
||
pub fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
|
||
if params.len() != self.n_parameters() {
|
||
return Err(PinnError::DimensionMismatch(format!(
|
||
"Expected {} parameters, got {}",
|
||
self.n_parameters(),
|
||
params.len()
|
||
)));
|
||
}
|
||
|
||
let mut idx = 0;
|
||
|
||
for layer in &mut self.layers {
|
||
let w_len = layer.weights().len();
|
||
let b_len = layer.biases().len();
|
||
|
||
for (i, w) in layer.weights_mut().iter_mut().enumerate() {
|
||
*w = params[idx + i];
|
||
}
|
||
idx += w_len;
|
||
|
||
for (i, b) in layer.biases_mut().iter_mut().enumerate() {
|
||
*b = params[idx + i];
|
||
}
|
||
idx += b_len;
|
||
}
|
||
|
||
// Current output layer
|
||
let w_len = self.current_output.weights().len();
|
||
let b_len = self.current_output.biases().len();
|
||
|
||
for (i, w) in self.current_output.weights_mut().iter_mut().enumerate() {
|
||
*w = params[idx + i];
|
||
}
|
||
idx += w_len;
|
||
|
||
for (i, b) in self.current_output.biases_mut().iter_mut().enumerate() {
|
||
*b = params[idx + i];
|
||
}
|
||
idx += b_len;
|
||
|
||
// Conductivity layer (if present)
|
||
if let Some(ref mut layer) = self.conductivity_output {
|
||
let w_len = layer.weights().len();
|
||
|
||
for (i, w) in layer.weights_mut().iter_mut().enumerate() {
|
||
*w = params[idx + i];
|
||
}
|
||
idx += w_len;
|
||
|
||
for b in layer.biases_mut().iter_mut() {
|
||
*b = params[idx];
|
||
idx += 1;
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
/// Sensor network for mapping sensor data to latent representation
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SensorEncoder {
|
||
/// Encoder layers
|
||
layers: Vec<DenseLayer>,
|
||
/// Output dimension
|
||
latent_dim: usize,
|
||
}
|
||
|
||
impl SensorEncoder {
|
||
/// Create a new sensor encoder
|
||
pub fn new(n_sensors: usize, hidden_dims: &[usize], latent_dim: usize) -> Self {
|
||
let mut layers = Vec::new();
|
||
let mut prev_dim = n_sensors;
|
||
|
||
for &dim in hidden_dims {
|
||
layers.push(DenseLayer::new(prev_dim, dim, Some(Activation::Tanh)));
|
||
prev_dim = dim;
|
||
}
|
||
|
||
// Final layer to latent space
|
||
layers.push(DenseLayer::new(prev_dim, latent_dim, None));
|
||
|
||
Self { layers, latent_dim }
|
||
}
|
||
|
||
/// Encode sensor measurements
|
||
pub fn encode(&self, measurements: &Array1<f64>) -> Array1<f64> {
|
||
let mut x = measurements.clone();
|
||
for layer in &self.layers {
|
||
x = layer.forward_single(&x);
|
||
}
|
||
x
|
||
}
|
||
|
||
/// Encode batch of measurements
|
||
pub fn encode_batch(&self, measurements: &Array2<f64>) -> Array2<f64> {
|
||
let mut x = measurements.clone();
|
||
for layer in &self.layers {
|
||
x = layer.forward(&x);
|
||
}
|
||
x
|
||
}
|
||
|
||
/// Get latent dimension
|
||
pub fn latent_dim(&self) -> usize {
|
||
self.latent_dim
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use ndarray::array;
|
||
|
||
#[test]
|
||
fn test_fourier_features() {
|
||
let ff = FourierFeatures::new(16, 10.0);
|
||
assert_eq!(ff.output_dim(), 32);
|
||
|
||
let coords = Array2::from_shape_vec((2, 3), vec![0.0, 0.0, 0.05, 0.01, 0.0, 0.05]).unwrap();
|
||
|
||
let encoded = ff.encode(&coords);
|
||
assert_eq!(encoded.shape(), &[2, 32]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_activation_functions() {
|
||
assert!((Activation::Tanh.apply(0.0)).abs() < 1e-10);
|
||
assert!((Activation::ReLU.apply(-1.0)).abs() < 1e-10);
|
||
assert!((Activation::Sigmoid.apply(0.0) - 0.5).abs() < 1e-10);
|
||
}
|
||
|
||
#[test]
|
||
fn test_dense_layer() {
|
||
let layer = DenseLayer::new(4, 8, Some(Activation::Tanh));
|
||
let input = Array2::zeros((3, 4));
|
||
let output = layer.forward(&input);
|
||
assert_eq!(output.shape(), &[3, 8]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_source_network() {
|
||
let config = SourceNetworkConfig {
|
||
hidden_layers: vec![32, 64, 32],
|
||
n_fourier_features: 16,
|
||
fourier_scale: 10.0,
|
||
activation: Activation::Tanh,
|
||
output_dim: 3,
|
||
output_conductivity: true,
|
||
};
|
||
|
||
let network = SourceNetwork::new(config);
|
||
|
||
let coords = Array2::from_shape_vec(
|
||
(5, 3),
|
||
vec![
|
||
0.0, 0.0, 0.05, 0.01, 0.0, 0.05, -0.01, 0.0, 0.05, 0.0, 0.01, 0.05, 0.0, -0.01,
|
||
0.05,
|
||
],
|
||
)
|
||
.unwrap();
|
||
|
||
let (current, conductivity) = network.forward(&coords);
|
||
|
||
assert_eq!(current.shape(), &[5, 3]);
|
||
assert!(conductivity.is_some());
|
||
assert_eq!(conductivity.unwrap().len(), 5);
|
||
}
|
||
|
||
#[test]
|
||
fn test_parameter_io() {
|
||
let config = SourceNetworkConfig::default();
|
||
let mut network = SourceNetwork::new(config);
|
||
|
||
let params = network.get_parameters();
|
||
let n_params = network.n_parameters();
|
||
|
||
assert_eq!(params.len(), n_params);
|
||
|
||
// Modify and restore
|
||
let mut new_params = params.clone();
|
||
new_params[0] = 999.0;
|
||
|
||
network.set_parameters(&new_params).unwrap();
|
||
let retrieved = network.get_parameters();
|
||
|
||
assert!((retrieved[0] - 999.0).abs() < 1e-10);
|
||
}
|
||
|
||
#[test]
|
||
fn test_sensor_encoder() {
|
||
let encoder = SensorEncoder::new(64, &[32, 16], 8);
|
||
|
||
let measurements = Array1::zeros(64);
|
||
let latent = encoder.encode(&measurements);
|
||
|
||
assert_eq!(latent.len(), 8);
|
||
}
|
||
}
|