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