297 lines
8.0 KiB
Rust
297 lines
8.0 KiB
Rust
//! Ring buffer for streaming MEG/EEG data with timestamp tracking.
|
|
//!
|
|
//! Provides lock-free concurrent access for real-time applications.
|
|
|
|
use parking_lot::RwLock;
|
|
use std::collections::VecDeque;
|
|
use std::sync::Arc;
|
|
|
|
/// A timestamped sample
|
|
#[derive(Debug, Clone)]
|
|
pub struct TimestampedSample {
|
|
/// Sample data [n_channels]
|
|
pub data: Vec<f64>,
|
|
/// Timestamp in seconds (LSL time)
|
|
pub timestamp: f64,
|
|
/// Sample index
|
|
pub index: u64,
|
|
}
|
|
|
|
/// Thread-safe ring buffer for streaming data
|
|
#[derive(Debug)]
|
|
pub struct RingBuffer {
|
|
/// Data storage
|
|
samples: RwLock<VecDeque<TimestampedSample>>,
|
|
/// Maximum capacity in samples
|
|
capacity: usize,
|
|
/// Number of channels
|
|
n_channels: usize,
|
|
/// Sample rate
|
|
sample_rate: f64,
|
|
/// Total samples received
|
|
total_samples: RwLock<u64>,
|
|
}
|
|
|
|
impl RingBuffer {
|
|
/// Create a new ring buffer
|
|
///
|
|
/// # Arguments
|
|
/// * `duration_sec` - Buffer duration in seconds
|
|
/// * `sample_rate` - Sample rate in Hz
|
|
/// * `n_channels` - Number of channels
|
|
pub fn new(duration_sec: f64, sample_rate: f64, n_channels: usize) -> Self {
|
|
let capacity = (duration_sec * sample_rate).ceil() as usize;
|
|
Self {
|
|
samples: RwLock::new(VecDeque::with_capacity(capacity)),
|
|
capacity,
|
|
n_channels,
|
|
sample_rate,
|
|
total_samples: RwLock::new(0),
|
|
}
|
|
}
|
|
|
|
/// Push a single sample
|
|
pub fn push(&self, data: Vec<f64>, timestamp: f64) {
|
|
let mut samples = self.samples.write();
|
|
let mut total = self.total_samples.write();
|
|
|
|
let sample = TimestampedSample {
|
|
data,
|
|
timestamp,
|
|
index: *total,
|
|
};
|
|
|
|
if samples.len() >= self.capacity {
|
|
samples.pop_front();
|
|
}
|
|
samples.push_back(sample);
|
|
*total += 1;
|
|
}
|
|
|
|
/// Push multiple samples (batch)
|
|
pub fn push_batch(&self, data: &[Vec<f64>], timestamps: &[f64]) {
|
|
let mut samples = self.samples.write();
|
|
let mut total = self.total_samples.write();
|
|
|
|
for (d, &t) in data.iter().zip(timestamps.iter()) {
|
|
let sample = TimestampedSample {
|
|
data: d.clone(),
|
|
timestamp: t,
|
|
index: *total,
|
|
};
|
|
|
|
if samples.len() >= self.capacity {
|
|
samples.pop_front();
|
|
}
|
|
samples.push_back(sample);
|
|
*total += 1;
|
|
}
|
|
}
|
|
|
|
/// Get the most recent samples
|
|
///
|
|
/// # Arguments
|
|
/// * `duration_sec` - How many seconds of data to retrieve
|
|
///
|
|
/// # Returns
|
|
/// (data [n_samples x n_channels], timestamps)
|
|
pub fn get_recent(&self, duration_sec: f64) -> (Vec<Vec<f64>>, Vec<f64>) {
|
|
let samples = self.samples.read();
|
|
let n_samples = (duration_sec * self.sample_rate).ceil() as usize;
|
|
let n_available = samples.len();
|
|
let n_to_get = n_samples.min(n_available);
|
|
|
|
let start_idx = n_available.saturating_sub(n_to_get);
|
|
|
|
let data: Vec<Vec<f64>> = samples
|
|
.iter()
|
|
.skip(start_idx)
|
|
.map(|s| s.data.clone())
|
|
.collect();
|
|
|
|
let timestamps: Vec<f64> = samples
|
|
.iter()
|
|
.skip(start_idx)
|
|
.map(|s| s.timestamp)
|
|
.collect();
|
|
|
|
(data, timestamps)
|
|
}
|
|
|
|
/// Get data as a matrix [n_channels x n_times]
|
|
pub fn get_recent_matrix(&self, duration_sec: f64) -> Vec<Vec<f64>> {
|
|
let (samples, _) = self.get_recent(duration_sec);
|
|
if samples.is_empty() {
|
|
return vec![vec![]; self.n_channels];
|
|
}
|
|
|
|
// Transpose: [n_samples x n_channels] -> [n_channels x n_samples]
|
|
let n_times = samples.len();
|
|
let mut matrix = vec![vec![0.0; n_times]; self.n_channels];
|
|
|
|
for (t, sample) in samples.iter().enumerate() {
|
|
for (ch, &val) in sample.iter().enumerate() {
|
|
if ch < self.n_channels {
|
|
matrix[ch][t] = val;
|
|
}
|
|
}
|
|
}
|
|
|
|
matrix
|
|
}
|
|
|
|
/// Get the latest timestamp
|
|
pub fn latest_timestamp(&self) -> Option<f64> {
|
|
let samples = self.samples.read();
|
|
samples.back().map(|s| s.timestamp)
|
|
}
|
|
|
|
/// Get the oldest timestamp
|
|
pub fn oldest_timestamp(&self) -> Option<f64> {
|
|
let samples = self.samples.read();
|
|
samples.front().map(|s| s.timestamp)
|
|
}
|
|
|
|
/// Get current buffer fill level (0.0 - 1.0)
|
|
pub fn fill_level(&self) -> f64 {
|
|
let samples = self.samples.read();
|
|
samples.len() as f64 / self.capacity as f64
|
|
}
|
|
|
|
/// Get number of samples currently in buffer
|
|
pub fn len(&self) -> usize {
|
|
self.samples.read().len()
|
|
}
|
|
|
|
/// Check if buffer is empty
|
|
pub fn is_empty(&self) -> bool {
|
|
self.samples.read().is_empty()
|
|
}
|
|
|
|
/// Clear the buffer
|
|
pub fn clear(&self) {
|
|
self.samples.write().clear();
|
|
}
|
|
|
|
/// Get total samples received (including dropped)
|
|
pub fn total_samples(&self) -> u64 {
|
|
*self.total_samples.read()
|
|
}
|
|
|
|
/// Get buffer duration in seconds
|
|
pub fn duration(&self) -> f64 {
|
|
self.capacity as f64 / self.sample_rate
|
|
}
|
|
|
|
/// Get number of channels
|
|
pub fn n_channels(&self) -> usize {
|
|
self.n_channels
|
|
}
|
|
|
|
/// Get sample rate
|
|
pub fn sample_rate(&self) -> f64 {
|
|
self.sample_rate
|
|
}
|
|
|
|
/// Get capacity in samples
|
|
pub fn capacity(&self) -> usize {
|
|
self.capacity
|
|
}
|
|
}
|
|
|
|
/// Shared ring buffer handle
|
|
pub type SharedRingBuffer = Arc<RingBuffer>;
|
|
|
|
/// Create a new shared ring buffer
|
|
pub fn create_shared_buffer(
|
|
duration_sec: f64,
|
|
sample_rate: f64,
|
|
n_channels: usize,
|
|
) -> SharedRingBuffer {
|
|
Arc::new(RingBuffer::new(duration_sec, sample_rate, n_channels))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_ring_buffer_push() {
|
|
let buffer = RingBuffer::new(1.0, 100.0, 4);
|
|
|
|
buffer.push(vec![1.0, 2.0, 3.0, 4.0], 0.0);
|
|
assert_eq!(buffer.len(), 1);
|
|
assert_eq!(buffer.total_samples(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ring_buffer_overflow() {
|
|
let buffer = RingBuffer::new(0.1, 10.0, 2); // 1 sample capacity
|
|
|
|
// Push 3 samples
|
|
buffer.push(vec![1.0, 2.0], 0.0);
|
|
buffer.push(vec![3.0, 4.0], 0.1);
|
|
buffer.push(vec![5.0, 6.0], 0.2);
|
|
|
|
// Should only keep capacity worth
|
|
assert!(buffer.len() <= buffer.capacity());
|
|
assert_eq!(buffer.total_samples(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_recent() {
|
|
let buffer = RingBuffer::new(1.0, 100.0, 2);
|
|
|
|
for i in 0..50 {
|
|
buffer.push(vec![i as f64, (i * 2) as f64], i as f64 * 0.01);
|
|
}
|
|
|
|
let (data, timestamps) = buffer.get_recent(0.2); // 20 samples
|
|
assert_eq!(data.len(), 20);
|
|
assert_eq!(timestamps.len(), 20);
|
|
|
|
// Check we got the most recent samples
|
|
assert_eq!(data[19][0], 49.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_recent_matrix() {
|
|
let buffer = RingBuffer::new(1.0, 100.0, 3);
|
|
|
|
buffer.push(vec![1.0, 2.0, 3.0], 0.0);
|
|
buffer.push(vec![4.0, 5.0, 6.0], 0.01);
|
|
|
|
let matrix = buffer.get_recent_matrix(0.1);
|
|
|
|
assert_eq!(matrix.len(), 3); // 3 channels
|
|
assert_eq!(matrix[0].len(), 2); // 2 time points
|
|
assert_eq!(matrix[0][0], 1.0);
|
|
assert_eq!(matrix[0][1], 4.0);
|
|
assert_eq!(matrix[1][0], 2.0);
|
|
assert_eq!(matrix[1][1], 5.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_batch_push() {
|
|
let buffer = RingBuffer::new(1.0, 100.0, 2);
|
|
|
|
let data = vec![vec![1.0, 2.0], vec![3.0, 4.0], vec![5.0, 6.0]];
|
|
let timestamps = vec![0.0, 0.01, 0.02];
|
|
|
|
buffer.push_batch(&data, ×tamps);
|
|
|
|
assert_eq!(buffer.len(), 3);
|
|
assert_eq!(buffer.total_samples(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_fill_level() {
|
|
let buffer = RingBuffer::new(0.1, 10.0, 2); // 1 sample capacity
|
|
|
|
assert_eq!(buffer.fill_level(), 0.0);
|
|
|
|
buffer.push(vec![1.0, 2.0], 0.0);
|
|
assert!(buffer.fill_level() > 0.0);
|
|
}
|
|
}
|