1046 lines
31 KiB
Rust
1046 lines
31 KiB
Rust
//! Beamformer methods for source localization.
|
|
//!
|
|
//! This module provides beamformer inverse solutions:
|
|
//! - **LCMV** (Linearly Constrained Minimum Variance): Time-domain beamformer
|
|
//! - **DICS** (Dynamic Imaging of Coherent Sources): Frequency-domain beamformer
|
|
//!
|
|
//! ## Mathematical Background
|
|
//!
|
|
//! The beamformer filter is: w = (C^-1 * g) / (g^T * C^-1 * g)
|
|
//! where C is covariance/CSD, g is lead field for the target location.
|
|
//!
|
|
//! LCMV uses time-domain covariance, while DICS uses cross-spectral density.
|
|
|
|
use crate::{Covariance, InverseError, InverseResult, SourceEstimate};
|
|
use nalgebra::{Complex, DMatrix, DVector};
|
|
use rtx_neuro_forward::GainMatrix;
|
|
use std::f64::consts::PI;
|
|
|
|
/// LCMV beamformer
|
|
#[derive(Debug, Clone)]
|
|
pub struct LcmvBeamformer {
|
|
/// Spatial filters [n_sources x n_channels]
|
|
filters: DMatrix<f64>,
|
|
/// Number of sources
|
|
n_sources: usize,
|
|
/// Number of channels
|
|
n_channels: usize,
|
|
/// Regularization parameter
|
|
reg: f64,
|
|
/// Source indices
|
|
source_indices: Vec<usize>,
|
|
}
|
|
|
|
impl LcmvBeamformer {
|
|
/// Create an LCMV beamformer from forward model and data covariance
|
|
///
|
|
/// # Arguments
|
|
/// * `gain` - Forward model gain matrix [n_channels x n_sources]
|
|
/// * `data_cov` - Data covariance matrix
|
|
/// * `noise_cov` - Optional noise covariance for regularization
|
|
/// * `reg` - Regularization parameter (fraction of trace)
|
|
pub fn make_lcmv(
|
|
gain: &GainMatrix,
|
|
data_cov: &Covariance,
|
|
noise_cov: Option<&Covariance>,
|
|
reg: f64,
|
|
) -> InverseResult<Self> {
|
|
let n_channels = gain.n_sensors();
|
|
let n_source_cols = gain.n_source_columns();
|
|
|
|
if data_cov.n_channels() != n_channels {
|
|
return Err(InverseError::DimensionMismatch(format!(
|
|
"Data covariance has {} channels, gain has {}",
|
|
data_cov.n_channels(),
|
|
n_channels
|
|
)));
|
|
}
|
|
|
|
// Get data covariance matrix
|
|
let mut c = data_cov.data().clone();
|
|
|
|
// Apply regularization
|
|
if reg > 0.0 {
|
|
let trace: f64 = (0..n_channels).map(|i| c[(i, i)]).sum();
|
|
let reg_val = reg * trace / n_channels as f64;
|
|
for i in 0..n_channels {
|
|
c[(i, i)] += reg_val;
|
|
}
|
|
}
|
|
|
|
// Add noise covariance if provided
|
|
if let Some(nc) = noise_cov {
|
|
if nc.n_channels() != n_channels {
|
|
return Err(InverseError::DimensionMismatch(
|
|
"Noise covariance dimension mismatch".to_string(),
|
|
));
|
|
}
|
|
c += nc.data();
|
|
}
|
|
|
|
// Invert covariance matrix
|
|
let c_inv = Self::invert_covariance(&c)?;
|
|
|
|
// Convert gain to matrix
|
|
let g = DMatrix::from_fn(n_channels, n_source_cols, |i, j| gain.data()[i][j]);
|
|
|
|
// Compute filters: w_i = C^-1 * g_i / (g_i^T * C^-1 * g_i)
|
|
let mut filters = DMatrix::zeros(n_source_cols, n_channels);
|
|
|
|
for src in 0..n_source_cols {
|
|
// Extract lead field for this source
|
|
let g_i = g.column(src);
|
|
|
|
// C^-1 * g
|
|
let c_inv_g = &c_inv * g_i;
|
|
|
|
// g^T * C^-1 * g (normalization factor)
|
|
let norm = g_i.dot(&c_inv_g);
|
|
|
|
if norm > 1e-30 {
|
|
// w = C^-1 * g / norm
|
|
for ch in 0..n_channels {
|
|
filters[(src, ch)] = c_inv_g[ch] / norm;
|
|
}
|
|
}
|
|
}
|
|
|
|
// For free orientation, n_sources = n_source_cols / 3
|
|
let n_sources = if gain.is_free_orientation() {
|
|
n_source_cols / 3
|
|
} else {
|
|
n_source_cols
|
|
};
|
|
|
|
Ok(Self {
|
|
filters,
|
|
n_sources,
|
|
n_channels,
|
|
reg,
|
|
source_indices: (0..n_sources).collect(),
|
|
})
|
|
}
|
|
|
|
/// Apply the 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>]) -> InverseResult<SourceEstimate> {
|
|
if data.len() != self.n_channels {
|
|
return Err(InverseError::DimensionMismatch(format!(
|
|
"Expected {} channels, got {}",
|
|
self.n_channels,
|
|
data.len()
|
|
)));
|
|
}
|
|
|
|
let n_times = data[0].len();
|
|
let n_source_cols = self.filters.nrows();
|
|
|
|
// Convert data to matrix
|
|
let data_mat = DMatrix::from_fn(self.n_channels, n_times, |i, j| data[i][j]);
|
|
|
|
// Apply filters: S = W * M
|
|
let source_mat = &self.filters * &data_mat;
|
|
|
|
// Convert to Vec<Vec<f64>>
|
|
let source_data: Vec<Vec<f64>> = (0..n_source_cols)
|
|
.map(|i| (0..n_times).map(|j| source_mat[(i, j)]).collect())
|
|
.collect();
|
|
|
|
// Placeholder times
|
|
let times: Vec<f64> = (0..n_times).map(|i| i as f64).collect();
|
|
|
|
let free_orientation = n_source_cols != self.n_sources;
|
|
|
|
Ok(SourceEstimate::new(
|
|
source_data,
|
|
times,
|
|
self.source_indices.clone(),
|
|
free_orientation,
|
|
))
|
|
}
|
|
|
|
/// Get the spatial filters
|
|
pub fn filters(&self) -> &DMatrix<f64> {
|
|
&self.filters
|
|
}
|
|
|
|
/// Get number of sources
|
|
pub fn n_sources(&self) -> usize {
|
|
self.n_sources
|
|
}
|
|
|
|
/// Compute beamformer output power (pseudo-Z statistic)
|
|
///
|
|
/// Returns the neural activity index (NAI) for each source
|
|
pub fn compute_nai(
|
|
&self,
|
|
data_cov: &Covariance,
|
|
noise_cov: &Covariance,
|
|
) -> InverseResult<Vec<f64>> {
|
|
let n_source_cols = self.filters.nrows();
|
|
let mut nai = Vec::with_capacity(n_source_cols);
|
|
|
|
let c_data = data_cov.data();
|
|
let c_noise = noise_cov.data();
|
|
|
|
for src in 0..n_source_cols {
|
|
// Extract filter row
|
|
let w = self.filters.row(src);
|
|
|
|
// Signal power: w^T * C_data * w
|
|
let signal_power: f64 = (0..self.n_channels)
|
|
.map(|i| {
|
|
(0..self.n_channels)
|
|
.map(|j| w[i] * c_data[(i, j)] * w[j])
|
|
.sum::<f64>()
|
|
})
|
|
.sum();
|
|
|
|
// Noise power: w^T * C_noise * w
|
|
let noise_power: f64 = (0..self.n_channels)
|
|
.map(|i| {
|
|
(0..self.n_channels)
|
|
.map(|j| w[i] * c_noise[(i, j)] * w[j])
|
|
.sum::<f64>()
|
|
})
|
|
.sum();
|
|
|
|
// NAI = signal / noise
|
|
if noise_power > 1e-30 {
|
|
nai.push(signal_power / noise_power);
|
|
} else {
|
|
nai.push(0.0);
|
|
}
|
|
}
|
|
|
|
Ok(nai)
|
|
}
|
|
|
|
/// Invert covariance matrix using SVD
|
|
fn invert_covariance(c: &DMatrix<f64>) -> InverseResult<DMatrix<f64>> {
|
|
let n = c.nrows();
|
|
|
|
let svd = c.clone().svd(true, true);
|
|
let u = svd
|
|
.u
|
|
.ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
|
|
let vt = svd
|
|
.v_t
|
|
.ok_or_else(|| InverseError::ComputationError("SVD failed".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())
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// DICS Beamformer (Dynamic Imaging of Coherent Sources)
|
|
// ============================================================================
|
|
|
|
/// Configuration for DICS beamformer
|
|
#[derive(Debug, Clone)]
|
|
pub struct DicsConfig {
|
|
/// Target frequencies (Hz)
|
|
pub frequencies: Vec<f64>,
|
|
/// Frequency bandwidth for smoothing (Hz)
|
|
pub bandwidth: f64,
|
|
/// Regularization parameter (fraction of trace)
|
|
pub reg: f64,
|
|
/// Pick orientation method
|
|
pub pick_ori: PickOrientation,
|
|
/// Whether to compute power (true) or complex values (false)
|
|
pub real_filter: bool,
|
|
/// Sampling frequency
|
|
pub sfreq: f64,
|
|
}
|
|
|
|
impl Default for DicsConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
frequencies: vec![10.0], // Default: alpha band
|
|
bandwidth: 4.0,
|
|
reg: 0.05,
|
|
pick_ori: PickOrientation::MaxPower,
|
|
real_filter: true,
|
|
sfreq: 1000.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl DicsConfig {
|
|
/// Create DICS config for alpha band (8-12 Hz)
|
|
pub fn alpha(sfreq: f64) -> Self {
|
|
Self {
|
|
frequencies: vec![10.0],
|
|
bandwidth: 4.0,
|
|
sfreq,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Create DICS config for beta band (13-30 Hz)
|
|
pub fn beta(sfreq: f64) -> Self {
|
|
Self {
|
|
frequencies: vec![20.0],
|
|
bandwidth: 17.0,
|
|
sfreq,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Create DICS config for gamma band (30-100 Hz)
|
|
pub fn gamma(sfreq: f64) -> Self {
|
|
Self {
|
|
frequencies: vec![50.0],
|
|
bandwidth: 40.0,
|
|
sfreq,
|
|
..Self::default()
|
|
}
|
|
}
|
|
|
|
/// Set target frequencies
|
|
pub fn with_frequencies(mut self, freqs: Vec<f64>) -> Self {
|
|
self.frequencies = freqs;
|
|
self
|
|
}
|
|
|
|
/// Set bandwidth
|
|
pub fn with_bandwidth(mut self, bw: f64) -> Self {
|
|
self.bandwidth = bw;
|
|
self
|
|
}
|
|
|
|
/// Set regularization
|
|
pub fn with_reg(mut self, reg: f64) -> Self {
|
|
self.reg = reg;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Orientation picking method for free-orientation sources
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum PickOrientation {
|
|
/// Pick orientation with maximum power
|
|
MaxPower,
|
|
/// Use surface normal (requires fixed orientation in gain)
|
|
Normal,
|
|
/// Return all orientations (no projection)
|
|
None,
|
|
}
|
|
|
|
/// Cross-spectral density matrix
|
|
#[derive(Debug, Clone)]
|
|
pub struct CrossSpectralDensity {
|
|
/// CSD matrix [n_channels x n_channels] (complex)
|
|
data: Vec<Vec<Complex<f64>>>,
|
|
/// Number of channels
|
|
n_channels: usize,
|
|
/// Frequencies at which CSD was computed
|
|
frequencies: Vec<f64>,
|
|
/// Number of epochs/segments used
|
|
n_epochs: usize,
|
|
}
|
|
|
|
impl CrossSpectralDensity {
|
|
/// Create a new CSD matrix
|
|
pub fn new(data: Vec<Vec<Complex<f64>>>, frequencies: Vec<f64>, n_epochs: usize) -> Self {
|
|
let n_channels = data.len();
|
|
Self {
|
|
data,
|
|
n_channels,
|
|
frequencies,
|
|
n_epochs,
|
|
}
|
|
}
|
|
|
|
/// Compute CSD from epochs of data
|
|
///
|
|
/// # Arguments
|
|
/// * `epochs` - Data epochs [n_epochs][n_channels][n_samples]
|
|
/// * `sfreq` - Sampling frequency
|
|
/// * `fmin` - Minimum frequency
|
|
/// * `fmax` - Maximum frequency
|
|
/// * `n_fft` - FFT size (None = use epoch length)
|
|
pub fn from_epochs(
|
|
epochs: &[Vec<Vec<f64>>],
|
|
sfreq: f64,
|
|
fmin: f64,
|
|
fmax: f64,
|
|
n_fft: Option<usize>,
|
|
) -> InverseResult<Self> {
|
|
if epochs.is_empty() {
|
|
return Err(InverseError::InvalidParameter(
|
|
"No epochs provided".to_string(),
|
|
));
|
|
}
|
|
|
|
let n_epochs = epochs.len();
|
|
let n_channels = epochs[0].len();
|
|
let n_samples = epochs[0][0].len();
|
|
let n_fft = n_fft.unwrap_or(n_samples);
|
|
|
|
// Compute frequency bins
|
|
let df = sfreq / n_fft as f64;
|
|
let freq_min_idx = (fmin / df).ceil() as usize;
|
|
let freq_max_idx = (fmax / df).floor() as usize;
|
|
|
|
let frequencies: Vec<f64> = (freq_min_idx..=freq_max_idx)
|
|
.map(|i| i as f64 * df)
|
|
.collect();
|
|
|
|
if frequencies.is_empty() {
|
|
return Err(InverseError::InvalidParameter(
|
|
"No frequencies in specified range".to_string(),
|
|
));
|
|
}
|
|
|
|
// Initialize CSD matrix
|
|
let mut csd: Vec<Vec<Complex<f64>>> =
|
|
vec![vec![Complex::new(0.0, 0.0); n_channels]; n_channels];
|
|
|
|
// Compute FFT for each epoch and accumulate CSD
|
|
for epoch in epochs {
|
|
// Compute FFT for each channel
|
|
let ffts: Vec<Vec<Complex<f64>>> = epoch
|
|
.iter()
|
|
.map(|ch| Self::compute_fft(ch, n_fft))
|
|
.collect();
|
|
|
|
// Accumulate cross-spectra for frequency range
|
|
for i in 0..n_channels {
|
|
for j in i..n_channels {
|
|
let mut sum = Complex::new(0.0, 0.0);
|
|
for &freq_idx in &(freq_min_idx..=freq_max_idx).collect::<Vec<_>>() {
|
|
if freq_idx < ffts[0].len() {
|
|
// CSD_ij = X_i * conj(X_j)
|
|
sum += ffts[i][freq_idx] * ffts[j][freq_idx].conj();
|
|
}
|
|
}
|
|
csd[i][j] += sum;
|
|
if i != j {
|
|
csd[j][i] += sum.conj();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Normalize by number of epochs and frequency bins
|
|
let n_freqs = frequencies.len() as f64;
|
|
let norm = n_epochs as f64 * n_freqs;
|
|
for i in 0..n_channels {
|
|
for j in 0..n_channels {
|
|
csd[i][j] /= norm;
|
|
}
|
|
}
|
|
|
|
Ok(Self::new(csd, frequencies, n_epochs))
|
|
}
|
|
|
|
/// Get CSD data
|
|
pub fn data(&self) -> &Vec<Vec<Complex<f64>>> {
|
|
&self.data
|
|
}
|
|
|
|
/// Get number of channels
|
|
pub fn n_channels(&self) -> usize {
|
|
self.n_channels
|
|
}
|
|
|
|
/// Get frequencies
|
|
pub fn frequencies(&self) -> &[f64] {
|
|
&self.frequencies
|
|
}
|
|
|
|
/// Convert to real matrix (magnitude)
|
|
pub fn to_real(&self) -> DMatrix<f64> {
|
|
DMatrix::from_fn(self.n_channels, self.n_channels, |i, j| {
|
|
self.data[i][j].norm()
|
|
})
|
|
}
|
|
|
|
/// Convert to real matrix (real part only)
|
|
pub fn real_part(&self) -> DMatrix<f64> {
|
|
DMatrix::from_fn(self.n_channels, self.n_channels, |i, j| self.data[i][j].re)
|
|
}
|
|
|
|
/// Simple DFT implementation
|
|
fn compute_fft(signal: &[f64], n_fft: usize) -> Vec<Complex<f64>> {
|
|
let n = signal.len().min(n_fft);
|
|
let mut result = vec![Complex::new(0.0, 0.0); n_fft / 2 + 1];
|
|
|
|
for (k, res) in result.iter_mut().enumerate() {
|
|
let mut sum = Complex::new(0.0, 0.0);
|
|
for (t, &x) in signal.iter().take(n).enumerate() {
|
|
let angle = -2.0 * PI * k as f64 * t as f64 / n_fft as f64;
|
|
sum += Complex::new(x * angle.cos(), x * angle.sin());
|
|
}
|
|
*res = sum / (n as f64).sqrt();
|
|
}
|
|
|
|
result
|
|
}
|
|
}
|
|
|
|
/// DICS beamformer
|
|
#[derive(Debug, Clone)]
|
|
pub struct DicsBeamformer {
|
|
/// Spatial filters [n_sources x n_channels] (complex)
|
|
filters: Vec<Vec<Complex<f64>>>,
|
|
/// Number of sources
|
|
n_sources: usize,
|
|
/// Number of channels
|
|
n_channels: usize,
|
|
/// Whether sources have free orientation
|
|
free_orientation: bool,
|
|
/// Configuration used
|
|
config: DicsConfig,
|
|
/// Source indices
|
|
source_indices: Vec<usize>,
|
|
}
|
|
|
|
impl DicsBeamformer {
|
|
/// Create a DICS beamformer from forward model and CSD
|
|
///
|
|
/// # Arguments
|
|
/// * `gain` - Forward model gain matrix
|
|
/// * `csd` - Cross-spectral density matrix
|
|
/// * `noise_csd` - Optional noise CSD for regularization
|
|
/// * `config` - DICS configuration
|
|
pub fn make_dics(
|
|
gain: &GainMatrix,
|
|
csd: &CrossSpectralDensity,
|
|
noise_csd: Option<&CrossSpectralDensity>,
|
|
config: DicsConfig,
|
|
) -> InverseResult<Self> {
|
|
let n_channels = gain.n_sensors();
|
|
let n_source_cols = gain.n_source_columns();
|
|
let free_orientation = gain.is_free_orientation();
|
|
|
|
if csd.n_channels() != n_channels {
|
|
return Err(InverseError::DimensionMismatch(format!(
|
|
"CSD has {} channels, gain has {}",
|
|
csd.n_channels(),
|
|
n_channels
|
|
)));
|
|
}
|
|
|
|
// Convert CSD to matrix form (use real part for real filter)
|
|
let csd_mat = if config.real_filter {
|
|
csd.real_part()
|
|
} else {
|
|
csd.to_real()
|
|
};
|
|
|
|
// Apply regularization
|
|
let mut c = csd_mat.clone();
|
|
if config.reg > 0.0 {
|
|
let trace: f64 = (0..n_channels).map(|i| c[(i, i)]).sum();
|
|
let reg_val = config.reg * trace / n_channels as f64;
|
|
for i in 0..n_channels {
|
|
c[(i, i)] += reg_val;
|
|
}
|
|
}
|
|
|
|
// Add noise CSD if provided
|
|
if let Some(nc) = noise_csd {
|
|
let noise_mat = if config.real_filter {
|
|
nc.real_part()
|
|
} else {
|
|
nc.to_real()
|
|
};
|
|
c += &noise_mat;
|
|
}
|
|
|
|
// Invert CSD matrix
|
|
let c_inv = Self::invert_matrix(&c)?;
|
|
|
|
// Convert gain to matrix
|
|
let g = DMatrix::from_fn(n_channels, n_source_cols, |i, j| gain.data()[i][j]);
|
|
|
|
// Compute filters
|
|
let filters = if free_orientation {
|
|
Self::compute_filters_free(&g, &c_inv, n_channels, n_source_cols, &config)?
|
|
} else {
|
|
Self::compute_filters_fixed(&g, &c_inv, n_channels, n_source_cols)?
|
|
};
|
|
|
|
let n_sources = if free_orientation {
|
|
n_source_cols / 3
|
|
} else {
|
|
n_source_cols
|
|
};
|
|
|
|
Ok(Self {
|
|
filters,
|
|
n_sources,
|
|
n_channels,
|
|
free_orientation,
|
|
config,
|
|
source_indices: (0..n_sources).collect(),
|
|
})
|
|
}
|
|
|
|
/// Compute source power from CSD
|
|
///
|
|
/// Returns power at each source location
|
|
pub fn compute_power(&self, csd: &CrossSpectralDensity) -> InverseResult<Vec<f64>> {
|
|
if csd.n_channels() != self.n_channels {
|
|
return Err(InverseError::DimensionMismatch(
|
|
"CSD channel mismatch".to_string(),
|
|
));
|
|
}
|
|
|
|
let n_source_cols = self.filters.len();
|
|
let mut power = Vec::with_capacity(n_source_cols);
|
|
|
|
for src in 0..n_source_cols {
|
|
let w = &self.filters[src];
|
|
|
|
// Power = w^H * CSD * w
|
|
let mut p = Complex::new(0.0, 0.0);
|
|
for i in 0..self.n_channels {
|
|
for j in 0..self.n_channels {
|
|
p += w[i].conj() * csd.data()[i][j] * w[j];
|
|
}
|
|
}
|
|
|
|
power.push(p.re.max(0.0)); // Power is real and non-negative
|
|
}
|
|
|
|
// For free orientation, combine orientations
|
|
if self.free_orientation && n_source_cols > self.n_sources {
|
|
let mut combined = Vec::with_capacity(self.n_sources);
|
|
for src in 0..self.n_sources {
|
|
let p = power[3 * src] + power[3 * src + 1] + power[3 * src + 2];
|
|
combined.push(p);
|
|
}
|
|
return Ok(combined);
|
|
}
|
|
|
|
Ok(power)
|
|
}
|
|
|
|
/// Apply beamformer to time-domain data
|
|
///
|
|
/// First computes CSD from data, then returns source power
|
|
pub fn apply_to_epochs(&self, epochs: &[Vec<Vec<f64>>]) -> InverseResult<Vec<f64>> {
|
|
// Compute CSD from epochs
|
|
let fmin =
|
|
self.config.frequencies.first().copied().unwrap_or(1.0) - self.config.bandwidth / 2.0;
|
|
let fmax =
|
|
self.config.frequencies.last().copied().unwrap_or(100.0) + self.config.bandwidth / 2.0;
|
|
|
|
let csd = CrossSpectralDensity::from_epochs(
|
|
epochs,
|
|
self.config.sfreq,
|
|
fmin.max(0.1),
|
|
fmax,
|
|
None,
|
|
)?;
|
|
|
|
self.compute_power(&csd)
|
|
}
|
|
|
|
/// Get the spatial filters
|
|
pub fn filters(&self) -> &Vec<Vec<Complex<f64>>> {
|
|
&self.filters
|
|
}
|
|
|
|
/// Get number of sources
|
|
pub fn n_sources(&self) -> usize {
|
|
self.n_sources
|
|
}
|
|
|
|
/// Get frequencies
|
|
pub fn frequencies(&self) -> &[f64] {
|
|
&self.config.frequencies
|
|
}
|
|
|
|
// ========== Private methods ==========
|
|
|
|
/// Compute filters for fixed orientation
|
|
fn compute_filters_fixed(
|
|
g: &DMatrix<f64>,
|
|
c_inv: &DMatrix<f64>,
|
|
n_channels: usize,
|
|
n_sources: usize,
|
|
) -> InverseResult<Vec<Vec<Complex<f64>>>> {
|
|
let mut filters = Vec::with_capacity(n_sources);
|
|
|
|
for src in 0..n_sources {
|
|
let g_i = g.column(src);
|
|
let c_inv_g = c_inv * g_i;
|
|
let norm = g_i.dot(&c_inv_g);
|
|
|
|
let filter: Vec<Complex<f64>> = if norm > 1e-30 {
|
|
(0..n_channels)
|
|
.map(|ch| Complex::new(c_inv_g[ch] / norm, 0.0))
|
|
.collect()
|
|
} else {
|
|
vec![Complex::new(0.0, 0.0); n_channels]
|
|
};
|
|
|
|
filters.push(filter);
|
|
}
|
|
|
|
Ok(filters)
|
|
}
|
|
|
|
/// Compute filters for free orientation
|
|
fn compute_filters_free(
|
|
g: &DMatrix<f64>,
|
|
c_inv: &DMatrix<f64>,
|
|
n_channels: usize,
|
|
n_source_cols: usize,
|
|
config: &DicsConfig,
|
|
) -> InverseResult<Vec<Vec<Complex<f64>>>> {
|
|
let n_sources = n_source_cols / 3;
|
|
let mut filters = Vec::with_capacity(n_source_cols);
|
|
|
|
match config.pick_ori {
|
|
PickOrientation::MaxPower => {
|
|
// Compute filter for each orientation, then pick max power
|
|
for src in 0..n_sources {
|
|
// Compute 3x3 matrix: G^T * C^(-1) * G
|
|
let mut gtcg = [[0.0; 3]; 3];
|
|
for ori1 in 0..3 {
|
|
let col1 = 3 * src + ori1;
|
|
for ori2 in 0..3 {
|
|
let col2 = 3 * src + ori2;
|
|
for i in 0..n_channels {
|
|
for j in 0..n_channels {
|
|
gtcg[ori1][ori2] += g[(i, col1)] * c_inv[(i, j)] * g[(j, col2)];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Find eigenvector with largest eigenvalue
|
|
let (max_ori, _) = Self::max_eigenvector_3x3(>cg);
|
|
|
|
// Create combined leadfield
|
|
let g_combined: Vec<f64> = (0..n_channels)
|
|
.map(|ch| {
|
|
max_ori[0] * g[(ch, 3 * src)]
|
|
+ max_ori[1] * g[(ch, 3 * src + 1)]
|
|
+ max_ori[2] * g[(ch, 3 * src + 2)]
|
|
})
|
|
.collect();
|
|
|
|
// Compute filter for combined orientation
|
|
let g_vec = DVector::from_column_slice(&g_combined);
|
|
let c_inv_g = c_inv * &g_vec;
|
|
let norm = g_vec.dot(&c_inv_g);
|
|
|
|
let filter: Vec<Complex<f64>> = if norm > 1e-30 {
|
|
(0..n_channels)
|
|
.map(|ch| Complex::new(c_inv_g[ch] / norm, 0.0))
|
|
.collect()
|
|
} else {
|
|
vec![Complex::new(0.0, 0.0); n_channels]
|
|
};
|
|
|
|
// Store same filter for all orientations (max power direction)
|
|
filters.push(filter.clone());
|
|
filters.push(filter.clone());
|
|
filters.push(filter);
|
|
}
|
|
}
|
|
PickOrientation::None | PickOrientation::Normal => {
|
|
// Compute separate filter for each orientation
|
|
for src_col in 0..n_source_cols {
|
|
let g_i = g.column(src_col);
|
|
let c_inv_g = c_inv * g_i;
|
|
let norm = g_i.dot(&c_inv_g);
|
|
|
|
let filter: Vec<Complex<f64>> = if norm > 1e-30 {
|
|
(0..n_channels)
|
|
.map(|ch| Complex::new(c_inv_g[ch] / norm, 0.0))
|
|
.collect()
|
|
} else {
|
|
vec![Complex::new(0.0, 0.0); n_channels]
|
|
};
|
|
|
|
filters.push(filter);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(filters)
|
|
}
|
|
|
|
/// Find eigenvector with maximum eigenvalue for 3x3 symmetric matrix
|
|
fn max_eigenvector_3x3(a: &[[f64; 3]; 3]) -> ([f64; 3], f64) {
|
|
// Power iteration for dominant eigenvector
|
|
let mut v = [1.0 / 3.0_f64.sqrt(); 3];
|
|
|
|
for _ in 0..20 {
|
|
// Av
|
|
let mut av = [0.0; 3];
|
|
for i in 0..3 {
|
|
for j in 0..3 {
|
|
av[i] += a[i][j] * v[j];
|
|
}
|
|
}
|
|
|
|
// Normalize
|
|
let norm = (av[0] * av[0] + av[1] * av[1] + av[2] * av[2]).sqrt();
|
|
if norm > 1e-15 {
|
|
for i in 0..3 {
|
|
v[i] = av[i] / norm;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute eigenvalue
|
|
let mut av = [0.0; 3];
|
|
for i in 0..3 {
|
|
for j in 0..3 {
|
|
av[i] += a[i][j] * v[j];
|
|
}
|
|
}
|
|
let eigenvalue = v[0] * av[0] + v[1] * av[1] + v[2] * av[2];
|
|
|
|
(v, eigenvalue)
|
|
}
|
|
|
|
/// Invert matrix using SVD
|
|
fn invert_matrix(c: &DMatrix<f64>) -> InverseResult<DMatrix<f64>> {
|
|
let n = c.nrows();
|
|
|
|
let svd = c.clone().svd(true, true);
|
|
let u = svd
|
|
.u
|
|
.ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
|
|
let vt = svd
|
|
.v_t
|
|
.ok_or_else(|| InverseError::ComputationError("SVD failed".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())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::covariance::CovarianceType;
|
|
|
|
fn create_simple_gain() -> GainMatrix {
|
|
let data = vec![
|
|
vec![1.0, 0.5, 0.2],
|
|
vec![0.5, 1.0, 0.5],
|
|
vec![0.2, 0.5, 1.0],
|
|
vec![0.1, 0.3, 0.6],
|
|
];
|
|
let names = vec![
|
|
"S1".to_string(),
|
|
"S2".to_string(),
|
|
"S3".to_string(),
|
|
"S4".to_string(),
|
|
];
|
|
GainMatrix::new(data, false, names).unwrap()
|
|
}
|
|
|
|
fn create_free_gain() -> GainMatrix {
|
|
// 4 sensors, 2 sources with free orientation (6 columns)
|
|
let data = vec![
|
|
vec![1.0, 0.2, 0.1, 0.3, 0.5, 0.2],
|
|
vec![0.2, 1.0, 0.3, 0.5, 0.3, 0.1],
|
|
vec![0.1, 0.3, 1.0, 0.2, 0.1, 0.5],
|
|
vec![0.3, 0.1, 0.2, 1.0, 0.2, 0.3],
|
|
];
|
|
let names = vec![
|
|
"S1".to_string(),
|
|
"S2".to_string(),
|
|
"S3".to_string(),
|
|
"S4".to_string(),
|
|
];
|
|
GainMatrix::new(data, true, names).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn test_make_lcmv() {
|
|
let gain = create_simple_gain();
|
|
let data_cov = Covariance::identity(4, CovarianceType::Data);
|
|
|
|
let bf = LcmvBeamformer::make_lcmv(&gain, &data_cov, None, 0.05).unwrap();
|
|
|
|
assert_eq!(bf.n_sources(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_lcmv() {
|
|
let gain = create_simple_gain();
|
|
let data_cov = Covariance::identity(4, CovarianceType::Data);
|
|
|
|
let bf = LcmvBeamformer::make_lcmv(&gain, &data_cov, None, 0.05).unwrap();
|
|
|
|
let data = vec![
|
|
vec![1.0, 0.0],
|
|
vec![0.0, 1.0],
|
|
vec![0.5, 0.5],
|
|
vec![0.2, 0.8],
|
|
];
|
|
|
|
let stc = bf.apply(&data).unwrap();
|
|
assert_eq!(stc.n_sources(), 3);
|
|
assert_eq!(stc.n_times(), 2);
|
|
}
|
|
|
|
// DICS tests
|
|
|
|
#[test]
|
|
fn test_dics_config() {
|
|
let config = DicsConfig::default();
|
|
assert_eq!(config.frequencies, vec![10.0]);
|
|
assert!((config.bandwidth - 4.0).abs() < 1e-10);
|
|
|
|
let config = DicsConfig::alpha(1000.0).with_reg(0.1);
|
|
assert!((config.reg - 0.1).abs() < 1e-10);
|
|
assert!((config.sfreq - 1000.0).abs() < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_csd_from_epochs() {
|
|
// Create simple test epochs
|
|
let n_epochs = 5;
|
|
let n_channels = 4;
|
|
let n_samples = 100;
|
|
|
|
let epochs: Vec<Vec<Vec<f64>>> = (0..n_epochs)
|
|
.map(|_| {
|
|
(0..n_channels)
|
|
.map(|ch| {
|
|
(0..n_samples)
|
|
.map(|t| {
|
|
// Simple sinusoid + noise
|
|
let freq = 10.0;
|
|
let sfreq = 100.0;
|
|
(2.0 * PI * freq * t as f64 / sfreq).sin()
|
|
+ 0.1 * (ch as f64 * t as f64 * 0.01).sin()
|
|
})
|
|
.collect()
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let csd = CrossSpectralDensity::from_epochs(
|
|
&epochs, 100.0, // sfreq
|
|
5.0, // fmin
|
|
15.0, // fmax
|
|
None,
|
|
)
|
|
.unwrap();
|
|
|
|
assert_eq!(csd.n_channels(), 4);
|
|
assert!(!csd.frequencies().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_make_dics_fixed() {
|
|
let gain = create_simple_gain();
|
|
|
|
// Create simple CSD (identity-like for testing)
|
|
let n_channels = 4;
|
|
let csd_data: Vec<Vec<Complex<f64>>> = (0..n_channels)
|
|
.map(|i| {
|
|
(0..n_channels)
|
|
.map(|j| {
|
|
if i == j {
|
|
Complex::new(1.0, 0.0)
|
|
} else {
|
|
Complex::new(0.1, 0.0)
|
|
}
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let csd = CrossSpectralDensity::new(csd_data, vec![10.0], 10);
|
|
|
|
let config = DicsConfig::default();
|
|
let bf = DicsBeamformer::make_dics(&gain, &csd, None, config).unwrap();
|
|
|
|
assert_eq!(bf.n_sources(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_make_dics_free() {
|
|
let gain = create_free_gain();
|
|
|
|
let n_channels = 4;
|
|
let csd_data: Vec<Vec<Complex<f64>>> = (0..n_channels)
|
|
.map(|i| {
|
|
(0..n_channels)
|
|
.map(|j| {
|
|
if i == j {
|
|
Complex::new(1.0, 0.0)
|
|
} else {
|
|
Complex::new(0.1, 0.0)
|
|
}
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let csd = CrossSpectralDensity::new(csd_data, vec![10.0], 10);
|
|
|
|
let config = DicsConfig::default();
|
|
let bf = DicsBeamformer::make_dics(&gain, &csd, None, config).unwrap();
|
|
|
|
assert_eq!(bf.n_sources(), 2); // 6 columns / 3 orientations
|
|
}
|
|
|
|
#[test]
|
|
fn test_dics_power() {
|
|
let gain = create_simple_gain();
|
|
|
|
let n_channels = 4;
|
|
let csd_data: Vec<Vec<Complex<f64>>> = (0..n_channels)
|
|
.map(|i| {
|
|
(0..n_channels)
|
|
.map(|j| {
|
|
if i == j {
|
|
Complex::new(1.0, 0.0)
|
|
} else {
|
|
Complex::new(0.1, 0.01)
|
|
}
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let csd = CrossSpectralDensity::new(csd_data, vec![10.0], 10);
|
|
|
|
let config = DicsConfig::default();
|
|
let bf = DicsBeamformer::make_dics(&gain, &csd, None, config).unwrap();
|
|
|
|
let power = bf.compute_power(&csd).unwrap();
|
|
assert_eq!(power.len(), 3);
|
|
|
|
// All power values should be non-negative
|
|
for p in &power {
|
|
assert!(*p >= 0.0);
|
|
}
|
|
}
|
|
}
|