rust-NMF

A modern Rust port of R's NMF package plus 2019-2024 SOTA algorithms (HALS, E-HALS, RcppML-style diagonalised NMF). Six algorithms are bit-equivalent to R within f64 round-off; four are modern single-cell-tuned methods. Parallelised with rayon + matrixmultiply, on omicverse pbmc8k it's 76-280× faster than R depending on configuration.

Quick decision table

Key empirical finding (from examples/benchmark_scaling.ipynb init-pairing section, real pbmc8k 3000 cells × 1500 HVG, rank 8): NNDSVD init is the magic, not the algorithm. With NNDSVD seeding, every algorithm reaches its biological-signal ceiling in ~25 iterations. Random init takes 100 iters and still ends ~30% worse on ARI. Pair NNDSVD with lee/brunet/dnmf and you get the best of both worlds (R-equivalence + best biology, or speed + best biology).

Four one-liner recipes

import numpy as np
from nmf_rs import nmf, nndsvd_init

# 1. EXPLORATORY single-cell, best ARI vs cell types (0.89), still bit-eq R
W0, H0 = nndsvd_init(V, rank=10, fill='mean')
res = nmf(V, rank=10, method='lee',             # 'brunet' also reaches ARI 0.86
          W0=W0, H0=H0, max_iter=25, num_threads=16)

# 2. SPEED + biology, modern algorithm (ARI 0.85, 26% faster than recipe 1)
W0, H0 = nndsvd_init(V, rank=10, fill='mean')
res = nmf(V, rank=10, method='dnmf',
          W0=W0, H0=H0, max_iter=25, num_threads=16)

# 3. REPRODUCE R analysis exactly — bit-eq, max|ΔW| < 1e-12
res = nmf(V, rank=10, method='brunet',           # or 'lee', 'offset', 'nsNMF', 'snmf/r', 'snmf/l'
          W0=W0_from_R, H0=H0_from_R, max_iter=200, num_threads=16)

# 4. PUREST SPEED at atlas scale (>500k cells), ARI tradeoff
W0, H0 = nndsvd_init(V, rank=20, fill='mean')
res = nmf(V, rank=20, method='hals',
          W0=W0, H0=H0, max_iter=25, num_threads=16)

Algorithm catalogue

The six "✅ 1e-12" rows are bitwise-identical to R NMF within f64 round-off on every parity test (pytest tests/).

Speed (omicverse pbmc8k, HVG=2000, cells=7750, rank=10)

See examples/benchmark_pbmc8k.ipynb and examples/comparison_vs_cnmf.ipynb.

Objective accuracy (anchor-free, no algorithm bias)

The standard NMF evaluation metrics from Brunet 2004 and Kim & Park 2007 on real pbmc8k (3000 cells × 1500 HVG, rank 8, K=5 random inits per algorithm):

ARI vs cell types (real pbmc8k, 3k cells × 1.5k HVG, rank 8)

Pivoted by (algorithm × init) to expose the NNDSVD effect:

Wall-clock at the same configuration (16 threads)

Three takeaways

1. NNDSVD init dominates the choice of algorithm. Every algorithm gets higher ARI with NNDSVD than with random — typically 30-60% higher. So the question "which algorithm" is downstream of "use NNDSVD".

2. lee + NNDSVD is the gold standard for biological signal. On real single-cell data with predicted_celltype labels, lee + NNDSVD@25 reaches ARI = 0.89 — the best across all (algorithm × init) combinations. And lee is bit-equivalent to R NMF::nmf(method="lee") so you can reproduce R analyses exactly.

3. dnmf + NNDSVD is the fast alternative. Same NNDSVD trick on the modern (DeBruine 2024) diagonalised algorithm: 0.14 s, ARI 0.85. Trades 0.04 ARI for 26% speed.

Caveat: on synthetic data without label structure, the same recipes undertrain in 25 iters — for ground-truth recovery prefer lee/brunet with 100+ iters (medium tier in examples/benchmark_scaling.ipynb).

Tutorial roadmap

Suggested reading order for new users:

1. Start with tutorial_quickstart.ipynb for API.
2. Then benchmark_scaling.ipynb to pick the right algorithm for your scale + use case.
3. If reproducing an R analysis, benchmark_vs_R.ipynb.
4. If you've used cNMF before, comparison_vs_cnmf.ipynb.
5. For deep-dive: tutorial_full.ipynb.

Install

# Sherlock or any host with rust toolchain:
module load rust/1.90.0
maturin develop --release --manifest-path rust/Cargo.toml

# Or via pip into an active venv (release wheels coming to PyPI as `nmf-rs`):
VIRTUAL_ENV=/path/to/venv pip install -e .

Reproducing an R NMF analysis

# 1. Generate W0, H0 in R with set.seed(s); export as TSV.
# 2. Load same V, W0, H0 in Python:
import numpy as np
from nmf_rs import nmf

V  = np.loadtxt("V.tsv")
W0 = np.loadtxt("W0.tsv")
H0 = np.loadtxt("H0.tsv")
res = nmf(V, rank=W0.shape[1], W0=W0, H0=H0, method="brunet", max_iter=200)

W_R = np.loadtxt("W_from_R.tsv")
H_R = np.loadtxt("H_from_R.tsv")
np.testing.assert_allclose(res.W, W_R, atol=1e-12)   # passes exactly
np.testing.assert_allclose(res.H, H_R, atol=1e-12)

See tests/test_parity.py and tests/test_snmf_parity.py for the auto-regenerated reference fixtures + R-parity asserts.

License

GPL ≥ 2 (inherited from R NMF).