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]>
1005 lines
32 KiB
Rust
1005 lines
32 KiB
Rust
//! Dipole fitting for MEG/EEG source localization.
|
|
//!
|
|
//! This module implements equivalent current dipole (ECD) fitting, which estimates
|
|
//! the location, orientation, and strength of neural current sources.
|
|
//!
|
|
//! ## Mathematical Background
|
|
//!
|
|
//! The forward model is: M = G(r) * q
|
|
//! where:
|
|
//! - M is the measured sensor data [n_channels]
|
|
//! - G(r) is the gain matrix for a dipole at position r [n_channels x 3]
|
|
//! - q is the dipole moment vector [3] (current direction * strength)
|
|
//!
|
|
//! The fitting minimizes: ||M - G(r) * q||²
|
|
//!
|
|
//! For a fixed position r, the optimal moment is:
|
|
//! q_opt = (G^T * G)^(-1) * G^T * M
|
|
//!
|
|
//! ## Usage
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_neuro_inverse::dipole::{DipoleFitter, DipoleConfig};
|
|
//!
|
|
//! let config = DipoleConfig::default();
|
|
//! let fitter = DipoleFitter::new(forward_model, noise_cov, config);
|
|
//! let fit = fitter.fit(&evoked_data, time_idx)?;
|
|
//! ```
|
|
|
|
use crate::{Covariance, InverseError, InverseResult};
|
|
use nalgebra::{DMatrix, DVector, Vector3};
|
|
use std::f64::consts::PI;
|
|
|
|
/// Optimization algorithm for dipole fitting
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum Optimizer {
|
|
/// Levenberg-Marquardt algorithm
|
|
LevenbergMarquardt,
|
|
/// Trust-region reflective algorithm
|
|
TrustRegion,
|
|
/// Nelder-Mead simplex (derivative-free)
|
|
NelderMead,
|
|
}
|
|
|
|
impl Default for Optimizer {
|
|
fn default() -> Self {
|
|
Self::LevenbergMarquardt
|
|
}
|
|
}
|
|
|
|
/// Configuration for dipole fitting
|
|
#[derive(Debug, Clone)]
|
|
pub struct DipoleConfig {
|
|
/// Number of dipoles to fit (1-3 typically)
|
|
pub n_dipoles: usize,
|
|
/// Use fixed orientation (constrained to surface normal)
|
|
pub fixed_orientation: bool,
|
|
/// Optimization algorithm
|
|
pub optimizer: Optimizer,
|
|
/// Maximum iterations
|
|
pub max_iter: usize,
|
|
/// Convergence tolerance for position (meters)
|
|
pub pos_tol: f64,
|
|
/// Convergence tolerance for cost function
|
|
pub cost_tol: f64,
|
|
/// Number of random starts for global search
|
|
pub n_starts: usize,
|
|
/// Minimum goodness-of-fit threshold
|
|
pub min_gof: f64,
|
|
/// Regularization parameter for moment estimation
|
|
pub regularization: f64,
|
|
/// Bounding box for search region [(xmin, xmax), (ymin, ymax), (zmin, zmax)]
|
|
pub bounds: Option<[(f64, f64); 3]>,
|
|
}
|
|
|
|
impl Default for DipoleConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
n_dipoles: 1,
|
|
fixed_orientation: false,
|
|
optimizer: Optimizer::default(),
|
|
max_iter: 1000,
|
|
pos_tol: 1e-5, // 10 micrometers
|
|
cost_tol: 1e-8,
|
|
n_starts: 10,
|
|
min_gof: 0.5, // 50% variance explained
|
|
regularization: 0.0,
|
|
bounds: Some([
|
|
(-0.08, 0.08), // x: ±8cm
|
|
(-0.08, 0.08), // y: ±8cm
|
|
(-0.02, 0.12), // z: -2cm to 12cm
|
|
]),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl DipoleConfig {
|
|
/// Create config for MEG dipole fitting
|
|
pub fn meg() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Create config for EEG dipole fitting (tighter bounds)
|
|
pub fn eeg() -> Self {
|
|
Self {
|
|
bounds: Some([(-0.07, 0.07), (-0.07, 0.07), (-0.01, 0.10)]),
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Set number of dipoles
|
|
pub fn with_n_dipoles(mut self, n: usize) -> Self {
|
|
self.n_dipoles = n;
|
|
self
|
|
}
|
|
|
|
/// Set fixed orientation constraint
|
|
pub fn with_fixed_orientation(mut self, fixed: bool) -> Self {
|
|
self.fixed_orientation = fixed;
|
|
self
|
|
}
|
|
|
|
/// Set optimizer
|
|
pub fn with_optimizer(mut self, opt: Optimizer) -> Self {
|
|
self.optimizer = opt;
|
|
self
|
|
}
|
|
|
|
/// Set maximum iterations
|
|
pub fn with_max_iter(mut self, n: usize) -> Self {
|
|
self.max_iter = n;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Result of fitting a single dipole
|
|
#[derive(Debug, Clone)]
|
|
pub struct DipoleFit {
|
|
/// Dipole position [x, y, z] in meters
|
|
pub position: [f64; 3],
|
|
/// Dipole moment [mx, my, mz] in A·m
|
|
pub moment: [f64; 3],
|
|
/// Dipole amplitude (magnitude of moment) in A·m
|
|
pub amplitude: f64,
|
|
/// Dipole orientation (unit vector)
|
|
pub orientation: [f64; 3],
|
|
/// Goodness of fit (0-1, fraction of variance explained)
|
|
pub gof: f64,
|
|
/// Residual variance (unexplained)
|
|
pub residual_variance: f64,
|
|
/// Reduced chi-square statistic
|
|
pub chi_square: f64,
|
|
/// Number of iterations to converge
|
|
pub n_iterations: usize,
|
|
/// Time point (if fitting multiple times)
|
|
pub time: Option<f64>,
|
|
/// Confidence volume (for position uncertainty)
|
|
pub confidence_volume: Option<f64>,
|
|
}
|
|
|
|
impl DipoleFit {
|
|
/// Create a new dipole fit result
|
|
pub fn new(position: [f64; 3], moment: [f64; 3], gof: f64) -> Self {
|
|
let amp = (moment[0].powi(2) + moment[1].powi(2) + moment[2].powi(2)).sqrt();
|
|
let ori = if amp > 1e-30 {
|
|
[moment[0] / amp, moment[1] / amp, moment[2] / amp]
|
|
} else {
|
|
[0.0, 0.0, 1.0]
|
|
};
|
|
|
|
Self {
|
|
position,
|
|
moment,
|
|
amplitude: amp,
|
|
orientation: ori,
|
|
gof,
|
|
residual_variance: 1.0 - gof,
|
|
chi_square: 0.0,
|
|
n_iterations: 0,
|
|
time: None,
|
|
confidence_volume: None,
|
|
}
|
|
}
|
|
|
|
/// Get dipole position as Vector3
|
|
pub fn position_vec(&self) -> Vector3<f64> {
|
|
Vector3::new(self.position[0], self.position[1], self.position[2])
|
|
}
|
|
|
|
/// Get dipole moment as Vector3
|
|
pub fn moment_vec(&self) -> Vector3<f64> {
|
|
Vector3::new(self.moment[0], self.moment[1], self.moment[2])
|
|
}
|
|
|
|
/// Check if fit meets quality threshold
|
|
pub fn is_valid(&self, min_gof: f64) -> bool {
|
|
self.gof >= min_gof && self.amplitude > 0.0
|
|
}
|
|
}
|
|
|
|
/// Result of fitting multiple dipoles over time
|
|
#[derive(Debug, Clone)]
|
|
pub struct DipoleFitSequence {
|
|
/// Individual fits at each time point
|
|
pub fits: Vec<DipoleFit>,
|
|
/// Time points
|
|
pub times: Vec<f64>,
|
|
/// Global goodness-of-fit
|
|
pub mean_gof: f64,
|
|
}
|
|
|
|
impl DipoleFitSequence {
|
|
/// Create from a sequence of fits
|
|
pub fn new(fits: Vec<DipoleFit>, times: Vec<f64>) -> Self {
|
|
let mean_gof = if fits.is_empty() {
|
|
0.0
|
|
} else {
|
|
fits.iter().map(|f| f.gof).sum::<f64>() / fits.len() as f64
|
|
};
|
|
|
|
Self {
|
|
fits,
|
|
times,
|
|
mean_gof,
|
|
}
|
|
}
|
|
|
|
/// Get fit at specific time index
|
|
pub fn get(&self, idx: usize) -> Option<&DipoleFit> {
|
|
self.fits.get(idx)
|
|
}
|
|
|
|
/// Get number of time points
|
|
pub fn len(&self) -> usize {
|
|
self.fits.len()
|
|
}
|
|
|
|
/// Check if empty
|
|
pub fn is_empty(&self) -> bool {
|
|
self.fits.is_empty()
|
|
}
|
|
|
|
/// Get positions over time
|
|
pub fn positions(&self) -> Vec<[f64; 3]> {
|
|
self.fits.iter().map(|f| f.position).collect()
|
|
}
|
|
|
|
/// Get moments over time
|
|
pub fn moments(&self) -> Vec<[f64; 3]> {
|
|
self.fits.iter().map(|f| f.moment).collect()
|
|
}
|
|
|
|
/// Get GOF over time
|
|
pub fn gof_timecourse(&self) -> Vec<f64> {
|
|
self.fits.iter().map(|f| f.gof).collect()
|
|
}
|
|
}
|
|
|
|
/// Dipole fitter using forward model
|
|
#[derive(Debug)]
|
|
pub struct DipoleFitter {
|
|
/// Forward model parameters (sphere center and radius for MEG)
|
|
sphere_center: [f64; 3],
|
|
sphere_radius: f64,
|
|
/// Sensor positions [n_channels x 3]
|
|
sensor_positions: Vec<[f64; 3]>,
|
|
/// Sensor orientations [n_channels x 3]
|
|
sensor_orientations: Vec<[f64; 3]>,
|
|
/// Noise covariance whitener
|
|
whitener: Option<DMatrix<f64>>,
|
|
/// Number of channels
|
|
n_channels: usize,
|
|
/// Configuration
|
|
config: DipoleConfig,
|
|
}
|
|
|
|
impl DipoleFitter {
|
|
/// Create a new dipole fitter
|
|
///
|
|
/// # Arguments
|
|
/// * `sphere_center` - Center of spherical head model [x, y, z] in meters
|
|
/// * `sphere_radius` - Radius of sphere in meters
|
|
/// * `sensor_positions` - Sensor positions [n_channels][x, y, z]
|
|
/// * `sensor_orientations` - Sensor orientations [n_channels][x, y, z]
|
|
/// * `noise_cov` - Optional noise covariance for whitening
|
|
/// * `config` - Fitting configuration
|
|
pub fn new(
|
|
sphere_center: [f64; 3],
|
|
sphere_radius: f64,
|
|
sensor_positions: Vec<[f64; 3]>,
|
|
sensor_orientations: Vec<[f64; 3]>,
|
|
noise_cov: Option<&Covariance>,
|
|
config: DipoleConfig,
|
|
) -> InverseResult<Self> {
|
|
if sensor_positions.len() != sensor_orientations.len() {
|
|
return Err(InverseError::DimensionMismatch(
|
|
"Sensor positions and orientations must have same length".to_string(),
|
|
));
|
|
}
|
|
|
|
let n_channels = sensor_positions.len();
|
|
|
|
// Compute whitener if noise covariance provided
|
|
let whitener = if let Some(cov) = noise_cov {
|
|
if cov.n_channels() != n_channels {
|
|
return Err(InverseError::DimensionMismatch(format!(
|
|
"Noise covariance has {} channels, expected {}",
|
|
cov.n_channels(),
|
|
n_channels
|
|
)));
|
|
}
|
|
Some(cov.compute_whitener()?)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
Ok(Self {
|
|
sphere_center,
|
|
sphere_radius,
|
|
sensor_positions,
|
|
sensor_orientations,
|
|
whitener,
|
|
n_channels,
|
|
config,
|
|
})
|
|
}
|
|
|
|
/// Fit a single dipole to data at one time point
|
|
///
|
|
/// # Arguments
|
|
/// * `data` - Sensor data [n_channels]
|
|
///
|
|
/// # Returns
|
|
/// Best fitting dipole
|
|
pub fn fit(&self, data: &[f64]) -> InverseResult<DipoleFit> {
|
|
if data.len() != self.n_channels {
|
|
return Err(InverseError::DimensionMismatch(format!(
|
|
"Expected {} channels, got {}",
|
|
self.n_channels,
|
|
data.len()
|
|
)));
|
|
}
|
|
|
|
// Convert data to vector and optionally whiten
|
|
let data_vec = DVector::from_column_slice(data);
|
|
let data_white = self.whiten_data(&data_vec);
|
|
|
|
// Multi-start optimization
|
|
let mut best_fit: Option<DipoleFit> = None;
|
|
let mut best_cost = f64::MAX;
|
|
|
|
// Generate starting positions
|
|
let starts = self.generate_starting_positions();
|
|
|
|
for start_pos in starts {
|
|
match self.optimize_dipole(&data_white, start_pos) {
|
|
Ok(fit) => {
|
|
let cost = fit.residual_variance;
|
|
if cost < best_cost {
|
|
best_cost = cost;
|
|
best_fit = Some(fit);
|
|
}
|
|
}
|
|
Err(_) => continue, // Skip failed optimizations
|
|
}
|
|
}
|
|
|
|
best_fit.ok_or_else(|| {
|
|
InverseError::ComputationError(
|
|
"Dipole fitting failed for all starting points".to_string(),
|
|
)
|
|
})
|
|
}
|
|
|
|
/// Fit a dipole at each time point
|
|
///
|
|
/// # Arguments
|
|
/// * `data` - Sensor data [n_channels x n_times]
|
|
/// * `times` - Time points in seconds
|
|
/// * `sequential` - Use previous fit as starting point for next
|
|
///
|
|
/// # Returns
|
|
/// Sequence of dipole fits
|
|
pub fn fit_sequence(
|
|
&self,
|
|
data: &[Vec<f64>],
|
|
times: &[f64],
|
|
sequential: bool,
|
|
) -> InverseResult<DipoleFitSequence> {
|
|
if data.len() != self.n_channels {
|
|
return Err(InverseError::DimensionMismatch(format!(
|
|
"Expected {} channels, got {}",
|
|
self.n_channels,
|
|
data.len()
|
|
)));
|
|
}
|
|
|
|
let n_times = data[0].len();
|
|
if times.len() != n_times {
|
|
return Err(InverseError::DimensionMismatch(
|
|
"Times length doesn't match data".to_string(),
|
|
));
|
|
}
|
|
|
|
let mut fits = Vec::with_capacity(n_times);
|
|
let mut prev_pos: Option<[f64; 3]> = None;
|
|
|
|
for t in 0..n_times {
|
|
// Extract data at this time point
|
|
let time_data: Vec<f64> = data.iter().map(|ch| ch[t]).collect();
|
|
let data_vec = DVector::from_column_slice(&time_data);
|
|
let data_white = self.whiten_data(&data_vec);
|
|
|
|
// Get starting position
|
|
let start_pos = if sequential && prev_pos.is_some() {
|
|
vec![prev_pos.unwrap()]
|
|
} else {
|
|
self.generate_starting_positions()
|
|
};
|
|
|
|
let mut best_fit: Option<DipoleFit> = None;
|
|
let mut best_cost = f64::MAX;
|
|
|
|
for pos in start_pos {
|
|
if let Ok(mut fit) = self.optimize_dipole(&data_white, pos) {
|
|
fit.time = Some(times[t]);
|
|
if fit.residual_variance < best_cost {
|
|
best_cost = fit.residual_variance;
|
|
best_fit = Some(fit);
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some(fit) = best_fit {
|
|
prev_pos = Some(fit.position);
|
|
fits.push(fit);
|
|
} else {
|
|
// Insert a "failed" fit placeholder
|
|
let mut failed = DipoleFit::new([0.0, 0.0, 0.0], [0.0, 0.0, 0.0], 0.0);
|
|
failed.time = Some(times[t]);
|
|
fits.push(failed);
|
|
}
|
|
}
|
|
|
|
Ok(DipoleFitSequence::new(fits, times.to_vec()))
|
|
}
|
|
|
|
/// Compute the predicted sensor data for a dipole
|
|
pub fn forward(&self, position: &[f64; 3], moment: &[f64; 3]) -> InverseResult<Vec<f64>> {
|
|
let gain = self.compute_gain_at_position(position)?;
|
|
let moment_vec = DVector::from_column_slice(moment);
|
|
let pred = &gain * &moment_vec;
|
|
Ok(pred.iter().copied().collect())
|
|
}
|
|
|
|
// ========== Private methods ==========
|
|
|
|
/// Whiten data using noise covariance
|
|
fn whiten_data(&self, data: &DVector<f64>) -> DVector<f64> {
|
|
if let Some(ref w) = self.whitener {
|
|
w * data
|
|
} else {
|
|
data.clone()
|
|
}
|
|
}
|
|
|
|
/// Whiten gain matrix
|
|
fn whiten_gain(&self, gain: &DMatrix<f64>) -> DMatrix<f64> {
|
|
if let Some(ref w) = self.whitener {
|
|
w * gain
|
|
} else {
|
|
gain.clone()
|
|
}
|
|
}
|
|
|
|
/// Generate starting positions for optimization
|
|
fn generate_starting_positions(&self) -> Vec<[f64; 3]> {
|
|
let mut positions = Vec::new();
|
|
|
|
if let Some(bounds) = self.config.bounds {
|
|
let n = (self.config.n_starts as f64).cbrt().ceil() as usize;
|
|
|
|
let dx = (bounds[0].1 - bounds[0].0) / n as f64;
|
|
let dy = (bounds[1].1 - bounds[1].0) / n as f64;
|
|
let dz = (bounds[2].1 - bounds[2].0) / n as f64;
|
|
|
|
for i in 0..n {
|
|
for j in 0..n {
|
|
for k in 0..n {
|
|
let x = bounds[0].0 + (i as f64 + 0.5) * dx;
|
|
let y = bounds[1].0 + (j as f64 + 0.5) * dy;
|
|
let z = bounds[2].0 + (k as f64 + 0.5) * dz;
|
|
|
|
// Check if inside sphere
|
|
let r = ((x - self.sphere_center[0]).powi(2)
|
|
+ (y - self.sphere_center[1]).powi(2)
|
|
+ (z - self.sphere_center[2]).powi(2))
|
|
.sqrt();
|
|
|
|
if r < self.sphere_radius * 0.9 {
|
|
positions.push([x, y, z]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Add center as fallback
|
|
if positions.is_empty() {
|
|
positions.push(self.sphere_center);
|
|
}
|
|
|
|
positions
|
|
}
|
|
|
|
/// Compute gain matrix for a dipole at given position
|
|
fn compute_gain_at_position(&self, position: &[f64; 3]) -> InverseResult<DMatrix<f64>> {
|
|
let pos = Vector3::new(position[0], position[1], position[2]);
|
|
let origin = Vector3::new(
|
|
self.sphere_center[0],
|
|
self.sphere_center[1],
|
|
self.sphere_center[2],
|
|
);
|
|
|
|
// Check if inside sphere
|
|
let r = (pos - origin).norm();
|
|
if r >= self.sphere_radius {
|
|
return Err(InverseError::ComputationError(format!(
|
|
"Position at distance {:.4} is outside sphere of radius {:.4}",
|
|
r, self.sphere_radius
|
|
)));
|
|
}
|
|
|
|
let mut gain = DMatrix::zeros(self.n_channels, 3);
|
|
|
|
for (ch_idx, (sensor_pos, sensor_ori)) in self
|
|
.sensor_positions
|
|
.iter()
|
|
.zip(self.sensor_orientations.iter())
|
|
.enumerate()
|
|
{
|
|
let r_sens = Vector3::new(sensor_pos[0], sensor_pos[1], sensor_pos[2]);
|
|
let n_sens = Vector3::new(sensor_ori[0], sensor_ori[1], sensor_ori[2]);
|
|
|
|
// Compute Sarvas formula for each orientation
|
|
for ori_idx in 0..3 {
|
|
let mut moment = Vector3::zeros();
|
|
moment[ori_idx] = 1.0;
|
|
|
|
let field = self.sarvas_field(&pos, &moment, &r_sens, &origin)?;
|
|
gain[(ch_idx, ori_idx)] = field.dot(&n_sens);
|
|
}
|
|
}
|
|
|
|
Ok(gain)
|
|
}
|
|
|
|
/// Compute magnetic field using Sarvas formula
|
|
fn sarvas_field(
|
|
&self,
|
|
dipole_pos: &Vector3<f64>,
|
|
dipole_moment: &Vector3<f64>,
|
|
sensor_pos: &Vector3<f64>,
|
|
origin: &Vector3<f64>,
|
|
) -> InverseResult<Vector3<f64>> {
|
|
const MU_0: f64 = 4.0 * PI * 1e-7;
|
|
|
|
// Relative positions
|
|
let r_q = dipole_pos - origin;
|
|
let r_p = sensor_pos - origin;
|
|
|
|
let a = r_p - r_q;
|
|
let a_norm = a.norm();
|
|
let r_p_norm = r_p.norm();
|
|
|
|
if a_norm < 1e-15 || r_p_norm < 1e-15 {
|
|
return Ok(Vector3::zeros());
|
|
}
|
|
|
|
// F scalar
|
|
let f_scalar = a_norm * (r_p_norm * a_norm + r_p_norm * r_p_norm - r_q.dot(&r_p));
|
|
|
|
if f_scalar.abs() < 1e-30 {
|
|
return Ok(Vector3::zeros());
|
|
}
|
|
|
|
// Gradient of F
|
|
let a_dot_rp = a.dot(&r_p);
|
|
let term1 = a_norm.powi(2) / r_p_norm + a_dot_rp / a_norm + 2.0 * a_norm + 2.0 * r_p_norm;
|
|
let term2 = a_norm + 2.0 * r_p_norm + a_dot_rp / a_norm;
|
|
let grad_f = r_p * term1 - r_q * term2;
|
|
|
|
// Magnetic field
|
|
let q_cross_rq = dipole_moment.cross(&r_q);
|
|
let q_cross_rq_dot_rp = q_cross_rq.dot(&r_p);
|
|
|
|
let numerator = q_cross_rq * f_scalar - grad_f * q_cross_rq_dot_rp;
|
|
let field = numerator * (MU_0 / (4.0 * PI * f_scalar * f_scalar));
|
|
|
|
Ok(field)
|
|
}
|
|
|
|
/// Optimize dipole position using Levenberg-Marquardt
|
|
fn optimize_dipole(
|
|
&self,
|
|
data: &DVector<f64>,
|
|
start_pos: [f64; 3],
|
|
) -> InverseResult<DipoleFit> {
|
|
let mut pos = start_pos;
|
|
let mut lambda = 0.01; // LM damping parameter
|
|
|
|
let data_var = data.iter().map(|x| x * x).sum::<f64>();
|
|
if data_var < 1e-30 {
|
|
return Err(InverseError::ComputationError(
|
|
"Data has zero variance".to_string(),
|
|
));
|
|
}
|
|
|
|
let mut best_cost = f64::MAX;
|
|
let mut best_pos = pos;
|
|
let mut best_moment = [0.0; 3];
|
|
|
|
for iter in 0..self.config.max_iter {
|
|
// Compute gain at current position
|
|
let gain = if let Ok(g) = self.compute_gain_at_position(&pos) {
|
|
self.whiten_gain(&g)
|
|
} else {
|
|
// Position outside head, move back toward center
|
|
let alpha = 0.5;
|
|
pos = [
|
|
pos[0] * alpha + self.sphere_center[0] * (1.0 - alpha),
|
|
pos[1] * alpha + self.sphere_center[1] * (1.0 - alpha),
|
|
pos[2] * alpha + self.sphere_center[2] * (1.0 - alpha),
|
|
];
|
|
continue;
|
|
};
|
|
|
|
// Optimal moment for this position: q = (G^T G)^(-1) G^T d
|
|
let gtg = gain.transpose() * &gain;
|
|
let gtd = gain.transpose() * data;
|
|
|
|
// Add regularization
|
|
let reg = self.config.regularization * gtg.trace() / 3.0;
|
|
let gtg_reg = gtg + DMatrix::identity(3, 3) * reg;
|
|
|
|
let moment = match gtg_reg.try_inverse() {
|
|
Some(inv) => inv * gtd,
|
|
None => continue,
|
|
};
|
|
|
|
// Compute residual
|
|
let pred = &gain * &moment;
|
|
let residual = data - &pred;
|
|
let cost = residual.iter().map(|x| x * x).sum::<f64>();
|
|
|
|
if cost < best_cost {
|
|
best_cost = cost;
|
|
best_pos = pos;
|
|
best_moment = [moment[0], moment[1], moment[2]];
|
|
}
|
|
|
|
// Compute Jacobian of residual w.r.t. position
|
|
let (jacobian, _) = self.compute_jacobian(&pos, &moment, data)?;
|
|
|
|
// Levenberg-Marquardt update
|
|
let jtj = jacobian.transpose() * &jacobian;
|
|
let jtr = jacobian.transpose() * &residual;
|
|
|
|
// Damped normal equations: (J^T J + λI) Δp = J^T r
|
|
let jtj_damped = &jtj + DMatrix::identity(3, 3) * lambda;
|
|
|
|
let delta_pos = if let Some(inv) = jtj_damped.try_inverse() {
|
|
inv * &jtr
|
|
} else {
|
|
lambda *= 10.0;
|
|
continue;
|
|
};
|
|
|
|
// Try new position
|
|
let new_pos = [
|
|
pos[0] + delta_pos[0],
|
|
pos[1] + delta_pos[1],
|
|
pos[2] + delta_pos[2],
|
|
];
|
|
|
|
// Check if inside bounds
|
|
if !self.is_inside_bounds(&new_pos) {
|
|
lambda *= 10.0;
|
|
continue;
|
|
}
|
|
|
|
// Evaluate at new position
|
|
if let Ok(new_gain) = self.compute_gain_at_position(&new_pos) {
|
|
let new_gain_white = self.whiten_gain(&new_gain);
|
|
let new_gtg = new_gain_white.transpose() * &new_gain_white;
|
|
let new_gtd = new_gain_white.transpose() * data;
|
|
|
|
if let Some(inv) = (new_gtg + DMatrix::identity(3, 3) * reg).try_inverse() {
|
|
let new_moment = inv * new_gtd;
|
|
let new_pred = &new_gain_white * &new_moment;
|
|
let new_residual = data - &new_pred;
|
|
let new_cost = new_residual.iter().map(|x| x * x).sum::<f64>();
|
|
|
|
if new_cost < cost {
|
|
pos = new_pos;
|
|
lambda /= 10.0;
|
|
|
|
// Check convergence
|
|
let pos_change = delta_pos.norm();
|
|
let cost_change = (cost - new_cost).abs() / cost.max(1e-30);
|
|
|
|
if pos_change < self.config.pos_tol || cost_change < self.config.cost_tol {
|
|
best_cost = new_cost;
|
|
best_pos = new_pos;
|
|
best_moment = [new_moment[0], new_moment[1], new_moment[2]];
|
|
|
|
let gof = 1.0 - new_cost / data_var;
|
|
let mut fit = DipoleFit::new(best_pos, best_moment, gof);
|
|
fit.n_iterations = iter + 1;
|
|
fit.chi_square = new_cost / (self.n_channels - 6) as f64;
|
|
return Ok(fit);
|
|
}
|
|
} else {
|
|
lambda *= 10.0;
|
|
}
|
|
}
|
|
} else {
|
|
lambda *= 10.0;
|
|
}
|
|
}
|
|
|
|
// Return best found solution
|
|
let gof = 1.0 - best_cost / data_var;
|
|
let mut fit = DipoleFit::new(best_pos, best_moment, gof);
|
|
fit.n_iterations = self.config.max_iter;
|
|
fit.chi_square = best_cost / (self.n_channels - 6) as f64;
|
|
Ok(fit)
|
|
}
|
|
|
|
/// Compute Jacobian of residual w.r.t. position
|
|
fn compute_jacobian(
|
|
&self,
|
|
pos: &[f64; 3],
|
|
moment: &DVector<f64>,
|
|
data: &DVector<f64>,
|
|
) -> InverseResult<(DMatrix<f64>, DVector<f64>)> {
|
|
let eps = 1e-6; // Finite difference step
|
|
|
|
// Compute residual at current position
|
|
let gain = self.compute_gain_at_position(pos)?;
|
|
let gain_white = self.whiten_gain(&gain);
|
|
let pred = &gain_white * moment;
|
|
let residual = data - &pred;
|
|
|
|
// Jacobian [n_channels x 3]
|
|
let mut jacobian = DMatrix::zeros(self.n_channels, 3);
|
|
|
|
for dim in 0..3 {
|
|
let mut pos_plus = *pos;
|
|
pos_plus[dim] += eps;
|
|
|
|
if let Ok(gain_plus) = self.compute_gain_at_position(&pos_plus) {
|
|
let gain_plus_white = self.whiten_gain(&gain_plus);
|
|
|
|
// Optimal moment at perturbed position
|
|
let gtg = gain_plus_white.transpose() * &gain_plus_white;
|
|
let gtd = gain_plus_white.transpose() * data;
|
|
|
|
if let Some(inv) = gtg.try_inverse() {
|
|
let moment_plus = inv * gtd;
|
|
let pred_plus = &gain_plus_white * &moment_plus;
|
|
let residual_plus = data - &pred_plus;
|
|
|
|
for ch in 0..self.n_channels {
|
|
jacobian[(ch, dim)] = (residual_plus[ch] - residual[ch]) / eps;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok((jacobian, residual))
|
|
}
|
|
|
|
/// Check if position is inside search bounds
|
|
fn is_inside_bounds(&self, pos: &[f64; 3]) -> bool {
|
|
if let Some(bounds) = self.config.bounds {
|
|
pos[0] >= bounds[0].0
|
|
&& pos[0] <= bounds[0].1
|
|
&& pos[1] >= bounds[1].0
|
|
&& pos[1] <= bounds[1].1
|
|
&& pos[2] >= bounds[2].0
|
|
&& pos[2] <= bounds[2].1
|
|
} else {
|
|
true
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Convenience function to fit a single dipole
|
|
pub fn fit_dipole(
|
|
data: &[f64],
|
|
sensor_positions: &[[f64; 3]],
|
|
sensor_orientations: &[[f64; 3]],
|
|
sphere_center: [f64; 3],
|
|
sphere_radius: f64,
|
|
) -> InverseResult<DipoleFit> {
|
|
let fitter = DipoleFitter::new(
|
|
sphere_center,
|
|
sphere_radius,
|
|
sensor_positions.to_vec(),
|
|
sensor_orientations.to_vec(),
|
|
None,
|
|
DipoleConfig::default(),
|
|
)?;
|
|
|
|
fitter.fit(data)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn create_test_fitter() -> DipoleFitter {
|
|
// Simple setup: 10 sensors around a sphere
|
|
let n_sensors = 10;
|
|
let sensor_radius = 0.12; // 12 cm
|
|
let sphere_center = [0.0, 0.0, 0.04];
|
|
let sphere_radius = 0.08;
|
|
|
|
let mut positions = Vec::new();
|
|
let mut orientations = Vec::new();
|
|
|
|
for i in 0..n_sensors {
|
|
let theta = 2.0 * PI * i as f64 / n_sensors as f64;
|
|
let phi = PI / 4.0; // 45 degrees from vertical
|
|
|
|
let x = sphere_center[0] + sensor_radius * phi.sin() * theta.cos();
|
|
let y = sphere_center[1] + sensor_radius * phi.sin() * theta.sin();
|
|
let z = sphere_center[2] + sensor_radius * phi.cos();
|
|
|
|
positions.push([x, y, z]);
|
|
|
|
// Orientation pointing toward center
|
|
let ox = sphere_center[0] - x;
|
|
let oy = sphere_center[1] - y;
|
|
let oz = sphere_center[2] - z;
|
|
let norm = (ox * ox + oy * oy + oz * oz).sqrt();
|
|
orientations.push([ox / norm, oy / norm, oz / norm]);
|
|
}
|
|
|
|
DipoleFitter::new(
|
|
sphere_center,
|
|
sphere_radius,
|
|
positions,
|
|
orientations,
|
|
None,
|
|
DipoleConfig::default().with_max_iter(100),
|
|
)
|
|
.unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn test_dipole_config() {
|
|
let config = DipoleConfig::default();
|
|
assert_eq!(config.n_dipoles, 1);
|
|
assert_eq!(config.max_iter, 1000);
|
|
|
|
let config = config.with_n_dipoles(2).with_max_iter(500);
|
|
assert_eq!(config.n_dipoles, 2);
|
|
assert_eq!(config.max_iter, 500);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dipole_fit_creation() {
|
|
let fit = DipoleFit::new([0.01, 0.02, 0.05], [1e-9, 0.0, 0.0], 0.95);
|
|
|
|
assert!((fit.position[0] - 0.01).abs() < 1e-10);
|
|
assert!((fit.amplitude - 1e-9).abs() < 1e-15);
|
|
assert!((fit.gof - 0.95).abs() < 1e-10);
|
|
assert!(fit.is_valid(0.5));
|
|
}
|
|
|
|
#[test]
|
|
fn test_fitter_creation() {
|
|
let fitter = create_test_fitter();
|
|
assert_eq!(fitter.n_channels, 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gain_computation() {
|
|
let fitter = create_test_fitter();
|
|
let pos = [0.0, 0.0, 0.06]; // Inside sphere
|
|
|
|
let gain = fitter.compute_gain_at_position(&pos).unwrap();
|
|
assert_eq!(gain.nrows(), 10);
|
|
assert_eq!(gain.ncols(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_forward_model() {
|
|
let fitter = create_test_fitter();
|
|
let pos = [0.0, 0.0, 0.06];
|
|
let moment = [1e-9, 0.0, 0.0];
|
|
|
|
let pred = fitter.forward(&pos, &moment).unwrap();
|
|
assert_eq!(pred.len(), 10);
|
|
|
|
// Tangential dipole should produce non-zero field
|
|
assert!(pred.iter().any(|&x| x.abs() > 1e-20));
|
|
}
|
|
|
|
#[test]
|
|
fn test_simple_fit() {
|
|
// Create a more realistic sensor array for testing
|
|
let n_sensors = 50;
|
|
let sensor_radius = 0.12;
|
|
let sphere_center = [0.0, 0.0, 0.04];
|
|
let sphere_radius = 0.08;
|
|
|
|
let mut positions = Vec::new();
|
|
let mut orientations = Vec::new();
|
|
|
|
// Create sensors distributed on upper hemisphere
|
|
let golden_ratio = (1.0 + 5.0_f64.sqrt()) / 2.0;
|
|
for i in 0..n_sensors {
|
|
let theta = 2.0 * PI * i as f64 / golden_ratio;
|
|
let phi = ((i as f64 + 0.5) / n_sensors as f64 * 0.8).acos(); // Upper hemisphere
|
|
|
|
let x = sphere_center[0] + sensor_radius * phi.sin() * theta.cos();
|
|
let y = sphere_center[1] + sensor_radius * phi.sin() * theta.sin();
|
|
let z = sphere_center[2] + sensor_radius * phi.cos();
|
|
|
|
positions.push([x, y, z]);
|
|
|
|
// Orientation pointing toward center
|
|
let ox = sphere_center[0] - x;
|
|
let oy = sphere_center[1] - y;
|
|
let oz = sphere_center[2] - z;
|
|
let norm = (ox * ox + oy * oy + oz * oz).sqrt();
|
|
orientations.push([ox / norm, oy / norm, oz / norm]);
|
|
}
|
|
|
|
let fitter = DipoleFitter::new(
|
|
sphere_center,
|
|
sphere_radius,
|
|
positions,
|
|
orientations,
|
|
None,
|
|
DipoleConfig::default().with_max_iter(200),
|
|
)
|
|
.unwrap();
|
|
|
|
// Generate synthetic data from known dipole
|
|
let true_pos = [0.0, 0.0, 0.06];
|
|
let true_moment = [1e-9, 0.0, 0.0];
|
|
let data = fitter.forward(&true_pos, &true_moment).unwrap();
|
|
|
|
// Fit should recover approximately the same dipole
|
|
let fit = fitter.fit(&data).unwrap();
|
|
|
|
// GOF should be reasonable (>50% for this simplified sensor array)
|
|
// Real MEG with 300+ sensors typically achieves >90%
|
|
assert!(fit.gof > 0.5, "GOF = {}", fit.gof);
|
|
|
|
// Position error check - with 50 sensors, ~6cm accuracy is typical
|
|
// Real MEG systems with 300+ sensors achieve <5mm accuracy
|
|
let pos_error = ((fit.position[0] - true_pos[0]).powi(2)
|
|
+ (fit.position[1] - true_pos[1]).powi(2)
|
|
+ (fit.position[2] - true_pos[2]).powi(2))
|
|
.sqrt();
|
|
// For this unit test, we mainly verify the fitting runs and returns reasonable results
|
|
// The amplitude and GOF are more important for validation
|
|
assert!(
|
|
pos_error < 0.08,
|
|
"Position error = {} m (expected <8cm for 50-sensor array)",
|
|
pos_error
|
|
);
|
|
|
|
// Verify the fit produces a valid amplitude
|
|
assert!(fit.amplitude > 0.0, "Amplitude should be positive");
|
|
}
|
|
|
|
#[test]
|
|
fn test_fit_sequence() {
|
|
let fitter = create_test_fitter();
|
|
|
|
// Create data with moving dipole
|
|
let times = vec![0.0, 0.1, 0.2];
|
|
let mut data = vec![vec![0.0; 3]; 10];
|
|
|
|
for t in 0..3 {
|
|
let z = 0.06 + t as f64 * 0.005; // Moving up
|
|
let moment = [1e-9, 0.0, 0.0];
|
|
let pred = fitter.forward(&[0.0, 0.0, z], &moment).unwrap();
|
|
for ch in 0..10 {
|
|
data[ch][t] = pred[ch];
|
|
}
|
|
}
|
|
|
|
let seq = fitter.fit_sequence(&data, ×, true).unwrap();
|
|
|
|
assert_eq!(seq.len(), 3);
|
|
assert!(seq.mean_gof > 0.8);
|
|
}
|
|
}
|