567 lines
16 KiB
Rust
567 lines
16 KiB
Rust
//! Epoched data for event-related analysis.
|
|
|
|
use crate::{ChannelInfo, Event, Events, NeuroData, NeuroResult, SampleRate};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
|
|
/// Configuration for creating epochs
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct EpochsConfig {
|
|
/// Start time relative to event (negative = before event)
|
|
pub tmin: f64,
|
|
/// End time relative to event
|
|
pub tmax: f64,
|
|
/// Baseline correction window (None = no baseline correction)
|
|
pub baseline: Option<(f64, f64)>,
|
|
/// Reject epochs where any channel exceeds this threshold
|
|
pub reject: Option<HashMap<String, f64>>,
|
|
/// Decimate by this factor
|
|
pub decim: usize,
|
|
}
|
|
|
|
impl Default for EpochsConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
tmin: -0.2,
|
|
tmax: 0.5,
|
|
baseline: Some((-0.2, 0.0)),
|
|
reject: None,
|
|
decim: 1,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl EpochsConfig {
|
|
/// Create a new epochs configuration
|
|
#[must_use]
|
|
pub fn new(tmin: f64, tmax: f64) -> Self {
|
|
Self {
|
|
tmin,
|
|
tmax,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
/// Set baseline correction window
|
|
#[must_use]
|
|
pub fn with_baseline(mut self, start: f64, end: f64) -> Self {
|
|
self.baseline = Some((start, end));
|
|
self
|
|
}
|
|
|
|
/// Disable baseline correction
|
|
#[must_use]
|
|
pub fn no_baseline(mut self) -> Self {
|
|
self.baseline = None;
|
|
self
|
|
}
|
|
|
|
/// Set rejection threshold for a channel type
|
|
pub fn reject_by_type(&mut self, ch_type: &str, threshold: f64) {
|
|
self.reject
|
|
.get_or_insert_with(HashMap::new)
|
|
.insert(ch_type.to_string(), threshold);
|
|
}
|
|
}
|
|
|
|
/// Collection of epochs (segmented data around events)
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Epochs {
|
|
/// Data: [n_epochs x n_channels x n_times]
|
|
data: Vec<f64>,
|
|
/// Number of epochs
|
|
n_epochs: usize,
|
|
/// Number of channels
|
|
n_channels: usize,
|
|
/// Number of time points per epoch
|
|
n_times: usize,
|
|
/// Sampling frequency
|
|
sfreq: SampleRate,
|
|
/// Start time relative to event
|
|
tmin: f64,
|
|
/// End time relative to event
|
|
tmax: f64,
|
|
/// Channel information
|
|
channels: ChannelInfo,
|
|
/// Events that define each epoch
|
|
events: Vec<Event>,
|
|
/// Indices of dropped epochs
|
|
dropped: Vec<usize>,
|
|
/// Selection (which epochs are currently active)
|
|
selection: Vec<usize>,
|
|
}
|
|
|
|
impl Epochs {
|
|
/// Create epochs from continuous data
|
|
///
|
|
/// # Arguments
|
|
/// * `raw` - Continuous data
|
|
/// * `events` - Events to epoch around
|
|
/// * `event_ids` - Which event values to include (None = all)
|
|
/// * `config` - Epoching configuration
|
|
pub fn from_data(
|
|
raw: &NeuroData,
|
|
events: &Events,
|
|
event_ids: Option<&[i32]>,
|
|
config: &EpochsConfig,
|
|
) -> NeuroResult<Self> {
|
|
let sfreq = raw.sfreq();
|
|
let n_channels = raw.n_channels();
|
|
|
|
// Calculate sample indices
|
|
let tmin_samples = (config.tmin * sfreq).round() as i64;
|
|
let tmax_samples = (config.tmax * sfreq).round() as i64;
|
|
let n_times = (tmax_samples - tmin_samples) as usize;
|
|
|
|
// Filter events by event_ids if specified
|
|
let selected_events: Vec<&Event> = events
|
|
.iter()
|
|
.filter(|e| event_ids.is_none_or(|ids| ids.contains(&e.value)))
|
|
.collect();
|
|
|
|
let n_epochs = selected_events.len();
|
|
if n_epochs == 0 {
|
|
return Err(crate::NeuroError::Event(
|
|
"No events found matching the criteria".to_string(),
|
|
));
|
|
}
|
|
|
|
// Allocate data
|
|
let mut data = vec![0.0; n_epochs * n_channels * n_times];
|
|
let mut epoch_events = Vec::with_capacity(n_epochs);
|
|
let mut dropped = Vec::new();
|
|
|
|
// Extract epochs
|
|
for (epoch_idx, event) in selected_events.iter().enumerate() {
|
|
let event_sample = event.sample as i64;
|
|
let start_sample = event_sample + tmin_samples;
|
|
let end_sample = event_sample + tmax_samples;
|
|
|
|
// Check if epoch is within data bounds
|
|
if start_sample < 0 || end_sample as usize > raw.n_samples() {
|
|
dropped.push(epoch_idx);
|
|
continue;
|
|
}
|
|
|
|
let start_sample = start_sample as usize;
|
|
|
|
// Copy data for this epoch
|
|
for ch in 0..n_channels {
|
|
let raw_ch = raw.get_channel(ch).unwrap();
|
|
let epoch_offset = epoch_idx * n_channels * n_times + ch * n_times;
|
|
|
|
for t in 0..n_times {
|
|
data[epoch_offset + t] = raw_ch[start_sample + t];
|
|
}
|
|
}
|
|
|
|
epoch_events.push(**event);
|
|
}
|
|
|
|
let mut epochs = Self {
|
|
data,
|
|
n_epochs,
|
|
n_channels,
|
|
n_times,
|
|
sfreq,
|
|
tmin: config.tmin,
|
|
tmax: config.tmax,
|
|
channels: raw.channels().clone(),
|
|
events: epoch_events,
|
|
dropped,
|
|
selection: (0..n_epochs).collect(),
|
|
};
|
|
|
|
// Apply baseline correction
|
|
if let Some((bmin, bmax)) = config.baseline {
|
|
epochs.apply_baseline(bmin, bmax)?;
|
|
}
|
|
|
|
Ok(epochs)
|
|
}
|
|
|
|
/// Number of epochs
|
|
#[must_use]
|
|
pub fn n_epochs(&self) -> usize {
|
|
self.selection.len()
|
|
}
|
|
|
|
/// Total number of epochs (including dropped)
|
|
#[must_use]
|
|
pub fn n_epochs_total(&self) -> usize {
|
|
self.n_epochs
|
|
}
|
|
|
|
/// Number of channels
|
|
#[must_use]
|
|
pub fn n_channels(&self) -> usize {
|
|
self.n_channels
|
|
}
|
|
|
|
/// Number of time points per epoch
|
|
#[must_use]
|
|
pub fn n_times(&self) -> usize {
|
|
self.n_times
|
|
}
|
|
|
|
/// Sampling frequency
|
|
#[must_use]
|
|
pub fn sfreq(&self) -> SampleRate {
|
|
self.sfreq
|
|
}
|
|
|
|
/// Start time relative to event
|
|
#[must_use]
|
|
pub fn tmin(&self) -> f64 {
|
|
self.tmin
|
|
}
|
|
|
|
/// End time relative to event
|
|
#[must_use]
|
|
pub fn tmax(&self) -> f64 {
|
|
self.tmax
|
|
}
|
|
|
|
/// Get time vector
|
|
#[must_use]
|
|
pub fn times(&self) -> Vec<f64> {
|
|
(0..self.n_times)
|
|
.map(|i| self.tmin + i as f64 / self.sfreq)
|
|
.collect()
|
|
}
|
|
|
|
/// Reference to channel information
|
|
#[must_use]
|
|
pub fn channels(&self) -> &ChannelInfo {
|
|
&self.channels
|
|
}
|
|
|
|
/// Get events
|
|
#[must_use]
|
|
pub fn events(&self) -> &[Event] {
|
|
&self.events
|
|
}
|
|
|
|
/// Get indices of dropped epochs
|
|
#[must_use]
|
|
pub fn dropped(&self) -> &[usize] {
|
|
&self.dropped
|
|
}
|
|
|
|
/// Get data for a single epoch [n_channels x n_times]
|
|
#[must_use]
|
|
pub fn get_epoch(&self, epoch: usize) -> Option<Vec<Vec<f64>>> {
|
|
if epoch >= self.n_epochs {
|
|
return None;
|
|
}
|
|
|
|
let epoch_offset = epoch * self.n_channels * self.n_times;
|
|
let mut result = Vec::with_capacity(self.n_channels);
|
|
|
|
for ch in 0..self.n_channels {
|
|
let ch_offset = epoch_offset + ch * self.n_times;
|
|
result.push(self.data[ch_offset..ch_offset + self.n_times].to_vec());
|
|
}
|
|
|
|
Some(result)
|
|
}
|
|
|
|
/// Get data for a single epoch and channel
|
|
#[must_use]
|
|
pub fn get_epoch_channel(&self, epoch: usize, channel: usize) -> Option<&[f64]> {
|
|
if epoch >= self.n_epochs || channel >= self.n_channels {
|
|
return None;
|
|
}
|
|
|
|
let offset = epoch * self.n_channels * self.n_times + channel * self.n_times;
|
|
Some(&self.data[offset..offset + self.n_times])
|
|
}
|
|
|
|
/// Apply baseline correction in-place
|
|
pub fn apply_baseline(&mut self, bmin: f64, bmax: f64) -> NeuroResult<()> {
|
|
// Convert time to samples
|
|
let bmin_sample = ((bmin - self.tmin) * self.sfreq).round() as usize;
|
|
let bmax_sample = ((bmax - self.tmin) * self.sfreq).round() as usize;
|
|
|
|
if bmin_sample >= self.n_times || bmax_sample > self.n_times {
|
|
return Err(crate::NeuroError::InvalidParameter(format!(
|
|
"Baseline window [{bmin}, {bmax}] outside epoch range [{}, {}]",
|
|
self.tmin, self.tmax
|
|
)));
|
|
}
|
|
|
|
let baseline_len = bmax_sample - bmin_sample;
|
|
|
|
for epoch in 0..self.n_epochs {
|
|
let epoch_offset = epoch * self.n_channels * self.n_times;
|
|
|
|
for ch in 0..self.n_channels {
|
|
let ch_offset = epoch_offset + ch * self.n_times;
|
|
|
|
// Compute baseline mean
|
|
let baseline_sum: f64 = self.data[ch_offset + bmin_sample..ch_offset + bmax_sample]
|
|
.iter()
|
|
.sum();
|
|
let baseline_mean = baseline_sum / baseline_len as f64;
|
|
|
|
// Subtract baseline
|
|
for t in 0..self.n_times {
|
|
self.data[ch_offset + t] -= baseline_mean;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Compute evoked response (average across epochs)
|
|
#[must_use]
|
|
pub fn average(&self) -> Evoked {
|
|
let n_active = self.selection.len();
|
|
let mut averaged = vec![0.0; self.n_channels * self.n_times];
|
|
|
|
for &epoch_idx in &self.selection {
|
|
let epoch_offset = epoch_idx * self.n_channels * self.n_times;
|
|
|
|
for ch in 0..self.n_channels {
|
|
let ch_offset = ch * self.n_times;
|
|
let epoch_ch_offset = epoch_offset + ch_offset;
|
|
|
|
for t in 0..self.n_times {
|
|
averaged[ch_offset + t] += self.data[epoch_ch_offset + t];
|
|
}
|
|
}
|
|
}
|
|
|
|
// Divide by number of epochs
|
|
for x in &mut averaged {
|
|
*x /= n_active as f64;
|
|
}
|
|
|
|
Evoked {
|
|
data: averaged,
|
|
n_channels: self.n_channels,
|
|
n_times: self.n_times,
|
|
sfreq: self.sfreq,
|
|
tmin: self.tmin,
|
|
tmax: self.tmax,
|
|
channels: self.channels.clone(),
|
|
n_averaged: n_active,
|
|
comment: String::new(),
|
|
}
|
|
}
|
|
|
|
/// Drop epochs where any channel exceeds threshold
|
|
pub fn drop_bad(&mut self, threshold: f64) {
|
|
let mut new_selection = Vec::new();
|
|
|
|
for &epoch_idx in &self.selection {
|
|
let epoch_offset = epoch_idx * self.n_channels * self.n_times;
|
|
let mut is_bad = false;
|
|
|
|
'outer: for ch in 0..self.n_channels {
|
|
let ch_offset = epoch_offset + ch * self.n_times;
|
|
for t in 0..self.n_times {
|
|
if self.data[ch_offset + t].abs() > threshold {
|
|
is_bad = true;
|
|
break 'outer;
|
|
}
|
|
}
|
|
}
|
|
|
|
if !is_bad {
|
|
new_selection.push(epoch_idx);
|
|
} else {
|
|
self.dropped.push(epoch_idx);
|
|
}
|
|
}
|
|
|
|
self.selection = new_selection;
|
|
}
|
|
}
|
|
|
|
/// Evoked response (averaged epochs)
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Evoked {
|
|
/// Data: [n_channels x n_times]
|
|
data: Vec<f64>,
|
|
/// Number of channels
|
|
n_channels: usize,
|
|
/// Number of time points
|
|
n_times: usize,
|
|
/// Sampling frequency
|
|
sfreq: SampleRate,
|
|
/// Start time relative to event
|
|
tmin: f64,
|
|
/// End time relative to event
|
|
tmax: f64,
|
|
/// Channel information
|
|
channels: ChannelInfo,
|
|
/// Number of epochs averaged
|
|
n_averaged: usize,
|
|
/// Comment/description
|
|
comment: String,
|
|
}
|
|
|
|
impl Evoked {
|
|
/// Number of channels
|
|
#[must_use]
|
|
pub fn n_channels(&self) -> usize {
|
|
self.n_channels
|
|
}
|
|
|
|
/// Number of time points
|
|
#[must_use]
|
|
pub fn n_times(&self) -> usize {
|
|
self.n_times
|
|
}
|
|
|
|
/// Sampling frequency
|
|
#[must_use]
|
|
pub fn sfreq(&self) -> SampleRate {
|
|
self.sfreq
|
|
}
|
|
|
|
/// Get time vector
|
|
#[must_use]
|
|
pub fn times(&self) -> Vec<f64> {
|
|
(0..self.n_times)
|
|
.map(|i| self.tmin + i as f64 / self.sfreq)
|
|
.collect()
|
|
}
|
|
|
|
/// Reference to channel information
|
|
#[must_use]
|
|
pub fn channels(&self) -> &ChannelInfo {
|
|
&self.channels
|
|
}
|
|
|
|
/// Number of epochs averaged
|
|
#[must_use]
|
|
pub fn n_averaged(&self) -> usize {
|
|
self.n_averaged
|
|
}
|
|
|
|
/// Get raw data as slice
|
|
#[must_use]
|
|
pub fn data(&self) -> &[f64] {
|
|
&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_times;
|
|
Some(&self.data[start..start + self.n_times])
|
|
}
|
|
|
|
/// Get data as 2D vector
|
|
#[must_use]
|
|
pub fn to_2d(&self) -> Vec<Vec<f64>> {
|
|
(0..self.n_channels)
|
|
.map(|ch| self.get_channel(ch).unwrap().to_vec())
|
|
.collect()
|
|
}
|
|
|
|
/// Find peak amplitude and latency for a channel
|
|
#[must_use]
|
|
pub fn peak(&self, ch: usize, tmin: Option<f64>, tmax: Option<f64>) -> Option<(f64, f64)> {
|
|
let ch_data = self.get_channel(ch)?;
|
|
let times = self.times();
|
|
|
|
let tmin = tmin.unwrap_or(self.tmin);
|
|
let tmax = tmax.unwrap_or(self.tmax);
|
|
|
|
let mut max_val = f64::NEG_INFINITY;
|
|
let mut max_time = 0.0;
|
|
|
|
for (i, (&val, &t)) in ch_data.iter().zip(times.iter()).enumerate() {
|
|
if t >= tmin && t <= tmax {
|
|
let abs_val = val.abs();
|
|
if abs_val > max_val {
|
|
max_val = abs_val;
|
|
max_time = t;
|
|
if val < 0.0 {
|
|
max_val = -max_val;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if max_val.is_finite() {
|
|
Some((max_val, max_time))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::{Channel, ChannelType, Event};
|
|
|
|
fn create_test_epochs() -> Epochs {
|
|
// Create simple test data: 2 epochs, 2 channels, 100 samples each
|
|
let n_epochs = 2;
|
|
let n_channels = 2;
|
|
let n_times = 100;
|
|
let sfreq = 1000.0;
|
|
|
|
let mut data = Vec::with_capacity(n_epochs * n_channels * n_times);
|
|
for epoch in 0..n_epochs {
|
|
for ch in 0..n_channels {
|
|
for t in 0..n_times {
|
|
data.push((epoch * 100 + ch * 10 + t) as f64);
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut channels = ChannelInfo::new();
|
|
channels.add_channel(Channel::new("Ch1", ChannelType::EegScalp));
|
|
channels.add_channel(Channel::new("Ch2", ChannelType::EegScalp));
|
|
|
|
Epochs {
|
|
data,
|
|
n_epochs,
|
|
n_channels,
|
|
n_times,
|
|
sfreq,
|
|
tmin: -0.05,
|
|
tmax: 0.05,
|
|
channels,
|
|
events: vec![Event::new(500, 1), Event::new(1500, 1)],
|
|
dropped: Vec::new(),
|
|
selection: vec![0, 1],
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_epochs_average() {
|
|
let epochs = create_test_epochs();
|
|
let evoked = epochs.average();
|
|
|
|
assert_eq!(evoked.n_channels(), 2);
|
|
assert_eq!(evoked.n_times(), 100);
|
|
assert_eq!(evoked.n_averaged(), 2);
|
|
|
|
// Average of ch0 at t0: (0 + 100) / 2 = 50
|
|
// Epoch 0: 0*100 + 0*10 + 0 = 0
|
|
// Epoch 1: 1*100 + 0*10 + 0 = 100
|
|
let ch0 = evoked.get_channel(0).unwrap();
|
|
assert!((ch0[0] - 50.0).abs() < 1e-10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_epochs_config() {
|
|
let config = EpochsConfig::new(-0.1, 0.4).with_baseline(-0.1, 0.0);
|
|
|
|
assert_eq!(config.tmin, -0.1);
|
|
assert_eq!(config.tmax, 0.4);
|
|
assert_eq!(config.baseline, Some((-0.1, 0.0)));
|
|
}
|
|
}
|