Files
rustytorch/crates/specialized/rtx-neuro-inverse/src/source_estimate.rs
T
2026-03-04 00:08:42 +00:00

245 lines
6.6 KiB
Rust

//! Source estimate data structure.
//!
//! Represents brain source activity estimated from MEG/EEG data.
use crate::InverseResult;
/// Source time course estimate
#[derive(Debug, Clone)]
pub struct SourceEstimate {
/// Source data [n_source_columns x n_times]
data: Vec<Vec<f64>>,
/// Time points in seconds
times: Vec<f64>,
/// Source indices (vertex indices)
vertices: Vec<usize>,
/// Whether sources have free orientation
free_orientation: bool,
}
impl SourceEstimate {
/// Create a new source estimate
pub fn new(
data: Vec<Vec<f64>>,
times: Vec<f64>,
vertices: Vec<usize>,
free_orientation: bool,
) -> Self {
Self {
data,
times,
vertices,
free_orientation,
}
}
/// Get number of sources
pub fn n_sources(&self) -> usize {
if self.free_orientation {
self.data.len() / 3
} else {
self.data.len()
}
}
/// Get number of time points
pub fn n_times(&self) -> usize {
self.times.len()
}
/// Get the source data
pub fn data(&self) -> &Vec<Vec<f64>> {
&self.data
}
/// Get the time points
pub fn times(&self) -> &[f64] {
&self.times
}
/// Get vertex indices
pub fn vertices(&self) -> &[usize] {
&self.vertices
}
/// Check if sources have free orientation
pub fn is_free_orientation(&self) -> bool {
self.free_orientation
}
/// Get source activity at a specific time
pub fn at_time(&self, time_idx: usize) -> Option<Vec<f64>> {
if time_idx >= self.n_times() {
return None;
}
Some(self.data.iter().map(|src| src[time_idx]).collect())
}
/// Get time course for a specific source
pub fn get_source(&self, source_idx: usize) -> Option<&[f64]> {
if self.free_orientation {
// Return combined magnitude
None // For free orientation, use get_source_vector
} else {
self.data.get(source_idx).map(std::vec::Vec::as_slice)
}
}
/// Get vector time course for a free-orientation source
pub fn get_source_vector(&self, source_idx: usize) -> Option<[&[f64]; 3]> {
if !self.free_orientation {
return None;
}
let idx = source_idx * 3;
if idx + 2 >= self.data.len() {
return None;
}
Some([&self.data[idx], &self.data[idx + 1], &self.data[idx + 2]])
}
/// Compute the magnitude time course for each source
///
/// For fixed orientation: absolute value
/// For free orientation: sqrt(x^2 + y^2 + z^2)
pub fn magnitude(&self) -> Vec<Vec<f64>> {
let n_sources = self.n_sources();
let n_times = self.n_times();
if self.free_orientation {
(0..n_sources)
.map(|src| {
let idx = src * 3;
(0..n_times)
.map(|t| {
let x = self.data[idx][t];
let y = self.data[idx + 1][t];
let z = self.data[idx + 2][t];
(x * x + y * y + z * z).sqrt()
})
.collect()
})
.collect()
} else {
self.data
.iter()
.map(|src| src.iter().map(|&v| v.abs()).collect())
.collect()
}
}
/// Get the mean activity across time
pub fn mean(&self) -> Vec<f64> {
self.data
.iter()
.map(|src| src.iter().sum::<f64>() / src.len() as f64)
.collect()
}
/// Get the peak activity for each source
pub fn peak(&self) -> Vec<(f64, usize)> {
self.data
.iter()
.map(|src| {
src.iter()
.enumerate()
.map(|(i, &v)| (v.abs(), i))
.max_by(|a, b| a.0.partial_cmp(&b.0).unwrap())
.unwrap_or((0.0, 0))
})
.collect()
}
/// Crop to a time window
pub fn crop(&self, tmin: f64, tmax: f64) -> InverseResult<Self> {
let start_idx = self.times.iter().position(|&t| t >= tmin).unwrap_or(0);
let end_idx = self
.times
.iter()
.rposition(|&t| t <= tmax)
.unwrap_or(self.n_times() - 1)
+ 1;
let new_data: Vec<Vec<f64>> = self
.data
.iter()
.map(|src| src[start_idx..end_idx].to_vec())
.collect();
let new_times = self.times[start_idx..end_idx].to_vec();
Ok(Self {
data: new_data,
times: new_times,
vertices: self.vertices.clone(),
free_orientation: self.free_orientation,
})
}
/// Extract sources above a threshold
pub fn threshold(&self, thresh: f64) -> Vec<usize> {
let mag = self.magnitude();
mag.iter()
.enumerate()
.filter(|(_, src)| src.iter().any(|&v| v > thresh))
.map(|(i, _)| i)
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_source_estimate() {
let data = vec![vec![1.0, 2.0, 3.0], vec![0.5, 1.5, 2.5]];
let times = vec![0.0, 0.1, 0.2];
let vertices = vec![0, 1];
let stc = SourceEstimate::new(data, times, vertices, false);
assert_eq!(stc.n_sources(), 2);
assert_eq!(stc.n_times(), 3);
}
#[test]
fn test_magnitude_fixed() {
let data = vec![vec![-1.0, 2.0, -3.0], vec![0.5, -1.5, 2.5]];
let stc = SourceEstimate::new(data, vec![0.0, 0.1, 0.2], vec![0, 1], false);
let mag = stc.magnitude();
assert!((mag[0][0] - 1.0).abs() < 1e-10);
assert!((mag[0][2] - 3.0).abs() < 1e-10);
}
#[test]
fn test_magnitude_free() {
// 1 source with 3 orientations
let data = vec![
vec![3.0], // x
vec![4.0], // y
vec![0.0], // z
];
let stc = SourceEstimate::new(data, vec![0.0], vec![0], true);
let mag = stc.magnitude();
assert!((mag[0][0] - 5.0).abs() < 1e-10); // sqrt(9+16) = 5
}
#[test]
fn test_crop() {
let data = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0]];
let times = vec![0.0, 0.1, 0.2, 0.3, 0.4];
let stc = SourceEstimate::new(data, times, vec![0], false);
let cropped = stc.crop(0.1, 0.3).unwrap();
assert_eq!(cropped.n_times(), 3);
assert!((cropped.times()[0] - 0.1).abs() < 1e-10);
}
}