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