Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
420 lines
13 KiB
Rust
420 lines
13 KiB
Rust
//! 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(¤t, &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
|
||
}
|
||
}
|