Initial commit

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

@@ -0,0 +1,419 @@
+//! Maxwell equation residuals for bioelectromagnetic source localization.
+//!
+//! Implements the quasi-static Maxwell equations that govern neural current flow:
+//!
+//! ```text
+//! ∇·(σ∇Φ) = ∇·Jp
+//! ```
+//!
+//! where Φ is the electric potential, σ is the tissue conductivity,
+//! and Jp is the primary (neural) current density.
+
+use crate::error::{PinnError, PinnResult};
+use ndarray::{Array1, Array2, Array3};
+use serde::{Deserialize, Serialize};
+
+/// Current density representation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CurrentDensity {
+    /// Current density components [n_sources, 3] (Jx, Jy, Jz)
+    pub components: Array2<f64>,
+    /// Source positions [n_sources, 3]
+    pub positions: Array2<f64>,
+    /// Magnitude at each source
+    pub magnitudes: Array1<f64>,
+}
+
+impl CurrentDensity {
+    /// Create a new current density field
+    pub fn new(n_sources: usize) -> Self {
+        Self {
+            components: Array2::zeros((n_sources, 3)),
+            positions: Array2::zeros((n_sources, 3)),
+            magnitudes: Array1::zeros(n_sources),
+        }
+    }
+
+    /// Create from positions and orientations
+    pub fn from_dipoles(
+        positions: Array2<f64>,
+        orientations: Array2<f64>,
+        magnitudes: Array1<f64>,
+    ) -> PinnResult<Self> {
+        let n = positions.nrows();
+        if orientations.nrows() != n || magnitudes.len() != n {
+            return Err(PinnError::DimensionMismatch(
+                "Positions, orientations, and magnitudes must have same length".into(),
+            ));
+        }
+
+        // Scale orientations by magnitudes
+        let mut components = Array2::zeros((n, 3));
+        for i in 0..n {
+            for j in 0..3 {
+                components[[i, j]] = orientations[[i, j]] * magnitudes[i];
+            }
+        }
+
+        Ok(Self {
+            components,
+            positions,
+            magnitudes,
+        })
+    }
+
+    /// Number of sources
+    pub fn n_sources(&self) -> usize {
+        self.positions.nrows()
+    }
+
+    /// Compute divergence at a point (for physics residual)
+    pub fn divergence_at(&self, point: &[f64; 3], sigma: f64) -> f64 {
+        // Simplified: sum of contributions from all sources
+        // In reality, this would involve Green's functions
+        let mut div = 0.0;
+        let eps = 1e-10;
+
+        for i in 0..self.n_sources() {
+            let dx = point[0] - self.positions[[i, 0]];
+            let dy = point[1] - self.positions[[i, 1]];
+            let dz = point[2] - self.positions[[i, 2]];
+            let r2 = dx * dx + dy * dy + dz * dz + eps;
+            let r = r2.sqrt();
+
+            // Dipole-like contribution
+            let jx = self.components[[i, 0]];
+            let jy = self.components[[i, 1]];
+            let jz = self.components[[i, 2]];
+
+            // ∇·J contribution (simplified monopole approximation)
+            let dot = jx * dx + jy * dy + jz * dz;
+            div += dot / (r2 * r) / sigma;
+        }
+
+        div
+    }
+
+    /// Total current magnitude
+    pub fn total_magnitude(&self) -> f64 {
+        self.magnitudes.iter().map(|&m| m.abs()).sum()
+    }
+}
+
+/// Quasi-static Maxwell equation for bioelectric problems
+#[derive(Debug, Clone)]
+pub struct QuasiStaticMaxwell {
+    /// Default conductivity (S/m)
+    conductivity: f64,
+    /// Regularization parameter for sparsity
+    sparsity_weight: f64,
+    /// Regularization for smoothness
+    smoothness_weight: f64,
+}
+
+impl QuasiStaticMaxwell {
+    /// Create a new quasi-static Maxwell equation
+    pub fn new(conductivity: f64) -> Self {
+        Self {
+            conductivity,
+            sparsity_weight: 0.01,
+            smoothness_weight: 0.001,
+        }
+    }
+
+    /// Set sparsity regularization weight
+    pub fn with_sparsity(mut self, weight: f64) -> Self {
+        self.sparsity_weight = weight;
+        self
+    }
+
+    /// Set smoothness regularization weight
+    pub fn with_smoothness(mut self, weight: f64) -> Self {
+        self.smoothness_weight = weight;
+        self
+    }
+
+    /// Get conductivity
+    pub fn conductivity(&self) -> f64 {
+        self.conductivity
+    }
+
+    /// Compute the PDE residual at collocation points
+    ///
+    /// The residual is: ∇·(σ∇Φ) - ∇·Jp = 0
+    ///
+    /// # Arguments
+    /// * `potential` - Electric potential field [n_points]
+    /// * `potential_grad` - Gradient of potential [n_points, 3]
+    /// * `current_density` - Current density field
+    /// * `conductivity_field` - Spatially varying conductivity [n_points]
+    ///
+    /// # Returns
+    /// Residual at each collocation point
+    pub fn compute_residual(
+        &self,
+        potential: &Array1<f64>,
+        potential_grad: &Array2<f64>,
+        potential_hessian: &Array3<f64>,
+        current_density: &CurrentDensity,
+        conductivity_field: &Array1<f64>,
+        points: &Array2<f64>,
+    ) -> PinnResult<Array1<f64>> {
+        let n_points = potential.len();
+        if potential_grad.nrows() != n_points {
+            return Err(PinnError::DimensionMismatch(
+                "Potential and gradient dimensions don't match".into(),
+            ));
+        }
+
+        let mut residual = Array1::zeros(n_points);
+
+        for i in 0..n_points {
+            let sigma = conductivity_field[i];
+
+            // Laplacian term: ∇²Φ (trace of Hessian)
+            let laplacian = potential_hessian[[i, 0, 0]]
+                + potential_hessian[[i, 1, 1]]
+                + potential_hessian[[i, 2, 2]];
+
+            // σ∇²Φ term
+            let diffusion = sigma * laplacian;
+
+            // ∇σ·∇Φ term (if conductivity varies spatially)
+            // For now, assume locally constant conductivity
+            let advection = 0.0;
+
+            // Source term: ∇·Jp at this point
+            let point = [points[[i, 0]], points[[i, 1]], points[[i, 2]]];
+            let source = current_density.divergence_at(&point, sigma);
+
+            // Residual: ∇·(σ∇Φ) - ∇·Jp should be zero
+            residual[i] = diffusion + advection - source;
+        }
+
+        Ok(residual)
+    }
+
+    /// Compute physics loss from residuals
+    pub fn physics_loss(&self, residual: &Array1<f64>) -> f64 {
+        // Mean squared residual
+        
+        residual.iter().map(|&r| r * r).sum::<f64>() / residual.len() as f64
+    }
+
+    /// Compute sparsity loss for source estimation
+    pub fn sparsity_loss(&self, current_density: &CurrentDensity) -> f64 {
+        // L1 norm of source magnitudes
+        let l1 = current_density
+            .magnitudes
+            .iter()
+            .map(|&m| m.abs())
+            .sum::<f64>();
+        self.sparsity_weight * l1
+    }
+
+    /// Compute smoothness loss for source distribution
+    pub fn smoothness_loss(&self, current_density: &CurrentDensity) -> f64 {
+        if current_density.n_sources() < 2 {
+            return 0.0;
+        }
+
+        // Total variation of source magnitudes (simplified)
+        let mut tv = 0.0;
+        for i in 1..current_density.n_sources() {
+            let diff = current_density.magnitudes[i] - current_density.magnitudes[i - 1];
+            tv += diff.abs();
+        }
+
+        self.smoothness_weight * tv
+    }
+
+    /// Compute total loss
+    pub fn total_loss(&self, residual: &Array1<f64>, current_density: &CurrentDensity) -> f64 {
+        self.physics_loss(residual)
+            + self.sparsity_loss(current_density)
+            + self.smoothness_loss(current_density)
+    }
+}
+
+/// Maxwell residual trait for different formulations
+pub trait MaxwellResidual {
+    /// Compute residual at points
+    fn residual(
+        &self,
+        potential: &Array1<f64>,
+        gradients: &Array2<f64>,
+        hessians: &Array3<f64>,
+        sources: &CurrentDensity,
+        conductivity: &Array1<f64>,
+        points: &Array2<f64>,
+    ) -> PinnResult<Array1<f64>>;
+
+    /// Compute loss from residual
+    fn loss(&self, residual: &Array1<f64>) -> f64;
+}
+
+impl MaxwellResidual for QuasiStaticMaxwell {
+    fn residual(
+        &self,
+        potential: &Array1<f64>,
+        gradients: &Array2<f64>,
+        hessians: &Array3<f64>,
+        sources: &CurrentDensity,
+        conductivity: &Array1<f64>,
+        points: &Array2<f64>,
+    ) -> PinnResult<Array1<f64>> {
+        self.compute_residual(
+            potential,
+            gradients,
+            hessians,
+            sources,
+            conductivity,
+            points,
+        )
+    }
+
+    fn loss(&self, residual: &Array1<f64>) -> f64 {
+        self.physics_loss(residual)
+    }
+}
+
+/// MEG-specific Maxwell formulation (magnetic field)
+#[derive(Debug, Clone)]
+pub struct MagneticMaxwell {
+    /// Permeability of free space
+    mu0: f64,
+}
+
+impl MagneticMaxwell {
+    /// Create new magnetic Maxwell formulation
+    pub fn new() -> Self {
+        Self {
+            mu0: 4.0 * std::f64::consts::PI * 1e-7, // H/m
+        }
+    }
+
+    /// Compute magnetic field from current density using Biot-Savart
+    pub fn magnetic_field(
+        &self,
+        current_density: &CurrentDensity,
+        sensor_positions: &Array2<f64>,
+    ) -> Array2<f64> {
+        let n_sensors = sensor_positions.nrows();
+        let mut b_field = Array2::zeros((n_sensors, 3));
+
+        for s in 0..n_sensors {
+            let sensor = [
+                sensor_positions[[s, 0]],
+                sensor_positions[[s, 1]],
+                sensor_positions[[s, 2]],
+            ];
+
+            for i in 0..current_density.n_sources() {
+                let source = [
+                    current_density.positions[[i, 0]],
+                    current_density.positions[[i, 1]],
+                    current_density.positions[[i, 2]],
+                ];
+
+                let j = [
+                    current_density.components[[i, 0]],
+                    current_density.components[[i, 1]],
+                    current_density.components[[i, 2]],
+                ];
+
+                // r = sensor - source
+                let r = [
+                    sensor[0] - source[0],
+                    sensor[1] - source[1],
+                    sensor[2] - source[2],
+                ];
+                let r_mag = (r[0] * r[0] + r[1] * r[1] + r[2] * r[2]).sqrt();
+                let r_mag3 = r_mag * r_mag * r_mag;
+
+                if r_mag > 1e-10 {
+                    // B = (μ0/4π) * (J × r) / |r|³
+                    let cross = [
+                        j[1] * r[2] - j[2] * r[1],
+                        j[2] * r[0] - j[0] * r[2],
+                        j[0] * r[1] - j[1] * r[0],
+                    ];
+
+                    let coeff = self.mu0 / (4.0 * std::f64::consts::PI * r_mag3);
+                    b_field[[s, 0]] += coeff * cross[0];
+                    b_field[[s, 1]] += coeff * cross[1];
+                    b_field[[s, 2]] += coeff * cross[2];
+                }
+            }
+        }
+
+        b_field
+    }
+
+    /// Compute data fitting loss for MEG
+    pub fn data_loss(&self, predicted: &Array2<f64>, measured: &Array2<f64>) -> PinnResult<f64> {
+        if predicted.shape() != measured.shape() {
+            return Err(PinnError::DimensionMismatch(
+                "Predicted and measured have different shapes".into(),
+            ));
+        }
+
+        let mse = predicted
+            .iter()
+            .zip(measured.iter())
+            .map(|(&p, &m)| (p - m).powi(2))
+            .sum::<f64>()
+            / predicted.len() as f64;
+
+        Ok(mse)
+    }
+}
+
+impl Default for MagneticMaxwell {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_current_density_creation() {
+        let positions =
+            Array2::from_shape_vec((2, 3), vec![0.0, 0.0, 0.05, 0.01, 0.0, 0.05]).unwrap();
+        let orientations =
+            Array2::from_shape_vec((2, 3), vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0]).unwrap();
+        let magnitudes = Array1::from_vec(vec![1e-9, 2e-9]);
+
+        let current = CurrentDensity::from_dipoles(positions, orientations, magnitudes).unwrap();
+        assert_eq!(current.n_sources(), 2);
+    }
+
+    #[test]
+    fn test_quasi_static_maxwell() {
+        let maxwell = QuasiStaticMaxwell::new(0.33);
+        assert!((maxwell.conductivity() - 0.33).abs() < 1e-10);
+    }
+
+    #[test]
+    fn test_magnetic_field() {
+        let mag = MagneticMaxwell::new();
+
+        // Single dipole at origin pointing in z
+        let positions = Array2::from_shape_vec((1, 3), vec![0.0, 0.0, 0.0]).unwrap();
+        let orientations = Array2::from_shape_vec((1, 3), vec![0.0, 0.0, 1.0]).unwrap();
+        let magnitudes = Array1::from_vec(vec![1e-9]);
+        let current = CurrentDensity::from_dipoles(positions, orientations, magnitudes).unwrap();
+
+        // Sensor at (0.1, 0, 0)
+        let sensors = Array2::from_shape_vec((1, 3), vec![0.1, 0.0, 0.0]).unwrap();
+        let b_field = mag.magnetic_field(&current, &sensors);
+
+        // B should be in y direction (perpendicular to both z and x)
+        assert!(b_field[[0, 0]].abs() < 1e-20); // Bx ~ 0
+        assert!(b_field[[0, 1]].abs() > 0.0); // By != 0
+        assert!(b_field[[0, 2]].abs() < 1e-20); // Bz ~ 0
+    }
+}