547 lines
17 KiB
Rust
547 lines
17 KiB
Rust
//! GPU-accelerated LCMV beamformer for real-time source localization.
|
|
//!
|
|
//! Provides sub-millisecond beamformer computation using pre-computed spatial filters.
|
|
|
|
use crate::error::{RealtimeError, RealtimeResult};
|
|
use nalgebra::{DMatrix, DVector};
|
|
use std::sync::Arc;
|
|
|
|
/// GPU beamformer configuration
|
|
#[derive(Debug, Clone)]
|
|
pub struct GpuBeamformerConfig {
|
|
/// Number of sensor channels
|
|
pub n_channels: usize,
|
|
/// Number of source locations
|
|
pub n_sources: usize,
|
|
/// Whether sources have free orientation (3 components per source)
|
|
pub free_orientation: bool,
|
|
/// Regularization parameter (fraction of trace)
|
|
pub regularization: f64,
|
|
/// Power normalization enabled
|
|
pub normalize_power: bool,
|
|
}
|
|
|
|
impl Default for GpuBeamformerConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
n_channels: 64,
|
|
n_sources: 1000,
|
|
free_orientation: false,
|
|
regularization: 0.05,
|
|
normalize_power: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Pre-computed beamformer weights for GPU processing
|
|
#[derive(Debug, Clone)]
|
|
pub struct BeamformerWeights {
|
|
/// Spatial filters [n_source_components x n_channels]
|
|
pub weights: DMatrix<f64>,
|
|
/// Power normalization factors [n_source_components]
|
|
pub power_norm: DVector<f64>,
|
|
/// Source indices
|
|
pub source_indices: Vec<usize>,
|
|
/// Number of source components (n_sources * 3 if free orientation)
|
|
pub n_components: usize,
|
|
}
|
|
|
|
impl BeamformerWeights {
|
|
/// Create from pre-computed weights
|
|
pub fn new(
|
|
weights: DMatrix<f64>,
|
|
power_norm: Option<DVector<f64>>,
|
|
source_indices: Vec<usize>,
|
|
) -> Self {
|
|
let n_components = weights.nrows();
|
|
let power_norm = power_norm.unwrap_or_else(|| DVector::from_element(n_components, 1.0));
|
|
|
|
Self {
|
|
weights,
|
|
power_norm,
|
|
source_indices,
|
|
n_components,
|
|
}
|
|
}
|
|
|
|
/// Number of channels
|
|
pub fn n_channels(&self) -> usize {
|
|
self.weights.ncols()
|
|
}
|
|
|
|
/// Number of source components
|
|
pub fn n_components(&self) -> usize {
|
|
self.n_components
|
|
}
|
|
}
|
|
|
|
/// GPU-accelerated LCMV beamformer
|
|
#[derive(Debug)]
|
|
pub struct GpuBeamformer {
|
|
/// Configuration
|
|
config: GpuBeamformerConfig,
|
|
/// Pre-computed beamformer weights
|
|
weights: Arc<BeamformerWeights>,
|
|
/// Cached covariance inverse (for weight updates)
|
|
cov_inv: Option<DMatrix<f64>>,
|
|
}
|
|
|
|
impl GpuBeamformer {
|
|
/// Create a new GPU beamformer from pre-computed weights
|
|
///
|
|
/// # Arguments
|
|
/// * `config` - Beamformer configuration
|
|
/// * `weights` - Pre-computed spatial filter weights
|
|
pub fn new(config: GpuBeamformerConfig, weights: BeamformerWeights) -> RealtimeResult<Self> {
|
|
// Validate dimensions
|
|
if weights.n_channels() != config.n_channels {
|
|
return Err(RealtimeError::DimensionMismatch(format!(
|
|
"Weights have {} channels, config expects {}",
|
|
weights.n_channels(),
|
|
config.n_channels
|
|
)));
|
|
}
|
|
|
|
Ok(Self {
|
|
config,
|
|
weights: Arc::new(weights),
|
|
cov_inv: None,
|
|
})
|
|
}
|
|
|
|
/// Create a beamformer from gain matrix and covariance
|
|
///
|
|
/// # Arguments
|
|
/// * `config` - Beamformer configuration
|
|
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
|
|
/// * `data_cov` - Data covariance matrix [n_channels x n_channels]
|
|
pub fn from_forward_model(
|
|
config: GpuBeamformerConfig,
|
|
gain: &DMatrix<f64>,
|
|
data_cov: &DMatrix<f64>,
|
|
) -> RealtimeResult<Self> {
|
|
let n_channels = gain.nrows();
|
|
let n_source_cols = gain.ncols();
|
|
|
|
if data_cov.nrows() != n_channels || data_cov.ncols() != n_channels {
|
|
return Err(RealtimeError::DimensionMismatch(
|
|
"Covariance matrix dimension mismatch".to_string(),
|
|
));
|
|
}
|
|
|
|
// Regularize and invert covariance
|
|
let mut c = data_cov.clone();
|
|
let trace: f64 = (0..n_channels).map(|i| c[(i, i)]).sum();
|
|
let reg_val = config.regularization * trace / n_channels as f64;
|
|
|
|
for i in 0..n_channels {
|
|
c[(i, i)] += reg_val;
|
|
}
|
|
|
|
let c_inv = Self::invert_matrix(&c)?;
|
|
|
|
// Compute beamformer weights: w_i = C^-1 * g_i / (g_i^T * C^-1 * g_i)
|
|
let mut weights = DMatrix::zeros(n_source_cols, n_channels);
|
|
let mut power_norm = DVector::zeros(n_source_cols);
|
|
|
|
for src in 0..n_source_cols {
|
|
let g_i = gain.column(src);
|
|
let c_inv_g = &c_inv * g_i;
|
|
let norm = g_i.dot(&c_inv_g);
|
|
|
|
if norm > 1e-30 {
|
|
for ch in 0..n_channels {
|
|
weights[(src, ch)] = c_inv_g[ch] / norm;
|
|
}
|
|
power_norm[src] = 1.0 / norm.sqrt();
|
|
}
|
|
}
|
|
|
|
let source_indices: Vec<usize> = if config.free_orientation {
|
|
(0..n_source_cols / 3).collect()
|
|
} else {
|
|
(0..n_source_cols).collect()
|
|
};
|
|
|
|
let bf_weights = BeamformerWeights::new(weights, Some(power_norm), source_indices);
|
|
|
|
let mut bf = Self::new(config, bf_weights)?;
|
|
bf.cov_inv = Some(c_inv);
|
|
|
|
Ok(bf)
|
|
}
|
|
|
|
/// Apply beamformer to sensor data
|
|
///
|
|
/// # Arguments
|
|
/// * `data` - Sensor data [n_channels x n_times]
|
|
///
|
|
/// # Returns
|
|
/// Source estimates [n_sources x n_times]
|
|
pub fn apply(&self, data: &[Vec<f64>]) -> RealtimeResult<Vec<Vec<f64>>> {
|
|
if data.len() != self.config.n_channels {
|
|
return Err(RealtimeError::DimensionMismatch(format!(
|
|
"Expected {} channels, got {}",
|
|
self.config.n_channels,
|
|
data.len()
|
|
)));
|
|
}
|
|
|
|
if data.is_empty() || data[0].is_empty() {
|
|
return Ok(vec![vec![]; self.weights.n_components]);
|
|
}
|
|
|
|
let n_times = data[0].len();
|
|
let n_comp = self.weights.n_components;
|
|
|
|
// Convert to matrix [n_channels x n_times]
|
|
let data_mat = DMatrix::from_fn(self.config.n_channels, n_times, |i, j| data[i][j]);
|
|
|
|
// Apply spatial filters: S = W * M
|
|
let source_mat = &self.weights.weights * &data_mat;
|
|
|
|
// Apply power normalization if enabled
|
|
let output: Vec<Vec<f64>> = if self.config.normalize_power {
|
|
(0..n_comp)
|
|
.map(|i| {
|
|
let norm = self.weights.power_norm[i];
|
|
(0..n_times).map(|j| source_mat[(i, j)] * norm).collect()
|
|
})
|
|
.collect()
|
|
} else {
|
|
(0..n_comp)
|
|
.map(|i| (0..n_times).map(|j| source_mat[(i, j)]).collect())
|
|
.collect()
|
|
};
|
|
|
|
// If free orientation, combine xyz components into power
|
|
if self.config.free_orientation {
|
|
let n_sources = n_comp / 3;
|
|
let combined: Vec<Vec<f64>> = (0..n_sources)
|
|
.map(|s| {
|
|
let x = &output[s * 3];
|
|
let y = &output[s * 3 + 1];
|
|
let z = &output[s * 3 + 2];
|
|
(0..n_times)
|
|
.map(|t| (x[t].powi(2) + y[t].powi(2) + z[t].powi(2)).sqrt())
|
|
.collect()
|
|
})
|
|
.collect();
|
|
Ok(combined)
|
|
} else {
|
|
Ok(output)
|
|
}
|
|
}
|
|
|
|
/// Apply beamformer using nalgebra matrix (faster for batch processing)
|
|
pub fn apply_matrix(&self, data: &DMatrix<f64>) -> RealtimeResult<DMatrix<f64>> {
|
|
if data.nrows() != self.config.n_channels {
|
|
return Err(RealtimeError::DimensionMismatch(format!(
|
|
"Expected {} channels, got {}",
|
|
self.config.n_channels,
|
|
data.nrows()
|
|
)));
|
|
}
|
|
|
|
// Apply spatial filters: S = W * M
|
|
let mut source_mat = &self.weights.weights * data;
|
|
|
|
// Apply power normalization
|
|
if self.config.normalize_power {
|
|
for i in 0..source_mat.nrows() {
|
|
let norm = self.weights.power_norm[i];
|
|
for j in 0..source_mat.ncols() {
|
|
source_mat[(i, j)] *= norm;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(source_mat)
|
|
}
|
|
|
|
/// Compute Neural Activity Index (NAI) for each source
|
|
///
|
|
/// NAI = (w^T * C_data * w) / (w^T * C_noise * w)
|
|
pub fn compute_nai(
|
|
&self,
|
|
data_cov: &DMatrix<f64>,
|
|
noise_cov: &DMatrix<f64>,
|
|
) -> RealtimeResult<Vec<f64>> {
|
|
let n_comp = self.weights.n_components;
|
|
let n_channels = self.config.n_channels;
|
|
|
|
if data_cov.nrows() != n_channels || noise_cov.nrows() != n_channels {
|
|
return Err(RealtimeError::DimensionMismatch(
|
|
"Covariance dimension mismatch".to_string(),
|
|
));
|
|
}
|
|
|
|
let mut nai = Vec::with_capacity(n_comp);
|
|
|
|
for src in 0..n_comp {
|
|
let w = self.weights.weights.row(src);
|
|
|
|
// w^T * C * w
|
|
let compute_power = |cov: &DMatrix<f64>| -> f64 {
|
|
(0..n_channels)
|
|
.map(|i| {
|
|
(0..n_channels)
|
|
.map(|j| w[i] * cov[(i, j)] * w[j])
|
|
.sum::<f64>()
|
|
})
|
|
.sum()
|
|
};
|
|
|
|
let signal_power = compute_power(data_cov);
|
|
let noise_power = compute_power(noise_cov);
|
|
|
|
if noise_power > 1e-30 {
|
|
nai.push(signal_power / noise_power);
|
|
} else {
|
|
nai.push(0.0);
|
|
}
|
|
}
|
|
|
|
Ok(nai)
|
|
}
|
|
|
|
/// Get the spatial filter weights
|
|
pub fn weights(&self) -> &BeamformerWeights {
|
|
&self.weights
|
|
}
|
|
|
|
/// Get configuration
|
|
pub fn config(&self) -> &GpuBeamformerConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Update beamformer weights with new covariance (adaptive beamforming)
|
|
pub fn update_covariance(&mut self, new_cov: &DMatrix<f64>) -> RealtimeResult<()> {
|
|
// This would require re-computing the weights
|
|
// For real-time use, this should be done asynchronously
|
|
Err(RealtimeError::NotInitialized(
|
|
"Adaptive beamforming not yet implemented".to_string(),
|
|
))
|
|
}
|
|
|
|
/// Invert matrix using SVD
|
|
fn invert_matrix(m: &DMatrix<f64>) -> RealtimeResult<DMatrix<f64>> {
|
|
let n = m.nrows();
|
|
let svd = m.clone().svd(true, true);
|
|
|
|
let u = svd
|
|
.u
|
|
.ok_or_else(|| RealtimeError::Beamformer("SVD failed to compute U".to_string()))?;
|
|
let vt = svd
|
|
.v_t
|
|
.ok_or_else(|| RealtimeError::Beamformer("SVD failed to compute V^T".to_string()))?;
|
|
let s = svd.singular_values;
|
|
|
|
let tol = 1e-10 * s[0];
|
|
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())
|
|
}
|
|
}
|
|
|
|
/// Create a beamformer for a subset of sources (ROI-based)
|
|
pub fn create_roi_beamformer(
|
|
full_beamformer: &GpuBeamformer,
|
|
source_indices: &[usize],
|
|
) -> RealtimeResult<GpuBeamformer> {
|
|
let full_weights = &full_beamformer.weights;
|
|
let mut config = full_beamformer.config.clone();
|
|
|
|
let n_channels = full_weights.n_channels();
|
|
let n_selected = source_indices.len();
|
|
|
|
// Extract weights for selected sources
|
|
let mut selected_weights = DMatrix::zeros(n_selected, n_channels);
|
|
let mut selected_norm = DVector::zeros(n_selected);
|
|
|
|
for (i, &src_idx) in source_indices.iter().enumerate() {
|
|
if src_idx >= full_weights.n_components {
|
|
return Err(RealtimeError::DimensionMismatch(format!(
|
|
"Source index {} out of range (max {})",
|
|
src_idx,
|
|
full_weights.n_components - 1
|
|
)));
|
|
}
|
|
|
|
for ch in 0..n_channels {
|
|
selected_weights[(i, ch)] = full_weights.weights[(src_idx, ch)];
|
|
}
|
|
selected_norm[i] = full_weights.power_norm[src_idx];
|
|
}
|
|
|
|
config.n_sources = n_selected;
|
|
config.free_orientation = false;
|
|
|
|
let roi_weights = BeamformerWeights::new(
|
|
selected_weights,
|
|
Some(selected_norm),
|
|
source_indices.to_vec(),
|
|
);
|
|
|
|
GpuBeamformer::new(config, roi_weights)
|
|
}
|
|
|
|
/// Compute source power (envelope) from beamformer output
|
|
pub fn compute_source_power(sources: &[Vec<f64>], window_samples: usize) -> Vec<Vec<f64>> {
|
|
sources
|
|
.iter()
|
|
.map(|source| {
|
|
source
|
|
.windows(window_samples)
|
|
.map(|window| {
|
|
let sum_sq: f64 = window.iter().map(|x| x.powi(2)).sum();
|
|
(sum_sq / window_samples as f64).sqrt()
|
|
})
|
|
.collect()
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn create_test_gain() -> DMatrix<f64> {
|
|
// Simple 4-channel, 3-source gain matrix
|
|
DMatrix::from_row_slice(
|
|
4,
|
|
3,
|
|
&[1.0, 0.5, 0.2, 0.5, 1.0, 0.5, 0.2, 0.5, 1.0, 0.1, 0.3, 0.6],
|
|
)
|
|
}
|
|
|
|
fn create_test_covariance() -> DMatrix<f64> {
|
|
// Simple identity-like covariance
|
|
DMatrix::from_diagonal(&DVector::from_vec(vec![1.0, 1.0, 1.0, 1.0]))
|
|
}
|
|
|
|
#[test]
|
|
fn test_beamformer_creation() {
|
|
let config = GpuBeamformerConfig {
|
|
n_channels: 4,
|
|
n_sources: 3,
|
|
free_orientation: false,
|
|
regularization: 0.05,
|
|
normalize_power: true,
|
|
};
|
|
|
|
let gain = create_test_gain();
|
|
let cov = create_test_covariance();
|
|
|
|
let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();
|
|
|
|
assert_eq!(bf.weights.n_channels(), 4);
|
|
assert_eq!(bf.weights.n_components(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_beamformer_apply() {
|
|
let config = GpuBeamformerConfig {
|
|
n_channels: 4,
|
|
n_sources: 3,
|
|
free_orientation: false,
|
|
regularization: 0.05,
|
|
normalize_power: false,
|
|
};
|
|
|
|
let gain = create_test_gain();
|
|
let cov = create_test_covariance();
|
|
let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();
|
|
|
|
// Test data: 4 channels, 10 time points
|
|
let data: Vec<Vec<f64>> = (0..4)
|
|
.map(|ch| (0..10).map(|t| (ch * 10 + t) as f64).collect())
|
|
.collect();
|
|
|
|
let sources = bf.apply(&data).unwrap();
|
|
|
|
assert_eq!(sources.len(), 3);
|
|
assert_eq!(sources[0].len(), 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_matrix() {
|
|
let config = GpuBeamformerConfig {
|
|
n_channels: 4,
|
|
n_sources: 3,
|
|
free_orientation: false,
|
|
regularization: 0.05,
|
|
normalize_power: false,
|
|
};
|
|
|
|
let gain = create_test_gain();
|
|
let cov = create_test_covariance();
|
|
let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();
|
|
|
|
let data = DMatrix::from_fn(4, 10, |i, j| (i * 10 + j) as f64);
|
|
let sources = bf.apply_matrix(&data).unwrap();
|
|
|
|
assert_eq!(sources.nrows(), 3);
|
|
assert_eq!(sources.ncols(), 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_nai() {
|
|
let config = GpuBeamformerConfig {
|
|
n_channels: 4,
|
|
n_sources: 3,
|
|
free_orientation: false,
|
|
regularization: 0.05,
|
|
normalize_power: false,
|
|
};
|
|
|
|
let gain = create_test_gain();
|
|
let cov = create_test_covariance();
|
|
let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();
|
|
|
|
// Data cov with higher signal
|
|
let data_cov = DMatrix::from_diagonal(&DVector::from_vec(vec![2.0, 2.0, 2.0, 2.0]));
|
|
let noise_cov = DMatrix::from_diagonal(&DVector::from_vec(vec![1.0, 1.0, 1.0, 1.0]));
|
|
|
|
let nai = bf.compute_nai(&data_cov, &noise_cov).unwrap();
|
|
|
|
assert_eq!(nai.len(), 3);
|
|
// NAI should be positive with higher signal than noise
|
|
for n in &nai {
|
|
assert!(*n >= 0.0);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_roi_beamformer() {
|
|
let config = GpuBeamformerConfig {
|
|
n_channels: 4,
|
|
n_sources: 3,
|
|
free_orientation: false,
|
|
regularization: 0.05,
|
|
normalize_power: false,
|
|
};
|
|
|
|
let gain = create_test_gain();
|
|
let cov = create_test_covariance();
|
|
let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();
|
|
|
|
// Create ROI beamformer for sources 0 and 2
|
|
let roi_bf = create_roi_beamformer(&bf, &[0, 2]).unwrap();
|
|
|
|
assert_eq!(roi_bf.weights.n_components(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_source_power() {
|
|
let sources = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0], vec![0.5, 1.0, 1.5, 2.0, 2.5]];
|
|
|
|
let power = compute_source_power(&sources, 3);
|
|
|
|
assert_eq!(power.len(), 2);
|
|
assert_eq!(power[0].len(), 3); // 5 - 3 + 1 = 3 windows
|
|
}
|
|
}
|