rustytorch/demos/slidescope-shared/src/nmf.rs

//! NMF (Non-negative Matrix Factorization) configuration and result types

use serde::{Deserialize, Serialize};

/// Configuration for NMF stain unmixing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NmfConfig {
    /// Number of stain components to extract (typically 2-4)
    pub n_components: usize,
    /// Maximum iterations for convergence
    pub max_iterations: usize,
    /// Convergence tolerance
    pub tolerance: f32,
    /// Use GPU acceleration if available
    pub use_gpu: bool,
    /// Random seed for reproducibility
    pub random_seed: Option<u64>,
    /// Divergence type for optimization
    pub divergence: DivergenceType,
}

impl Default for NmfConfig {
    fn default() -> Self {
        Self {
            n_components: 2, // H&E has 2 stains
            max_iterations: 200,
            tolerance: 1e-4,
            use_gpu: true,
            random_seed: Some(42),
            divergence: DivergenceType::Frobenius,
        }
    }
}

impl NmfConfig {
    /// Create config for H&E staining (2 components)
    pub fn he() -> Self {
        Self {
            n_components: 2,
            ..Default::default()
        }
    }

    /// Create config for IHC-DAB staining (3 components)
    pub fn ihc_dab() -> Self {
        Self {
            n_components: 3,
            ..Default::default()
        }
    }

    /// Fast config for testing
    pub fn fast() -> Self {
        Self {
            n_components: 2,
            max_iterations: 50,
            tolerance: 1e-3,
            use_gpu: false,
            random_seed: Some(42),
            divergence: DivergenceType::Frobenius,
        }
    }
}

/// Divergence type for NMF optimization
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum DivergenceType {
    /// Frobenius norm (Euclidean distance)
    #[default]
    Frobenius,
    /// Kullback-Leibler divergence
    KullbackLeibler,
    /// Itakura-Saito divergence
    ItakuraSaito,
    /// Beta-divergence with parameter
    Beta(f32),
}

/// Result of NMF stain unmixing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NmfResult {
    /// Slide ID this result belongs to
    pub slide_id: String,
    /// Paths to concentration map images for each stain
    pub stain_images: Vec<String>,
    /// Path to composite/reconstructed image
    pub composite_path: Option<String>,
    /// Final reconstruction error
    pub reconstruction_error: f32,
    /// Number of iterations performed
    pub iterations: usize,
    /// Processing time in milliseconds
    pub processing_time_ms: u64,
    /// GPU memory used in MB (if GPU processing)
    pub gpu_memory_mb: Option<f32>,
    /// Estimated stain vectors (W matrix columns as RGB)
    pub stain_vectors: Option<Vec<[f32; 3]>>,
}

/// Stain vector matrix (W matrix from NMF)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StainMatrix {
    /// Number of stain components
    pub n_components: usize,
    /// Stain vectors as RGB triplets (column-major)
    pub vectors: Vec<f32>,
}

impl StainMatrix {
    /// Create from raw data
    pub fn new(n_components: usize, vectors: Vec<f32>) -> Self {
        assert_eq!(vectors.len(), n_components * 3);
        Self {
            n_components,
            vectors,
        }
    }

    /// Get stain vector for component i as [R, G, B]
    pub fn get_vector(&self, i: usize) -> Option<[f32; 3]> {
        if i >= self.n_components {
            return None;
        }
        Some([
            self.vectors[i * 3],
            self.vectors[i * 3 + 1],
            self.vectors[i * 3 + 2],
        ])
    }

    /// Standard H&E stain vectors (Ruifrok & Johnston)
    pub fn he_ruifrok() -> Self {
        Self {
            n_components: 2,
            vectors: vec![
                0.644, 0.717, 0.267, // Hematoxylin
                0.093, 0.954, 0.283, // Eosin
            ],
        }
    }

    /// Standard IHC-DAB stain vectors
    pub fn ihc_dab_standard() -> Self {
        Self {
            n_components: 3,
            vectors: vec![
                0.650, 0.704, 0.286, // Hematoxylin
                0.268, 0.570, 0.776, // DAB
                0.0, 0.0, 0.0, // Residual
            ],
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_nmf_config_defaults() {
        let config = NmfConfig::default();
        assert_eq!(config.n_components, 2);
        assert_eq!(config.max_iterations, 200);
        assert!(config.use_gpu);
    }

    #[test]
    fn test_nmf_config_serialization() {
        let config = NmfConfig::he();
        let json = serde_json::to_string(&config).unwrap();
        let deserialized: NmfConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.n_components, config.n_components);
    }

    #[test]
    fn test_stain_matrix() {
        let matrix = StainMatrix::he_ruifrok();
        assert_eq!(matrix.n_components, 2);

        let hematoxylin = matrix.get_vector(0).unwrap();
        assert!((hematoxylin[0] - 0.644).abs() < 1e-6);
    }

    #[test]
    fn test_divergence_serialization() {
        let div = DivergenceType::Beta(0.5);
        let json = serde_json::to_string(&div).unwrap();
        assert!(json.contains("beta"));
    }
}