608 lines
18 KiB
Rust
608 lines
18 KiB
Rust
//! Real-time streaming pipeline for MEG/EEG source localization.
|
|
//!
|
|
//! Orchestrates the data flow from LSL input through filtering, beamforming,
|
|
//! and source estimation with sub-10ms latency.
|
|
|
|
use crate::error::{RealtimeError, RealtimeResult};
|
|
use crate::gpu_beamformer::GpuBeamformer;
|
|
use crate::gpu_filter::{FilterType, GpuFilter, GpuFilterConfig};
|
|
use crate::latency::{LatencyMonitor, LatencyStage, LatencyStats};
|
|
use crate::ring_buffer::{SharedRingBuffer, create_shared_buffer};
|
|
|
|
use crossbeam_channel::{Receiver, Sender, TryRecvError, bounded};
|
|
use parking_lot::RwLock;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
use std::thread::{self, JoinHandle};
|
|
use std::time::{Duration, Instant};
|
|
|
|
/// Pipeline state
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum PipelineState {
|
|
/// Pipeline is created but not running
|
|
Idle,
|
|
/// Pipeline is starting up
|
|
Starting,
|
|
/// Pipeline is running and processing data
|
|
Running,
|
|
/// Pipeline is paused (can resume)
|
|
Paused,
|
|
/// Pipeline is stopping
|
|
Stopping,
|
|
/// Pipeline has stopped
|
|
Stopped,
|
|
/// Pipeline encountered an error
|
|
Error,
|
|
}
|
|
|
|
impl std::fmt::Display for PipelineState {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
Self::Idle => write!(f, "Idle"),
|
|
Self::Starting => write!(f, "Starting"),
|
|
Self::Running => write!(f, "Running"),
|
|
Self::Paused => write!(f, "Paused"),
|
|
Self::Stopping => write!(f, "Stopping"),
|
|
Self::Stopped => write!(f, "Stopped"),
|
|
Self::Error => write!(f, "Error"),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Pipeline configuration
|
|
#[derive(Debug, Clone)]
|
|
pub struct PipelineConfig {
|
|
/// Number of channels
|
|
pub n_channels: usize,
|
|
/// Sample rate in Hz
|
|
pub sample_rate: f64,
|
|
/// Buffer duration in milliseconds
|
|
pub buffer_duration_ms: f64,
|
|
/// Processing window size in milliseconds
|
|
pub window_size_ms: f64,
|
|
/// Window stride in milliseconds
|
|
pub window_stride_ms: f64,
|
|
/// Filter low cutoff (Hz)
|
|
pub filter_low: f64,
|
|
/// Filter high cutoff (Hz)
|
|
pub filter_high: f64,
|
|
/// Enable filtering
|
|
pub enable_filtering: bool,
|
|
/// Enable beamforming
|
|
pub enable_beamforming: bool,
|
|
/// Target latency in milliseconds
|
|
pub target_latency_ms: f64,
|
|
/// Output queue size
|
|
pub output_queue_size: usize,
|
|
}
|
|
|
|
impl Default for PipelineConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
n_channels: 64,
|
|
sample_rate: 1000.0,
|
|
buffer_duration_ms: 500.0,
|
|
window_size_ms: 100.0,
|
|
window_stride_ms: 50.0,
|
|
filter_low: 1.0,
|
|
filter_high: 40.0,
|
|
enable_filtering: true,
|
|
enable_beamforming: true,
|
|
target_latency_ms: 10.0,
|
|
output_queue_size: 100,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Source estimate output
|
|
#[derive(Debug, Clone)]
|
|
pub struct SourceOutput {
|
|
/// Source data [n_sources x n_times]
|
|
pub sources: Vec<Vec<f64>>,
|
|
/// Timestamps for source estimates
|
|
pub timestamps: Vec<f64>,
|
|
/// Processing latency in microseconds
|
|
pub latency_us: u64,
|
|
/// Window index
|
|
pub window_index: u64,
|
|
}
|
|
|
|
/// Real-time processing pipeline
|
|
pub struct RealtimePipeline {
|
|
/// Configuration
|
|
config: PipelineConfig,
|
|
/// Pipeline state
|
|
state: Arc<RwLock<PipelineState>>,
|
|
/// Input ring buffer
|
|
input_buffer: SharedRingBuffer,
|
|
/// GPU filter (if enabled)
|
|
filter: Option<Arc<RwLock<GpuFilter>>>,
|
|
/// GPU beamformer (if enabled)
|
|
beamformer: Option<Arc<GpuBeamformer>>,
|
|
/// Output channel sender
|
|
output_tx: Option<Sender<SourceOutput>>,
|
|
/// Output channel receiver
|
|
output_rx: Option<Receiver<SourceOutput>>,
|
|
/// Latency monitor
|
|
latency_monitor: Arc<LatencyMonitor>,
|
|
/// Running flag
|
|
running: Arc<AtomicBool>,
|
|
/// Processed windows counter
|
|
windows_processed: Arc<AtomicU64>,
|
|
/// Processing thread handle
|
|
processing_thread: Option<JoinHandle<()>>,
|
|
/// Last error message
|
|
last_error: Arc<RwLock<Option<String>>>,
|
|
}
|
|
|
|
impl RealtimePipeline {
|
|
/// Create a new pipeline
|
|
pub fn new(config: PipelineConfig) -> RealtimeResult<Self> {
|
|
// Create input buffer
|
|
let buffer_duration = config.buffer_duration_ms / 1000.0;
|
|
let input_buffer =
|
|
create_shared_buffer(buffer_duration, config.sample_rate, config.n_channels);
|
|
|
|
// Create filter if enabled
|
|
let filter = if config.enable_filtering {
|
|
let filter_config = GpuFilterConfig {
|
|
filter_type: FilterType::Bandpass,
|
|
low_freq: config.filter_low,
|
|
high_freq: config.filter_high,
|
|
sample_rate: config.sample_rate,
|
|
order: 101,
|
|
transition_bw: 2.0,
|
|
};
|
|
Some(Arc::new(RwLock::new(GpuFilter::new(
|
|
filter_config,
|
|
config.n_channels,
|
|
)?)))
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Create output channel
|
|
let (output_tx, output_rx) = bounded(config.output_queue_size);
|
|
|
|
// Create latency monitor
|
|
let latency_monitor = Arc::new(LatencyMonitor::new(1000));
|
|
|
|
Ok(Self {
|
|
config,
|
|
state: Arc::new(RwLock::new(PipelineState::Idle)),
|
|
input_buffer,
|
|
filter,
|
|
beamformer: None,
|
|
output_tx: Some(output_tx),
|
|
output_rx: Some(output_rx),
|
|
latency_monitor,
|
|
running: Arc::new(AtomicBool::new(false)),
|
|
windows_processed: Arc::new(AtomicU64::new(0)),
|
|
processing_thread: None,
|
|
last_error: Arc::new(RwLock::new(None)),
|
|
})
|
|
}
|
|
|
|
/// Set the beamformer
|
|
pub fn set_beamformer(&mut self, beamformer: GpuBeamformer) {
|
|
self.beamformer = Some(Arc::new(beamformer));
|
|
}
|
|
|
|
/// Push data into the pipeline
|
|
///
|
|
/// # Arguments
|
|
/// * `data` - Sample data [n_channels]
|
|
/// * `timestamp` - LSL timestamp
|
|
pub fn push_sample(&self, data: Vec<f64>, timestamp: f64) {
|
|
self.input_buffer.push(data, timestamp);
|
|
}
|
|
|
|
/// Push a batch of data
|
|
pub fn push_batch(&self, data: &[Vec<f64>], timestamps: &[f64]) {
|
|
self.input_buffer.push_batch(data, timestamps);
|
|
}
|
|
|
|
/// Start the processing pipeline
|
|
pub fn start(&mut self) -> RealtimeResult<()> {
|
|
if *self.state.read() == PipelineState::Running {
|
|
return Ok(());
|
|
}
|
|
|
|
*self.state.write() = PipelineState::Starting;
|
|
self.running.store(true, Ordering::SeqCst);
|
|
|
|
// Clone shared resources for the processing thread
|
|
let config = self.config.clone();
|
|
let state = Arc::clone(&self.state);
|
|
let input_buffer = Arc::clone(&self.input_buffer);
|
|
let filter = self.filter.clone();
|
|
let beamformer = self.beamformer.clone();
|
|
let output_tx = self.output_tx.clone().ok_or_else(|| {
|
|
RealtimeError::PipelineState("Output channel not available".to_string())
|
|
})?;
|
|
let latency_monitor = Arc::clone(&self.latency_monitor);
|
|
let running = Arc::clone(&self.running);
|
|
let windows_processed = Arc::clone(&self.windows_processed);
|
|
let last_error = Arc::clone(&self.last_error);
|
|
|
|
// Spawn processing thread
|
|
let handle = thread::spawn(move || {
|
|
*state.write() = PipelineState::Running;
|
|
|
|
let window_samples = (config.window_size_ms / 1000.0 * config.sample_rate) as usize;
|
|
let stride_samples = (config.window_stride_ms / 1000.0 * config.sample_rate) as usize;
|
|
let mut last_process_time = Instant::now();
|
|
let stride_duration = Duration::from_secs_f64(config.window_stride_ms / 1000.0);
|
|
|
|
while running.load(Ordering::SeqCst) {
|
|
let start = Instant::now();
|
|
|
|
// Check if we have enough data
|
|
if input_buffer.len() < window_samples {
|
|
thread::sleep(Duration::from_millis(1));
|
|
continue;
|
|
}
|
|
|
|
// Wait for stride interval
|
|
let elapsed = last_process_time.elapsed();
|
|
if elapsed < stride_duration {
|
|
thread::sleep(stride_duration.checked_sub(elapsed).unwrap());
|
|
}
|
|
last_process_time = Instant::now();
|
|
|
|
// Get data window
|
|
latency_monitor.start(LatencyStage::LslReceive);
|
|
let window_sec = config.window_size_ms / 1000.0;
|
|
let data_matrix = input_buffer.get_recent_matrix(window_sec);
|
|
let (_, timestamps) = input_buffer.get_recent(window_sec);
|
|
latency_monitor.stop(LatencyStage::LslReceive);
|
|
|
|
// Apply filter if enabled
|
|
let filtered_data = if let Some(ref filter) = filter {
|
|
latency_monitor.start(LatencyStage::GpuFilter);
|
|
let result = {
|
|
let mut f = filter.write();
|
|
f.apply(&data_matrix)
|
|
};
|
|
latency_monitor.stop(LatencyStage::GpuFilter);
|
|
|
|
match result {
|
|
Ok(d) => d,
|
|
Err(e) => {
|
|
*last_error.write() = Some(e.to_string());
|
|
continue;
|
|
}
|
|
}
|
|
} else {
|
|
data_matrix
|
|
};
|
|
|
|
// Apply beamformer if enabled
|
|
let source_data = if let Some(ref bf) = beamformer {
|
|
latency_monitor.start(LatencyStage::GpuBeamformer);
|
|
let result = bf.apply(&filtered_data);
|
|
latency_monitor.stop(LatencyStage::GpuBeamformer);
|
|
|
|
match result {
|
|
Ok(d) => d,
|
|
Err(e) => {
|
|
*last_error.write() = Some(e.to_string());
|
|
continue;
|
|
}
|
|
}
|
|
} else {
|
|
filtered_data
|
|
};
|
|
|
|
// Record total latency
|
|
let total_latency = start.elapsed();
|
|
latency_monitor.record(LatencyStage::Total, total_latency);
|
|
|
|
// Create output
|
|
let output = SourceOutput {
|
|
sources: source_data,
|
|
timestamps: timestamps.clone(),
|
|
latency_us: total_latency.as_micros() as u64,
|
|
window_index: windows_processed.fetch_add(1, Ordering::SeqCst),
|
|
};
|
|
|
|
// Send output (non-blocking)
|
|
if output_tx.try_send(output).is_err() {
|
|
// Queue full, drop oldest
|
|
}
|
|
}
|
|
|
|
*state.write() = PipelineState::Stopped;
|
|
});
|
|
|
|
self.processing_thread = Some(handle);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Stop the pipeline
|
|
pub fn stop(&mut self) -> RealtimeResult<()> {
|
|
*self.state.write() = PipelineState::Stopping;
|
|
self.running.store(false, Ordering::SeqCst);
|
|
|
|
if let Some(handle) = self.processing_thread.take() {
|
|
handle.join().map_err(|_| {
|
|
RealtimeError::PipelineState("Failed to join processing thread".to_string())
|
|
})?;
|
|
}
|
|
|
|
*self.state.write() = PipelineState::Stopped;
|
|
Ok(())
|
|
}
|
|
|
|
/// Pause the pipeline
|
|
pub fn pause(&mut self) -> RealtimeResult<()> {
|
|
*self.state.write() = PipelineState::Paused;
|
|
self.running.store(false, Ordering::SeqCst);
|
|
Ok(())
|
|
}
|
|
|
|
/// Resume the pipeline
|
|
pub fn resume(&mut self) -> RealtimeResult<()> {
|
|
if *self.state.read() == PipelineState::Paused {
|
|
self.start()
|
|
} else {
|
|
Err(RealtimeError::PipelineState(
|
|
"Pipeline is not paused".to_string(),
|
|
))
|
|
}
|
|
}
|
|
|
|
/// Get the next source output (blocking)
|
|
pub fn recv(&self) -> Option<SourceOutput> {
|
|
self.output_rx.as_ref()?.recv().ok()
|
|
}
|
|
|
|
/// Get the next source output (non-blocking)
|
|
pub fn try_recv(&self) -> Option<SourceOutput> {
|
|
match self.output_rx.as_ref()?.try_recv() {
|
|
Ok(output) => Some(output),
|
|
Err(TryRecvError::Empty) => None,
|
|
Err(TryRecvError::Disconnected) => None,
|
|
}
|
|
}
|
|
|
|
/// Get the next source output with timeout
|
|
pub fn recv_timeout(&self, timeout: Duration) -> Option<SourceOutput> {
|
|
self.output_rx.as_ref()?.recv_timeout(timeout).ok()
|
|
}
|
|
|
|
/// Get current pipeline state
|
|
pub fn state(&self) -> PipelineState {
|
|
*self.state.read()
|
|
}
|
|
|
|
/// Get latency statistics
|
|
pub fn latency_stats(&self) -> Vec<LatencyStats> {
|
|
self.latency_monitor.all_stats()
|
|
}
|
|
|
|
/// Get latency report
|
|
pub fn latency_report(&self) -> String {
|
|
self.latency_monitor.report()
|
|
}
|
|
|
|
/// Check if latency targets are met
|
|
pub fn meets_latency_targets(&self) -> bool {
|
|
self.latency_monitor.all_meet_targets()
|
|
}
|
|
|
|
/// Get number of processed windows
|
|
pub fn windows_processed(&self) -> u64 {
|
|
self.windows_processed.load(Ordering::SeqCst)
|
|
}
|
|
|
|
/// Get input buffer fill level (0.0 - 1.0)
|
|
pub fn buffer_fill_level(&self) -> f64 {
|
|
self.input_buffer.fill_level()
|
|
}
|
|
|
|
/// Get last error message
|
|
pub fn last_error(&self) -> Option<String> {
|
|
self.last_error.read().clone()
|
|
}
|
|
|
|
/// Clear last error
|
|
pub fn clear_error(&self) {
|
|
*self.last_error.write() = None;
|
|
}
|
|
|
|
/// Get configuration
|
|
pub fn config(&self) -> &PipelineConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Reset the pipeline (clears buffers and statistics)
|
|
pub fn reset(&mut self) {
|
|
self.input_buffer.clear();
|
|
self.latency_monitor.clear();
|
|
self.windows_processed.store(0, Ordering::SeqCst);
|
|
self.clear_error();
|
|
}
|
|
}
|
|
|
|
impl Drop for RealtimePipeline {
|
|
fn drop(&mut self) {
|
|
self.running.store(false, Ordering::SeqCst);
|
|
if let Some(handle) = self.processing_thread.take() {
|
|
let _ = handle.join();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Pipeline builder for fluent configuration
|
|
pub struct PipelineBuilder {
|
|
config: PipelineConfig,
|
|
}
|
|
|
|
impl PipelineBuilder {
|
|
/// Create a new builder with default configuration
|
|
pub fn new() -> Self {
|
|
Self {
|
|
config: PipelineConfig::default(),
|
|
}
|
|
}
|
|
|
|
/// Set number of channels
|
|
pub fn channels(mut self, n_channels: usize) -> Self {
|
|
self.config.n_channels = n_channels;
|
|
self
|
|
}
|
|
|
|
/// Set sample rate
|
|
pub fn sample_rate(mut self, rate: f64) -> Self {
|
|
self.config.sample_rate = rate;
|
|
self
|
|
}
|
|
|
|
/// Set buffer duration in milliseconds
|
|
pub fn buffer_duration_ms(mut self, duration: f64) -> Self {
|
|
self.config.buffer_duration_ms = duration;
|
|
self
|
|
}
|
|
|
|
/// Set processing window size in milliseconds
|
|
pub fn window_size_ms(mut self, size: f64) -> Self {
|
|
self.config.window_size_ms = size;
|
|
self
|
|
}
|
|
|
|
/// Set window stride in milliseconds
|
|
pub fn window_stride_ms(mut self, stride: f64) -> Self {
|
|
self.config.window_stride_ms = stride;
|
|
self
|
|
}
|
|
|
|
/// Set filter band (low, high in Hz)
|
|
pub fn filter_band(mut self, low: f64, high: f64) -> Self {
|
|
self.config.filter_low = low;
|
|
self.config.filter_high = high;
|
|
self
|
|
}
|
|
|
|
/// Enable or disable filtering
|
|
pub fn enable_filtering(mut self, enable: bool) -> Self {
|
|
self.config.enable_filtering = enable;
|
|
self
|
|
}
|
|
|
|
/// Enable or disable beamforming
|
|
pub fn enable_beamforming(mut self, enable: bool) -> Self {
|
|
self.config.enable_beamforming = enable;
|
|
self
|
|
}
|
|
|
|
/// Set target latency in milliseconds
|
|
pub fn target_latency_ms(mut self, latency: f64) -> Self {
|
|
self.config.target_latency_ms = latency;
|
|
self
|
|
}
|
|
|
|
/// Build the pipeline
|
|
pub fn build(self) -> RealtimeResult<RealtimePipeline> {
|
|
RealtimePipeline::new(self.config)
|
|
}
|
|
}
|
|
|
|
impl Default for PipelineBuilder {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_pipeline_creation() {
|
|
let config = PipelineConfig::default();
|
|
let pipeline = RealtimePipeline::new(config).unwrap();
|
|
|
|
assert_eq!(pipeline.state(), PipelineState::Idle);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_builder() {
|
|
let pipeline = PipelineBuilder::new()
|
|
.channels(32)
|
|
.sample_rate(500.0)
|
|
.filter_band(8.0, 30.0)
|
|
.window_size_ms(50.0)
|
|
.build()
|
|
.unwrap();
|
|
|
|
assert_eq!(pipeline.config().n_channels, 32);
|
|
assert_eq!(pipeline.config().sample_rate, 500.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_push_sample() {
|
|
let config = PipelineConfig {
|
|
n_channels: 4,
|
|
sample_rate: 100.0,
|
|
..Default::default()
|
|
};
|
|
let pipeline = RealtimePipeline::new(config).unwrap();
|
|
|
|
pipeline.push_sample(vec![1.0, 2.0, 3.0, 4.0], 0.0);
|
|
pipeline.push_sample(vec![5.0, 6.0, 7.0, 8.0], 0.01);
|
|
|
|
assert_eq!(pipeline.input_buffer.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_push_batch() {
|
|
let config = PipelineConfig {
|
|
n_channels: 2,
|
|
sample_rate: 100.0,
|
|
..Default::default()
|
|
};
|
|
let pipeline = RealtimePipeline::new(config).unwrap();
|
|
|
|
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];
|
|
|
|
pipeline.push_batch(&data, ×tamps);
|
|
|
|
assert_eq!(pipeline.input_buffer.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipeline_lifecycle() {
|
|
let config = PipelineConfig {
|
|
n_channels: 4,
|
|
sample_rate: 100.0,
|
|
window_size_ms: 100.0,
|
|
enable_beamforming: false,
|
|
..Default::default()
|
|
};
|
|
let mut pipeline = RealtimePipeline::new(config).unwrap();
|
|
|
|
// Start
|
|
pipeline.start().unwrap();
|
|
assert_eq!(pipeline.state(), PipelineState::Starting);
|
|
|
|
// Give it time to start
|
|
thread::sleep(Duration::from_millis(10));
|
|
|
|
// Stop
|
|
pipeline.stop().unwrap();
|
|
assert_eq!(pipeline.state(), PipelineState::Stopped);
|
|
}
|
|
|
|
#[test]
|
|
fn test_latency_monitor() {
|
|
let config = PipelineConfig::default();
|
|
let pipeline = RealtimePipeline::new(config).unwrap();
|
|
|
|
let stats = pipeline.latency_stats();
|
|
assert_eq!(stats.len(), 5);
|
|
}
|
|
}
|