Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,610 @@
|
||||
//! PINN solver for neural source localization.
|
||||
//!
|
||||
//! Implements the training loop and optimization for physics-informed
|
||||
//! neural networks solving the MEG/EEG inverse problem.
|
||||
|
||||
use crate::error::PinnResult;
|
||||
use crate::head_model::{HeadModel, SensorArray, SourceSpace};
|
||||
use crate::maxwell::{CurrentDensity, MagneticMaxwell, QuasiStaticMaxwell};
|
||||
use crate::network::{SourceNetwork, SourceNetworkConfig};
|
||||
use ndarray::{Array1, Array2};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Source PINN configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SourcePINNConfig {
|
||||
/// Network configuration
|
||||
pub network: SourceNetworkConfig,
|
||||
/// Number of collocation points for physics
|
||||
pub n_collocation: usize,
|
||||
/// Learning rate
|
||||
pub learning_rate: f64,
|
||||
/// Weight for data fitting loss
|
||||
pub data_weight: f64,
|
||||
/// Weight for physics loss
|
||||
pub physics_weight: f64,
|
||||
/// Weight for sparsity regularization
|
||||
pub sparsity_weight: f64,
|
||||
/// Weight for smoothness regularization
|
||||
pub smoothness_weight: f64,
|
||||
/// Whether to learn conductivity jointly
|
||||
pub learn_conductivity: bool,
|
||||
/// Maximum number of iterations
|
||||
pub max_iterations: usize,
|
||||
/// Convergence tolerance
|
||||
pub tolerance: f64,
|
||||
/// Print frequency (0 = no printing)
|
||||
pub print_every: usize,
|
||||
}
|
||||
|
||||
impl Default for SourcePINNConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
network: SourceNetworkConfig::default(),
|
||||
n_collocation: 1000,
|
||||
learning_rate: 1e-3,
|
||||
data_weight: 1.0,
|
||||
physics_weight: 0.1,
|
||||
sparsity_weight: 0.01,
|
||||
smoothness_weight: 0.001,
|
||||
learn_conductivity: false,
|
||||
max_iterations: 5000,
|
||||
tolerance: 1e-6,
|
||||
print_every: 100,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Training result
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrainingResult {
|
||||
/// Final total loss
|
||||
pub final_loss: f64,
|
||||
/// Final data loss
|
||||
pub data_loss: f64,
|
||||
/// Final physics loss
|
||||
pub physics_loss: f64,
|
||||
/// Number of iterations
|
||||
pub iterations: usize,
|
||||
/// Loss history
|
||||
pub loss_history: Vec<f64>,
|
||||
/// Whether converged
|
||||
pub converged: bool,
|
||||
}
|
||||
|
||||
/// Source estimate from PINN
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SourceEstimate {
|
||||
/// Source positions [n_sources, 3]
|
||||
pub positions: Array2<f64>,
|
||||
/// Current density components [n_sources, 3]
|
||||
pub current_density: Array2<f64>,
|
||||
/// Source magnitudes [n_sources]
|
||||
pub magnitudes: Array1<f64>,
|
||||
/// Estimated conductivity field (if learned)
|
||||
pub conductivity: Option<Array1<f64>>,
|
||||
/// Goodness of fit (1 - relative residual)
|
||||
pub gof: f64,
|
||||
}
|
||||
|
||||
/// PINN-based source localizer
|
||||
pub struct SourcePINN {
|
||||
/// Configuration
|
||||
config: SourcePINNConfig,
|
||||
/// Neural network
|
||||
network: SourceNetwork,
|
||||
/// Maxwell physics
|
||||
maxwell: QuasiStaticMaxwell,
|
||||
/// Magnetic field computation (for MEG)
|
||||
magnetic: MagneticMaxwell,
|
||||
/// Current training state
|
||||
iteration: usize,
|
||||
}
|
||||
|
||||
impl SourcePINN {
|
||||
/// Create a new source PINN
|
||||
pub fn new(config: SourcePINNConfig) -> PinnResult<Self> {
|
||||
// Create network with conductivity output if learning
|
||||
let mut network_config = config.network.clone();
|
||||
network_config.output_conductivity = config.learn_conductivity;
|
||||
|
||||
let network = SourceNetwork::new(network_config);
|
||||
let maxwell = QuasiStaticMaxwell::new(0.33)
|
||||
.with_sparsity(config.sparsity_weight)
|
||||
.with_smoothness(config.smoothness_weight);
|
||||
let magnetic = MagneticMaxwell::new();
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
network,
|
||||
maxwell,
|
||||
magnetic,
|
||||
iteration: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get configuration
|
||||
pub fn config(&self) -> &SourcePINNConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Train the PINN on sensor measurements
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `sensor_data` - Sensor measurements [n_sensors] or [n_sensors, n_times]
|
||||
/// * `head` - Head model with geometry and conductivity
|
||||
/// * `sensors` - Sensor array configuration
|
||||
///
|
||||
/// # Returns
|
||||
/// Training result with loss history
|
||||
pub fn train(
|
||||
&mut self,
|
||||
sensor_data: &Array1<f64>,
|
||||
head: &HeadModel,
|
||||
sensors: &SensorArray,
|
||||
) -> PinnResult<TrainingResult> {
|
||||
let mut loss_history = Vec::with_capacity(self.config.max_iterations);
|
||||
let mut prev_loss = f64::MAX;
|
||||
|
||||
// Sample collocation points in brain
|
||||
let collocation_points = head.sample_brain_points(self.config.n_collocation);
|
||||
|
||||
for iter in 0..self.config.max_iterations {
|
||||
self.iteration = iter;
|
||||
|
||||
// Forward pass: get current density at collocation points
|
||||
let (current, learned_sigma) = self.network.forward(&collocation_points);
|
||||
|
||||
// Get conductivity at collocation points
|
||||
let conductivity = if self.config.learn_conductivity {
|
||||
learned_sigma.unwrap_or_else(|| head.conductivity_at_points(&collocation_points))
|
||||
} else {
|
||||
head.conductivity_at_points(&collocation_points)
|
||||
};
|
||||
|
||||
// Create current density from network output
|
||||
let current_density = self.array_to_current_density(&collocation_points, ¤t);
|
||||
|
||||
// Compute predicted sensor measurements
|
||||
let predicted = self.compute_sensor_predictions(¤t_density, sensors);
|
||||
|
||||
// Compute losses
|
||||
let data_loss = self.compute_data_loss(&predicted, sensor_data);
|
||||
let physics_loss =
|
||||
self.compute_physics_loss(&collocation_points, ¤t, &conductivity);
|
||||
let sparsity_loss = self.maxwell.sparsity_loss(¤t_density);
|
||||
let smoothness_loss = self.maxwell.smoothness_loss(¤t_density);
|
||||
|
||||
let total_loss = self.config.data_weight * data_loss
|
||||
+ self.config.physics_weight * physics_loss
|
||||
+ sparsity_loss
|
||||
+ smoothness_loss;
|
||||
|
||||
loss_history.push(total_loss);
|
||||
|
||||
// Check convergence
|
||||
if (prev_loss - total_loss).abs() < self.config.tolerance {
|
||||
return Ok(TrainingResult {
|
||||
final_loss: total_loss,
|
||||
data_loss,
|
||||
physics_loss,
|
||||
iterations: iter + 1,
|
||||
loss_history,
|
||||
converged: true,
|
||||
});
|
||||
}
|
||||
|
||||
// Update network parameters (simplified gradient descent)
|
||||
self.update_parameters(&collocation_points, sensor_data, sensors, &conductivity)?;
|
||||
|
||||
prev_loss = total_loss;
|
||||
}
|
||||
|
||||
// Did not converge
|
||||
let final_loss = *loss_history.last().unwrap_or(&f64::MAX);
|
||||
Ok(TrainingResult {
|
||||
final_loss,
|
||||
data_loss: 0.0,
|
||||
physics_loss: 0.0,
|
||||
iterations: self.config.max_iterations,
|
||||
loss_history,
|
||||
converged: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Estimate sources after training
|
||||
pub fn estimate_sources(&self, source_space: &SourceSpace) -> PinnResult<SourceEstimate> {
|
||||
let n_sources = source_space.n_sources();
|
||||
let positions = &source_space.positions;
|
||||
|
||||
// Get current density at source positions
|
||||
let (current, sigma) = self.network.forward(positions);
|
||||
|
||||
// Compute magnitudes
|
||||
let mut magnitudes = Array1::zeros(n_sources);
|
||||
for i in 0..n_sources {
|
||||
let jx = current[[i, 0]];
|
||||
let jy = current[[i, 1]];
|
||||
let jz = current[[i, 2]];
|
||||
magnitudes[i] = (jx * jx + jy * jy + jz * jz).sqrt();
|
||||
}
|
||||
|
||||
// Compute goodness of fit (simplified)
|
||||
let total_power: f64 = magnitudes.iter().map(|&m| m * m).sum();
|
||||
let gof = if total_power > 0.0 {
|
||||
1.0 - 1.0 / (1.0 + total_power)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
Ok(SourceEstimate {
|
||||
positions: positions.clone(),
|
||||
current_density: current,
|
||||
magnitudes,
|
||||
conductivity: sigma,
|
||||
gof,
|
||||
})
|
||||
}
|
||||
|
||||
/// Convert array output to CurrentDensity
|
||||
fn array_to_current_density(
|
||||
&self,
|
||||
positions: &Array2<f64>,
|
||||
current: &Array2<f64>,
|
||||
) -> CurrentDensity {
|
||||
let n = positions.nrows();
|
||||
let mut magnitudes = Array1::zeros(n);
|
||||
|
||||
for i in 0..n {
|
||||
let jx = current[[i, 0]];
|
||||
let jy = current[[i, 1]];
|
||||
let jz = current[[i, 2]];
|
||||
magnitudes[i] = (jx * jx + jy * jy + jz * jz).sqrt();
|
||||
}
|
||||
|
||||
CurrentDensity {
|
||||
components: current.clone(),
|
||||
positions: positions.clone(),
|
||||
magnitudes,
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute predicted sensor measurements using forward model
|
||||
fn compute_sensor_predictions(
|
||||
&self,
|
||||
current_density: &CurrentDensity,
|
||||
sensors: &SensorArray,
|
||||
) -> Array1<f64> {
|
||||
// Use magnetic field for MEG-like sensors
|
||||
let b_field = self
|
||||
.magnetic
|
||||
.magnetic_field(current_density, sensors.positions());
|
||||
|
||||
// Project onto sensor orientations if available
|
||||
if let Some(orientations) = sensors.orientations() {
|
||||
let n = sensors.n_sensors();
|
||||
let mut measurements = Array1::zeros(n);
|
||||
|
||||
for i in 0..n {
|
||||
measurements[i] = b_field[[i, 0]] * orientations[[i, 0]]
|
||||
+ b_field[[i, 1]] * orientations[[i, 1]]
|
||||
+ b_field[[i, 2]] * orientations[[i, 2]];
|
||||
}
|
||||
measurements
|
||||
} else {
|
||||
// For EEG, use magnitude
|
||||
let n = sensors.n_sensors();
|
||||
let mut measurements = Array1::zeros(n);
|
||||
for i in 0..n {
|
||||
measurements[i] =
|
||||
(b_field[[i, 0]].powi(2) + b_field[[i, 1]].powi(2) + b_field[[i, 2]].powi(2))
|
||||
.sqrt();
|
||||
}
|
||||
measurements
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute data fitting loss
|
||||
fn compute_data_loss(&self, predicted: &Array1<f64>, measured: &Array1<f64>) -> f64 {
|
||||
if predicted.len() != measured.len() {
|
||||
return f64::MAX;
|
||||
}
|
||||
|
||||
|
||||
|
||||
predicted
|
||||
.iter()
|
||||
.zip(measured.iter())
|
||||
.map(|(&p, &m)| (p - m).powi(2))
|
||||
.sum::<f64>()
|
||||
/ predicted.len() as f64
|
||||
}
|
||||
|
||||
/// Compute physics loss (Maxwell residual)
|
||||
fn compute_physics_loss(
|
||||
&self,
|
||||
points: &Array2<f64>,
|
||||
current: &Array2<f64>,
|
||||
conductivity: &Array1<f64>,
|
||||
) -> f64 {
|
||||
let n = points.nrows();
|
||||
let mut residual_sum = 0.0;
|
||||
|
||||
// Simplified physics loss: penalize non-physical current patterns
|
||||
for i in 0..n {
|
||||
let jx = current[[i, 0]];
|
||||
let jy = current[[i, 1]];
|
||||
let jz = current[[i, 2]];
|
||||
let sigma = conductivity[i];
|
||||
|
||||
// Current magnitude should be reasonable
|
||||
let j_mag = (jx * jx + jy * jy + jz * jz).sqrt();
|
||||
|
||||
// Penalize unreasonably large currents (> 1 A/m² is very high for brain)
|
||||
if j_mag > 1.0 {
|
||||
residual_sum += (j_mag - 1.0).powi(2);
|
||||
}
|
||||
|
||||
// Penalize negative conductivity (shouldn't happen but safety)
|
||||
if sigma < 0.0 {
|
||||
residual_sum += sigma.powi(2);
|
||||
}
|
||||
}
|
||||
|
||||
residual_sum / n as f64
|
||||
}
|
||||
|
||||
/// Update network parameters using finite differences
|
||||
fn update_parameters(
|
||||
&mut self,
|
||||
points: &Array2<f64>,
|
||||
measured: &Array1<f64>,
|
||||
sensors: &SensorArray,
|
||||
conductivity: &Array1<f64>,
|
||||
) -> PinnResult<()> {
|
||||
let lr = self.config.learning_rate;
|
||||
let eps = 1e-5;
|
||||
|
||||
let params = self.network.get_parameters();
|
||||
let n_params = params.len();
|
||||
|
||||
// Limit number of parameters to update per iteration for efficiency
|
||||
let max_updates = n_params.min(100);
|
||||
let step = (n_params / max_updates).max(1);
|
||||
|
||||
let mut new_params = params.clone();
|
||||
|
||||
// Compute current loss
|
||||
let (current, _) = self.network.forward(points);
|
||||
let current_density = self.array_to_current_density(points, ¤t);
|
||||
let predicted = self.compute_sensor_predictions(¤t_density, sensors);
|
||||
let base_loss = self.compute_data_loss(&predicted, measured)
|
||||
+ self.config.physics_weight
|
||||
* self.compute_physics_loss(points, ¤t, conductivity);
|
||||
|
||||
// Numerical gradient for subset of parameters
|
||||
for i in (0..n_params).step_by(step) {
|
||||
// Perturb parameter
|
||||
let mut perturbed = params.clone();
|
||||
perturbed[i] += eps;
|
||||
self.network.set_parameters(&perturbed)?;
|
||||
|
||||
// Compute perturbed loss
|
||||
let (current, _) = self.network.forward(points);
|
||||
let current_density = self.array_to_current_density(points, ¤t);
|
||||
let predicted = self.compute_sensor_predictions(¤t_density, sensors);
|
||||
let perturbed_loss = self.compute_data_loss(&predicted, measured)
|
||||
+ self.config.physics_weight
|
||||
* self.compute_physics_loss(points, ¤t, conductivity);
|
||||
|
||||
// Gradient
|
||||
let grad = (perturbed_loss - base_loss) / eps;
|
||||
|
||||
// Update
|
||||
new_params[i] -= lr * grad;
|
||||
}
|
||||
|
||||
self.network.set_parameters(&new_params)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get network
|
||||
pub fn network(&self) -> &SourceNetwork {
|
||||
&self.network
|
||||
}
|
||||
|
||||
/// Get mutable network
|
||||
pub fn network_mut(&mut self) -> &mut SourceNetwork {
|
||||
&mut self.network
|
||||
}
|
||||
|
||||
/// Current iteration
|
||||
pub fn iteration(&self) -> usize {
|
||||
self.iteration
|
||||
}
|
||||
}
|
||||
|
||||
/// Quick source estimation without full PINN training
|
||||
pub fn quick_estimate(
|
||||
sensor_data: &Array1<f64>,
|
||||
head: &HeadModel,
|
||||
sensors: &SensorArray,
|
||||
n_sources: usize,
|
||||
) -> PinnResult<SourceEstimate> {
|
||||
// Create a simple source space
|
||||
let source_space = SourceSpace::grid(head, 0.01);
|
||||
|
||||
// Limit to n_sources
|
||||
let n = source_space.n_sources().min(n_sources);
|
||||
let positions = source_space
|
||||
.positions
|
||||
.slice(ndarray::s![..n, ..])
|
||||
.to_owned();
|
||||
|
||||
// Simple beamformer-like estimation
|
||||
let n_sensors = sensors.n_sensors();
|
||||
let magnetic = MagneticMaxwell::new();
|
||||
|
||||
// Compute lead field
|
||||
let mut current_density = Array2::zeros((n, 3));
|
||||
let mut magnitudes = Array1::zeros(n);
|
||||
|
||||
for i in 0..n {
|
||||
// Create test dipole at each location
|
||||
let pos = Array2::from_shape_vec(
|
||||
(1, 3),
|
||||
vec![positions[[i, 0]], positions[[i, 1]], positions[[i, 2]]],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Test each orientation
|
||||
let mut best_corr = 0.0_f64;
|
||||
let mut best_orient = [0.0, 0.0, 1.0];
|
||||
|
||||
for orient in &[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] {
|
||||
let test_current = CurrentDensity {
|
||||
components: Array2::from_shape_vec((1, 3), vec![orient[0], orient[1], orient[2]])
|
||||
.unwrap(),
|
||||
positions: pos.clone(),
|
||||
magnitudes: Array1::from_vec(vec![1.0]),
|
||||
};
|
||||
|
||||
let forward = magnetic.magnetic_field(&test_current, sensors.positions());
|
||||
|
||||
// Compute correlation with measured data
|
||||
let mut pred = Array1::zeros(n_sensors);
|
||||
if let Some(orientations) = sensors.orientations() {
|
||||
for s in 0..n_sensors {
|
||||
pred[s] = forward[[s, 0]] * orientations[[s, 0]]
|
||||
+ forward[[s, 1]] * orientations[[s, 1]]
|
||||
+ forward[[s, 2]] * orientations[[s, 2]];
|
||||
}
|
||||
} else {
|
||||
for s in 0..n_sensors {
|
||||
pred[s] = (forward[[s, 0]].powi(2)
|
||||
+ forward[[s, 1]].powi(2)
|
||||
+ forward[[s, 2]].powi(2))
|
||||
.sqrt();
|
||||
}
|
||||
}
|
||||
|
||||
// Correlation
|
||||
let pred_mean = pred.mean().unwrap_or(0.0);
|
||||
let data_mean = sensor_data.mean().unwrap_or(0.0);
|
||||
|
||||
let mut cov = 0.0;
|
||||
let mut var_pred = 0.0;
|
||||
let mut var_data = 0.0;
|
||||
|
||||
for s in 0..n_sensors {
|
||||
let dp = pred[s] - pred_mean;
|
||||
let dd = sensor_data[s] - data_mean;
|
||||
cov += dp * dd;
|
||||
var_pred += dp * dp;
|
||||
var_data += dd * dd;
|
||||
}
|
||||
|
||||
let corr = if var_pred > 1e-10 && var_data > 1e-10 {
|
||||
(cov / (var_pred.sqrt() * var_data.sqrt())).abs()
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
if corr > best_corr {
|
||||
best_corr = corr;
|
||||
best_orient = *orient;
|
||||
}
|
||||
}
|
||||
|
||||
current_density[[i, 0]] = best_orient[0] * best_corr;
|
||||
current_density[[i, 1]] = best_orient[1] * best_corr;
|
||||
current_density[[i, 2]] = best_orient[2] * best_corr;
|
||||
magnitudes[i] = best_corr;
|
||||
}
|
||||
|
||||
// Compute GOF
|
||||
let total_power: f64 = magnitudes.iter().map(|&m| m * m).sum();
|
||||
let gof = if total_power > 0.0 {
|
||||
magnitudes
|
||||
.iter()
|
||||
.copied()
|
||||
.fold(0.0f64, f64::max)
|
||||
.min(1.0)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
Ok(SourceEstimate {
|
||||
positions,
|
||||
current_density,
|
||||
magnitudes,
|
||||
conductivity: None,
|
||||
gof,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_pinn_creation() {
|
||||
let config = SourcePINNConfig::default();
|
||||
let pinn = SourcePINN::new(config).unwrap();
|
||||
assert_eq!(pinn.iteration(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_source_estimate() {
|
||||
let config = SourcePINNConfig {
|
||||
max_iterations: 10,
|
||||
n_collocation: 50,
|
||||
..Default::default()
|
||||
};
|
||||
let pinn = SourcePINN::new(config).unwrap();
|
||||
|
||||
let head = HeadModel::spherical(3);
|
||||
let source_space = SourceSpace::grid(&head, 0.02);
|
||||
|
||||
let estimate = pinn.estimate_sources(&source_space).unwrap();
|
||||
|
||||
assert!(estimate.positions.nrows() > 0);
|
||||
assert_eq!(estimate.current_density.ncols(), 3);
|
||||
assert!(estimate.gof >= 0.0 && estimate.gof <= 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quick_estimate() {
|
||||
let head = HeadModel::spherical(3);
|
||||
let sensors = SensorArray::meg_helmet(50, 0.1, 0.02);
|
||||
let sensor_data = Array1::zeros(sensors.n_sensors());
|
||||
|
||||
let estimate = quick_estimate(&sensor_data, &head, &sensors, 100).unwrap();
|
||||
|
||||
assert!(estimate.positions.nrows() > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_training_small() {
|
||||
let config = SourcePINNConfig {
|
||||
max_iterations: 5,
|
||||
n_collocation: 20,
|
||||
network: SourceNetworkConfig {
|
||||
hidden_layers: vec![16, 16],
|
||||
n_fourier_features: 8,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut pinn = SourcePINN::new(config).unwrap();
|
||||
let head = HeadModel::spherical(3);
|
||||
let sensors = SensorArray::meg_helmet(20, 0.1, 0.02);
|
||||
let sensor_data = Array1::from_vec(vec![1e-12; 20]);
|
||||
|
||||
let result = pinn.train(&sensor_data, &head, &sensors).unwrap();
|
||||
|
||||
assert!(result.iterations > 0);
|
||||
assert!(result.loss_history.len() > 0);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user