Files
rustytorch/crates/specialized/rtx-neuro-inverse/src/loreta.rs
T
2026-03-04 00:08:42 +00:00

724 lines
23 KiB
Rust

//! eLORETA (Exact Low Resolution Electromagnetic Tomography) inverse solution.
//!
//! eLORETA is an inverse method with the unique property of exact (zero-error)
//! localization for point sources. Unlike sLORETA which only standardizes the
//! current density, eLORETA computes weights that ensure the resolution matrix
//! has specific properties guaranteeing zero localization error.
//!
//! ## Mathematical Background
//!
//! The eLORETA solution is: J = W * M
//!
//! where W is computed iteratively such that the resulting resolution matrix
//! R = W * G satisfies specific properties.
//!
//! The weight matrix W_i for source i is:
//! W_i = [G_i^T * C^(-1) * G_i]^(-1/2) for each source orientation block
//!
//! This ensures that the point-spread function is centered on the true source.
//!
//! ## References
//!
//! - Pascual-Marqui, R.D. (2007). Discrete, 3D distributed, linear imaging methods
//! of electric neuronal activity. Part 1: exact, zero error localization.
//! - Pascual-Marqui, R.D. (2009). Theory of the EEG inverse problem.
//!
//! ## Usage
//!
//! ```rust,ignore
//! use rtx_neuro_inverse::loreta::{EloretaInverse, EloretaConfig};
//!
//! let config = EloretaConfig::default();
//! let inverse = EloretaInverse::make_inverse(&gain, &noise_cov, config)?;
//! let stc = inverse.apply(&evoked_data)?;
//! ```
use crate::{Covariance, InverseError, InverseResult, SourceEstimate};
use nalgebra::{DMatrix, DVector, Matrix3};
use rtx_neuro_forward::GainMatrix;
/// Configuration for eLORETA inverse
#[derive(Debug, Clone)]
pub struct EloretaConfig {
/// Regularization parameter (lambda^2)
pub lambda2: f64,
/// Maximum iterations for weight computation
pub max_iter: usize,
/// Convergence tolerance
pub tol: f64,
/// Depth weighting exponent (0 = none, typically 0.5-0.9)
pub depth: f64,
/// Whether sources have free orientation
pub free_orientation: bool,
}
impl Default for EloretaConfig {
fn default() -> Self {
Self {
lambda2: 1.0 / 9.0, // SNR^2 = 9
max_iter: 100,
tol: 1e-6,
depth: 0.5,
free_orientation: true,
}
}
}
impl EloretaConfig {
/// Create config for fixed orientation sources
pub fn fixed() -> Self {
Self {
free_orientation: false,
..Self::default()
}
}
/// Set regularization parameter
pub fn with_lambda2(mut self, lambda2: f64) -> Self {
self.lambda2 = lambda2;
self
}
/// Set depth weighting
pub fn with_depth(mut self, depth: f64) -> Self {
self.depth = depth;
self
}
}
/// eLORETA inverse operator
#[derive(Debug, Clone)]
pub struct EloretaInverse {
/// Inverse kernel [n_source_columns x n_channels]
kernel: DMatrix<f64>,
/// Weight matrix for each source (for normalization)
weights: Vec<DMatrix<f64>>,
/// Number of sources
n_sources: usize,
/// Number of channels
n_channels: usize,
/// Whether sources have free orientation
free_orientation: bool,
/// Configuration used
config: EloretaConfig,
/// Source indices
source_indices: Vec<usize>,
}
impl EloretaInverse {
/// Create an eLORETA inverse operator
///
/// # Arguments
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
/// * `noise_cov` - Noise covariance matrix
/// * `config` - eLORETA configuration
pub fn make_inverse(
gain: &GainMatrix,
noise_cov: &Covariance,
config: EloretaConfig,
) -> InverseResult<Self> {
let n_channels = gain.n_sensors();
let n_source_cols = gain.n_source_columns();
let free_orientation = gain.is_free_orientation();
if noise_cov.n_channels() != n_channels {
return Err(InverseError::DimensionMismatch(format!(
"Noise covariance has {} channels, gain has {}",
noise_cov.n_channels(),
n_channels
)));
}
// Convert gain to nalgebra matrix
let g = Self::gain_to_matrix(gain);
// Compute inverse of regularized noise covariance
let c_inv = Self::compute_regularized_inv(noise_cov, config.lambda2)?;
// Compute depth weights
let depth_weights = Self::compute_depth_weights(&g, config.depth, free_orientation);
// Compute eLORETA weights iteratively
let (weights, kernel) = if free_orientation {
Self::compute_eloreta_weights_free(&g, &c_inv, &depth_weights, &config)?
} else {
Self::compute_eloreta_weights_fixed(&g, &c_inv, &depth_weights, &config)?
};
let n_sources = if free_orientation {
n_source_cols / 3
} else {
n_source_cols
};
Ok(Self {
kernel,
weights,
n_sources,
n_channels,
free_orientation,
config,
source_indices: (0..n_sources).collect(),
})
}
/// Apply the inverse operator to sensor data
///
/// # Arguments
/// * `data` - Sensor data [n_channels x n_times]
///
/// # Returns
/// Source estimates
pub fn apply(&self, data: &[Vec<f64>]) -> InverseResult<SourceEstimate> {
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();
// Convert to matrix
let data_mat = DMatrix::from_fn(self.n_channels, n_times, |i, j| data[i][j]);
// Apply inverse kernel
let source_mat = &self.kernel * &data_mat;
// Convert to output format
let n_rows = source_mat.nrows();
let source_data: Vec<Vec<f64>> = (0..n_rows)
.map(|i| (0..n_times).map(|j| source_mat[(i, j)]).collect())
.collect();
let times: Vec<f64> = (0..n_times).map(|i| i as f64).collect();
Ok(SourceEstimate::new(
source_data,
times,
self.source_indices.clone(),
self.free_orientation,
))
}
/// Get the inverse kernel matrix
pub fn kernel(&self) -> &DMatrix<f64> {
&self.kernel
}
/// Get number of sources
pub fn n_sources(&self) -> usize {
self.n_sources
}
/// Get the weight matrices
pub fn weights(&self) -> &[DMatrix<f64>] {
&self.weights
}
// ========== Private methods ==========
/// Convert GainMatrix to nalgebra DMatrix
fn gain_to_matrix(gain: &GainMatrix) -> DMatrix<f64> {
let n_rows = gain.n_sensors();
let n_cols = gain.n_source_columns();
let data = gain.data();
DMatrix::from_fn(n_rows, n_cols, |i, j| data[i][j])
}
/// Compute regularized inverse of noise covariance
fn compute_regularized_inv(
noise_cov: &Covariance,
lambda2: f64,
) -> InverseResult<DMatrix<f64>> {
let n = noise_cov.n_channels();
let c = noise_cov.data();
// Add regularization: C_reg = C + lambda^2 * trace(C)/n * I
let trace: f64 = (0..n).map(|i| c[(i, i)]).sum();
let reg_val = lambda2 * trace / n as f64;
let mut c_reg = c.clone();
for i in 0..n {
c_reg[(i, i)] += reg_val;
}
// Compute inverse via eigendecomposition for stability
let eigen = c_reg.symmetric_eigen();
let eigenvalues = eigen.eigenvalues;
let eigenvectors = eigen.eigenvectors;
let min_eig = eigenvalues.iter().copied().fold(f64::INFINITY, f64::min);
if min_eig <= 0.0 {
return Err(InverseError::ComputationError(format!(
"Regularized covariance has non-positive eigenvalue: {:.2e}",
min_eig
)));
}
let d_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| 1.0 / eigenvalues[i]));
Ok(&eigenvectors * &d_inv * eigenvectors.transpose())
}
/// Compute depth weights
fn compute_depth_weights(gain: &DMatrix<f64>, depth: f64, free_orientation: bool) -> Vec<f64> {
let n_cols = gain.ncols();
if depth == 0.0 {
return vec![1.0; n_cols];
}
if free_orientation {
let n_sources = n_cols / 3;
let mut weights = Vec::with_capacity(n_cols);
for src in 0..n_sources {
// Compute norm across all orientations
let mut sum_sq = 0.0;
for ori in 0..3 {
let col_idx = 3 * src + ori;
for row in 0..gain.nrows() {
sum_sq += gain[(row, col_idx)].powi(2);
}
}
let norm = sum_sq.sqrt();
let w = if norm > 1e-15 { norm.powf(-depth) } else { 1.0 };
// Same weight for all orientations
weights.push(w);
weights.push(w);
weights.push(w);
}
weights
} else {
(0..n_cols)
.map(|col| {
let sum_sq: f64 = (0..gain.nrows()).map(|row| gain[(row, col)].powi(2)).sum();
let norm = sum_sq.sqrt();
if norm > 1e-15 { norm.powf(-depth) } else { 1.0 }
})
.collect()
}
}
/// Compute eLORETA weights for fixed orientation sources
fn compute_eloreta_weights_fixed(
gain: &DMatrix<f64>,
c_inv: &DMatrix<f64>,
depth_weights: &[f64],
config: &EloretaConfig,
) -> InverseResult<(Vec<DMatrix<f64>>, DMatrix<f64>)> {
let n_channels = gain.nrows();
let n_sources = gain.ncols();
// Initialize weights to identity (scaled by depth)
let mut source_weights: Vec<f64> = depth_weights.to_vec();
// Iterative weight computation
for _iter in 0..config.max_iter {
let old_weights = source_weights.clone();
// Compute H = sum_i (w_i^2 * g_i * g_i^T)
let mut h = DMatrix::zeros(n_channels, n_channels);
for i in 0..n_sources {
let g_i = gain.column(i);
let w_sq = source_weights[i].powi(2);
for row in 0..n_channels {
for col in 0..n_channels {
h[(row, col)] += w_sq * g_i[row] * g_i[col];
}
}
}
// Regularize H
let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
let reg = config.lambda2 * trace / n_channels as f64;
for i in 0..n_channels {
h[(i, i)] += reg;
}
// Compute T = (H + λC)^(-1)
let h_reg = &h + c_inv * config.lambda2;
let t = match h_reg.clone().try_inverse() {
Some(inv) => inv,
None => {
return Err(InverseError::ComputationError(
"Failed to invert H matrix".to_string(),
));
}
};
// Update weights: w_i = (g_i^T * T * g_i)^(-1/4)
for i in 0..n_sources {
let g_i = gain.column(i);
let tg = &t * g_i;
let gtg = g_i.dot(&tg);
if gtg > 1e-30 {
source_weights[i] = depth_weights[i] * gtg.powf(-0.25);
}
}
// Check convergence
let change: f64 = source_weights
.iter()
.zip(old_weights.iter())
.map(|(a, b)| (a - b).abs())
.sum::<f64>()
/ n_sources as f64;
if change < config.tol {
break;
}
}
// Compute final kernel: K = W * G^T * T
// where W is diagonal weight matrix
let mut h = DMatrix::zeros(n_channels, n_channels);
for i in 0..n_sources {
let g_i = gain.column(i);
let w_sq = source_weights[i].powi(2);
for row in 0..n_channels {
for col in 0..n_channels {
h[(row, col)] += w_sq * g_i[row] * g_i[col];
}
}
}
let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
let reg = config.lambda2 * trace / n_channels as f64;
for i in 0..n_channels {
h[(i, i)] += reg;
}
let h_reg = &h + c_inv * config.lambda2;
let t = h_reg.try_inverse().ok_or_else(|| {
InverseError::ComputationError("Failed to compute final inverse".to_string())
})?;
// Kernel [n_sources x n_channels]
let mut kernel = DMatrix::zeros(n_sources, n_channels);
for i in 0..n_sources {
let g_i = gain.column(i);
let w_sq = source_weights[i].powi(2);
let tg = &t * g_i;
for ch in 0..n_channels {
kernel[(i, ch)] = w_sq * tg[ch];
}
}
// Convert weights to matrix format
let weight_matrices: Vec<DMatrix<f64>> = source_weights
.iter()
.map(|&w| DMatrix::from_element(1, 1, w))
.collect();
Ok((weight_matrices, kernel))
}
/// Compute eLORETA weights for free orientation sources
fn compute_eloreta_weights_free(
gain: &DMatrix<f64>,
c_inv: &DMatrix<f64>,
depth_weights: &[f64],
config: &EloretaConfig,
) -> InverseResult<(Vec<DMatrix<f64>>, DMatrix<f64>)> {
let n_channels = gain.nrows();
let n_sources = gain.ncols() / 3;
// Initialize weight matrices (3x3 for each source)
let mut source_weights: Vec<Matrix3<f64>> = (0..n_sources)
.map(|i| {
let d = depth_weights[3 * i];
Matrix3::identity() * d
})
.collect();
// Iterative weight computation
for _iter in 0..config.max_iter {
let old_weights: Vec<Matrix3<f64>> = source_weights.clone();
// Compute H = sum_i (G_i * W_i^2 * G_i^T)
let mut h = DMatrix::zeros(n_channels, n_channels);
for i in 0..n_sources {
let g_i = Self::extract_source_gain(gain, i);
let w_sq = source_weights[i] * source_weights[i];
// H += G_i * W_i^2 * G_i^T
for ch1 in 0..n_channels {
for ch2 in 0..n_channels {
let mut sum = 0.0;
for ori1 in 0..3 {
for ori2 in 0..3 {
sum += g_i[(ch1, ori1)] * w_sq[(ori1, ori2)] * g_i[(ch2, ori2)];
}
}
h[(ch1, ch2)] += sum;
}
}
}
// Regularize H
let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
let reg = config.lambda2 * trace / n_channels as f64;
for i in 0..n_channels {
h[(i, i)] += reg;
}
// Compute T = (H + λC)^(-1)
let h_reg = &h + c_inv * config.lambda2;
let t = match h_reg.clone().try_inverse() {
Some(inv) => inv,
None => {
return Err(InverseError::ComputationError(
"Failed to invert H matrix".to_string(),
));
}
};
// Update weights: W_i = D_i * (G_i^T * T * G_i)^(-1/4)
for i in 0..n_sources {
let g_i = Self::extract_source_gain(gain, i);
// Compute G_i^T * T * G_i (3x3 matrix)
let mut gtg = Matrix3::zeros();
for ori1 in 0..3 {
for ori2 in 0..3 {
let mut sum = 0.0;
for ch1 in 0..n_channels {
for ch2 in 0..n_channels {
sum += g_i[(ch1, ori1)] * t[(ch1, ch2)] * g_i[(ch2, ori2)];
}
}
gtg[(ori1, ori2)] = sum;
}
}
// Compute matrix power -1/4 via eigendecomposition
if let Some(new_w) = Self::matrix_power_quarter_inv(&gtg) {
let d = depth_weights[3 * i];
source_weights[i] = new_w * d;
}
}
// Check convergence
let change: f64 = source_weights
.iter()
.zip(old_weights.iter())
.map(|(a, b)| (a - b).norm())
.sum::<f64>()
/ n_sources as f64;
if change < config.tol {
break;
}
}
// Compute final kernel
let mut h = DMatrix::zeros(n_channels, n_channels);
for i in 0..n_sources {
let g_i = Self::extract_source_gain(gain, i);
let w_sq = source_weights[i] * source_weights[i];
for ch1 in 0..n_channels {
for ch2 in 0..n_channels {
let mut sum = 0.0;
for ori1 in 0..3 {
for ori2 in 0..3 {
sum += g_i[(ch1, ori1)] * w_sq[(ori1, ori2)] * g_i[(ch2, ori2)];
}
}
h[(ch1, ch2)] += sum;
}
}
}
let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
let reg = config.lambda2 * trace / n_channels as f64;
for i in 0..n_channels {
h[(i, i)] += reg;
}
let h_reg = &h + c_inv * config.lambda2;
let t = h_reg.try_inverse().ok_or_else(|| {
InverseError::ComputationError("Failed to compute final inverse".to_string())
})?;
// Kernel [3*n_sources x n_channels]
let mut kernel = DMatrix::zeros(3 * n_sources, n_channels);
for i in 0..n_sources {
let g_i = Self::extract_source_gain(gain, i);
let w_sq = source_weights[i] * source_weights[i];
for ori in 0..3 {
for ch in 0..n_channels {
let mut sum = 0.0;
for ori2 in 0..3 {
for ch2 in 0..n_channels {
sum += w_sq[(ori, ori2)] * g_i[(ch2, ori2)] * t[(ch2, ch)];
}
}
kernel[(3 * i + ori, ch)] = sum;
}
}
}
// Convert to output format
let weight_matrices: Vec<DMatrix<f64>> = source_weights
.iter()
.map(|w| {
let mut m = DMatrix::zeros(3, 3);
for i in 0..3 {
for j in 0..3 {
m[(i, j)] = w[(i, j)];
}
}
m
})
.collect();
Ok((weight_matrices, kernel))
}
/// Extract 3-column gain matrix for a single source
fn extract_source_gain(gain: &DMatrix<f64>, source_idx: usize) -> DMatrix<f64> {
let n_channels = gain.nrows();
let start_col = 3 * source_idx;
DMatrix::from_fn(n_channels, 3, |row, col| gain[(row, start_col + col)])
}
/// Compute A^(-1/4) for a 3x3 symmetric positive definite matrix
fn matrix_power_quarter_inv(a: &Matrix3<f64>) -> Option<Matrix3<f64>> {
// Convert to nalgebra for eigendecomposition
let eigen = a.symmetric_eigen();
// Check all eigenvalues are positive
for &ev in eigen.eigenvalues.iter() {
if ev <= 1e-30 {
return None;
}
}
// Compute D^(-1/4)
let d_inv_quarter = Matrix3::from_diagonal(&nalgebra::Vector3::new(
eigen.eigenvalues[0].powf(-0.25),
eigen.eigenvalues[1].powf(-0.25),
eigen.eigenvalues[2].powf(-0.25),
));
// A^(-1/4) = V * D^(-1/4) * V^T
Some(eigen.eigenvectors * d_inv_quarter * eigen.eigenvectors.transpose())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::covariance::CovarianceType;
fn create_simple_gain() -> GainMatrix {
// 4 sensors, 2 sources with free orientation (6 columns)
let data = vec![
vec![1.0, 0.2, 0.1, 0.3, 0.5, 0.2],
vec![0.2, 1.0, 0.3, 0.5, 0.3, 0.1],
vec![0.1, 0.3, 1.0, 0.2, 0.1, 0.5],
vec![0.3, 0.1, 0.2, 1.0, 0.2, 0.3],
];
let names = vec![
"S1".to_string(),
"S2".to_string(),
"S3".to_string(),
"S4".to_string(),
];
GainMatrix::new(data, true, names).unwrap()
}
fn create_fixed_gain() -> GainMatrix {
// 4 sensors, 3 fixed sources
let data = vec![
vec![1.0, 0.5, 0.2],
vec![0.5, 1.0, 0.5],
vec![0.2, 0.5, 1.0],
vec![0.1, 0.3, 0.6],
];
let names = vec![
"S1".to_string(),
"S2".to_string(),
"S3".to_string(),
"S4".to_string(),
];
GainMatrix::new(data, false, names).unwrap()
}
#[test]
fn test_eloreta_config() {
let config = EloretaConfig::default();
assert!(config.free_orientation);
assert!((config.lambda2 - 1.0 / 9.0).abs() < 1e-10);
let config = EloretaConfig::fixed().with_lambda2(0.1);
assert!(!config.free_orientation);
assert!((config.lambda2 - 0.1).abs() < 1e-10);
}
#[test]
fn test_make_eloreta_fixed() {
let gain = create_fixed_gain();
let noise_cov = Covariance::identity(4, CovarianceType::Noise);
let config = EloretaConfig::fixed();
let inv = EloretaInverse::make_inverse(&gain, &noise_cov, config).unwrap();
assert_eq!(inv.n_sources(), 3);
assert!(!inv.free_orientation);
}
#[test]
fn test_make_eloreta_free() {
let gain = create_simple_gain();
let noise_cov = Covariance::identity(4, CovarianceType::Noise);
let config = EloretaConfig::default();
let inv = EloretaInverse::make_inverse(&gain, &noise_cov, config).unwrap();
assert_eq!(inv.n_sources(), 2); // 6 columns / 3 orientations
assert!(inv.free_orientation);
}
#[test]
fn test_apply_eloreta() {
let gain = create_fixed_gain();
let noise_cov = Covariance::identity(4, CovarianceType::Noise);
let config = EloretaConfig::fixed();
let inv = EloretaInverse::make_inverse(&gain, &noise_cov, config).unwrap();
let data = vec![
vec![1.0, 0.0],
vec![0.0, 1.0],
vec![0.5, 0.5],
vec![0.2, 0.8],
];
let stc = inv.apply(&data).unwrap();
assert_eq!(stc.n_sources(), 3);
assert_eq!(stc.n_times(), 2);
}
#[test]
fn test_eloreta_kernel_shape() {
let gain = create_simple_gain();
let noise_cov = Covariance::identity(4, CovarianceType::Noise);
let inv =
EloretaInverse::make_inverse(&gain, &noise_cov, EloretaConfig::default()).unwrap();
let kernel = inv.kernel();
assert_eq!(kernel.nrows(), 6); // 3 orientations * 2 sources
assert_eq!(kernel.ncols(), 4); // 4 channels
}
}