339 lines
9.9 KiB
Rust
339 lines
9.9 KiB
Rust
//! Noise and data covariance estimation.
|
|
//!
|
|
//! Covariance matrices are essential for inverse solutions:
|
|
//! - **Noise covariance**: Estimated from baseline or empty room recordings
|
|
//! - **Data covariance**: Used for beamforming approaches
|
|
|
|
use crate::{InverseError, InverseResult};
|
|
use nalgebra::{DMatrix, DVector};
|
|
|
|
/// Type of covariance matrix
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum CovarianceType {
|
|
/// Noise covariance (from baseline)
|
|
Noise,
|
|
/// Data covariance (from epochs)
|
|
Data,
|
|
/// Empty room noise
|
|
EmptyRoom,
|
|
}
|
|
|
|
/// Covariance matrix with metadata
|
|
#[derive(Debug, Clone)]
|
|
pub struct Covariance {
|
|
/// Covariance matrix [n_channels x n_channels]
|
|
data: DMatrix<f64>,
|
|
/// Channel names
|
|
channel_names: Vec<String>,
|
|
/// Type of covariance
|
|
cov_type: CovarianceType,
|
|
/// Number of samples used
|
|
n_samples: usize,
|
|
/// Regularization applied
|
|
reg: f64,
|
|
}
|
|
|
|
impl Covariance {
|
|
/// Create a covariance from a matrix
|
|
pub fn new(
|
|
data: Vec<Vec<f64>>,
|
|
channel_names: Vec<String>,
|
|
cov_type: CovarianceType,
|
|
n_samples: usize,
|
|
) -> InverseResult<Self> {
|
|
let n = data.len();
|
|
if n == 0 {
|
|
return Err(InverseError::InvalidParameter(
|
|
"Covariance matrix cannot be empty".to_string(),
|
|
));
|
|
}
|
|
|
|
if data[0].len() != n {
|
|
return Err(InverseError::DimensionMismatch(
|
|
"Covariance matrix must be square".to_string(),
|
|
));
|
|
}
|
|
|
|
if channel_names.len() != n {
|
|
return Err(InverseError::DimensionMismatch(format!(
|
|
"Expected {} channel names, got {}",
|
|
n,
|
|
channel_names.len()
|
|
)));
|
|
}
|
|
|
|
// Convert to nalgebra matrix
|
|
let mat = DMatrix::from_fn(n, n, |i, j| data[i][j]);
|
|
|
|
Ok(Self {
|
|
data: mat,
|
|
channel_names,
|
|
cov_type,
|
|
n_samples,
|
|
reg: 0.0,
|
|
})
|
|
}
|
|
|
|
/// Create an identity covariance (for testing)
|
|
pub fn identity(n_channels: usize, cov_type: CovarianceType) -> Self {
|
|
let names: Vec<String> = (0..n_channels).map(|i| format!("CH{:03}", i + 1)).collect();
|
|
|
|
Self {
|
|
data: DMatrix::identity(n_channels, n_channels),
|
|
channel_names: names,
|
|
cov_type,
|
|
n_samples: 0,
|
|
reg: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Compute covariance from data epochs
|
|
///
|
|
/// # Arguments
|
|
/// * `epochs` - Data epochs [n_epochs][n_channels][n_times]
|
|
/// * `channel_names` - Channel names
|
|
/// * `cov_type` - Type of covariance
|
|
/// * `tmin` - Start time index for baseline
|
|
/// * `tmax` - End time index for baseline
|
|
pub fn from_epochs(
|
|
epochs: &[Vec<Vec<f64>>],
|
|
channel_names: Vec<String>,
|
|
cov_type: CovarianceType,
|
|
tmin: Option<usize>,
|
|
tmax: Option<usize>,
|
|
) -> InverseResult<Self> {
|
|
if epochs.is_empty() {
|
|
return Err(InverseError::InvalidParameter(
|
|
"No epochs provided".to_string(),
|
|
));
|
|
}
|
|
|
|
let n_channels = epochs[0].len();
|
|
let n_times_total = epochs[0][0].len();
|
|
|
|
// Determine time range
|
|
let t_start = tmin.unwrap_or(0);
|
|
let t_end = tmax.unwrap_or(n_times_total);
|
|
|
|
// Collect all samples across epochs and time
|
|
let mut all_samples: Vec<Vec<f64>> = Vec::new();
|
|
|
|
for epoch in epochs {
|
|
for t in t_start..t_end {
|
|
let sample: Vec<f64> = epoch.iter().map(|ch| ch[t]).collect();
|
|
all_samples.push(sample);
|
|
}
|
|
}
|
|
|
|
let n_samples = all_samples.len();
|
|
|
|
// Compute mean
|
|
let mut mean = vec![0.0; n_channels];
|
|
for sample in &all_samples {
|
|
for (i, &val) in sample.iter().enumerate() {
|
|
mean[i] += val;
|
|
}
|
|
}
|
|
for m in &mut mean {
|
|
*m /= n_samples as f64;
|
|
}
|
|
|
|
// Compute covariance
|
|
let mut cov = vec![vec![0.0; n_channels]; n_channels];
|
|
for sample in &all_samples {
|
|
for i in 0..n_channels {
|
|
for j in i..n_channels {
|
|
cov[i][j] += (sample[i] - mean[i]) * (sample[j] - mean[j]);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Normalize and symmetrize
|
|
for i in 0..n_channels {
|
|
for j in i..n_channels {
|
|
cov[i][j] /= (n_samples - 1) as f64;
|
|
cov[j][i] = cov[i][j];
|
|
}
|
|
}
|
|
|
|
Self::new(cov, channel_names, cov_type, n_samples)
|
|
}
|
|
|
|
/// Get number of channels
|
|
pub fn n_channels(&self) -> usize {
|
|
self.data.nrows()
|
|
}
|
|
|
|
/// Get channel names
|
|
pub fn channel_names(&self) -> &[String] {
|
|
&self.channel_names
|
|
}
|
|
|
|
/// Get covariance type
|
|
pub fn cov_type(&self) -> CovarianceType {
|
|
self.cov_type
|
|
}
|
|
|
|
/// Get number of samples used
|
|
pub fn n_samples(&self) -> usize {
|
|
self.n_samples
|
|
}
|
|
|
|
/// Get the covariance matrix
|
|
pub fn data(&self) -> &DMatrix<f64> {
|
|
&self.data
|
|
}
|
|
|
|
/// Get a copy of the data as nested vectors
|
|
pub fn to_vec(&self) -> Vec<Vec<f64>> {
|
|
let n = self.n_channels();
|
|
(0..n)
|
|
.map(|i| (0..n).map(|j| self.data[(i, j)]).collect())
|
|
.collect()
|
|
}
|
|
|
|
/// Apply Tikhonov regularization
|
|
///
|
|
/// C_reg = C + reg * trace(C)/n * I
|
|
pub fn regularize(&mut self, reg: f64) -> InverseResult<()> {
|
|
if !(0.0..=1.0).contains(®) {
|
|
return Err(InverseError::InvalidParameter(
|
|
"Regularization must be between 0 and 1".to_string(),
|
|
));
|
|
}
|
|
|
|
let n = self.n_channels();
|
|
let trace: f64 = (0..n).map(|i| self.data[(i, i)]).sum();
|
|
let reg_value = reg * trace / n as f64;
|
|
|
|
for i in 0..n {
|
|
self.data[(i, i)] += reg_value;
|
|
}
|
|
|
|
self.reg = reg;
|
|
Ok(())
|
|
}
|
|
|
|
/// Compute the whitening matrix W = C^(-1/2)
|
|
///
|
|
/// Such that W * C * W^T = I
|
|
pub fn compute_whitener(&self) -> InverseResult<DMatrix<f64>> {
|
|
let n = self.n_channels();
|
|
|
|
// Eigendecomposition: C = V * D * V^T
|
|
let eigen = self.data.clone().symmetric_eigen();
|
|
let eigenvalues = eigen.eigenvalues;
|
|
let eigenvectors = eigen.eigenvectors;
|
|
|
|
// Check for negative eigenvalues
|
|
let min_eig = eigenvalues.min();
|
|
if min_eig <= 0.0 {
|
|
return Err(InverseError::ComputationError(format!(
|
|
"Covariance has non-positive eigenvalue: {:.2e}. Consider regularization.",
|
|
min_eig
|
|
)));
|
|
}
|
|
|
|
// W = V * D^(-1/2) * V^T
|
|
let d_inv_sqrt =
|
|
DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| 1.0 / eigenvalues[i].sqrt()));
|
|
|
|
let whitener = &eigenvectors * &d_inv_sqrt * eigenvectors.transpose();
|
|
Ok(whitener)
|
|
}
|
|
|
|
/// Compute the pseudo-inverse using SVD with rank control
|
|
pub fn compute_pseudo_inverse(&self, rank: Option<usize>) -> InverseResult<DMatrix<f64>> {
|
|
let n = self.n_channels();
|
|
|
|
let svd = self.data.clone().svd(true, true);
|
|
let u = svd
|
|
.u
|
|
.ok_or_else(|| InverseError::ComputationError("SVD failed to compute U".to_string()))?;
|
|
let vt = svd.v_t.ok_or_else(|| {
|
|
InverseError::ComputationError("SVD failed to compute V^T".to_string())
|
|
})?;
|
|
let s = svd.singular_values;
|
|
|
|
// Determine rank
|
|
let effective_rank = rank.unwrap_or_else(|| {
|
|
let tol = 1e-10 * s[0];
|
|
s.iter().filter(|&&v| v > tol).count()
|
|
});
|
|
|
|
// Compute pseudo-inverse: V * S^(-1) * U^T (only for non-zero singular values)
|
|
let s_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| {
|
|
if i < effective_rank && s[i] > 1e-15 {
|
|
1.0 / s[i]
|
|
} else {
|
|
0.0
|
|
}
|
|
}));
|
|
|
|
let pinv = vt.transpose() * &s_inv * u.transpose();
|
|
Ok(pinv)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_identity_covariance() {
|
|
let cov = Covariance::identity(10, CovarianceType::Noise);
|
|
assert_eq!(cov.n_channels(), 10);
|
|
assert_eq!(cov.cov_type(), CovarianceType::Noise);
|
|
}
|
|
|
|
#[test]
|
|
fn test_from_epochs() {
|
|
// Simple test data: 2 epochs, 3 channels, 5 time points
|
|
let epochs = vec![
|
|
vec![
|
|
vec![1.0, 2.0, 3.0, 4.0, 5.0],
|
|
vec![2.0, 3.0, 4.0, 5.0, 6.0],
|
|
vec![3.0, 4.0, 5.0, 6.0, 7.0],
|
|
],
|
|
vec![
|
|
vec![1.5, 2.5, 3.5, 4.5, 5.5],
|
|
vec![2.5, 3.5, 4.5, 5.5, 6.5],
|
|
vec![3.5, 4.5, 5.5, 6.5, 7.5],
|
|
],
|
|
];
|
|
let names = vec!["CH1".to_string(), "CH2".to_string(), "CH3".to_string()];
|
|
|
|
let cov =
|
|
Covariance::from_epochs(&epochs, names, CovarianceType::Data, None, None).unwrap();
|
|
|
|
assert_eq!(cov.n_channels(), 3);
|
|
assert_eq!(cov.n_samples(), 10); // 2 epochs * 5 time points
|
|
}
|
|
|
|
#[test]
|
|
fn test_regularization() {
|
|
let mut cov = Covariance::identity(5, CovarianceType::Noise);
|
|
cov.regularize(0.1).unwrap();
|
|
|
|
// Diagonal should be > 1 after regularization
|
|
let data = cov.data();
|
|
assert!(data[(0, 0)] > 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_whitener() {
|
|
// 2x2 covariance matrix
|
|
let data = vec![vec![4.0, 2.0], vec![2.0, 3.0]];
|
|
let names = vec!["A".to_string(), "B".to_string()];
|
|
let cov = Covariance::new(data, names, CovarianceType::Noise, 100).unwrap();
|
|
|
|
let whitener = cov.compute_whitener().unwrap();
|
|
|
|
// W * C * W^T should be close to identity
|
|
let result = &whitener * cov.data() * whitener.transpose();
|
|
assert!((result[(0, 0)] - 1.0).abs() < 1e-10);
|
|
assert!((result[(1, 1)] - 1.0).abs() < 1e-10);
|
|
assert!((result[(0, 1)]).abs() < 1e-10);
|
|
}
|
|
}
|