Initial commit

crates/specialized/rtx-neuro-realtime/src/pipeline.rs

@@ -0,0 +1,607 @@
+//! 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, &timestamps);
+
+        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);
+    }
+}