Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,338 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user