Initial commit

crates/specialized/rtx-neuro-pinn/src/network.rs

@@ -0,0 +1,617 @@
+//! 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);
+    }
+}