Files
rustytorch/crates/specialized/rtx-neuro-pinn/src/maxwell.rs
T
osobhandClaude Opus 4.6 02d382d5f6 style: apply rustfmt across all crates and demos
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]>
2026-04-12 07:01:58 -07:00

420 lines
13 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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
}
}