296 lines
9.5 KiB
Rust
296 lines
9.5 KiB
Rust
//! GPU-accelerated Non-negative Matrix Factorization
|
|
//!
|
|
//! Provides NMF decomposition using GPU acceleration when available,
|
|
//! with automatic fallback to CPU computation.
|
|
|
|
use crate::config::NmfProcessingConfig;
|
|
use crate::gpu::{BackendType, GpuBackend, select_backend};
|
|
use crate::nmf::NmfOutput;
|
|
use anyhow::{Result, bail};
|
|
use ndarray::Array2;
|
|
use tracing::{debug, info};
|
|
|
|
/// GPU-accelerated NMF processor
|
|
///
|
|
/// Uses CUDA on NVIDIA GPUs, Metal on Apple Silicon, or falls back to CPU.
|
|
pub struct GpuNmfProcessor {
|
|
config: NmfProcessingConfig,
|
|
backend: Box<dyn GpuBackend>,
|
|
}
|
|
|
|
impl GpuNmfProcessor {
|
|
/// Create a new GPU NMF processor with the given configuration
|
|
///
|
|
/// Automatically selects the best available GPU backend.
|
|
pub fn new(config: NmfProcessingConfig) -> Result<Self> {
|
|
let backend = select_backend()?;
|
|
let device = backend.device_info();
|
|
|
|
info!(
|
|
"GpuNmfProcessor initialized with {} backend: {}",
|
|
device.backend, device.name
|
|
);
|
|
|
|
Ok(Self { config, backend })
|
|
}
|
|
|
|
/// Create a processor with a specific backend
|
|
pub fn with_backend(config: NmfProcessingConfig, backend: Box<dyn GpuBackend>) -> Self {
|
|
Self { config, backend }
|
|
}
|
|
|
|
/// Get the backend type being used
|
|
pub fn backend_type(&self) -> BackendType {
|
|
self.backend.device_info().backend
|
|
}
|
|
|
|
/// Check if GPU acceleration is active (not CPU fallback)
|
|
pub fn is_gpu_accelerated(&self) -> bool {
|
|
matches!(
|
|
self.backend.device_info().backend,
|
|
BackendType::Cuda | BackendType::Metal
|
|
)
|
|
}
|
|
|
|
/// Perform NMF decomposition: V ≈ W * H
|
|
///
|
|
/// # Arguments
|
|
/// * `v` - Input matrix (n_features x n_samples), e.g., (3 x n_pixels) for RGB/OD
|
|
/// * `w_init` - Optional initial W matrix (stain vectors)
|
|
///
|
|
/// # Returns
|
|
/// NMF output containing W and H matrices
|
|
pub fn decompose(&self, v: &Array2<f32>, w_init: Option<&Array2<f32>>) -> Result<NmfOutput> {
|
|
let (n_features, n_samples) = v.dim();
|
|
let k = self.config.n_components;
|
|
|
|
info!(
|
|
"GPU NMF: {}x{} matrix, {} components, max {} iterations ({})",
|
|
n_features,
|
|
n_samples,
|
|
k,
|
|
self.config.max_iterations,
|
|
self.backend.name()
|
|
);
|
|
|
|
if n_features == 0 || n_samples == 0 {
|
|
bail!("Input matrix must not be empty");
|
|
}
|
|
if k > n_features || k > n_samples {
|
|
bail!("n_components must be <= min(n_features, n_samples)");
|
|
}
|
|
|
|
// Initialize W and H
|
|
let seed = self.config.random_seed.unwrap_or(42);
|
|
let w = if let Some(w_init) = w_init {
|
|
if w_init.dim() != (n_features, k) {
|
|
bail!("w_init must have shape ({}, {})", n_features, k);
|
|
}
|
|
w_init.clone()
|
|
} else {
|
|
random_nonneg_matrix(n_features, k, seed)
|
|
};
|
|
|
|
let h = random_nonneg_matrix(k, n_samples, seed + 1);
|
|
|
|
// Upload matrices to GPU
|
|
let gpu_v = self.backend.upload(v)?;
|
|
let mut gpu_w = self.backend.upload(&w)?;
|
|
let mut gpu_h = self.backend.upload(&h)?;
|
|
|
|
// Pre-compute V transpose for efficient operations (reserved for future optimizations)
|
|
let _gpu_v_t = self.backend.transpose(&gpu_v)?;
|
|
|
|
let mut prev_error = f32::INFINITY;
|
|
let mut converged = false;
|
|
let mut iterations = 0;
|
|
|
|
const EPSILON: f32 = 1e-10;
|
|
|
|
for iter in 0..self.config.max_iterations {
|
|
iterations = iter + 1;
|
|
|
|
// Update H: H = H .* (W^T @ V) ./ (W^T @ W @ H + eps)
|
|
let gpu_w_t = self.backend.transpose(&gpu_w)?;
|
|
let numerator_h = self.backend.matmul(&gpu_w_t, &gpu_v)?;
|
|
let wtw = self.backend.matmul(&gpu_w_t, &gpu_w)?;
|
|
let denominator_h = self.backend.matmul(&wtw, &gpu_h)?;
|
|
self.backend
|
|
.nmf_update_inplace(&mut gpu_h, &numerator_h, &denominator_h, EPSILON)?;
|
|
|
|
// Update W: W = W .* (V @ H^T) ./ (W @ H @ H^T + eps)
|
|
let gpu_h_t = self.backend.transpose(&gpu_h)?;
|
|
let numerator_w = self.backend.matmul(&gpu_v, &gpu_h_t)?;
|
|
let wh = self.backend.matmul(&gpu_w, &gpu_h)?;
|
|
let denominator_w = self.backend.matmul(&wh, &gpu_h_t)?;
|
|
self.backend
|
|
.nmf_update_inplace(&mut gpu_w, &numerator_w, &denominator_w, EPSILON)?;
|
|
|
|
// Check convergence every 10 iterations
|
|
if iter % 10 == 0 || iter == self.config.max_iterations - 1 {
|
|
// Compute reconstruction: approx = W @ H
|
|
let approx = self.backend.matmul(&gpu_w, &gpu_h)?;
|
|
let error = self.backend.frobenius_diff(&gpu_v, &approx)?;
|
|
|
|
debug!("Iteration {}: error = {:.6}", iter, error);
|
|
|
|
if (prev_error - error).abs() < self.config.tolerance {
|
|
info!("Converged at iteration {}", iter);
|
|
converged = true;
|
|
break;
|
|
}
|
|
|
|
prev_error = error;
|
|
}
|
|
}
|
|
|
|
// Download results
|
|
self.backend.synchronize()?;
|
|
let w_matrix = self.backend.download(&gpu_w)?;
|
|
let h_matrix = self.backend.download(&gpu_h)?;
|
|
|
|
info!(
|
|
"GPU NMF completed: {} iterations, error = {:.6}, converged = {}",
|
|
iterations, prev_error, converged
|
|
);
|
|
|
|
Ok(NmfOutput {
|
|
w_matrix,
|
|
h_matrix,
|
|
iterations,
|
|
reconstruction_error: prev_error,
|
|
converged,
|
|
})
|
|
}
|
|
|
|
/// Apply NMF to separate stains from optical density image
|
|
///
|
|
/// # Arguments
|
|
/// * `od_pixels` - Optical density values, shape (n_pixels, 3)
|
|
/// * `stain_matrix` - Known stain vectors, shape (3, n_components)
|
|
///
|
|
/// # Returns
|
|
/// Concentration values for each stain, shape (n_pixels, n_components)
|
|
pub fn unmix_stains(
|
|
&self,
|
|
od_pixels: &Array2<f32>,
|
|
stain_matrix: &Array2<f32>,
|
|
) -> Result<Array2<f32>> {
|
|
let (_n_pixels, n_channels) = od_pixels.dim();
|
|
let _n_components = stain_matrix.ncols();
|
|
|
|
if n_channels != 3 {
|
|
bail!("Expected 3 channels (RGB), got {}", n_channels);
|
|
}
|
|
if stain_matrix.nrows() != 3 {
|
|
bail!("Stain matrix must have 3 rows (RGB)");
|
|
}
|
|
|
|
// Transpose to (channels x pixels) for NMF
|
|
let v = od_pixels.t().to_owned();
|
|
|
|
// Use stain matrix as initial W
|
|
let result = self.decompose(&v, Some(stain_matrix))?;
|
|
|
|
// H is (n_components x n_pixels), transpose to (n_pixels x n_components)
|
|
Ok(result.h_matrix.t().to_owned())
|
|
}
|
|
}
|
|
|
|
/// Generate a random non-negative matrix
|
|
fn random_nonneg_matrix(rows: usize, cols: usize, seed: u64) -> Array2<f32> {
|
|
use rand::Rng;
|
|
use rand::SeedableRng;
|
|
use rand_chacha::ChaCha8Rng;
|
|
|
|
let mut rng = ChaCha8Rng::seed_from_u64(seed);
|
|
Array2::from_shape_fn((rows, cols), |_| rng.r#gen::<f32>() * 0.1 + 0.01)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn test_config() -> NmfProcessingConfig {
|
|
NmfProcessingConfig {
|
|
n_components: 2,
|
|
max_iterations: 50,
|
|
tolerance: 1e-3,
|
|
use_gpu: true,
|
|
random_seed: Some(42),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_nmf_basic() {
|
|
let config = test_config();
|
|
let processor = GpuNmfProcessor::new(config).unwrap();
|
|
|
|
// Create a simple test matrix
|
|
let v = Array2::from_shape_vec(
|
|
(3, 100),
|
|
(0..300).map(|i| (i as f32 / 300.0) + 0.1).collect(),
|
|
)
|
|
.unwrap();
|
|
|
|
let result = processor.decompose(&v, None).unwrap();
|
|
|
|
assert_eq!(result.w_matrix.shape(), &[3, 2]);
|
|
assert_eq!(result.h_matrix.shape(), &[2, 100]);
|
|
assert!(result.reconstruction_error.is_finite());
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_nmf_with_init() {
|
|
let config = test_config();
|
|
let processor = GpuNmfProcessor::new(config).unwrap();
|
|
|
|
// Create test matrix and initial W
|
|
let v = Array2::from_shape_fn((3, 50), |(i, j)| ((i + j) as f32 / 10.0) + 0.1);
|
|
let w_init = Array2::from_shape_fn((3, 2), |_| 0.5);
|
|
|
|
let result = processor.decompose(&v, Some(&w_init)).unwrap();
|
|
|
|
assert_eq!(result.w_matrix.shape(), &[3, 2]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_unmix_stains() {
|
|
let config = test_config();
|
|
let processor = GpuNmfProcessor::new(config).unwrap();
|
|
|
|
// Create synthetic OD pixels
|
|
let od_pixels =
|
|
Array2::from_shape_fn((100, 3), |(i, j)| ((i * 3 + j) as f32 / 300.0) + 0.1);
|
|
|
|
// Standard H&E stain matrix (simplified)
|
|
let stain_matrix =
|
|
Array2::from_shape_vec((3, 2), vec![0.65, 0.07, 0.70, 0.99, 0.29, 0.11]).unwrap();
|
|
|
|
let concentrations = processor.unmix_stains(&od_pixels, &stain_matrix).unwrap();
|
|
|
|
assert_eq!(concentrations.shape(), &[100, 2]);
|
|
|
|
// Concentrations should be non-negative
|
|
for &c in concentrations.iter() {
|
|
assert!(c >= 0.0);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_backend_detection() {
|
|
let config = test_config();
|
|
let processor = GpuNmfProcessor::new(config).unwrap();
|
|
|
|
// Should have some backend (at minimum CPU)
|
|
let backend_type = processor.backend_type();
|
|
println!("Using backend: {:?}", backend_type);
|
|
|
|
// backend_type should be one of the valid types
|
|
assert!(matches!(
|
|
backend_type,
|
|
BackendType::Cuda | BackendType::Metal | BackendType::Cpu
|
|
));
|
|
}
|
|
}
|