Initial commit

crates/specialized/rtx-neuro-artifacts/src/labels.rs

@@ -0,0 +1,457 @@
+//! Artifact type definitions and labels.
+//!
+//! Defines the types of artifacts that can be detected in MEG/EEG recordings.
+
+use serde::{Deserialize, Serialize};
+
+/// Types of artifacts that can be detected
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum ArtifactType {
+    /// Eye blink artifact (EOG)
+    /// Characterized by large frontal deflections lasting 100-400ms
+    EyeBlink,
+
+    /// Eye movement artifact (saccades, smooth pursuit)
+    /// Horizontal or vertical eye movements
+    EyeMovement,
+
+    /// Muscle artifact (EMG)
+    /// High-frequency (>20Hz) contamination from facial/neck muscles
+    Muscle,
+
+    /// Heartbeat artifact (ECG/BCG)
+    /// Cardiac-related artifact, especially prominent in MEG
+    Heartbeat,
+
+    /// Power line noise (50/60 Hz)
+    /// Narrow-band interference from electrical mains
+    LineNoise,
+
+    /// Movement artifact
+    /// Head or body movement causing signal distortion
+    Movement,
+
+    /// Electrode pop/jump
+    /// Sudden amplitude changes from electrode issues
+    ElectrodePop,
+
+    /// Channel noise/bad channel
+    /// Persistently noisy or disconnected channel
+    ChannelNoise,
+
+    /// Environmental interference
+    /// External electromagnetic interference
+    Environmental,
+
+    /// Unknown artifact
+    /// Detected anomaly of unknown origin
+    Unknown,
+}
+
+impl ArtifactType {
+    /// Get the index for this artifact type (for model output mapping)
+    pub fn index(&self) -> usize {
+        match self {
+            Self::EyeBlink => 0,
+            Self::EyeMovement => 1,
+            Self::Muscle => 2,
+            Self::Heartbeat => 3,
+            Self::LineNoise => 4,
+            Self::Movement => 5,
+            Self::ElectrodePop => 6,
+            Self::ChannelNoise => 7,
+            Self::Environmental => 8,
+            Self::Unknown => 9,
+        }
+    }
+
+    /// Create from index
+    pub fn from_index(index: usize) -> Option<Self> {
+        match index {
+            0 => Some(Self::EyeBlink),
+            1 => Some(Self::EyeMovement),
+            2 => Some(Self::Muscle),
+            3 => Some(Self::Heartbeat),
+            4 => Some(Self::LineNoise),
+            5 => Some(Self::Movement),
+            6 => Some(Self::ElectrodePop),
+            7 => Some(Self::ChannelNoise),
+            8 => Some(Self::Environmental),
+            9 => Some(Self::Unknown),
+            _ => None,
+        }
+    }
+
+    /// Number of artifact types
+    pub const fn count() -> usize {
+        10
+    }
+
+    /// Get all artifact types
+    pub fn all() -> Vec<Self> {
+        vec![
+            Self::EyeBlink,
+            Self::EyeMovement,
+            Self::Muscle,
+            Self::Heartbeat,
+            Self::LineNoise,
+            Self::Movement,
+            Self::ElectrodePop,
+            Self::ChannelNoise,
+            Self::Environmental,
+            Self::Unknown,
+        ]
+    }
+
+    /// Get human-readable name
+    pub fn name(&self) -> &'static str {
+        match self {
+            Self::EyeBlink => "Eye Blink",
+            Self::EyeMovement => "Eye Movement",
+            Self::Muscle => "Muscle",
+            Self::Heartbeat => "Heartbeat",
+            Self::LineNoise => "Line Noise",
+            Self::Movement => "Movement",
+            Self::ElectrodePop => "Electrode Pop",
+            Self::ChannelNoise => "Channel Noise",
+            Self::Environmental => "Environmental",
+            Self::Unknown => "Unknown",
+        }
+    }
+
+    /// Get short code
+    pub fn code(&self) -> &'static str {
+        match self {
+            Self::EyeBlink => "EOG_B",
+            Self::EyeMovement => "EOG_M",
+            Self::Muscle => "EMG",
+            Self::Heartbeat => "ECG",
+            Self::LineNoise => "LINE",
+            Self::Movement => "MOV",
+            Self::ElectrodePop => "POP",
+            Self::ChannelNoise => "BAD",
+            Self::Environmental => "ENV",
+            Self::Unknown => "UNK",
+        }
+    }
+
+    /// Get typical frequency range for this artifact type
+    pub fn frequency_range(&self) -> (f64, f64) {
+        match self {
+            Self::EyeBlink => (0.5, 4.0),
+            Self::EyeMovement => (0.5, 10.0),
+            Self::Muscle => (20.0, 200.0),
+            Self::Heartbeat => (0.5, 40.0),
+            Self::LineNoise => (49.0, 61.0), // Covers both 50 and 60 Hz
+            Self::Movement => (0.1, 5.0),
+            Self::ElectrodePop => (0.1, 100.0),  // Broadband
+            Self::ChannelNoise => (0.1, 200.0),  // Broadband
+            Self::Environmental => (0.1, 200.0), // Variable
+            Self::Unknown => (0.1, 200.0),
+        }
+    }
+
+    /// Get typical duration range in milliseconds
+    pub fn duration_range_ms(&self) -> (f64, f64) {
+        match self {
+            Self::EyeBlink => (100.0, 400.0),
+            Self::EyeMovement => (50.0, 500.0),
+            Self::Muscle => (10.0, 1000.0),    // Can be sustained
+            Self::Heartbeat => (200.0, 400.0), // QRS complex
+            Self::LineNoise => (0.0, f64::INFINITY), // Continuous
+            Self::Movement => (100.0, 5000.0),
+            Self::ElectrodePop => (1.0, 50.0), // Very brief
+            Self::ChannelNoise => (0.0, f64::INFINITY), // Continuous
+            Self::Environmental => (0.0, f64::INFINITY), // Variable
+            Self::Unknown => (0.0, f64::INFINITY),
+        }
+    }
+
+    /// Get color for visualization (RGB)
+    pub fn color(&self) -> (u8, u8, u8) {
+        match self {
+            Self::EyeBlink => (255, 165, 0),     // Orange
+            Self::EyeMovement => (255, 200, 0),  // Gold
+            Self::Muscle => (255, 0, 0),         // Red
+            Self::Heartbeat => (255, 0, 128),    // Pink
+            Self::LineNoise => (128, 128, 128),  // Gray
+            Self::Movement => (0, 128, 255),     // Blue
+            Self::ElectrodePop => (255, 255, 0), // Yellow
+            Self::ChannelNoise => (128, 0, 128), // Purple
+            Self::Environmental => (0, 128, 0),  // Green
+            Self::Unknown => (64, 64, 64),       // Dark gray
+        }
+    }
+}
+
+impl std::fmt::Display for ArtifactType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name())
+    }
+}
+
+/// A detected artifact label with probability
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ArtifactLabel {
+    /// Type of artifact
+    pub artifact_type: ArtifactType,
+    /// Detection probability (0.0 - 1.0)
+    pub probability: f64,
+    /// Confidence score (may differ from probability for some models)
+    pub confidence: f64,
+}
+
+impl ArtifactLabel {
+    /// Create a new artifact label
+    pub fn new(artifact_type: ArtifactType, probability: f64) -> Self {
+        Self {
+            artifact_type,
+            probability,
+            confidence: probability,
+        }
+    }
+
+    /// Create with separate confidence score
+    pub fn with_confidence(artifact_type: ArtifactType, probability: f64, confidence: f64) -> Self {
+        Self {
+            artifact_type,
+            probability,
+            confidence,
+        }
+    }
+
+    /// Check if this artifact is considered detected (probability > threshold)
+    pub fn is_detected(&self, threshold: f64) -> bool {
+        self.probability > threshold
+    }
+}
+
+/// A region in the signal where an artifact was detected
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ArtifactRegion {
+    /// Artifact type
+    pub artifact_type: ArtifactType,
+    /// Start time in seconds
+    pub start_time: f64,
+    /// End time in seconds
+    pub end_time: f64,
+    /// Affected channel indices
+    pub channels: Vec<usize>,
+    /// Detection probability
+    pub probability: f64,
+    /// Peak location (time of maximum artifact influence)
+    pub peak_time: Option<f64>,
+    /// Severity score (0.0 - 1.0)
+    pub severity: f64,
+}
+
+impl ArtifactRegion {
+    /// Create a new artifact region
+    pub fn new(
+        artifact_type: ArtifactType,
+        start_time: f64,
+        end_time: f64,
+        channels: Vec<usize>,
+        probability: f64,
+    ) -> Self {
+        Self {
+            artifact_type,
+            start_time,
+            end_time,
+            channels,
+            probability,
+            peak_time: None,
+            severity: probability,
+        }
+    }
+
+    /// Duration in seconds
+    pub fn duration(&self) -> f64 {
+        self.end_time - self.start_time
+    }
+
+    /// Duration in milliseconds
+    pub fn duration_ms(&self) -> f64 {
+        self.duration() * 1000.0
+    }
+
+    /// Check if this region overlaps with a time range
+    pub fn overlaps(&self, start: f64, end: f64) -> bool {
+        self.start_time < end && self.end_time > start
+    }
+
+    /// Check if this region affects a specific channel
+    pub fn affects_channel(&self, channel: usize) -> bool {
+        self.channels.contains(&channel)
+    }
+
+    /// Merge with another overlapping region of the same type
+    pub fn merge(&self, other: &Self) -> Option<Self> {
+        if self.artifact_type != other.artifact_type {
+            return None;
+        }
+
+        if !self.overlaps(other.start_time, other.end_time) {
+            return None;
+        }
+
+        let mut channels: Vec<usize> = self.channels.clone();
+        for ch in &other.channels {
+            if !channels.contains(ch) {
+                channels.push(*ch);
+            }
+        }
+        channels.sort_unstable();
+
+        Some(Self {
+            artifact_type: self.artifact_type,
+            start_time: self.start_time.min(other.start_time),
+            end_time: self.end_time.max(other.end_time),
+            channels,
+            probability: self.probability.max(other.probability),
+            peak_time: match (self.peak_time, other.peak_time) {
+                (Some(a), Some(b)) => Some(if self.probability > other.probability {
+                    a
+                } else {
+                    b
+                }),
+                (Some(a), None) => Some(a),
+                (None, Some(b)) => Some(b),
+                (None, None) => None,
+            },
+            severity: self.severity.max(other.severity),
+        })
+    }
+}
+
+/// Summary statistics for artifact detection
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ArtifactSummary {
+    /// Total number of artifacts detected
+    pub total_artifacts: usize,
+    /// Count per artifact type
+    pub counts: Vec<(ArtifactType, usize)>,
+    /// Total contaminated duration (seconds)
+    pub total_duration: f64,
+    /// Percentage of signal contaminated
+    pub contamination_percent: f64,
+    /// Most common artifact type
+    pub most_common: Option<ArtifactType>,
+    /// Channels most affected
+    pub most_affected_channels: Vec<usize>,
+}
+
+impl ArtifactSummary {
+    /// Create from a list of artifact regions
+    pub fn from_regions(regions: &[ArtifactRegion], total_signal_duration: f64) -> Self {
+        let mut counts: std::collections::HashMap<ArtifactType, usize> =
+            std::collections::HashMap::new();
+        let mut total_duration = 0.0;
+        let mut channel_counts: std::collections::HashMap<usize, usize> =
+            std::collections::HashMap::new();
+
+        for region in regions {
+            *counts.entry(region.artifact_type).or_insert(0) += 1;
+            total_duration += region.duration();
+
+            for &ch in &region.channels {
+                *channel_counts.entry(ch).or_insert(0) += 1;
+            }
+        }
+
+        let counts_vec: Vec<(ArtifactType, usize)> = counts.into_iter().collect();
+        let most_common = counts_vec.iter().max_by_key(|(_, c)| *c).map(|(t, _)| *t);
+
+        let mut channel_vec: Vec<(usize, usize)> = channel_counts.into_iter().collect();
+        channel_vec.sort_by(|a, b| b.1.cmp(&a.1));
+        let most_affected_channels: Vec<usize> =
+            channel_vec.into_iter().take(5).map(|(ch, _)| ch).collect();
+
+        let contamination_percent = if total_signal_duration > 0.0 {
+            (total_duration / total_signal_duration) * 100.0
+        } else {
+            0.0
+        };
+
+        Self {
+            total_artifacts: regions.len(),
+            counts: counts_vec,
+            total_duration,
+            contamination_percent,
+            most_common,
+            most_affected_channels,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_artifact_type_index() {
+        assert_eq!(ArtifactType::EyeBlink.index(), 0);
+        assert_eq!(ArtifactType::Muscle.index(), 2);
+        assert_eq!(ArtifactType::from_index(2), Some(ArtifactType::Muscle));
+    }
+
+    #[test]
+    fn test_artifact_type_all() {
+        let all = ArtifactType::all();
+        assert_eq!(all.len(), ArtifactType::count());
+    }
+
+    #[test]
+    fn test_artifact_label() {
+        let label = ArtifactLabel::new(ArtifactType::EyeBlink, 0.85);
+        assert!(label.is_detected(0.5));
+        assert!(!label.is_detected(0.9));
+    }
+
+    #[test]
+    fn test_artifact_region() {
+        let region = ArtifactRegion::new(ArtifactType::Muscle, 1.0, 1.5, vec![10, 11, 12], 0.9);
+
+        assert_eq!(region.duration(), 0.5);
+        assert_eq!(region.duration_ms(), 500.0);
+        assert!(region.affects_channel(10));
+        assert!(!region.affects_channel(5));
+    }
+
+    #[test]
+    fn test_region_overlap() {
+        let region = ArtifactRegion::new(ArtifactType::Muscle, 1.0, 2.0, vec![0], 0.9);
+
+        assert!(region.overlaps(0.5, 1.5));
+        assert!(region.overlaps(1.5, 2.5));
+        assert!(!region.overlaps(2.5, 3.0));
+        assert!(!region.overlaps(0.0, 0.5));
+    }
+
+    #[test]
+    fn test_region_merge() {
+        let r1 = ArtifactRegion::new(ArtifactType::Muscle, 1.0, 2.0, vec![0, 1], 0.8);
+        let r2 = ArtifactRegion::new(ArtifactType::Muscle, 1.5, 2.5, vec![1, 2], 0.9);
+
+        let merged = r1.merge(&r2).unwrap();
+        assert_eq!(merged.start_time, 1.0);
+        assert_eq!(merged.end_time, 2.5);
+        assert_eq!(merged.channels, vec![0, 1, 2]);
+        assert_eq!(merged.probability, 0.9);
+    }
+
+    #[test]
+    fn test_artifact_summary() {
+        let regions = vec![
+            ArtifactRegion::new(ArtifactType::EyeBlink, 0.0, 0.3, vec![0, 1], 0.9),
+            ArtifactRegion::new(ArtifactType::EyeBlink, 1.0, 1.3, vec![0, 1], 0.85),
+            ArtifactRegion::new(ArtifactType::Muscle, 2.0, 2.5, vec![10], 0.7),
+        ];
+
+        let summary = ArtifactSummary::from_regions(&regions, 10.0);
+
+        assert_eq!(summary.total_artifacts, 3);
+        assert!((summary.total_duration - 1.1).abs() < 0.01);
+        assert_eq!(summary.most_common, Some(ArtifactType::EyeBlink));
+    }
+}