123 lines
4.1 KiB
Rust
123 lines
4.1 KiB
Rust
//! # RTX-NMF: GPU-Accelerated Non-negative Matrix Factorization
|
||
//!
|
||
//! This crate provides GPU-accelerated Non-negative Matrix Factorization (NMF)
|
||
//! algorithms optimized for the RustyTorch++ ecosystem.
|
||
//!
|
||
//! ## Features
|
||
//! - **GPU-Native**: Direct integration with rtx-tensor and cudarc 0.17.2
|
||
//! - **High Performance**: Custom CUDA kernels for element-wise operations
|
||
//! - **Memory Efficient**: Zero-copy operations when possible
|
||
//! - **Image Processing**: Specialized support for image decomposition
|
||
//! - **Demo Capabilities**: Built-in visualization and testing utilities
|
||
//!
|
||
//! ## Quick Start
|
||
//!
|
||
//! ```rust
|
||
//! use rtx_nmf::{NMFDecomposer, NMFConfig};
|
||
//! use rtx_tensor::{Tensor, Device};
|
||
//!
|
||
//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||
//! // Create GPU device
|
||
//! let device = Device::cuda(0)?;
|
||
//!
|
||
//! // Create input matrix (non-negative)
|
||
//! let matrix = Tensor::rand([100, 50], &device)?;
|
||
//!
|
||
//! // Configure NMF
|
||
//! let config = NMFConfig::new()
|
||
//! .with_components(10)
|
||
//! .with_max_iterations(100)
|
||
//! .with_tolerance(1e-4);
|
||
//!
|
||
//! // Run decomposition
|
||
//! let mut nmf = NMFDecomposer::new(config);
|
||
//! let (w, h) = nmf.fit_transform(&matrix)?;
|
||
//!
|
||
//! println!("Decomposed {} x {} matrix into W({} x {}) and H({} x {})",
|
||
//! matrix.shape()[0], matrix.shape()[1],
|
||
//! w.shape()[0], w.shape()[1],
|
||
//! h.shape()[0], h.shape()[1]);
|
||
//! # Ok(())
|
||
//! # }
|
||
//! ```
|
||
//!
|
||
//! ## NMF Algorithm
|
||
//!
|
||
//! NMF decomposes a non-negative matrix V into two non-negative factors:
|
||
//! V ≈ W × H
|
||
//!
|
||
//! Where:
|
||
//! - V: Input matrix (m × n)
|
||
//! - W: Basis matrix (m × k)
|
||
//! - H: Coefficient matrix (k × n)
|
||
//! - k: Number of components (rank)
|
||
//!
|
||
//! ## GPU Acceleration
|
||
//!
|
||
//! This implementation supports multiple GPU backends:
|
||
//!
|
||
//! - **NVIDIA CUDA**: Custom CUDA kernels for NVIDIA GPUs
|
||
//! - **Apple Metal**: Custom Metal shaders for Apple Silicon (M1/M2/M3)
|
||
//!
|
||
//! The backend is selected automatically based on platform and available features:
|
||
//! - On macOS with Metal feature: Uses Metal compute shaders
|
||
//! - On Linux/Windows with CUDA feature: Uses CUDA kernels
|
||
//! - Otherwise: Falls back to CPU tensor operations
|
||
//!
|
||
//! Key GPU operations:
|
||
//! - Element-wise multiplication with broadcasting
|
||
//! - Element-wise division with epsilon stability
|
||
//! - Matrix multiplication via cuBLAS (CUDA) or MPS (Metal)
|
||
//! - Convergence checking with reduction operations
|
||
//! - Fused NMF multiplicative updates
|
||
|
||
pub mod algorithm;
|
||
pub mod color_unmixing; // Color blind unmixing for biomedical imaging
|
||
pub mod demo;
|
||
pub mod error;
|
||
pub mod gpu_ready_color_unmixing;
|
||
pub mod honest_nmf; // Honest, working implementation
|
||
pub mod image_processing; // Image processing for color unmixing
|
||
pub mod kernels; // GPU-ready implementation with CPU fallback
|
||
|
||
// Re-export main types
|
||
pub use algorithm::{NMFConfig, NMFDecomposer, NMFResult};
|
||
// Re-export kernel types
|
||
pub use demo::{DemoResult, NMFDemo};
|
||
pub use error::{NMFError, Result};
|
||
#[cfg(all(target_os = "macos", feature = "metal"))]
|
||
pub use kernels::MetalNMFKernels;
|
||
pub use kernels::{GpuBackend, current_backend, is_gpu_available};
|
||
// Re-export honest implementation
|
||
pub use honest_nmf::{
|
||
HonestDemoResult, HonestNMF, HonestNMFConfig, HonestNMFDemo, HonestNMFResult,
|
||
};
|
||
// Re-export image processing types
|
||
pub use image_processing::{Color, ImageData, ImageProcessor, ImageStats};
|
||
// Re-export color unmixing types
|
||
pub use color_unmixing::{
|
||
ColorUnmixer, ColorUnmixingConfig, ColorUnmixingDemo, ColorUnmixingResult, ColorVisualization,
|
||
};
|
||
// Re-export GPU-ready color unmixing (primary interface)
|
||
pub use gpu_ready_color_unmixing::{
|
||
ClinicalExportData, ExportImageFile, GPUColorUnmixer, GPUColorUnmixingConfig,
|
||
GPUColorUnmixingDemo, GPUColorUnmixingResult, VisualizationData,
|
||
};
|
||
|
||
// Re-export rtx-tensor types for convenience
|
||
pub use rtx_tensor::{Device, Tensor};
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_nmf_basic() {
|
||
// Basic smoke test that the API compiles
|
||
let config = NMFConfig::new().with_components(5).with_max_iterations(10);
|
||
|
||
assert_eq!(config.components(), 5);
|
||
assert_eq!(config.max_iterations(), 10);
|
||
}
|
||
}
|