Files
rustytorch/crates/specialized/rtx-neuro-signal/src/ssp.rs
T
osobhandClaude Opus 4.6 02d382d5f6 style: apply rustfmt across all crates and demos
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]>
2026-04-12 07:01:58 -07:00

695 lines
21 KiB
Rust

//! Signal Space Projection (SSP) for artifact removal.
//!
//! SSP is a spatial filtering technique that removes artifacts by projecting
//! out directions in sensor space that contain artifact activity.
//!
//! ## Mathematical Background
//!
//! Given artifact data matrix A [n_channels x n_samples], we compute the
//! covariance C = A * A^T and find its principal components via SVD.
//!
//! The projection operator is: P = I - Σ(u_i * u_i^T)
//! where u_i are the artifact subspace vectors to remove.
//!
//! Clean data is obtained by: M_clean = P * M
use crate::{SignalError, SignalResult};
/// SSP projector for artifact removal
#[derive(Debug, Clone)]
pub struct SspProjector {
/// Projection vectors [n_projectors x n_channels]
/// Each row is a unit vector defining an artifact direction
vectors: Vec<Vec<f64>>,
/// Whether each projector is active
active: Vec<bool>,
/// Explained variance by each projector
explained_var: Vec<f64>,
/// Names/descriptions of projectors
names: Vec<String>,
/// Number of channels
n_channels: usize,
}
impl SspProjector {
/// Create an SSP projector from artifact epochs
///
/// # Arguments
/// * `epochs` - Artifact epochs [n_epochs][n_channels][n_samples]
/// * `n_components` - Number of projectors to compute
/// * `name` - Name for this projector set (e.g., "EOG", "ECG")
///
/// # Returns
/// SSP projector with the requested number of components
pub fn from_epochs(
epochs: &[Vec<Vec<f64>>],
n_components: usize,
name: &str,
) -> SignalResult<Self> {
if epochs.is_empty() {
return Err(SignalError::InvalidLength("No epochs provided".to_string()));
}
let n_channels = epochs[0].len();
if n_channels == 0 {
return Err(SignalError::InvalidLength("Empty channels".to_string()));
}
// Concatenate all epochs into a data matrix [n_channels x total_samples]
let total_samples: usize = epochs.iter().map(|e| e[0].len()).sum();
let mut data = vec![vec![0.0; total_samples]; n_channels];
let mut offset = 0;
for epoch in epochs {
let n_times = epoch[0].len();
for ch in 0..n_channels {
for t in 0..n_times {
data[ch][offset + t] = epoch[ch][t];
}
}
offset += n_times;
}
// Compute covariance matrix: C = (1/n) * X * X^T
let mut cov = vec![vec![0.0; n_channels]; n_channels];
for i in 0..n_channels {
for j in i..n_channels {
let mut sum = 0.0;
for t in 0..total_samples {
sum += data[i][t] * data[j][t];
}
cov[i][j] = sum / total_samples as f64;
cov[j][i] = cov[i][j];
}
}
// Compute SVD of covariance matrix
let (eigenvectors, eigenvalues) = symmetric_eigen(&cov)?;
// Total variance
let total_var: f64 = eigenvalues.iter().sum();
// Extract top n_components
let n_proj = n_components.min(n_channels);
let mut vectors = Vec::with_capacity(n_proj);
let mut explained_var = Vec::with_capacity(n_proj);
let mut names = Vec::with_capacity(n_proj);
let mut active = Vec::with_capacity(n_proj);
for i in 0..n_proj {
vectors.push(eigenvectors[i].clone());
explained_var.push(eigenvalues[i] / total_var);
names.push(format!("{}-{}", name, i + 1));
active.push(true);
}
Ok(Self {
vectors,
active,
explained_var,
names,
n_channels,
})
}
/// Create SSP projector from EOG data
///
/// # Arguments
/// * `data` - Continuous data [n_channels][n_samples]
/// * `eog_channels` - Indices of EOG channels
/// * `n_components` - Number of projectors (typically 1-2)
pub fn from_eog(
data: &[Vec<f64>],
eog_channels: &[usize],
n_components: usize,
) -> SignalResult<Self> {
if data.is_empty() {
return Err(SignalError::InvalidLength("Empty data".to_string()));
}
let n_channels = data.len();
let n_samples = data[0].len();
// Create artifact epochs from EOG threshold crossings
let mut artifact_epochs = Vec::new();
for &eog_ch in eog_channels {
if eog_ch >= n_channels {
return Err(SignalError::InvalidParameters(format!(
"EOG channel {eog_ch} out of range"
)));
}
// Find threshold (use robust estimate: 5 * median absolute deviation)
let eog_data = &data[eog_ch];
let threshold = compute_robust_threshold(eog_data);
// Find peaks above threshold
let peaks = find_peaks(eog_data, threshold);
// Extract epochs around each peak (±100 samples or available)
let half_window = 100.min(n_samples / 2);
for peak in peaks {
let start = peak.saturating_sub(half_window);
let end = (peak + half_window).min(n_samples);
let epoch: Vec<Vec<f64>> = data.iter().map(|ch| ch[start..end].to_vec()).collect();
artifact_epochs.push(epoch);
}
}
if artifact_epochs.is_empty() {
// Create a projector from the raw EOG channel correlation
return Self::from_eog_correlation(data, eog_channels, n_components);
}
Self::from_epochs(&artifact_epochs, n_components, "EOG")
}
/// Create SSP from EOG channel correlation (fallback method)
fn from_eog_correlation(
data: &[Vec<f64>],
eog_channels: &[usize],
n_components: usize,
) -> SignalResult<Self> {
let n_channels = data.len();
let n_samples = data[0].len();
// Compute correlation of each channel with EOG channels
let mut vectors = Vec::new();
let mut explained_var = Vec::new();
for (comp_idx, &eog_ch) in eog_channels.iter().take(n_components).enumerate() {
let eog_data = &data[eog_ch];
let eog_var: f64 = eog_data.iter().map(|&x| x * x).sum::<f64>() / n_samples as f64;
if eog_var < 1e-15 {
continue;
}
// Compute correlation vector
let mut corr_vec = vec![0.0; n_channels];
for ch in 0..n_channels {
let cov: f64 = (0..n_samples)
.map(|t| data[ch][t] * eog_data[t])
.sum::<f64>()
/ n_samples as f64;
corr_vec[ch] = cov / eog_var.sqrt();
}
// Normalize
let norm: f64 = corr_vec.iter().map(|&x| x * x).sum::<f64>().sqrt();
if norm > 1e-10 {
for v in &mut corr_vec {
*v /= norm;
}
vectors.push(corr_vec);
explained_var.push(1.0 / (comp_idx + 1) as f64);
}
}
let n_proj = vectors.len();
Ok(Self {
vectors,
active: vec![true; n_proj],
explained_var,
names: (0..n_proj).map(|i| format!("EOG-{}", i + 1)).collect(),
n_channels,
})
}
/// Create SSP projector from ECG data
///
/// # Arguments
/// * `data` - Continuous data [n_channels][n_samples]
/// * `ecg_channel` - Index of ECG channel
/// * `n_components` - Number of projectors (typically 1-3)
pub fn from_ecg(
data: &[Vec<f64>],
ecg_channel: usize,
n_components: usize,
) -> SignalResult<Self> {
if data.is_empty() {
return Err(SignalError::InvalidLength("Empty data".to_string()));
}
let n_channels = data.len();
let n_samples = data[0].len();
if ecg_channel >= n_channels {
return Err(SignalError::InvalidParameters(format!(
"ECG channel {ecg_channel} out of range"
)));
}
let ecg_data = &data[ecg_channel];
// Find R-peaks using simple threshold detection
let threshold = compute_robust_threshold(ecg_data);
let r_peaks = find_peaks(ecg_data, threshold);
if r_peaks.len() < 3 {
return Err(SignalError::InvalidLength(
"Not enough ECG peaks detected".to_string(),
));
}
// Extract epochs around each R-peak (±50 samples)
let half_window = 50.min(n_samples / 2);
let mut artifact_epochs = Vec::new();
for peak in r_peaks {
let start = peak.saturating_sub(half_window);
let end = (peak + half_window).min(n_samples);
let epoch: Vec<Vec<f64>> = data.iter().map(|ch| ch[start..end].to_vec()).collect();
artifact_epochs.push(epoch);
}
Self::from_epochs(&artifact_epochs, n_components, "ECG")
}
/// Apply the SSP projector to data
///
/// # Arguments
/// * `data` - Data to clean [n_channels][n_samples]
///
/// # Returns
/// Cleaned data with artifacts removed
pub fn apply(&self, data: &[Vec<f64>]) -> SignalResult<Vec<Vec<f64>>> {
if data.len() != self.n_channels {
return Err(SignalError::InvalidLength(format!(
"Expected {} channels, got {}",
self.n_channels,
data.len()
)));
}
let n_samples = data[0].len();
// Compute projection matrix: P = I - Σ(u_i * u_i^T)
// For efficiency, we apply it directly: y = x - Σ((u_i^T * x) * u_i)
let mut result = data.to_vec();
for (i, vec) in self.vectors.iter().enumerate() {
if !self.active[i] {
continue;
}
for t in 0..n_samples {
// Compute dot product: u^T * x
let dot: f64 = (0..self.n_channels).map(|ch| vec[ch] * data[ch][t]).sum();
// Subtract: x = x - (u^T * x) * u
for ch in 0..self.n_channels {
result[ch][t] -= dot * vec[ch];
}
}
}
Ok(result)
}
/// Apply SSP to epoched data
///
/// # Arguments
/// * `epochs` - Epoched data [n_epochs][n_channels][n_samples]
///
/// # Returns
/// Cleaned epochs
pub fn apply_epochs(&self, epochs: &[Vec<Vec<f64>>]) -> SignalResult<Vec<Vec<Vec<f64>>>> {
epochs.iter().map(|epoch| self.apply(epoch)).collect()
}
/// Activate or deactivate a projector
pub fn set_active(&mut self, index: usize, active: bool) -> SignalResult<()> {
if index >= self.vectors.len() {
return Err(SignalError::InvalidParameters(format!(
"Projector index {index} out of range"
)));
}
self.active[index] = active;
Ok(())
}
/// Get number of projectors
pub fn n_projectors(&self) -> usize {
self.vectors.len()
}
/// Get number of active projectors
pub fn n_active(&self) -> usize {
self.active.iter().filter(|&&a| a).count()
}
/// Get projector vectors
pub fn vectors(&self) -> &[Vec<f64>] {
&self.vectors
}
/// Get explained variance ratios
pub fn explained_var(&self) -> &[f64] {
&self.explained_var
}
/// Get projector names
pub fn names(&self) -> &[String] {
&self.names
}
/// Check if projector is active
pub fn is_active(&self, index: usize) -> bool {
self.active.get(index).copied().unwrap_or(false)
}
/// Compute the projection matrix P = I - Σ(u_i * u_i^T)
pub fn compute_projection_matrix(&self) -> Vec<Vec<f64>> {
let n = self.n_channels;
let mut proj = vec![vec![0.0; n]; n];
// Start with identity
for i in 0..n {
proj[i][i] = 1.0;
}
// Subtract outer products of active projectors
for (i, vec) in self.vectors.iter().enumerate() {
if !self.active[i] {
continue;
}
for row in 0..n {
for col in 0..n {
proj[row][col] -= vec[row] * vec[col];
}
}
}
proj
}
/// Merge multiple SSP projectors
pub fn merge(projectors: &[&SspProjector]) -> SignalResult<Self> {
if projectors.is_empty() {
return Err(SignalError::InvalidParameters(
"No projectors to merge".to_string(),
));
}
let n_channels = projectors[0].n_channels;
for p in projectors {
if p.n_channels != n_channels {
return Err(SignalError::InvalidParameters(
"Channel count mismatch".to_string(),
));
}
}
let mut vectors = Vec::new();
let mut active = Vec::new();
let mut explained_var = Vec::new();
let mut names = Vec::new();
for p in projectors {
vectors.extend(p.vectors.clone());
active.extend(p.active.clone());
explained_var.extend(p.explained_var.clone());
names.extend(p.names.clone());
}
Ok(Self {
vectors,
active,
explained_var,
names,
n_channels,
})
}
}
/// Compute robust threshold using median absolute deviation
fn compute_robust_threshold(data: &[f64]) -> f64 {
let mut sorted: Vec<f64> = data.iter().map(|&x| x.abs()).collect();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
let median = if sorted.len().is_multiple_of(2) {
f64::midpoint(sorted[sorted.len() / 2 - 1], sorted[sorted.len() / 2])
} else {
sorted[sorted.len() / 2]
};
// MAD = median(|x - median|)
let mut deviations: Vec<f64> = data.iter().map(|&x| (x.abs() - median).abs()).collect();
deviations.sort_by(|a, b| a.partial_cmp(b).unwrap());
let mad = if deviations.len().is_multiple_of(2) {
f64::midpoint(
deviations[deviations.len() / 2 - 1],
deviations[deviations.len() / 2],
)
} else {
deviations[deviations.len() / 2]
};
// 5 * MAD is a robust threshold
5.0 * mad.max(1e-10)
}
/// Find peaks above threshold with minimum separation
fn find_peaks(data: &[f64], threshold: f64) -> Vec<usize> {
let min_separation = 50; // Minimum samples between peaks
let mut peaks = Vec::new();
let mut last_peak: Option<usize> = None;
for i in 1..data.len().saturating_sub(1) {
if data[i] > threshold && data[i] > data[i - 1] && data[i] > data[i + 1] {
if let Some(last) = last_peak {
if i - last >= min_separation {
peaks.push(i);
last_peak = Some(i);
}
} else {
peaks.push(i);
last_peak = Some(i);
}
}
}
peaks
}
/// Compute eigendecomposition of symmetric matrix using power iteration
fn symmetric_eigen(matrix: &[Vec<f64>]) -> SignalResult<(Vec<Vec<f64>>, Vec<f64>)> {
let n = matrix.len();
if n == 0 {
return Ok((vec![], vec![]));
}
let mut eigenvectors = Vec::with_capacity(n);
let mut eigenvalues = Vec::with_capacity(n);
let mut work_matrix = matrix.to_vec();
for _ in 0..n {
// Power iteration to find largest eigenvalue/vector
let (eigval, eigvec) = power_iteration(&work_matrix, 100)?;
if eigval.abs() < 1e-15 {
break;
}
eigenvalues.push(eigval);
eigenvectors.push(eigvec.clone());
// Deflate: A = A - λ * v * v^T
for i in 0..n {
for j in 0..n {
work_matrix[i][j] -= eigval * eigvec[i] * eigvec[j];
}
}
}
Ok((eigenvectors, eigenvalues))
}
/// Power iteration to find dominant eigenvalue/vector
fn power_iteration(matrix: &[Vec<f64>], max_iter: usize) -> SignalResult<(f64, Vec<f64>)> {
let n = matrix.len();
if n == 0 {
return Err(SignalError::InvalidLength("Empty matrix".to_string()));
}
// Initialize with random-ish vector
let mut v: Vec<f64> = (0..n).map(|i| ((i + 1) as f64).sin()).collect();
// Normalize
let mut norm: f64 = v.iter().map(|&x| x * x).sum::<f64>().sqrt();
for x in &mut v {
*x /= norm;
}
let mut eigenvalue = 0.0;
for _ in 0..max_iter {
// w = A * v
let w: Vec<f64> = (0..n)
.map(|i| (0..n).map(|j| matrix[i][j] * v[j]).sum())
.collect();
// Rayleigh quotient: λ = v^T * A * v
eigenvalue = (0..n).map(|i| v[i] * w[i]).sum();
// Normalize w
norm = w.iter().map(|&x| x * x).sum::<f64>().sqrt();
if norm < 1e-15 {
break;
}
// Check convergence
let diff: f64 = (0..n).map(|i| (w[i] / norm - v[i]).abs()).sum::<f64>();
v = w.iter().map(|&x| x / norm).collect();
if diff < 1e-10 {
break;
}
}
Ok((eigenvalue, v))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ssp_from_epochs() {
// Create simple artifact epochs
let epochs = vec![
vec![
vec![1.0, 2.0, 3.0, 2.0, 1.0],
vec![0.5, 1.0, 1.5, 1.0, 0.5],
vec![0.2, 0.4, 0.6, 0.4, 0.2],
],
vec![
vec![1.1, 2.1, 3.1, 2.1, 1.1],
vec![0.55, 1.05, 1.55, 1.05, 0.55],
vec![0.22, 0.42, 0.62, 0.42, 0.22],
],
];
let ssp = SspProjector::from_epochs(&epochs, 2, "test").unwrap();
assert!(ssp.n_projectors() >= 1);
assert_eq!(ssp.n_active(), ssp.n_projectors());
}
#[test]
fn test_ssp_apply() {
// Create a known projector
let ssp = SspProjector {
vectors: vec![vec![1.0, 0.0, 0.0]], // Project out first channel
active: vec![true],
explained_var: vec![0.5],
names: vec!["test-1".to_string()],
n_channels: 3,
};
let data = vec![
vec![1.0, 2.0, 3.0],
vec![4.0, 5.0, 6.0],
vec![7.0, 8.0, 9.0],
];
let cleaned = ssp.apply(&data).unwrap();
// First channel should be zeroed
assert!(cleaned[0].iter().all(|&x| x.abs() < 1e-10));
// Other channels unchanged
assert_eq!(cleaned[1], vec![4.0, 5.0, 6.0]);
assert_eq!(cleaned[2], vec![7.0, 8.0, 9.0]);
}
#[test]
fn test_ssp_projection_matrix() {
let ssp = SspProjector {
vectors: vec![vec![1.0, 0.0], vec![0.0, 1.0]],
active: vec![true, false],
explained_var: vec![0.5, 0.3],
names: vec!["p1".to_string(), "p2".to_string()],
n_channels: 2,
};
let proj = ssp.compute_projection_matrix();
// Only first projector is active, so P = I - u1*u1^T
// u1 = [1, 0], so P = [[0, 0], [0, 1]]
assert!((proj[0][0] - 0.0).abs() < 1e-10);
assert!((proj[0][1] - 0.0).abs() < 1e-10);
assert!((proj[1][0] - 0.0).abs() < 1e-10);
assert!((proj[1][1] - 1.0).abs() < 1e-10);
}
#[test]
fn test_set_active() {
let mut ssp = SspProjector {
vectors: vec![vec![1.0, 0.0]],
active: vec![true],
explained_var: vec![0.5],
names: vec!["p1".to_string()],
n_channels: 2,
};
assert!(ssp.is_active(0));
ssp.set_active(0, false).unwrap();
assert!(!ssp.is_active(0));
assert_eq!(ssp.n_active(), 0);
}
#[test]
fn test_merge_projectors() {
let p1 = SspProjector {
vectors: vec![vec![1.0, 0.0, 0.0]],
active: vec![true],
explained_var: vec![0.5],
names: vec!["EOG-1".to_string()],
n_channels: 3,
};
let p2 = SspProjector {
vectors: vec![vec![0.0, 1.0, 0.0]],
active: vec![true],
explained_var: vec![0.3],
names: vec!["ECG-1".to_string()],
n_channels: 3,
};
let merged = SspProjector::merge(&[&p1, &p2]).unwrap();
assert_eq!(merged.n_projectors(), 2);
assert_eq!(merged.names()[0], "EOG-1");
assert_eq!(merged.names()[1], "ECG-1");
}
#[test]
fn test_power_iteration() {
// Diagonal matrix - eigenvalues are diagonal elements
let matrix = vec![vec![3.0, 0.0], vec![0.0, 1.0]];
let (eigval, eigvec) = power_iteration(&matrix, 100).unwrap();
// Largest eigenvalue should be 3
assert!((eigval - 3.0).abs() < 1e-6);
// Corresponding eigenvector should be [1, 0]
assert!((eigvec[0].abs() - 1.0).abs() < 1e-6);
assert!(eigvec[1].abs() < 1e-6);
}
#[test]
fn test_find_peaks() {
let data = vec![0.0, 1.0, 5.0, 2.0, 0.0, 0.0, 0.0];
let peaks = find_peaks(&data, 3.0);
assert_eq!(peaks.len(), 1);
assert_eq!(peaks[0], 2);
}
}