Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,411 @@
|
||||
//! Minimum Norm Estimate (MNE) and related inverse methods.
|
||||
//!
|
||||
//! This module implements:
|
||||
//! - **MNE**: Basic minimum-norm estimate
|
||||
//! - **dSPM**: Dynamic Statistical Parametric Mapping (noise-normalized)
|
||||
//! - **sLORETA**: Standardized Low Resolution Electromagnetic Tomography
|
||||
//!
|
||||
//! ## Mathematical Background
|
||||
//!
|
||||
//! The forward model is: M = G * J + n
|
||||
//! where M is measurements, G is gain matrix, J is source activity, n is noise.
|
||||
//!
|
||||
//! The MNE solution is: J = W * M
|
||||
//! where W = R * G^T * (G * R * G^T + λ * C)^(-1)
|
||||
//!
|
||||
//! - R is source covariance (with depth weighting)
|
||||
//! - C is noise covariance
|
||||
//! - λ is regularization parameter
|
||||
|
||||
use crate::{Covariance, InverseError, InverseResult, SourceEstimate};
|
||||
use nalgebra::{DMatrix, DVector};
|
||||
use rtx_neuro_forward::GainMatrix;
|
||||
|
||||
/// Inverse solution method
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum InverseMethod {
|
||||
/// Basic minimum norm estimate
|
||||
Mne,
|
||||
/// Dynamic SPM (noise-normalized)
|
||||
Dspm,
|
||||
/// Standardized LORETA (resolution-normalized)
|
||||
Sloreta,
|
||||
}
|
||||
|
||||
impl Default for InverseMethod {
|
||||
fn default() -> Self {
|
||||
Self::Dspm
|
||||
}
|
||||
}
|
||||
|
||||
/// MNE inverse operator
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MneInverse {
|
||||
/// Inverse kernel [n_sources x n_channels]
|
||||
kernel: DMatrix<f64>,
|
||||
/// Method used
|
||||
method: InverseMethod,
|
||||
/// Number of sources
|
||||
n_sources: usize,
|
||||
/// Number of channels
|
||||
n_channels: usize,
|
||||
/// Whether sources have free orientation
|
||||
free_orientation: bool,
|
||||
/// Regularization parameter (lambda^2)
|
||||
lambda2: f64,
|
||||
/// Depth weighting exponent
|
||||
depth: f64,
|
||||
/// Loose orientation constraint
|
||||
loose: f64,
|
||||
/// Noise normalization vector (for dSPM/sLORETA)
|
||||
noise_norm: Option<DVector<f64>>,
|
||||
/// Source names/indices
|
||||
source_indices: Vec<usize>,
|
||||
}
|
||||
|
||||
impl MneInverse {
|
||||
/// Create an MNE inverse operator from forward model and noise covariance
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
|
||||
/// * `noise_cov` - Noise covariance matrix
|
||||
/// * `method` - Inverse method (MNE, dSPM, sLORETA)
|
||||
/// * `loose` - Loose orientation constraint (0 = fixed, 1 = free)
|
||||
/// * `depth` - Depth weighting exponent (0 = none, typically 0.8)
|
||||
/// * `lambda2` - Regularization parameter (typically 1/SNR^2)
|
||||
pub fn make_inverse(
|
||||
gain: &GainMatrix,
|
||||
noise_cov: &Covariance,
|
||||
method: InverseMethod,
|
||||
loose: f64,
|
||||
depth: f64,
|
||||
lambda2: f64,
|
||||
) -> 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 whitener from noise covariance
|
||||
let whitener = noise_cov.compute_whitener()?;
|
||||
|
||||
// Whiten the gain matrix: G_w = W * G
|
||||
let g_white = &whitener * &g;
|
||||
|
||||
// Compute source covariance (with depth weighting)
|
||||
let source_cov = Self::compute_source_covariance(gain, depth, loose)?;
|
||||
|
||||
// Compute MNE kernel: W_mne = R * G^T * (G * R * G^T + λ * I)^(-1)
|
||||
// In whitened space: W_mne = R * G_w^T * (G_w * R * G_w^T + λ * I)^(-1) * W
|
||||
let grgt = &g_white * &source_cov * g_white.transpose();
|
||||
let grgt_reg = grgt + DMatrix::identity(n_channels, n_channels) * lambda2;
|
||||
|
||||
// Invert using SVD for stability
|
||||
let grgt_inv = Self::pseudo_inverse(&grgt_reg)?;
|
||||
|
||||
// MNE kernel (before normalization)
|
||||
let kernel_raw = &source_cov * g_white.transpose() * &grgt_inv * &whitener;
|
||||
|
||||
// Compute normalization based on method
|
||||
let (kernel, noise_norm) = match method {
|
||||
InverseMethod::Mne => (kernel_raw, None),
|
||||
InverseMethod::Dspm => {
|
||||
// dSPM: normalize by noise standard deviation
|
||||
let noise_norm = Self::compute_dspm_normalization(&kernel_raw, noise_cov)?;
|
||||
let kernel = Self::apply_row_normalization(&kernel_raw, &noise_norm);
|
||||
(kernel, Some(noise_norm))
|
||||
}
|
||||
InverseMethod::Sloreta => {
|
||||
// sLORETA: normalize by resolution matrix diagonal
|
||||
let noise_norm = Self::compute_sloreta_normalization(&kernel_raw, &g)?;
|
||||
let kernel = Self::apply_row_normalization(&kernel_raw, &noise_norm);
|
||||
(kernel, Some(noise_norm))
|
||||
}
|
||||
};
|
||||
|
||||
let n_sources = if free_orientation {
|
||||
n_source_cols / 3
|
||||
} else {
|
||||
n_source_cols
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
kernel,
|
||||
method,
|
||||
n_sources,
|
||||
n_channels,
|
||||
free_orientation,
|
||||
lambda2,
|
||||
depth,
|
||||
loose,
|
||||
noise_norm,
|
||||
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 [n_sources x n_times] (or [3*n_sources x n_times] for free)
|
||||
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: S = K * M
|
||||
let source_mat = &self.kernel * &data_mat;
|
||||
|
||||
// Convert back to Vec<Vec<f64>>
|
||||
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();
|
||||
|
||||
// Compute times (placeholder - should come from data)
|
||||
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 the inverse method
|
||||
pub fn method(&self) -> InverseMethod {
|
||||
self.method
|
||||
}
|
||||
|
||||
/// Get number of sources
|
||||
pub fn n_sources(&self) -> usize {
|
||||
self.n_sources
|
||||
}
|
||||
|
||||
/// Get regularization parameter
|
||||
pub fn lambda2(&self) -> f64 {
|
||||
self.lambda2
|
||||
}
|
||||
|
||||
// ========== Private helper 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 source covariance matrix with depth weighting
|
||||
fn compute_source_covariance(
|
||||
gain: &GainMatrix,
|
||||
depth: f64,
|
||||
_loose: f64,
|
||||
) -> InverseResult<DMatrix<f64>> {
|
||||
let n = gain.n_source_columns();
|
||||
|
||||
if depth == 0.0 {
|
||||
// No depth weighting: identity
|
||||
return Ok(DMatrix::identity(n, n));
|
||||
}
|
||||
|
||||
// Compute column norms of gain matrix for depth weighting
|
||||
let norms = gain.source_norms();
|
||||
|
||||
// Depth weighting: w_i = ||g_i||^(-depth)
|
||||
// Source covariance: R = diag(w_1^2, w_2^2, ...)
|
||||
let weights: Vec<f64> = if gain.is_free_orientation() {
|
||||
// For free orientation, apply same weight to all 3 orientations
|
||||
let n_sources = n / 3;
|
||||
let mut w = Vec::with_capacity(n);
|
||||
for i in 0..n_sources {
|
||||
let weight = if norms[i] > 1e-15 {
|
||||
norms[i].powf(-depth)
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
// Same weight for x, y, z
|
||||
w.push(weight * weight);
|
||||
w.push(weight * weight);
|
||||
w.push(weight * weight);
|
||||
}
|
||||
w
|
||||
} else {
|
||||
norms
|
||||
.iter()
|
||||
.map(|&norm| {
|
||||
let w = if norm > 1e-15 { norm.powf(-depth) } else { 1.0 };
|
||||
w * w
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
|
||||
Ok(DMatrix::from_diagonal(&DVector::from_vec(weights)))
|
||||
}
|
||||
|
||||
/// Compute pseudo-inverse using SVD
|
||||
fn pseudo_inverse(matrix: &DMatrix<f64>) -> InverseResult<DMatrix<f64>> {
|
||||
let svd = matrix.clone().svd(true, true);
|
||||
|
||||
let u = svd
|
||||
.u
|
||||
.ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
|
||||
let vt = svd
|
||||
.v_t
|
||||
.ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
|
||||
let s = svd.singular_values;
|
||||
|
||||
let tol = 1e-10 * s[0];
|
||||
let n = s.len();
|
||||
|
||||
let s_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| {
|
||||
if s[i] > tol { 1.0 / s[i] } else { 0.0 }
|
||||
}));
|
||||
|
||||
Ok(vt.transpose() * &s_inv * u.transpose())
|
||||
}
|
||||
|
||||
/// Compute dSPM normalization (noise standard deviation)
|
||||
fn compute_dspm_normalization(
|
||||
kernel: &DMatrix<f64>,
|
||||
noise_cov: &Covariance,
|
||||
) -> InverseResult<DVector<f64>> {
|
||||
let n_sources = kernel.nrows();
|
||||
let noise_cov_mat = noise_cov.data();
|
||||
|
||||
// Noise variance: diag(K * C * K^T)
|
||||
let kckt = kernel * noise_cov_mat * kernel.transpose();
|
||||
|
||||
// Extract diagonal and take sqrt
|
||||
let noise_norm = DVector::from_fn(n_sources, |i, _| {
|
||||
let var = kckt[(i, i)];
|
||||
if var > 1e-30 { var.sqrt() } else { 1.0 }
|
||||
});
|
||||
|
||||
Ok(noise_norm)
|
||||
}
|
||||
|
||||
/// Compute sLORETA normalization (resolution matrix diagonal)
|
||||
fn compute_sloreta_normalization(
|
||||
kernel: &DMatrix<f64>,
|
||||
gain: &DMatrix<f64>,
|
||||
) -> InverseResult<DVector<f64>> {
|
||||
let n_sources = kernel.nrows();
|
||||
|
||||
// Resolution matrix: R = K * G
|
||||
let resolution = kernel * gain;
|
||||
|
||||
// Take diagonal and sqrt
|
||||
let noise_norm = DVector::from_fn(n_sources, |i, _| {
|
||||
let val = resolution[(i, i)];
|
||||
if val > 1e-30 { val.sqrt() } else { 1.0 }
|
||||
});
|
||||
|
||||
Ok(noise_norm)
|
||||
}
|
||||
|
||||
/// Apply row-wise normalization to kernel
|
||||
fn apply_row_normalization(kernel: &DMatrix<f64>, norm: &DVector<f64>) -> DMatrix<f64> {
|
||||
let n_rows = kernel.nrows();
|
||||
let n_cols = kernel.ncols();
|
||||
|
||||
DMatrix::from_fn(n_rows, n_cols, |i, j| kernel[(i, j)] / norm[i])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::covariance::CovarianceType;
|
||||
|
||||
fn create_simple_gain() -> GainMatrix {
|
||||
// Simple 4x3 gain matrix (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_make_inverse_mne() {
|
||||
let gain = create_simple_gain();
|
||||
let noise_cov = Covariance::identity(4, CovarianceType::Noise);
|
||||
|
||||
let inv =
|
||||
MneInverse::make_inverse(&gain, &noise_cov, InverseMethod::Mne, 0.0, 0.0, 0.1).unwrap();
|
||||
|
||||
assert_eq!(inv.n_sources(), 3);
|
||||
assert_eq!(inv.method(), InverseMethod::Mne);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_make_inverse_dspm() {
|
||||
let gain = create_simple_gain();
|
||||
let noise_cov = Covariance::identity(4, CovarianceType::Noise);
|
||||
|
||||
let inv = MneInverse::make_inverse(&gain, &noise_cov, InverseMethod::Dspm, 0.0, 0.8, 0.1)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(inv.method(), InverseMethod::Dspm);
|
||||
assert!(inv.noise_norm.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_apply_inverse() {
|
||||
let gain = create_simple_gain();
|
||||
let noise_cov = Covariance::identity(4, CovarianceType::Noise);
|
||||
|
||||
let inv =
|
||||
MneInverse::make_inverse(&gain, &noise_cov, InverseMethod::Mne, 0.0, 0.0, 0.1).unwrap();
|
||||
|
||||
// Apply to simple data
|
||||
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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user