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

412 lines
12 KiB
Rust

//! Core data structures for neuroimaging data.
use crate::{Annotations, ChannelInfo, Events, NeuroResult, SampleRate};
use serde::{Deserialize, Serialize};
/// Main container for continuous neuroimaging data.
///
/// This structure holds the raw or processed signal data along with
/// all associated metadata (channels, events, sampling rate, etc.).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NeuroData {
/// Signal data as a 2D array [n_channels x n_samples]
/// Stored as row-major (each row is a channel's time series)
data: Vec<f64>,
/// Number of channels
n_channels: usize,
/// Number of samples per channel
n_samples: usize,
/// Sampling frequency in Hz
sfreq: SampleRate,
/// Channel information
channels: ChannelInfo,
/// Events/triggers
events: Events,
/// Annotations
annotations: Annotations,
/// First sample number (for time alignment with original recording)
first_sample: usize,
}
impl NeuroData {
/// Create new `NeuroData` from raw data.
///
/// # Arguments
/// * `data` - Flattened data array [n_channels x n_samples] in row-major order
/// * `n_channels` - Number of channels
/// * `sfreq` - Sampling frequency in Hz
/// * `channels` - Channel information
///
/// # Errors
/// Returns error if data length doesn't match n_channels * n_samples
pub fn new(
data: Vec<f64>,
n_channels: usize,
sfreq: SampleRate,
channels: ChannelInfo,
) -> NeuroResult<Self> {
if data.len() % n_channels != 0 {
return Err(crate::NeuroError::DimensionMismatch {
expected: format!("multiple of {n_channels}"),
got: format!("{}", data.len()),
});
}
let n_samples = data.len() / n_channels;
if channels.len() != n_channels {
return Err(crate::NeuroError::DimensionMismatch {
expected: format!("{n_channels} channels"),
got: format!("{} channel info entries", channels.len()),
});
}
Ok(Self {
data,
n_channels,
n_samples,
sfreq,
channels,
events: Events::new(),
annotations: Annotations::new(),
first_sample: 0,
})
}
/// Create empty `NeuroData` with specified dimensions
#[must_use]
pub fn zeros(n_channels: usize, n_samples: usize, sfreq: SampleRate) -> Self {
let mut channels = ChannelInfo::new();
for i in 0..n_channels {
channels.add_channel(crate::Channel::new(
format!("CH{i:03}"),
crate::ChannelType::EegScalp,
));
}
Self {
data: vec![0.0; n_channels * n_samples],
n_channels,
n_samples,
sfreq,
channels,
events: Events::new(),
annotations: Annotations::new(),
first_sample: 0,
}
}
/// Number of channels
#[must_use]
pub fn n_channels(&self) -> usize {
self.n_channels
}
/// Number of samples per channel
#[must_use]
pub fn n_samples(&self) -> usize {
self.n_samples
}
/// Sampling frequency in Hz
#[must_use]
pub fn sfreq(&self) -> SampleRate {
self.sfreq
}
/// Duration in seconds
#[must_use]
pub fn duration(&self) -> f64 {
self.n_samples as f64 / self.sfreq
}
/// First sample number
#[must_use]
pub fn first_sample(&self) -> usize {
self.first_sample
}
/// Last sample number (exclusive)
#[must_use]
pub fn last_sample(&self) -> usize {
self.first_sample + self.n_samples
}
/// Get time vector in seconds
#[must_use]
pub fn times(&self) -> Vec<f64> {
(0..self.n_samples).map(|i| i as f64 / self.sfreq).collect()
}
/// Reference to channel information
#[must_use]
pub fn channels(&self) -> &ChannelInfo {
&self.channels
}
/// Mutable reference to channel information
pub fn channels_mut(&mut self) -> &mut ChannelInfo {
&mut self.channels
}
/// Reference to events
#[must_use]
pub fn events(&self) -> &Events {
&self.events
}
/// Mutable reference to events
pub fn events_mut(&mut self) -> &mut Events {
&mut self.events
}
/// Set events
pub fn set_events(&mut self, events: Events) {
self.events = events;
}
/// Reference to annotations
#[must_use]
pub fn annotations(&self) -> &Annotations {
&self.annotations
}
/// Mutable reference to annotations
pub fn annotations_mut(&mut self) -> &mut Annotations {
&mut self.annotations
}
/// Get raw data as slice (row-major: [ch0_t0, ch0_t1, ..., ch1_t0, ch1_t1, ...])
#[must_use]
pub fn data(&self) -> &[f64] {
&self.data
}
/// Get mutable raw data
pub fn data_mut(&mut self) -> &mut [f64] {
&mut self.data
}
/// Get data for a single channel
#[must_use]
pub fn get_channel(&self, ch: usize) -> Option<&[f64]> {
if ch >= self.n_channels {
return None;
}
let start = ch * self.n_samples;
let end = start + self.n_samples;
Some(&self.data[start..end])
}
/// Get mutable data for a single channel
pub fn get_channel_mut(&mut self, ch: usize) -> Option<&mut [f64]> {
if ch >= self.n_channels {
return None;
}
let start = ch * self.n_samples;
let end = start + self.n_samples;
Some(&mut self.data[start..end])
}
/// Get a single sample value
#[must_use]
pub fn get(&self, ch: usize, sample: usize) -> Option<f64> {
if ch >= self.n_channels || sample >= self.n_samples {
return None;
}
Some(self.data[ch * self.n_samples + sample])
}
/// Set a single sample value
pub fn set(&mut self, ch: usize, sample: usize, value: f64) -> bool {
if ch >= self.n_channels || sample >= self.n_samples {
return false;
}
self.data[ch * self.n_samples + sample] = value;
true
}
/// Get data as 2D vector [n_channels][n_samples]
#[must_use]
pub fn to_2d(&self) -> Vec<Vec<f64>> {
(0..self.n_channels)
.map(|ch| self.get_channel(ch).unwrap().to_vec())
.collect()
}
/// Crop data to a time range
///
/// # Arguments
/// * `tmin` - Start time in seconds
/// * `tmax` - End time in seconds
#[must_use]
pub fn crop(&self, tmin: f64, tmax: f64) -> Self {
let start_sample = (tmin * self.sfreq).round() as usize;
let end_sample = (tmax * self.sfreq).round() as usize;
let start_sample = start_sample.min(self.n_samples);
let end_sample = end_sample.min(self.n_samples).max(start_sample);
let new_n_samples = end_sample - start_sample;
let mut new_data = Vec::with_capacity(self.n_channels * new_n_samples);
for ch in 0..self.n_channels {
let ch_data = self.get_channel(ch).unwrap();
new_data.extend_from_slice(&ch_data[start_sample..end_sample]);
}
Self {
data: new_data,
n_channels: self.n_channels,
n_samples: new_n_samples,
sfreq: self.sfreq,
channels: self.channels.clone(),
events: Events::new(), // Events would need to be adjusted
annotations: Annotations::new(),
first_sample: self.first_sample + start_sample,
}
}
/// Pick subset of channels
pub fn pick_channels(&self, indices: &[usize]) -> NeuroResult<Self> {
// Validate indices
for &idx in indices {
if idx >= self.n_channels {
return Err(crate::NeuroError::Channel(format!(
"Channel index {idx} out of range (max {})",
self.n_channels - 1
)));
}
}
let new_n_channels = indices.len();
let mut new_data = Vec::with_capacity(new_n_channels * self.n_samples);
let mut new_channels = ChannelInfo::new();
for &idx in indices {
new_data.extend_from_slice(self.get_channel(idx).unwrap());
new_channels.add_channel(self.channels[idx].clone());
}
Ok(Self {
data: new_data,
n_channels: new_n_channels,
n_samples: self.n_samples,
sfreq: self.sfreq,
channels: new_channels,
events: self.events.clone(),
annotations: self.annotations.clone(),
first_sample: self.first_sample,
})
}
/// Apply a function to each sample in-place
pub fn apply<F>(&mut self, f: F)
where
F: Fn(f64) -> f64,
{
for x in &mut self.data {
*x = f(*x);
}
}
/// Apply a function to each channel in-place
pub fn apply_per_channel<F>(&mut self, f: F)
where
F: Fn(&mut [f64]),
{
for ch in 0..self.n_channels {
let start = ch * self.n_samples;
let end = start + self.n_samples;
f(&mut self.data[start..end]);
}
}
/// Compute mean across all samples for each channel
#[must_use]
pub fn mean_per_channel(&self) -> Vec<f64> {
(0..self.n_channels)
.map(|ch| {
let data = self.get_channel(ch).unwrap();
data.iter().sum::<f64>() / data.len() as f64
})
.collect()
}
/// Compute standard deviation for each channel
#[must_use]
pub fn std_per_channel(&self) -> Vec<f64> {
let means = self.mean_per_channel();
(0..self.n_channels)
.map(|ch| {
let data = self.get_channel(ch).unwrap();
let mean = means[ch];
let variance =
data.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (data.len() - 1) as f64;
variance.sqrt()
})
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Channel, ChannelType};
fn create_test_data() -> NeuroData {
let n_channels = 3;
let n_samples = 100;
let sfreq = 1000.0;
// Create test data: each channel has values 0, 1, 2, ..., 99
let mut data = Vec::with_capacity(n_channels * n_samples);
for _ch in 0..n_channels {
for s in 0..n_samples {
data.push(s as f64);
}
}
let mut channels = ChannelInfo::new();
channels.add_channel(Channel::new("Ch1", ChannelType::EegScalp));
channels.add_channel(Channel::new("Ch2", ChannelType::EegScalp));
channels.add_channel(Channel::new("Ch3", ChannelType::EegScalp));
NeuroData::new(data, n_channels, sfreq, channels).unwrap()
}
#[test]
fn test_neuro_data_creation() {
let data = create_test_data();
assert_eq!(data.n_channels(), 3);
assert_eq!(data.n_samples(), 100);
assert_eq!(data.sfreq(), 1000.0);
assert!((data.duration() - 0.1).abs() < 1e-10);
}
#[test]
fn test_get_channel() {
let data = create_test_data();
let ch0 = data.get_channel(0).unwrap();
assert_eq!(ch0.len(), 100);
assert_eq!(ch0[0], 0.0);
assert_eq!(ch0[99], 99.0);
}
#[test]
fn test_crop() {
let data = create_test_data();
let cropped = data.crop(0.01, 0.05); // 10-50 samples at 1000 Hz
assert_eq!(cropped.n_samples(), 40);
assert_eq!(cropped.first_sample(), 10);
}
#[test]
fn test_pick_channels() {
let data = create_test_data();
let picked = data.pick_channels(&[0, 2]).unwrap();
assert_eq!(picked.n_channels(), 2);
assert_eq!(picked.channels().names(), vec!["Ch1", "Ch3"]);
}
}