diff --git a/src/cnmf/nmf_gpu.py b/src/cnmf/nmf_gpu.py index 9e4c2cb..fed8f50 100644 --- a/src/cnmf/nmf_gpu.py +++ b/src/cnmf/nmf_gpu.py @@ -1,64 +1,40 @@ #!/usr/bin/env python -"""GPU NMF kernel — raw-PyTorch Frobenius multiplicative-update (MU) factorization. - -Standalone (no cnmf dependency); exposes a core NMF factorization and a thin `cNMF._nmf` adapter. - -nmf_kwargs (from cnmf): reads n_components, max_iter, tol, random_state, init, and - consensus refit args H/update_H. - Ignores beta_loss/solver/alpha_W/alpha_H/l1_ratio by design — always Frobenius MU, no regularization - (matches cnmf defaults; would diverge if a run enables regularization). Defaults are centralized in - `DEFAULT_NMF` and `DEFAULT_GPU`. Init runs through sklearn's _initialize_nmf for parity - (random|nndsvd|nndsvda|nndsvdar; None uses `DEFAULT_NMF["init"]`; 'custom' unsupported). - -gpu_kwargs (from Nextflow config): precedence gpu_kwargs dict > `DEFAULT_GPU`. - device auto|cuda|cuda:N|mps|cpu - dtype auto|fp32|fp64|bf16 - allow_tf32 TF32 for CUDA fp32 matmul - compile torch.compile the MU step (CPU/CUDA; ignored on MPS) - eps MU denominator guard, stored in compute dtype - check_every iters/block when eager - compile_block iters/block when compiled - -Resolution — auto adapts silently; an explicit-but-unavailable value RAISES: - device : auto = CUDA->MPS->CPU; explicit cuda/mps unavailable -> error - dtype : auto = fp64 on CPU, fp32 on GPU; explicit fp64 on MPS -> error - bf16 is explicit CUDA-only experimental storage + matmul dtype - - device storage(default) - ------ ---------------- - CPU fp64 (stable CPU reference) - CUDA fp32 (bf16 is explicit experimental storage/matmul; TF32 is for fp32 matmul) - MPS fp32 (no fp64 in hardware) - -Convergence (relative drop in ||X-WH|| < tol) is checked once per block, between blocks. Eager runs -short blocks (check_every) for a responsive dynamic stop; compiled runs fixed compile_block-sized -blocks (the host sync that reads the error stays outside the compiled region) and so converges -tighter for the same tol. - -allow_tf32 is a CUDA fp32-matmul performance preference (not a storage dtype), scoped to the -factorization and restored on exit; it is silently ignored where unsupported. compile is allowed on -CPU/CUDA and ignored on MPS. TF32 under compile also flips float32_matmul_precision to 'high', the -knob the inductor backend honors. bf16 uses bf16 storage and bf16 matmul; it is not fp32 storage with -bf16 temporary casts. Input X must be non-negative & finite. -Output dtype contract: returned arrays are float64 for compatibility; numerical accuracy is set by -the compute dtype. +"""PyTorch Frobenius-MU NMF backend for cNMF. + +The kernel is cNMF-compatible: it returns `(spectra, usages) = (H, W)` as numpy +float64, while compute dtype/device are controlled by `gpu_kwargs`. + +Supported runtime options: + device: auto|cuda|cuda:N|mps|cpu auto = CUDA -> MPS -> CPU + dtype: auto|fp32|fp64|bf16 auto = fp64 on CPU, fp32 on GPU + allow_tf32, compile, eps, check_every, compile_block, batch + +Explicit unavailable GPUs raise instead of falling back to CPU. MPS is fp32-only +for this kernel; bf16 is explicit CUDA-only storage and matmul. The implementation +also supports batched same-k replicates for factorize and fixed-H consensus refits. """ import contextlib +import functools +from collections import namedtuple import numpy as np # --------------------------------------------------------------------- -# Defaults and option parsing +# Defaults # --------------------------------------------------------------------- + + _TRUTHY = {"1", "true", "yes", "on"} + DEFAULT_NMF = { "max_iter": 1000, "tol": 1e-4, "init": "random", } + DEFAULT_GPU = { "device": "auto", "dtype": "auto", @@ -67,9 +43,15 @@ "eps": 1e-9, "check_every": 10, "compile_block": 1, + "batch": 1, # factorize replicates per launch } +# --------------------------------------------------------------------- +# CLI argument parsing and option resolution +# --------------------------------------------------------------------- + + def parse_gpu_args(parser): """Register cNMF CLI flags for the optional PyTorch GPU NMF engine.""" group = parser.add_argument_group("NMF engine options") @@ -81,6 +63,7 @@ def parse_gpu_args(parser): group.add_argument("--gpu-eps", type=float, help="[factorize,consensus,gpu] Multiplicative-update denominator guard") group.add_argument("--gpu-check-every", type=int, help="[factorize,consensus,gpu] Eager-mode convergence check interval") group.add_argument("--gpu-compile-block", type=int, help="[factorize,consensus,gpu] Number of MU iterations per compiled block") + group.add_argument("--gpu-batch", type=int, help="[factorize] Replicates run per GPU launch (batched MU); 1 = single-replicate") return parser @@ -94,6 +77,7 @@ def gpu_kwargs_from_args(args): "eps": args.gpu_eps, "check_every": args.gpu_check_every, "compile_block": args.gpu_compile_block, + "batch": args.gpu_batch, } if args.engine != "gpu": if any(value is not None for value in raw.values()): @@ -117,29 +101,13 @@ def validate_engine_args_for_command(args, available_commands): "gpu_eps", "gpu_check_every", "gpu_compile_block", + "gpu_batch", ] if any(getattr(args, name) is not None for name in gpu_arg_names): commands = ", ".join(available_commands) raise ValueError(f"NMF engine/GPU options are only valid with: {commands}") -def configure_nmf_engine(cnmf_obj, engine="cpu", gpu_kwargs=None): - """Configure a cNMF instance with an optional GPU NMF adapter.""" - if engine not in ("cpu", "gpu"): - raise ValueError("engine must be 'cpu' or 'gpu'") - if engine == "cpu": - return cnmf_obj - - def _gpu_nmf(X, nmf_kwargs): - nmf_kwargs = dict(nmf_kwargs) - nmf_kwargs["engine"] = "gpu" - nmf_kwargs["gpu"] = gpu_kwargs or {} - return _nmf_gpu(cnmf_obj, X, nmf_kwargs) - - cnmf_obj._nmf = _gpu_nmf - return cnmf_obj - - def _resolve_gpu_opts(gpu_kwargs): """Merge Nextflow-provided gpu_kwargs over defaults into a typed opts dict.""" raw = dict(gpu_kwargs or {}) @@ -162,14 +130,17 @@ def parse_positive_int(value, default): eps = parse_typed(raw.get("eps"), DEFAULT_GPU["eps"], float), check_every = parse_positive_int(raw.get("check_every"), DEFAULT_GPU["check_every"]), compile_block = parse_positive_int(raw.get("compile_block"), DEFAULT_GPU["compile_block"]), + batch = parse_positive_int(raw.get("batch"), DEFAULT_GPU["batch"]), ) # --------------------------------------------------------------------- -# Runtime backend selection +# Runtime backend selection (device / dtype / TF32) # --------------------------------------------------------------------- + + def _select_device(torch, requested): - """auto = CUDA->MPS->CPU; an explicit GPU that's unavailable raises (never a silent slow-CPU).""" + """Resolve device, raising for explicit unavailable CUDA/MPS requests.""" gpu_availability = { "cuda": torch.cuda.is_available, "mps": torch.backends.mps.is_available, @@ -193,7 +164,7 @@ def _select_device(torch, requested): def _select_storage(torch, requested, device): - """auto = fp64 on CPU, fp32 on GPU; bf16 is explicit CUDA-only experimental storage.""" + """Resolve storage/matmul dtype for the selected device.""" base = device.split(":")[0] dtype_map = { "fp32": torch.float32, @@ -225,10 +196,7 @@ def _select_storage(torch, requested, device): @contextlib.contextmanager def _cuda_tf32(torch, enable, device): - """Scope CUDA TF32 fp32-matmul to `enable`, restoring the previous globals on exit (no-op off - CUDA). Flips BOTH backends.cuda.matmul.allow_tf32 and float32_matmul_precision — the latter is - what torch.compile's inductor backend honors. TF32 only affects fp32 matmul and is silently - ignored by pre-Ampere hardware.""" + """Temporarily set CUDA TF32 matmul flags; no-op off CUDA.""" if not device.startswith("cuda") or not torch.cuda.is_available(): yield return @@ -244,8 +212,10 @@ def _cuda_tf32(torch, enable, device): # --------------------------------------------------------------------- -# Input validation, loud imports, and initialization +# Input validation, lazy imports, and initialization # --------------------------------------------------------------------- + + def _to_checked_array(X): """Materialize X dense and enforce the NMF preconditions: finite and non-negative.""" # TODO: sparse RAM path. Sparse X is still densified by this prototype; add a dense-size @@ -262,6 +232,22 @@ def _to_checked_array(X): return np.ascontiguousarray(Xnp) +def _to_checked_fixed_h(H, k, n_features): + """Materialize and validate fixed H for update_H=False consensus refits.""" + if H is None: + raise ValueError("update_H=False requires a fixed H matrix") + Hnp = H.toarray() if hasattr(H, "toarray") else np.asarray(H) + if Hnp.ndim != 2: + raise ValueError("fixed H must be a 2D matrix") + if Hnp.shape != (k, n_features): + raise ValueError(f"fixed H shape must be ({k}, {n_features}); got {Hnp.shape}") + if not np.isfinite(Hnp).all(): + raise ValueError("fixed H contains NaN/inf") + if Hnp.size and Hnp.min() < 0: + raise ValueError(f"NMF requires non-negative fixed H; found min(H) = {float(Hnp.min()):.4g}") + return Hnp + + def _loud_import_torch(): """Import torch with an actionable environment error for pipeline users.""" try: @@ -300,12 +286,10 @@ def _loud_import_initialize_nmf(): def _init_wh(Xnp, k, seed, init): - """sklearn-parity init via _initialize_nmf -> (W0, H0) as numpy. `init` None uses the centralized - NMF default. Explicit random|nndsvd|nndsvda|nndsvdar pass through; 'custom' is unsupported - (supply W/H yourself).""" + """Initialize W/H with sklearn parity; `init=None` uses DEFAULT_NMF.""" if init == "custom": raise NotImplementedError("nmf_gpu does not support init='custom'") - _initialize_nmf = _loud_import_initialize_nmf() # private, but pinned by the cnmf-parity goal + _initialize_nmf = _loud_import_initialize_nmf() return _initialize_nmf( Xnp, n_components=k, @@ -315,75 +299,38 @@ def _init_wh(Xnp, k, seed, init): # --------------------------------------------------------------------- -# Multiplicative-update kernel and execution plan +# Multiplicative-update steppers (one MU iteration; batch-aware) # --------------------------------------------------------------------- -def _mu_step(W, H, Xg, eps): - """One functional Frobenius MU update — W first (old H), then H (new W), matching the update - order of sklearn's `_fit_multiplicative_update` (update W, then H using the just-updated W). - Out-of-place (no `*=`) so torch.compile can trace it; grouped as factor*(num/den) for stable - rounding. Σ_b over cell-blocks collapses to these dense matmuls.""" - W = W * ((Xg @ H.T) / (W @ (H @ H.T) + eps)) # W *= XHᵀ / (W·HHᵀ) (uses old H) - H = H * ((W.T @ Xg) / ((W.T @ W) @ H + eps)) # H *= WᵀX / (WᵀW·H) (uses new W) - return W, H - - -def _mu_step_fixed_h(W, H, Xg, eps): - """One Frobenius MU update for consensus refit: keep H fixed and update W only.""" - return W * ((Xg @ H.T) / (W @ (H @ H.T) + eps)) -def _execution_plan(torch, opt, device, step_fn=_mu_step): - """Resolve compile policy and convergence-check cadence in one place. +def _mu_step(W, H, Xg, eps): + """One MU update, sklearn order: update W from old H, then H from new W. - MPS does not support torch.compile reliably for this path, so compile requests are treated as - eager there. CPU/CUDA compile is allowed. Compiled runs use compile_block because the host sync - for convergence checks must stay outside the compiled step; eager runs use check_every for - responsive stopping. + Accepts either 2D single-replicate tensors or stacked `[R,...]` replicate + tensors with shared `Xg[1,n,g]`. Operations are out-of-place for compile. """ - use_compile = opt["compile"] and not device.startswith("mps") - if use_compile: - return torch.compile(step_fn), opt["compile_block"] - return step_fn, opt["check_every"] - - -def _to_device_factors(torch, W0, H0, dtype, device): - """Move initialized numpy factors to the selected runtime backend.""" - W = torch.as_tensor(np.ascontiguousarray(W0), dtype=dtype, device=device) # usages (cells x k) - H = torch.as_tensor(np.ascontiguousarray(H0), dtype=dtype, device=device) # spectra (k x genes) + Ht = H.transpose(-2, -1) # [g,k] or [R,g,k] + W = W * ((Xg @ Ht) / (W @ (H @ Ht) + eps)) # W *= XHᵀ / (W·HHᵀ) (uses old H) + Wt = W.transpose(-2, -1) # [k,n] or [R,k,n] + H = H * ((Wt @ Xg) / ((Wt @ W) @ H + eps)) # H *= WᵀX / (WᵀW·H) (uses new W) return W, H -def _to_device_eps(torch, eps, dtype, device): - """Move the MU denominator guard to the same dtype/device as X/W/H.""" - return torch.tensor(eps, dtype=dtype, device=device) - - -def _check_runtime_tensors(Xg, W, H, eps): - """Ensure storage dtype is also the matmul dtype for the MU loop.""" - dtypes = {Xg.dtype, W.dtype, H.dtype, eps.dtype} - if len(dtypes) != 1: - raise RuntimeError(f"NMF runtime tensors must share dtype; got {sorted(map(str, dtypes))}") +def _mu_step_fixed_h(W, H, Xg, eps): + """One fixed-H MU update; only W changes. Supports 2D or stacked W.""" + Ht = H.transpose(-2, -1) # [g,k] or [1,g,k] + return W * ((Xg @ Ht) / (W @ (H @ Ht) + eps)) -def _to_checked_fixed_h(H, k, n_features): - """Materialize and validate fixed H for update_H=False consensus refits.""" - if H is None: - raise ValueError("update_H=False requires a fixed H matrix") - Hnp = H.toarray() if hasattr(H, "toarray") else np.asarray(H) - if Hnp.ndim != 2: - raise ValueError("fixed H must be a 2D matrix") - if Hnp.shape != (k, n_features): - raise ValueError(f"fixed H shape must be ({k}, {n_features}); got {Hnp.shape}") - if not np.isfinite(Hnp).all(): - raise ValueError("fixed H contains NaN/inf") - if Hnp.size and Hnp.min() < 0: - raise ValueError(f"NMF requires non-negative fixed H; found min(H) = {float(Hnp.min()):.4g}") - return Hnp +# --------------------------------------------------------------------- +# MU fit loops (iterate to convergence; batch-aware) +# --------------------------------------------------------------------- def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): - """Run Frobenius MU updates until convergence or max_iter.""" + """Run full MU until `max_iter` or all replicate slices meet `tol`.""" _check_runtime_tensors(Xg, W, H, eps) + lead = W.shape[:-2] # () unbatched, (R,) batched err_init = prev_err = None with torch.no_grad(), _cuda_tf32(torch, tf32, device): it = 0 @@ -392,21 +339,19 @@ def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): for _ in range(n): # MU updates run inside the (compiled) step W, H = step(W, H, Xg, eps) it += n - # stop on relative drop in direct ||X-WH|| < tol (host sync stays outside the compiled step) - err = float(torch.linalg.norm(Xg - W @ H)) + err = torch.linalg.norm((Xg - W @ H).reshape(*lead, -1), dim=-1) if err_init is None: - err_init = err - if err_init == 0.0: # degenerate (e.g. all-zero X): nothing to factor, avoid 0/0 - break - elif prev_err is not None and (prev_err - err) / err_init < tol: + err_init = err.clamp_min(1e-30) # avoid 0/0 on a degenerate (all-zero) slice + elif prev_err is not None and bool((((prev_err - err) / err_init) < tol).all()): break prev_err = err return W, H def _fit_mu_fixed_h(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): - """Run Frobenius MU updates for W with H fixed.""" + """Run fixed-H MU until `max_iter` or all replicate slices meet `tol`.""" _check_runtime_tensors(Xg, W, H, eps) + lead = W.shape[:-2] # () unbatched, (R,) batched err_init = prev_err = None with torch.no_grad(), _cuda_tf32(torch, tf32, device): it = 0 @@ -415,90 +360,212 @@ def _fit_mu_fixed_h(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, devi for _ in range(n): W = step(W, H, Xg, eps) it += n - err = float(torch.linalg.norm(Xg - W @ H)) + err = torch.linalg.norm((Xg - W @ H).reshape(*lead, -1), dim=-1) if err_init is None: - err_init = err - if err_init == 0.0: - break - elif prev_err is not None and (prev_err - err) / err_init < tol: + err_init = err.clamp_min(1e-30) # avoid 0/0 on a degenerate (all-zero) slice + elif prev_err is not None and bool((((prev_err - err) / err_init) < tol).all()): break prev_err = err return W # --------------------------------------------------------------------- -# Public API and adapters +# Execution plan, run context, and device-tensor helpers # --------------------------------------------------------------------- + + +def _execution_plan(torch, opt, device, step_fn=_mu_step): + """Return `(step_fn, block_size)` for eager or compiled execution.""" + use_compile = opt["compile"] and not device.startswith("mps") + if use_compile: + return torch.compile(step_fn), opt["compile_block"] + return step_fn, opt["check_every"] + + +# Runtime context resolved once per factorize call. +_GpuRun = namedtuple("_GpuRun", "opt device dtype k max_iter tol eps Xnp") + + +def _gpu_setup(torch, X, nmf_kwargs, gpu_kwargs): + """Validate common inputs and resolve device, dtype, eps, k, max_iter, and tol.""" + opt = _resolve_gpu_opts(gpu_kwargs) + device = _select_device(torch, opt["device"]) + dtype = _select_storage(torch, opt["dtype"], device) + k = int(nmf_kwargs["n_components"]) + if k < 1: + raise ValueError("n_components must be >= 1") + max_iter = int(nmf_kwargs.get("max_iter", DEFAULT_NMF["max_iter"])) + tol = float(nmf_kwargs.get("tol", DEFAULT_NMF["tol"])) + eps = _to_device_eps(torch, opt["eps"], dtype, device) + Xnp = _to_checked_array(X) + return _GpuRun(opt, device, dtype, k, max_iter, tol, eps, Xnp) + + +def _want_tf32(torch, rc): + """TF32 applies only to explicitly-allowed CUDA fp32 matmul.""" + return rc.device.startswith("cuda") and rc.opt["allow_tf32"] and rc.dtype is torch.float32 + + +def _to_device_factors(torch, W0, H0, dtype, device): + """Move initialized factors to runtime dtype/device.""" + W = torch.as_tensor(np.ascontiguousarray(W0), dtype=dtype, device=device) # usages (cells x k) + H = torch.as_tensor(np.ascontiguousarray(H0), dtype=dtype, device=device) # spectra (k x genes) + return W, H + + +def _to_device_eps(torch, eps, dtype, device): + """Create the MU denominator guard on runtime dtype/device.""" + return torch.tensor(eps, dtype=dtype, device=device) + + def _to_nmf_output(H, W): """Return spectra/usages as numpy float64 for compatibility; compute precision is unchanged.""" return H.cpu().double().numpy(), W.cpu().double().numpy() -def _factorize_nmf_gpu_fixed_h(torch, Xnp, Xg, nmf_kwargs, opt, device, dtype, k, max_iter, tol, eps, seed): - """Consensus refit path: use fixed H and update W only.""" - Hnp = _to_checked_fixed_h(nmf_kwargs.get("H"), k, Xnp.shape[1]) +def _check_runtime_tensors(Xg, W, H, eps): + """Require one dtype across X/W/H/eps.""" + dtypes = {Xg.dtype, W.dtype, H.dtype, eps.dtype} + if len(dtypes) != 1: + raise RuntimeError(f"NMF runtime tensors must share dtype; got {sorted(map(str, dtypes))}") + - W0, _ = _init_wh(Xnp, k, seed, nmf_kwargs.get("init")) - W, H = _to_device_factors(torch, W0, Hnp, dtype, device) +# --------------------------------------------------------------------- +# NMF kernels (batch-aware; R>=1 replicates per launch, R=1 = single instance) +# --------------------------------------------------------------------- - step, block = _execution_plan(torch, opt, device, _mu_step_fixed_h) - tf32 = device.startswith("cuda") and opt["allow_tf32"] and dtype is torch.float32 - W = _fit_mu_fixed_h(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device) - return _to_nmf_output(H, W) +def _nmf_gpu_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + """Run same-X, same-k MU replicates in one launch; return one `(H, W)` per seed.""" + torch = _loud_import_torch() + seeds = [None if s is None else int(s) for s in seeds] + if not seeds: + raise ValueError("seeds must be a non-empty list of per-replicate random states") -def _factorize_nmf_gpu_full(torch, Xnp, Xg, nmf_kwargs, opt, device, dtype, k, max_iter, tol, eps, seed): - """Normal NMF path: initialize and update both W and H.""" - # sklearn-parity init on CPU (_initialize_nmf), then move to device (cnmf uses init='random') - W0, H0 = _init_wh(Xnp, k, seed, nmf_kwargs.get("init")) - W, H = _to_device_factors(torch, W0, H0, dtype, device) + rc = _gpu_setup(torch, X, nmf_kwargs, gpu_kwargs) + init = nmf_kwargs.get("init") + # TODO: stream sparse/row-blocked X instead of requiring full dense X in RAM/VRAM. + Xb = torch.as_tensor(rc.Xnp, dtype=rc.dtype, device=rc.device).unsqueeze(0) # [1, n, g] shared - step, block = _execution_plan(torch, opt, device) - tf32 = device.startswith("cuda") and opt["allow_tf32"] and dtype is torch.float32 + # Initialize each replicate independently, then stack. + Ws, Hs = [], [] + for s in seeds: + W0, H0 = _init_wh(rc.Xnp, rc.k, s, init) # (usages, spectra) + Ws.append(np.ascontiguousarray(W0)); Hs.append(np.ascontiguousarray(H0)) + W = torch.as_tensor(np.stack(Ws, 0), dtype=rc.dtype, device=rc.device) # [R, n, k] + H = torch.as_tensor(np.stack(Hs, 0), dtype=rc.dtype, device=rc.device) # [R, k, g] - W, H = _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device) - return _to_nmf_output(H, W) + step, block = _execution_plan(torch, rc.opt, rc.device) + W, H = _fit_mu(torch, Xb, W, H, rc.eps, rc.max_iter, rc.tol, step, block, _want_tf32(torch, rc), rc.device) + Hc = H.cpu().double().numpy(); Wc = W.cpu().double().numpy() + return [(Hc[r], Wc[r]) for r in range(len(seeds))] -def factorize_nmf_gpu(X, nmf_kwargs, gpu_kwargs=None): - """Standalone Frobenius MU NMF: X -> (spectra, usages) = (H, W). - CUDA is optional: CPU and MPS paths do not require CUDA. CUDA execution - requires an installed PyTorch build where `torch.cuda.is_available()` is - true; this function does not pin a CUDA toolkit/runtime version itself. - Explicit `dtype='bf16'` additionally requires PyTorch to report - `torch.cuda.is_bf16_supported()` for the selected CUDA device. - """ +def _nmf_gpu_fixed_h(X, seeds, nmf_kwargs, gpu_kwargs=None): + """Run fixed-H consensus refits in one launch; return one `(H, W)` per seed.""" torch = _loud_import_torch() - opt = _resolve_gpu_opts(gpu_kwargs) - device = _select_device(torch, opt["device"]) - dtype = _select_storage(torch, opt["dtype"], device) + seeds = [None if s is None else int(s) for s in seeds] + if not seeds: + raise ValueError("seeds must be a non-empty list of per-replicate random states") - k = int(nmf_kwargs["n_components"]) - if k < 1: - raise ValueError("n_components must be >= 1") - max_iter = int(nmf_kwargs.get("max_iter", DEFAULT_NMF["max_iter"])) - tol = float(nmf_kwargs.get("tol", DEFAULT_NMF["tol"])) - eps = _to_device_eps(torch, opt["eps"], dtype, device) - seed = nmf_kwargs.get("random_state", None) + rc = _gpu_setup(torch, X, nmf_kwargs, gpu_kwargs) + init = nmf_kwargs.get("init") + Hnp = _to_checked_fixed_h(nmf_kwargs.get("H"), rc.k, rc.Xnp.shape[1]) # fixed spectra [k, g] + Xb = torch.as_tensor(rc.Xnp, dtype=rc.dtype, device=rc.device).unsqueeze(0) # [1, n, g] shared - Xnp = _to_checked_array(X) - # TODO: sparse VRAM path. This prototype copies dense X to device; future row-blocked/sparse MU - # should stream X blocks to GPU/MPS instead of requiring full dense X in VRAM. - Xg = torch.as_tensor(Xnp, dtype=dtype, device=device) # cells x genes + # Initialize W independently per seed; share fixed H across slices. + Ws = [] + for s in seeds: + W0, _ = _init_wh(rc.Xnp, rc.k, s, init) + Ws.append(np.ascontiguousarray(W0)) + W = torch.as_tensor(np.stack(Ws, 0), dtype=rc.dtype, device=rc.device) # [R, n, k] + H = torch.as_tensor(np.ascontiguousarray(Hnp), dtype=rc.dtype, device=rc.device).unsqueeze(0) # [1, k, g] fixed + + step, block = _execution_plan(torch, rc.opt, rc.device, _mu_step_fixed_h) + W = _fit_mu_fixed_h(torch, Xb, W, H, rc.eps, rc.max_iter, rc.tol, step, block, _want_tf32(torch, rc), rc.device) + + Hc = H.squeeze(0).cpu().double().numpy() # shared fixed spectra [k, g] + Wc = W.cpu().double().numpy() + return [(Hc, Wc[r]) for r in range(len(seeds))] + + +def _nmf_gpu(X, nmf_kwargs, gpu_kwargs=None): + """Single-replicate NMF API; dispatches to full MU or fixed-H refit.""" + kernel = _nmf_gpu_fixed_h if nmf_kwargs.get("update_H", True) is False else _nmf_gpu_mu + (result,) = kernel(X, [nmf_kwargs.get("random_state")], nmf_kwargs, gpu_kwargs) + return result - if nmf_kwargs.get("update_H", True) is False: - return _factorize_nmf_gpu_fixed_h(torch, Xnp, Xg, nmf_kwargs, opt, device, dtype, - k, max_iter, tol, eps, seed) - return _factorize_nmf_gpu_full(torch, Xnp, Xg, nmf_kwargs, opt, device, dtype, - k, max_iter, tol, eps, seed) + +# --------------------------------------------------------------------- +# cNMF integration: engine wiring and adapters +# --------------------------------------------------------------------- + + +def configure_nmf_engine(cnmf_obj, engine="cpu", gpu_kwargs=None): + """Install GPU `_nmf` and batched factorize hooks on a cNMF instance.""" + if engine not in ("cpu", "gpu"): + raise ValueError("engine must be 'cpu' or 'gpu'") + if engine == "cpu": + return cnmf_obj + + def _gpu_nmf(X, nmf_kwargs): + nmf_kwargs = dict(nmf_kwargs) + nmf_kwargs["engine"] = "gpu" + nmf_kwargs["gpu"] = gpu_kwargs or {} + return nmf_gpu(cnmf_obj, X, nmf_kwargs) + + cnmf_obj._nmf = _gpu_nmf + + # Factorize packs same-k replicates according to gpu_kwargs["batch"]. + cnmf_obj.factorize = functools.partial(factorize_gpu, cnmf_obj, gpu_kwargs or {}) + return cnmf_obj -def _nmf_gpu(self, X, nmf_kwargs): - """cNMF adapter: same core NMF kernel, with `self` ignored for monkeypatch compatibility.""" +def nmf_gpu(self, X, nmf_kwargs): + """cNMF `_nmf` adapter; `self` is ignored for monkeypatch compatibility.""" nmf_kwargs = dict(nmf_kwargs) gpu_kwargs = nmf_kwargs.pop("gpu", None) nmf_kwargs.pop("engine", None) - return factorize_nmf_gpu(X, nmf_kwargs, gpu_kwargs) + return _nmf_gpu(X, nmf_kwargs, gpu_kwargs) + + +def factorize_gpu(cnmf_obj, gpu_kwargs, worker_i=0, total_workers=1, skip_completed_runs=False): + """GPU `factorize` drop-in: group worker jobs by k, batch seeds, write iter spectra.""" + import scanpy as sc + import yaml + import pandas as pd + from collections import OrderedDict + from .cnmf import load_df_from_npz, save_df_to_npz, worker_filter + + batch = _resolve_gpu_opts(gpu_kwargs)["batch"] + + run_params = load_df_from_npz(cnmf_obj.paths['nmf_replicate_parameters']) + norm_counts = sc.read(cnmf_obj.paths['normalized_counts']) + base_kwargs = yaml.load(open(cnmf_obj.paths['nmf_run_parameters']), Loader=yaml.FullLoader) + + if not skip_completed_runs: + job_idx = worker_filter(range(len(run_params)), worker_i, total_workers) + else: + job_idx = worker_filter(run_params.index[run_params['completed'] == False], worker_i, total_workers) + + genes = norm_counts.var.index + by_k = OrderedDict() + for idx in job_idx: + p = run_params.iloc[idx, :] + by_k.setdefault(int(p['n_components']), []).append((int(p['iter']), int(p['nmf_seed']))) + + for k, jobs in by_k.items(): + run_kwargs = dict(base_kwargs); run_kwargs['n_components'] = k + for start in range(0, len(jobs), batch): + chunk = jobs[start:start + batch] + iters = [it for it, _ in chunk] + seeds = [s for _, s in chunk] + print('[Worker %d]. k=%d: launching %d replicate(s), iters=%s.' + % (worker_i, k, len(chunk), iters)) + results = _nmf_gpu_mu(norm_counts.X, seeds, run_kwargs, gpu_kwargs) + for (spectra, _usages), it in zip(results, iters): + spectra = pd.DataFrame(spectra, index=np.arange(1, k + 1), columns=genes) + save_df_to_npz(spectra, cnmf_obj.paths['iter_spectra'] % (k, it)) diff --git a/tests/test_nmf_gpu.py b/tests/test_nmf_gpu.py index 0252c64..a66ce7e 100644 --- a/tests/test_nmf_gpu.py +++ b/tests/test_nmf_gpu.py @@ -1,4 +1,4 @@ -"""Reliability tests for the standalone NMF GPU kernel (`bin/nmf_gpu.py`). +"""Reliability tests for the NMF GPU kernel (`src/cnmf/nmf_gpu.py`). Scope ----- @@ -12,7 +12,7 @@ Public API and reconstruction contract Verifies that factorization actually reduces reconstruction error, returns the cNMF-compatible `(spectra, usages) = (H, W)` order, keeps float64 numpy - outputs for compatibility, and preserves the thin `_nmf_gpu` adapter shape. + outputs for compatibility, and preserves the thin `nmf_gpu` adapter shape. MU update order, convergence, and iteration bounds Pins the sklearn-style W-then-H multiplicative-update order, early-stop @@ -58,11 +58,12 @@ small_nonnegative_matrix, ) + # --------------------------------------------------------------------- # Test harness # --------------------------------------------------------------------- def test_kernel_loader_fails_when_kernel_file_is_missing(tmp_path): - """Fail the test harness clearly if the standalone kernel script is absent.""" + """Fail the test harness clearly if the kernel module is absent.""" missing_kernel = tmp_path / "missing_nmf_gpu.py" with pytest.raises(pytest.fail.Exception, match="Required NMF GPU kernel file is missing"): @@ -72,13 +73,13 @@ def test_kernel_loader_fails_when_kernel_file_is_missing(tmp_path): # --------------------------------------------------------------------- # Public API and reconstruction contract # --------------------------------------------------------------------- -def test_factorize_nmf_gpu_reconstructs_known_low_rank_matrix_with_small_relative_error(kernel): +def test_nmf_gpu_reconstructs_known_low_rank_matrix_with_small_relative_error(kernel): """Factorize an exact low-rank non-negative matrix and require low relative error.""" require_nmf_runtime() k = 3 X = low_rank_matrix(rank=k) - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": k, "max_iter": 600, "tol": 0, "random_state": 0}, {"device": "cpu", "check_every": 600}, @@ -89,12 +90,12 @@ def test_factorize_nmf_gpu_reconstructs_known_low_rank_matrix_with_small_relativ assert rel < 1e-3 -def test_factorize_nmf_gpu_returns_spectra_then_usages_with_cnmf_orientation(kernel): +def test_nmf_gpu_returns_spectra_then_usages_with_cnmf_orientation(kernel): """Pin the public return order as spectra H first, usages W second.""" require_nmf_runtime() X = small_nonnegative_matrix(cells=7, genes=5) - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 2, "max_iter": 2, "random_state": 0}, {"device": "cpu"}, @@ -104,11 +105,11 @@ def test_factorize_nmf_gpu_returns_spectra_then_usages_with_cnmf_orientation(ker assert W.shape == (7, 2) -def test_factorize_nmf_gpu_cpu_smoke_shapes_dtype_sign_and_finiteness(kernel): +def test_nmf_gpu_cpu_smoke_shapes_dtype_sign_and_finiteness(kernel): """Smoke-test CPU output shape, float64 compatibility dtype, finite values, and non-negativity.""" require_nmf_runtime() X = small_nonnegative_matrix() - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 3, "max_iter": 3, "random_state": 0}, {"device": "cpu"}, @@ -117,12 +118,12 @@ def test_factorize_nmf_gpu_cpu_smoke_shapes_dtype_sign_and_finiteness(kernel): assert_valid_nmf_output(X, H, W, 3) -def test_factorize_nmf_gpu_fp32_compute_still_returns_float64_numpy_outputs(kernel): +def test_nmf_gpu_fp32_compute_still_returns_float64_numpy_outputs(kernel): """Exercise fp32 compute while keeping the public numpy output contract as float64.""" require_nmf_runtime() X = small_nonnegative_matrix() - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 2, "max_iter": 1, "random_state": 0}, {"device": "cpu", "dtype": "fp32"}, @@ -132,7 +133,7 @@ def test_factorize_nmf_gpu_fp32_compute_still_returns_float64_numpy_outputs(kern assert W.dtype == np.float64 -def test_nmf_gpu_adapter_ignores_self_and_delegates_to_factorize_nmf_gpu(kernel, monkeypatch): +def test_nmf_gpu_adapter_ignores_self_and_delegates_to_nmf_gpu(kernel, monkeypatch): """Verify the cNMF adapter extracts embedded GPU args and ignores its bound `self`.""" calls = [] sentinel = (object(), object()) @@ -141,12 +142,12 @@ def fake_factorize(X, nmf_kwargs, gpu_kwargs=None): calls.append((X, nmf_kwargs, gpu_kwargs)) return sentinel - monkeypatch.setattr(kernel, "factorize_nmf_gpu", fake_factorize) + monkeypatch.setattr(kernel, "_nmf_gpu", fake_factorize) X = np.ones((3, 2)) gpu_kwargs = {"device": "cpu"} nmf_kwargs = {"engine": "gpu", "gpu": gpu_kwargs, "n_components": 1} - result = kernel._nmf_gpu(object(), X, nmf_kwargs) + result = kernel.nmf_gpu(object(), X, nmf_kwargs) assert result is sentinel assert calls == [(X, {"n_components": 1}, gpu_kwargs)] @@ -230,12 +231,12 @@ def test_compile_mode_matches_eager_output_for_same_seed_and_options(kernel, mon X = small_nonnegative_matrix(cells=8, genes=6) nmf_kwargs = {"n_components": 2, "max_iter": 4, "tol": -1.0, "random_state": 0} - eager_H, eager_W = kernel.factorize_nmf_gpu( + eager_H, eager_W = kernel._nmf_gpu( X, nmf_kwargs, {"device": "cpu", "dtype": "fp64", "compile": False, "check_every": 1}, ) - compiled_H, compiled_W = kernel.factorize_nmf_gpu( + compiled_H, compiled_W = kernel._nmf_gpu( X, nmf_kwargs, {"device": "cpu", "dtype": "fp64", "compile": True, "compile_block": 2}, @@ -262,28 +263,28 @@ def test_compile_mode_uses_explicit_multi_iteration_compile_block_when_requested # --------------------------------------------------------------------- # Initialization and reproducibility # --------------------------------------------------------------------- -def test_factorize_nmf_gpu_random_state_is_reproducible(kernel): +def test_nmf_gpu_random_state_is_reproducible(kernel): """The same random_state should produce identical initialization and final factors.""" require_nmf_runtime() X = small_nonnegative_matrix() kwargs = {"n_components": 3, "max_iter": 3, "random_state": 13} gpu = {"device": "cpu", "check_every": 3} - H1, W1 = kernel.factorize_nmf_gpu(X, kwargs, gpu) - H2, W2 = kernel.factorize_nmf_gpu(X, kwargs, gpu) + H1, W1 = kernel._nmf_gpu(X, kwargs, gpu) + H2, W2 = kernel._nmf_gpu(X, kwargs, gpu) assert np.allclose(H1, H2) assert np.allclose(W1, W2) -def test_factorize_nmf_gpu_different_random_state_changes_result(kernel): +def test_nmf_gpu_different_random_state_changes_result(kernel): """Different random_state values should produce different random initial factors.""" require_nmf_runtime() X = small_nonnegative_matrix() kwargs = {"n_components": 3, "max_iter": 0, "init": "random"} - H1, W1 = kernel.factorize_nmf_gpu(X, dict(kwargs, random_state=1), {"device": "cpu"}) - H2, W2 = kernel.factorize_nmf_gpu(X, dict(kwargs, random_state=2), {"device": "cpu"}) + H1, W1 = kernel._nmf_gpu(X, dict(kwargs, random_state=1), {"device": "cpu"}) + H2, W2 = kernel._nmf_gpu(X, dict(kwargs, random_state=2), {"device": "cpu"}) assert not np.allclose(H1, H2) assert not np.allclose(W1, W2) @@ -333,7 +334,7 @@ def fake_initialize(X, n_components, init, random_state): assert seen == ["nndsvd", "nndsvda", "nndsvdar"] -def test_factorize_nmf_gpu_custom_init_raises(kernel): +def test_nmf_gpu_custom_init_raises(kernel): """Document that custom W/H initialization is not implemented in this standalone API.""" with pytest.raises(NotImplementedError, match="custom"): kernel._init_wh(small_nonnegative_matrix(), 2, 0, "custom") @@ -342,30 +343,30 @@ def test_factorize_nmf_gpu_custom_init_raises(kernel): # --------------------------------------------------------------------- # Input validation and degenerate shapes # --------------------------------------------------------------------- -def test_factorize_nmf_gpu_rejects_negative_input(kernel): +def test_nmf_gpu_rejects_negative_input(kernel): """NMF input must be non-negative.""" with pytest.raises(ValueError, match="non-negative"): kernel._to_checked_array(np.array([[1.0, -0.1]])) -def test_factorize_nmf_gpu_rejects_nan_input(kernel): +def test_nmf_gpu_rejects_nan_input(kernel): """NaN input should fail before torch/sklearn runtime work begins.""" with pytest.raises(ValueError, match="NaN/inf"): kernel._to_checked_array(np.array([[1.0, np.nan]])) -def test_factorize_nmf_gpu_rejects_inf_input(kernel): +def test_nmf_gpu_rejects_inf_input(kernel): """Infinite input should fail before torch/sklearn runtime work begins.""" with pytest.raises(ValueError, match="NaN/inf"): kernel._to_checked_array(np.array([[1.0, np.inf]])) -def test_factorize_nmf_gpu_zero_matrix_does_not_crash_or_divide_by_zero(kernel): +def test_nmf_gpu_zero_matrix_does_not_crash_or_divide_by_zero(kernel): """All-zero input should take the degenerate path without crashing or producing invalid output.""" require_nmf_runtime() X = np.zeros((5, 4)) - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 2, "max_iter": 5, "random_state": 0}, {"device": "cpu"}, @@ -374,16 +375,16 @@ def test_factorize_nmf_gpu_zero_matrix_does_not_crash_or_divide_by_zero(kernel): assert_valid_nmf_output(X, H, W, 2) -def test_factorize_nmf_gpu_handles_single_row_and_single_column_inputs(kernel): +def test_nmf_gpu_handles_single_row_and_single_column_inputs(kernel): """Single-row and single-column matrices should keep valid H/W orientation.""" require_nmf_runtime() - H_row, W_row = kernel.factorize_nmf_gpu( + H_row, W_row = kernel._nmf_gpu( np.array([[1.0, 2.0, 3.0]]), {"n_components": 1, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, ) - H_col, W_col = kernel.factorize_nmf_gpu( + H_col, W_col = kernel._nmf_gpu( np.array([[1.0], [2.0], [3.0]]), {"n_components": 1, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, @@ -395,7 +396,7 @@ def test_factorize_nmf_gpu_handles_single_row_and_single_column_inputs(kernel): assert W_col.shape == (3, 1) -def test_factorize_nmf_gpu_rejects_empty_or_zero_dimensional_inputs(kernel): +def test_nmf_gpu_rejects_empty_or_zero_dimensional_inputs(kernel): """Reject empty matrices and non-2D arrays with clear validation errors.""" for X in (np.empty((0, 3)), np.empty((3, 0))): with pytest.raises(ValueError, match="at least one row"): @@ -405,23 +406,23 @@ def test_factorize_nmf_gpu_rejects_empty_or_zero_dimensional_inputs(kernel): kernel._to_checked_array(np.array([1.0, 2.0])) -def test_factorize_nmf_gpu_rejects_zero_components(kernel): +def test_nmf_gpu_rejects_zero_components(kernel): """Reject rank k=0 before reaching sklearn's initializer.""" require_nmf_runtime() with pytest.raises(ValueError, match="n_components"): - kernel.factorize_nmf_gpu( + kernel._nmf_gpu( np.ones((3, 3)), {"n_components": 0, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, ) -def test_factorize_nmf_gpu_defines_behavior_when_k_exceeds_min_dimension(kernel): +def test_nmf_gpu_defines_behavior_when_k_exceeds_min_dimension(kernel): """Allow sklearn-compatible overcomplete factorization when k exceeds matrix dimensions.""" require_nmf_runtime() X = small_nonnegative_matrix(cells=3, genes=2) - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 4, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, @@ -460,6 +461,7 @@ def test_resolve_gpu_opts_dict_values_override_defaults(kernel): "eps": 1e-8, "check_every": 7, "compile_block": 9, + "batch": 1, # not overridden here -> default } @@ -507,181 +509,6 @@ def test_resolve_gpu_opts_floors_check_every_and_compile_block_to_at_least_one(k assert opts["compile_block"] == 1 -# --------------------------------------------------------------------- -# GPU CLI and cNMF engine wiring -# --------------------------------------------------------------------- -def test_parse_gpu_args_defaults_to_none_until_user_sets_engine(kernel): - """CLI flags should default to None so absent options are distinguishable from explicit values.""" - parser = argparse.ArgumentParser() - kernel.parse_gpu_args(parser) - - args = parser.parse_args([]) - - assert args.engine is None - for name in ("gpu_device", "gpu_dtype", "gpu_allow_tf32", "gpu_compile", - "gpu_eps", "gpu_check_every", "gpu_compile_block"): - assert getattr(args, name) is None, f"{name} should default to None" - - -def test_gpu_kwargs_from_args_rejects_gpu_options_without_gpu_engine(kernel): - """GPU-specific CLI options should raise unless `--engine gpu` was explicitly selected.""" - parser = argparse.ArgumentParser() - kernel.parse_gpu_args(parser) - - for argv in (["--gpu-device", "cuda"], ["--engine", "cpu", "--gpu-device", "cuda"]): - with pytest.raises(ValueError, match="require --engine gpu"): - kernel.gpu_kwargs_from_args(parser.parse_args(argv)) - - -def test_gpu_kwargs_from_args_fills_defaults_when_gpu_engine_selected(kernel): - """`--engine gpu` alone should resolve missing GPU options from DEFAULT_GPU.""" - parser = argparse.ArgumentParser() - kernel.parse_gpu_args(parser) - - args = parser.parse_args(["--engine", "gpu"]) - - assert kernel.gpu_kwargs_from_args(args) == kernel.DEFAULT_GPU - - -def test_gpu_kwargs_from_args_normalizes_cli_overrides(kernel): - """GPU CLI override values should normalize through the same resolver as config dict values.""" - parser = argparse.ArgumentParser() - kernel.parse_gpu_args(parser) - - args = parser.parse_args([ - "--engine", "gpu", - "--gpu-device", "CUDA:0", # device is lower-cased by _resolve_gpu_opts - "--gpu-dtype", "FP32", - "--gpu-allow-tf32", # store_const flags -> True - "--gpu-compile", - "--gpu-eps", "1e-8", - "--gpu-check-every", "5", - "--gpu-compile-block", "100", - ]) - - assert kernel.gpu_kwargs_from_args(args) == { - "device": "cuda:0", - "dtype": "fp32", - "allow_tf32": True, - "compile": True, - "eps": 1e-8, - "check_every": 5, - "compile_block": 100, - } - - -def test_validate_engine_args_for_command_rejects_non_engine_command_gpu_options(kernel): - """Engine/GPU options should be accepted only for factorize and consensus.""" - parser = argparse.ArgumentParser() - parser.add_argument("command") - kernel.parse_gpu_args(parser) - supported = ("factorize", "consensus") - - # factorize and consensus accept engine/GPU options (no raise) - for argv in ( - ["factorize", "--engine", "gpu"], - ["consensus", "--engine", "gpu"], - ["consensus", "--gpu-device", "cuda"], - ): - kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) - - # non-engine commands carrying engine/GPU options are rejected - for argv in ( - ["prepare", "--engine", "gpu"], - ["combine", "--gpu-device", "cuda"], - ["k_selection_plot", "--gpu-dtype", "fp32"], - ): - with pytest.raises(ValueError, match="only valid with"): - kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) - - # non-engine commands without engine/GPU options are fine - kernel.validate_engine_args_for_command(parser.parse_args(["prepare"]), supported) - - -# --------------------------------------------------------------------- -# GPU CLI and cNMF consensus wiring -# --------------------------------------------------------------------- -def test_validate_engine_args_accepts_consensus_gpu_options(kernel): - """`consensus --engine gpu` and consensus GPU flags should be valid CLI input.""" - parser = argparse.ArgumentParser() - parser.add_argument("command") - kernel.parse_gpu_args(parser) - supported = ("factorize", "consensus") - - for argv in ( - ["consensus", "--engine", "gpu"], - ["consensus", "--engine", "gpu", "--gpu-device", "CUDA:0", "--gpu-dtype", "FP32"], - ["consensus", "--gpu-allow-tf32", "--gpu-compile"], - ): - kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) - - -def test_validate_engine_args_rejects_gpu_options_for_non_engine_commands(kernel): - """GPU flags should still be rejected for prepare/combine/k_selection_plot.""" - parser = argparse.ArgumentParser() - parser.add_argument("command") - kernel.parse_gpu_args(parser) - supported = ("factorize", "consensus") - - for argv in ( - ["prepare", "--engine", "gpu"], - ["combine", "--gpu-device", "cuda"], - ["k_selection_plot", "--gpu-check-every", "2"], - ): - with pytest.raises(ValueError, match="only valid with"): - kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) - - -def test_configure_nmf_engine_cpu_leaves_cnmf_instance_unchanged(kernel): - """The default CPU engine should be a no-op so existing sklearn behavior is preserved.""" - class DummyCNMF: - def _nmf(self, X, nmf_kwargs): - return "sklearn-path" - - obj = DummyCNMF() - result = kernel.configure_nmf_engine(obj, engine="cpu", gpu_kwargs={"device": "cuda"}) - - assert result is obj - assert "_nmf" not in vars(obj) # no instance override added - assert obj._nmf("X", {}) == "sklearn-path" # original class method intact - - -def test_configure_nmf_engine_rejects_unknown_engine(kernel): - """Unknown engine names should fail loudly instead of silently using CPU.""" - with pytest.raises(ValueError, match="engine must be 'cpu' or 'gpu'"): - kernel.configure_nmf_engine(object(), engine="tpu") - - -def test_configure_nmf_engine_gpu_patches_instance_nmf_with_adapter(kernel, monkeypatch): - """The GPU engine should replace the instance `_nmf` callable with the GPU adapter path.""" - captured = {} - - def fake_nmf_gpu(self, X, nmf_kwargs): - captured["self"], captured["X"], captured["nmf_kwargs"] = self, X, dict(nmf_kwargs) - return ("spectra", "usages") - - monkeypatch.setattr(kernel, "_nmf_gpu", fake_nmf_gpu) - - class DummyCNMF: - def _nmf(self, X, nmf_kwargs): - return "sklearn-path" - - obj = DummyCNMF() - gpu_kwargs = {"device": "cuda", "dtype": "fp32"} - result = kernel.configure_nmf_engine(obj, engine="gpu", gpu_kwargs=gpu_kwargs) - - assert result is obj - assert "_nmf" in vars(obj) # instance _nmf is now overridden - - out = obj._nmf("Xdata", {"n_components": 5}) - - assert out == ("spectra", "usages") # dispatched through the GPU adapter - assert captured["self"] is obj and captured["X"] == "Xdata" - assert captured["nmf_kwargs"]["engine"] == "gpu" # engine + gpu kwargs embedded first - assert captured["nmf_kwargs"]["gpu"] == gpu_kwargs - assert captured["nmf_kwargs"]["n_components"] == 5 # original kwargs preserved - - # --------------------------------------------------------------------- # Device selection # --------------------------------------------------------------------- @@ -806,7 +633,7 @@ def test_sparse_input_uses_densify_path_and_returns_valid_output(kernel): sparse = pytest.importorskip("scipy.sparse") X = sparse.csr_matrix(small_nonnegative_matrix(cells=6, genes=5)) - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 2, "max_iter": 2, "random_state": 0}, {"device": "cpu"}, @@ -821,7 +648,7 @@ def test_nndsvd_nndsvda_nndsvdar_initializers_return_valid_outputs(kernel): X = small_nonnegative_matrix(cells=8, genes=6) for init in ("nndsvd", "nndsvda", "nndsvdar"): - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 3, "max_iter": 1, "random_state": 0, "init": init}, {"device": "cpu"}, @@ -830,166 +657,74 @@ def test_nndsvd_nndsvda_nndsvdar_initializers_return_valid_outputs(kernel): # --------------------------------------------------------------------- -# Consensus fixed-H kernel behavior +# Backend-specific execution behavior # --------------------------------------------------------------------- -def test_factorize_nmf_gpu_update_h_false_keeps_fixed_h_and_updates_only_w(kernel): - """Consensus refit should preserve supplied H and optimize usages W only.""" +def test_execution_plan_ignores_compile_on_mps_and_uses_eager_path(kernel): + """MPS compile requests should resolve to eager execution with check_every cadence.""" + fake_torch = SimpleNamespace(compile=lambda fn: pytest.fail("compile should be ignored on MPS")) + opt = dict(kernel.DEFAULT_GPU, compile=True, check_every=4, compile_block=9) + + step, block = kernel._execution_plan(fake_torch, opt, "mps") + + assert step is kernel._mu_step + assert block == 4 + + +def test_tf32_flag_is_applied_only_for_cuda_float32_compute(kernel, monkeypatch): + """allow_tf32 should not be passed into the fit loop for non-CUDA execution.""" require_nmf_runtime() - fixed_h = np.array([[1.0, 0.3, 0.6], [0.2, 1.1, 0.4]], dtype=np.float64) - true_w = np.array([[1.0, 0.5], [0.4, 1.2], [1.5, 0.3], [0.7, 0.9]], dtype=np.float64) - X = true_w @ fixed_h + captured = [] + + def fake_fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): + captured.append((tf32, device, Xg.dtype)) + return W, H + + monkeypatch.setattr(kernel, "_fit_mu", fake_fit_mu) + X = small_nonnegative_matrix(cells=4, genes=3) - H, W = kernel.factorize_nmf_gpu( + kernel._nmf_gpu( X, - {"n_components": 2, "max_iter": 5, "random_state": 0, "update_H": False, "H": fixed_h}, - {"device": "cpu", "dtype": "fp64", "check_every": 5}, + {"n_components": 2, "max_iter": 1, "random_state": 0}, + {"device": "cpu", "dtype": "fp32", "allow_tf32": True}, ) - assert_valid_nmf_output(X, H, W, 2) - assert np.allclose(H, fixed_h) + assert captured[-1][0] is False -def test_to_checked_fixed_h_rejects_missing_invalid_or_incompatible_h(kernel): - """Fixed-H consensus refit should fail clearly for invalid supplied spectra.""" - with pytest.raises(ValueError, match="requires a fixed H"): - kernel._to_checked_fixed_h(None, 2, 3) +def test_tf32_scope_is_noop_off_cuda(kernel): + """The TF32 context manager should leave backend globals untouched off CUDA.""" + class FakeTorch: + def __init__(self): + self.cuda = SimpleNamespace(is_available=lambda: True) + self.backends = SimpleNamespace( + cuda=SimpleNamespace(matmul=SimpleNamespace(allow_tf32=False)) + ) + self.precision = "highest" - invalid_cases = [ - (np.array([1.0, 2.0, 3.0]), "2D"), - (np.ones((3, 3)), "shape"), - (np.array([[1.0, np.nan, 0.2], [0.4, 0.5, 0.6]]), "NaN/inf"), - (np.array([[1.0, -0.1, 0.2], [0.4, 0.5, 0.6]]), "non-negative"), - ] - for H, message in invalid_cases: - with pytest.raises(ValueError, match=message): - kernel._to_checked_fixed_h(H, 2, 3) + def get_float32_matmul_precision(self): + return self.precision + def set_float32_matmul_precision(self, value): + self.precision = value -def test_mu_step_fixed_h_matches_manual_w_only_update(kernel): - """One fixed-H MU step should match the manual W-only Frobenius update.""" - torch = require_nmf_runtime() - Xg = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64) - W0 = torch.tensor([[0.5, 0.7], [0.9, 1.1]], dtype=torch.float64) - H = torch.tensor([[0.6, 0.8], [1.0, 1.2]], dtype=torch.float64) - eps = torch.tensor(1e-9, dtype=torch.float64) - H_before = H.clone() + fake = FakeTorch() - expected_W = W0 * ((Xg @ H.T) / (W0 @ (H @ H.T) + eps)) - W = kernel._mu_step_fixed_h(W0, H, Xg, eps) + with kernel._cuda_tf32(fake, True, "cpu"): + assert fake.backends.cuda.matmul.allow_tf32 is False + assert fake.precision == "highest" - assert torch.allclose(W, expected_W) - assert torch.allclose(H, H_before) + assert fake.backends.cuda.matmul.allow_tf32 is False + assert fake.precision == "highest" -def test_fit_mu_fixed_h_respects_early_stop_and_max_iter_block_bounds(kernel): - """Fixed-H fit loop should share early-stop and no-overrun guarantees with full MU.""" - torch = require_nmf_runtime() - Xg = torch.full((3, 2), 2.0, dtype=torch.float64) - W = torch.ones((3, 1), dtype=torch.float64) - H = torch.ones((1, 2), dtype=torch.float64) - eps = torch.tensor(1e-9, dtype=torch.float64) - - early_calls = {"count": 0} - - def no_change_step_early(W, H, Xg, eps): - early_calls["count"] += 1 - return W - - kernel._fit_mu_fixed_h(torch, Xg, W, H, eps, 10, 1e-4, no_change_step_early, 1, False, "cpu") - assert early_calls["count"] == 2 - - max_iter_calls = {"count": 0} - - def no_change_step_max_iter(W, H, Xg, eps): - max_iter_calls["count"] += 1 - return W - - kernel._fit_mu_fixed_h(torch, Xg, W, H, eps, 6, -1.0, no_change_step_max_iter, 4, False, "cpu") - assert max_iter_calls["count"] == 6 - - -def test_execution_plan_for_fixed_h_compile_uses_fixed_h_step_and_compile_block(kernel): - """Compiled consensus refit should compile _mu_step_fixed_h and use compile_block.""" - calls = [] - fake_torch = SimpleNamespace(compile=lambda fn: calls.append(fn) or fn) - opt = dict(kernel.DEFAULT_GPU, compile=True, check_every=1, compile_block=3) - - step, block = kernel._execution_plan(fake_torch, opt, "cpu", kernel._mu_step_fixed_h) - - assert calls == [kernel._mu_step_fixed_h] - assert step is kernel._mu_step_fixed_h - assert block == 3 - - -# --------------------------------------------------------------------- -# Backend-specific execution behavior -# --------------------------------------------------------------------- -def test_execution_plan_ignores_compile_on_mps_and_uses_eager_path(kernel): - """MPS compile requests should resolve to eager execution with check_every cadence.""" - fake_torch = SimpleNamespace(compile=lambda fn: pytest.fail("compile should be ignored on MPS")) - opt = dict(kernel.DEFAULT_GPU, compile=True, check_every=4, compile_block=9) - - step, block = kernel._execution_plan(fake_torch, opt, "mps") - - assert step is kernel._mu_step - assert block == 4 - - -def test_tf32_flag_is_applied_only_for_cuda_float32_compute(kernel, monkeypatch): - """allow_tf32 should not be passed into the fit loop for non-CUDA execution.""" - require_nmf_runtime() - captured = [] - - def fake_fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): - captured.append((tf32, device, Xg.dtype)) - return W, H - - monkeypatch.setattr(kernel, "_fit_mu", fake_fit_mu) - X = small_nonnegative_matrix(cells=4, genes=3) - - kernel.factorize_nmf_gpu( - X, - {"n_components": 2, "max_iter": 1, "random_state": 0}, - {"device": "cpu", "dtype": "fp32", "allow_tf32": True}, - ) - - assert captured[-1][0] is False - - -def test_tf32_scope_is_noop_off_cuda(kernel): - """The TF32 context manager should leave backend globals untouched off CUDA.""" - class FakeTorch: - def __init__(self): - self.cuda = SimpleNamespace(is_available=lambda: True) - self.backends = SimpleNamespace( - cuda=SimpleNamespace(matmul=SimpleNamespace(allow_tf32=False)) - ) - self.precision = "highest" - - def get_float32_matmul_precision(self): - return self.precision - - def set_float32_matmul_precision(self, value): - self.precision = value - - fake = FakeTorch() - - with kernel._cuda_tf32(fake, True, "cpu"): - assert fake.backends.cuda.matmul.allow_tf32 is False - assert fake.precision == "highest" - - assert fake.backends.cuda.matmul.allow_tf32 is False - assert fake.precision == "highest" - - -def test_cuda_fp32_smoke_when_gpu_available(kernel): - """CUDA fp32 should run when PyTorch reports CUDA available; no CUDA version is pinned here.""" +def test_cuda_fp32_smoke_when_gpu_available(kernel): + """CUDA fp32 should run when PyTorch reports CUDA available; no CUDA version is pinned here.""" torch = require_nmf_runtime() if not torch.cuda.is_available(): pytest.skip("CUDA is not available") X = small_nonnegative_matrix(cells=6, genes=5) - H, W = kernel.factorize_nmf_gpu( + H, W = kernel._nmf_gpu( X, {"n_components": 2, "max_iter": 2, "random_state": 0}, {"device": "cuda", "dtype": "fp32"}, @@ -1014,7 +749,7 @@ def fake_fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): monkeypatch.setattr(kernel, "_fit_mu", fake_fit_mu) X = small_nonnegative_matrix(cells=4, genes=3) - kernel.factorize_nmf_gpu( + kernel._nmf_gpu( X, {"n_components": 2, "max_iter": 1, "random_state": 0}, {"device": "cuda", "dtype": "bf16"}, @@ -1037,3 +772,456 @@ def test_cuda_allow_tf32_scope_restores_previous_state_when_gpu_available(kernel assert torch.backends.cuda.matmul.allow_tf32 is prev_allow assert torch.get_float32_matmul_precision() == prev_precision + + +# --------------------------------------------------------------------- +# Batched-replicate factorize (--gpu-batch): batch-aware kernel parity +# --------------------------------------------------------------------- +def _rowwise_cosine(A, B): + """Cosine of each aligned program row (same seed -> same init -> same row order, no permutation).""" + num = (A * B).sum(axis=1) + den = np.linalg.norm(A, axis=1) * np.linalg.norm(B, axis=1) + 1e-30 + return num / den + + +def test_nmf_gpu_mu_matches_single_kernel_for_each_seed(kernel): + """Each batched replicate should match the single-replicate kernel at the same seed.""" + require_nmf_runtime() + k = 3 + X = low_rank_matrix(rank=k) + seeds = [7, 3, 101] + nmf_kwargs = {"n_components": k, "max_iter": 300, "tol": 0} + gpu_kwargs = {"device": "cpu"} + + batched = kernel._nmf_gpu_mu(X, seeds, nmf_kwargs, gpu_kwargs) + assert len(batched) == len(seeds) + + for (Hb, Wb), s in zip(batched, seeds): + Hs, Ws = kernel._nmf_gpu(X, dict(nmf_kwargs, random_state=s), gpu_kwargs) + assert Hb.shape == Hs.shape and Wb.shape == Ws.shape + assert _rowwise_cosine(Hb, Hs).min() > 0.9999 + rel_b = np.linalg.norm(X - Wb @ Hb) / np.linalg.norm(X) + rel_s = np.linalg.norm(X - Ws @ Hs) / np.linalg.norm(X) + assert abs(rel_b - rel_s) < 1e-4 + + +def test_nmf_gpu_mu_single_seed_reduces_to_single_kernel(kernel): + """A batch of one (R=1) must reproduce the single-replicate result at that seed.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=12, genes=6) + kw = {"n_components": 2, "max_iter": 80, "tol": 0} + + (Hb, Wb), = kernel._nmf_gpu_mu(X, [5], kw, {"device": "cpu"}) + Hs, Ws = kernel._nmf_gpu(X, dict(kw, random_state=5), {"device": "cpu"}) + + assert Hb.shape == Hs.shape and Wb.shape == Ws.shape + assert _rowwise_cosine(Hb, Hs).min() > 0.9999 + + +def test_nmf_gpu_mu_distinct_seeds_give_distinct_replicates(kernel): + """Distinct seeds should keep replicate outputs distinct.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=20, genes=8) + + out = kernel._nmf_gpu_mu(X, [1, 2], {"n_components": 3, "max_iter": 50}, {"device": "cpu"}) + + assert not np.allclose(out[0][0], out[1][0]) + + +def test_nmf_gpu_mu_rejects_empty_seeds(kernel): + """An empty seed list is a caller error (no replicates to run), not a silent no-op.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=6, genes=4) + + with pytest.raises(ValueError, match="non-empty"): + kernel._nmf_gpu_mu(X, [], {"n_components": 2, "max_iter": 1}, {"device": "cpu"}) + + +def test_nmf_gpu_mu_results_are_invariant_to_batch_grouping(kernel): + """Changing seed grouping should not change per-seed factors.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=16, genes=7) + seeds = [4, 11, 29] + kw = {"n_components": 3, "max_iter": 150, "tol": 0} + gpu = {"device": "cpu"} + + batched = kernel._nmf_gpu_mu(X, seeds, kw, gpu) # one launch, R=3 + assert len(batched) == len(seeds) + + for (Hb, Wb), s in zip(batched, seeds): + (Hs, Ws), = kernel._nmf_gpu_mu(X, [s], kw, gpu) # its own launch, R=1 + assert Hb.shape == Hs.shape and Wb.shape == Ws.shape + assert _rowwise_cosine(Hb, Hs).min() > 0.9999 + assert _rowwise_cosine(Wb.T, Ws.T).min() > 0.9999 + + +def test_fit_mu_is_batch_aware_and_stops_when_all_slices_converge(kernel): + """Batched `_fit_mu` should stop only after all slices meet tolerance.""" + torch = require_nmf_runtime() + R = 3 + Xb = torch.full((1, 3, 2), 2.0, dtype=torch.float64) + W = torch.ones((R, 3, 1), dtype=torch.float64) + H = torch.ones((R, 1, 2), dtype=torch.float64) + eps = torch.tensor(1e-9, dtype=torch.float64) + calls = {"count": 0} + + def no_change_step(W, H, Xg, eps): + calls["count"] += 1 + return W, H + + Wout, Hout = kernel._fit_mu(torch, Xb, W, H, eps, 10, 1e-4, no_change_step, 1, False, "cpu") + + assert calls["count"] == 2 + assert Wout.shape == (R, 3, 1) and Hout.shape == (R, 1, 2) + + +# --------------------------------------------------------------------- +# Fixed-H consensus refit (batch-aware): keep H fixed, update W only +# --------------------------------------------------------------------- +def _fixed_spectra(k, genes, seed=0): + """A valid non-negative fixed H (spectra) of shape (k, genes).""" + return np.abs(np.random.default_rng(seed).standard_normal((k, genes))) + + +def test_nmf_gpu_fixed_h_keeps_spectra_fixed_and_updates_usages(kernel): + """Fixed-H refit should return the supplied H unchanged.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=10, genes=6) + k, Hfix = 3, _fixed_spectra(3, 6) + kw = {"n_components": k, "max_iter": 50, "H": Hfix, "update_H": False} + + (H, W), = kernel._nmf_gpu_fixed_h(X, [7], kw, {"device": "cpu"}) + + assert H.shape == (k, 6) and W.shape == (10, k) + assert np.allclose(H, Hfix) + + +def test_nmf_gpu_fixed_h_batched_matches_single_refit_per_seed(kernel): + """Each batched fixed-H refit slice should match a single-seed refit.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=12, genes=5) + k, Hfix, seeds = 2, _fixed_spectra(2, 5, seed=1), [3, 9] + kw = {"n_components": k, "max_iter": 100, "tol": 0, "H": Hfix, "update_H": False} + + batched = kernel._nmf_gpu_fixed_h(X, seeds, kw, {"device": "cpu"}) + assert len(batched) == len(seeds) + + for (Hb, Wb), s in zip(batched, seeds): + (Hs, Ws), = kernel._nmf_gpu_fixed_h(X, [s], kw, {"device": "cpu"}) + assert np.allclose(Hb, Hfix) and np.allclose(Hs, Hfix) + assert _rowwise_cosine(Wb.T, Ws.T).min() > 0.9999 + + +def test_nmf_gpu_fixed_h_distinct_seeds_give_distinct_usages(kernel): + """Distinct W initializations should keep fixed-H usage outputs distinct.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=14, genes=6) + k, Hfix = 3, _fixed_spectra(3, 6, seed=2) + kw = {"n_components": k, "max_iter": 40, "H": Hfix, "update_H": False} + + out = kernel._nmf_gpu_fixed_h(X, [1, 2], kw, {"device": "cpu"}) + + assert not np.allclose(out[0][1], out[1][1]) + + +def test_nmf_gpu_update_H_false_dispatches_to_fixed_h_refit(kernel): + """`_nmf_gpu(update_H=False)` should dispatch to the fixed-H refit.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=8, genes=4) + k, Hfix = 2, _fixed_spectra(2, 4, seed=3) + + H, W = kernel._nmf_gpu(X, {"n_components": k, "max_iter": 20, "H": Hfix, "update_H": False}, {"device": "cpu"}) + + assert np.allclose(H, Hfix) and W.shape == (8, k) + + +# --------------------------------------------------------------------- +# --gpu-batch config plumbing +# --------------------------------------------------------------------- +def test_default_gpu_batch_is_single_replicate(kernel): + """The default batch is 1, so the single-replicate path is unchanged unless a user opts in.""" + assert kernel.DEFAULT_GPU["batch"] == 1 + assert kernel._resolve_gpu_opts(None)["batch"] == 1 + + +def test_resolve_gpu_opts_coerces_and_floors_batch_to_at_least_one(kernel): + """batch is a positive int: numeric strings coerce and non-positive values floor to 1.""" + assert kernel._resolve_gpu_opts({"batch": "4"})["batch"] == 4 + assert kernel._resolve_gpu_opts({"batch": 0})["batch"] == 1 + assert kernel._resolve_gpu_opts({"batch": -5})["batch"] == 1 + + +def test_parse_gpu_args_registers_batch_and_gpu_kwargs_carries_it(kernel): + """--gpu-batch parses under --engine gpu and flows into resolved gpu_kwargs; absent -> default 1.""" + import argparse + parser = kernel.parse_gpu_args(argparse.ArgumentParser()) + + args = parser.parse_args(["--engine", "gpu", "--gpu-batch", "8"]) + assert args.gpu_batch == 8 + assert kernel.gpu_kwargs_from_args(args)["batch"] == 8 + + default_args = parser.parse_args(["--engine", "gpu"]) + assert default_args.gpu_batch is None + assert kernel.gpu_kwargs_from_args(default_args)["batch"] == 1 + + +# --------------------------------------------------------------------- +# CLI parsing, engine wiring, and fixed-H consensus refit (integration) +# --------------------------------------------------------------------- +def test_parse_gpu_args_defaults_to_none_until_user_sets_engine(kernel): + """CLI flags should default to None so absent options are distinguishable from explicit values.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + args = parser.parse_args([]) + + assert args.engine is None + for name in ("gpu_device", "gpu_dtype", "gpu_allow_tf32", "gpu_compile", + "gpu_eps", "gpu_check_every", "gpu_compile_block"): + assert getattr(args, name) is None, f"{name} should default to None" + + +def test_gpu_kwargs_from_args_rejects_gpu_options_without_gpu_engine(kernel): + """GPU-specific CLI options should raise unless `--engine gpu` was explicitly selected.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + for argv in (["--gpu-device", "cuda"], ["--engine", "cpu", "--gpu-device", "cuda"]): + with pytest.raises(ValueError, match="require --engine gpu"): + kernel.gpu_kwargs_from_args(parser.parse_args(argv)) + + +def test_gpu_kwargs_from_args_fills_defaults_when_gpu_engine_selected(kernel): + """`--engine gpu` alone should resolve missing GPU options from DEFAULT_GPU.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + args = parser.parse_args(["--engine", "gpu"]) + + assert kernel.gpu_kwargs_from_args(args) == kernel.DEFAULT_GPU + + +def test_gpu_kwargs_from_args_normalizes_cli_overrides(kernel): + """GPU CLI override values should normalize through the same resolver as config dict values.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + args = parser.parse_args([ + "--engine", "gpu", + "--gpu-device", "CUDA:0", # device is lower-cased by _resolve_gpu_opts + "--gpu-dtype", "FP32", + "--gpu-allow-tf32", # store_const flags -> True + "--gpu-compile", + "--gpu-eps", "1e-8", + "--gpu-check-every", "5", + "--gpu-compile-block", "100", + ]) + + assert kernel.gpu_kwargs_from_args(args) == { + "device": "cuda:0", + "dtype": "fp32", + "allow_tf32": True, + "compile": True, + "eps": 1e-8, + "check_every": 5, + "compile_block": 100, + "batch": 1, # default added; not set on the CLI here + } + + +def test_validate_engine_args_for_command_rejects_non_engine_command_gpu_options(kernel): + """Engine/GPU options should be accepted only for factorize and consensus.""" + parser = argparse.ArgumentParser() + parser.add_argument("command") + kernel.parse_gpu_args(parser) + supported = ("factorize", "consensus") + + # factorize and consensus accept engine/GPU options (no raise) + for argv in ( + ["factorize", "--engine", "gpu"], + ["consensus", "--engine", "gpu"], + ["consensus", "--gpu-device", "cuda"], + ): + kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) + + # non-engine commands carrying engine/GPU options are rejected + for argv in ( + ["prepare", "--engine", "gpu"], + ["combine", "--gpu-device", "cuda"], + ["k_selection_plot", "--gpu-dtype", "fp32"], + ): + with pytest.raises(ValueError, match="only valid with"): + kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) + + # non-engine commands without engine/GPU options are fine + kernel.validate_engine_args_for_command(parser.parse_args(["prepare"]), supported) + + +def test_validate_engine_args_accepts_consensus_gpu_options(kernel): + """`consensus --engine gpu` and consensus GPU flags should be valid CLI input.""" + parser = argparse.ArgumentParser() + parser.add_argument("command") + kernel.parse_gpu_args(parser) + supported = ("factorize", "consensus") + + for argv in ( + ["consensus", "--engine", "gpu"], + ["consensus", "--engine", "gpu", "--gpu-device", "CUDA:0", "--gpu-dtype", "FP32"], + ["consensus", "--gpu-allow-tf32", "--gpu-compile"], + ): + kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) + + +def test_validate_engine_args_rejects_gpu_options_for_non_engine_commands(kernel): + """GPU flags should still be rejected for prepare/combine/k_selection_plot.""" + parser = argparse.ArgumentParser() + parser.add_argument("command") + kernel.parse_gpu_args(parser) + supported = ("factorize", "consensus") + + for argv in ( + ["prepare", "--engine", "gpu"], + ["combine", "--gpu-device", "cuda"], + ["k_selection_plot", "--gpu-check-every", "2"], + ): + with pytest.raises(ValueError, match="only valid with"): + kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) + + +def test_configure_nmf_engine_cpu_leaves_cnmf_instance_unchanged(kernel): + """The default CPU engine should be a no-op so existing sklearn behavior is preserved.""" + class DummyCNMF: + def _nmf(self, X, nmf_kwargs): + return "sklearn-path" + + obj = DummyCNMF() + result = kernel.configure_nmf_engine(obj, engine="cpu", gpu_kwargs={"device": "cuda"}) + + assert result is obj + assert "_nmf" not in vars(obj) # no instance override added + assert obj._nmf("X", {}) == "sklearn-path" + + +def test_configure_nmf_engine_rejects_unknown_engine(kernel): + """Unknown engine names should fail loudly instead of silently using CPU.""" + with pytest.raises(ValueError, match="engine must be 'cpu' or 'gpu'"): + kernel.configure_nmf_engine(object(), engine="tpu") + + +def test_configure_nmf_engine_gpu_patches_instance_nmf_with_adapter(kernel, monkeypatch): + """The GPU engine should replace the instance `_nmf` callable with the GPU adapter path.""" + captured = {} + + def fake_nmf_gpu(self, X, nmf_kwargs): + captured["self"], captured["X"], captured["nmf_kwargs"] = self, X, dict(nmf_kwargs) + return ("spectra", "usages") + + monkeypatch.setattr(kernel, "nmf_gpu", fake_nmf_gpu) + + class DummyCNMF: + def _nmf(self, X, nmf_kwargs): + return "sklearn-path" + + obj = DummyCNMF() + gpu_kwargs = {"device": "cuda", "dtype": "fp32"} + result = kernel.configure_nmf_engine(obj, engine="gpu", gpu_kwargs=gpu_kwargs) + + assert result is obj + assert "_nmf" in vars(obj) # instance _nmf is now overridden + + out = obj._nmf("Xdata", {"n_components": 5}) + + assert out == ("spectra", "usages") # dispatched through the GPU adapter + assert captured["self"] is obj and captured["X"] == "Xdata" + assert captured["nmf_kwargs"]["engine"] == "gpu" # engine + gpu kwargs embedded first + assert captured["nmf_kwargs"]["gpu"] == gpu_kwargs + assert captured["nmf_kwargs"]["n_components"] == 5 + + +def test_nmf_gpu_update_h_false_reconstructs_and_keeps_fixed_h(kernel): + """Consensus refit should preserve supplied H and optimize usages W only.""" + require_nmf_runtime() + fixed_h = np.array([[1.0, 0.3, 0.6], [0.2, 1.1, 0.4]], dtype=np.float64) + true_w = np.array([[1.0, 0.5], [0.4, 1.2], [1.5, 0.3], [0.7, 0.9]], dtype=np.float64) + X = true_w @ fixed_h + + H, W = kernel._nmf_gpu( + X, + {"n_components": 2, "max_iter": 5, "random_state": 0, "update_H": False, "H": fixed_h}, + {"device": "cpu", "dtype": "fp64", "check_every": 5}, + ) + + assert_valid_nmf_output(X, H, W, 2) + assert np.allclose(H, fixed_h) + + +def test_to_checked_fixed_h_rejects_missing_invalid_or_incompatible_h(kernel): + """Fixed-H consensus refit should fail clearly for invalid supplied spectra.""" + with pytest.raises(ValueError, match="requires a fixed H"): + kernel._to_checked_fixed_h(None, 2, 3) + + invalid_cases = [ + (np.array([1.0, 2.0, 3.0]), "2D"), + (np.ones((3, 3)), "shape"), + (np.array([[1.0, np.nan, 0.2], [0.4, 0.5, 0.6]]), "NaN/inf"), + (np.array([[1.0, -0.1, 0.2], [0.4, 0.5, 0.6]]), "non-negative"), + ] + for H, message in invalid_cases: + with pytest.raises(ValueError, match=message): + kernel._to_checked_fixed_h(H, 2, 3) + + +def test_mu_step_fixed_h_matches_manual_w_only_update(kernel): + """One fixed-H MU step should match the manual W-only Frobenius update.""" + torch = require_nmf_runtime() + Xg = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64) + W0 = torch.tensor([[0.5, 0.7], [0.9, 1.1]], dtype=torch.float64) + H = torch.tensor([[0.6, 0.8], [1.0, 1.2]], dtype=torch.float64) + eps = torch.tensor(1e-9, dtype=torch.float64) + H_before = H.clone() + + expected_W = W0 * ((Xg @ H.T) / (W0 @ (H @ H.T) + eps)) + W = kernel._mu_step_fixed_h(W0, H, Xg, eps) + + assert torch.allclose(W, expected_W) + assert torch.allclose(H, H_before) + + +def test_fit_mu_fixed_h_respects_early_stop_and_max_iter_block_bounds(kernel): + """Fixed-H fit loop should share early-stop and no-overrun guarantees with full MU.""" + torch = require_nmf_runtime() + Xg = torch.full((3, 2), 2.0, dtype=torch.float64) + W = torch.ones((3, 1), dtype=torch.float64) + H = torch.ones((1, 2), dtype=torch.float64) + eps = torch.tensor(1e-9, dtype=torch.float64) + + early_calls = {"count": 0} + + def no_change_step_early(W, H, Xg, eps): + early_calls["count"] += 1 + return W + + kernel._fit_mu_fixed_h(torch, Xg, W, H, eps, 10, 1e-4, no_change_step_early, 1, False, "cpu") + assert early_calls["count"] == 2 + + max_iter_calls = {"count": 0} + + def no_change_step_max_iter(W, H, Xg, eps): + max_iter_calls["count"] += 1 + return W + + kernel._fit_mu_fixed_h(torch, Xg, W, H, eps, 6, -1.0, no_change_step_max_iter, 4, False, "cpu") + assert max_iter_calls["count"] == 6 + + +def test_execution_plan_for_fixed_h_compile_uses_fixed_h_step_and_compile_block(kernel): + """Compiled consensus refit should compile _mu_step_fixed_h and use compile_block.""" + calls = [] + fake_torch = SimpleNamespace(compile=lambda fn: calls.append(fn) or fn) + opt = dict(kernel.DEFAULT_GPU, compile=True, check_every=1, compile_block=3) + + step, block = kernel._execution_plan(fake_torch, opt, "cpu", kernel._mu_step_fixed_h) + + assert calls == [kernel._mu_step_fixed_h] + assert step is kernel._mu_step_fixed_h + assert block == 3 diff --git a/tests/test_prepare.py b/tests/test_prepare.py index eefec72..83d1653 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -5,6 +5,7 @@ import os import scipy.sparse as sp from cnmf import cNMF, save_df_to_npz, load_df_from_npz +from cnmf import nmf_gpu # Global parameters for data simulation NUM_CELLS = 100 @@ -58,7 +59,105 @@ def generate_counts_file(tmp_path, file_format, dtype=np.int64, zero_count=False return str(counts_fn) +@pytest.mark.parametrize("file_format", ["txt", "npz", "h5ad"]) +@pytest.mark.parametrize("dtype", [np.int64, np.float32, np.float64]) +@pytest.mark.parametrize("densify", [True, False]) +def test_prepare(mock_cnmf, file_format, dtype, densify, tmp_path): + counts_fn = generate_counts_file(tmp_path, file_format, dtype) + + output_dir = tmp_path / "output" + os.makedirs(output_dir, exist_ok=True) + + mock_cnmf.prepare(counts_fn, components=[5, 10], n_iter=10, densify=densify) + + # Check if output files were created + expected_files = [ + mock_cnmf.paths['normalized_counts'], + mock_cnmf.paths['nmf_replicate_parameters'], + mock_cnmf.paths['nmf_run_parameters'], + mock_cnmf.paths['nmf_genes_list'], + mock_cnmf.paths['tpm'], + mock_cnmf.paths['tpm_stats'] + ] + + for file in expected_files: + assert os.path.exists(file), f"Expected output file {file} not found." + + # Clean up after test + for file in expected_files: + os.remove(file) + +@pytest.mark.parametrize("file_format", ["txt", "npz", "h5ad"]) +@pytest.mark.parametrize("dtype", [np.int64, np.float32, np.float64]) +@pytest.mark.parametrize("densify", [True, False]) +def test_prepare_raises_on_zero_count_cells(mock_cnmf, file_format, dtype, densify, tmp_path): + counts_fn = generate_counts_file(tmp_path, file_format, dtype, zero_count=True) + + with pytest.raises(Exception, match="Error: .* cells have zero counts of overdispersed genes.*"): + mock_cnmf.prepare(counts_fn, components=[5, 10], n_iter=10, densify=densify) + + +def test_configure_nmf_engine_gpu_factorize_groups_by_k_and_writes_each_replicate(mock_cnmf, monkeypatch, tmp_path): + """GPU factorize groups by k, batches seeds, and writes one spectra file per replicate.""" + counts_fn = generate_counts_file(tmp_path, "npz", np.float64) + mock_cnmf.prepare(counts_fn, components=[5, 7], n_iter=3, densify=True) + + calls = [] + + def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + k = int(nmf_kwargs["n_components"]) + calls.append((k, [int(s) for s in seeds])) + return [(np.zeros((k, X.shape[1])), np.zeros((X.shape[0], k))) for _ in seeds] + + monkeypatch.setattr(nmf_gpu, "_nmf_gpu_mu", fake_mu) + + nmf_gpu.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu", "batch": 2}) + mock_cnmf.factorize(worker_i=0, total_workers=1) + + # Two k-values x three replicates, chunked by batch=2 -> four launches. + assert len(calls) == 4 + assert sorted(len(seeds) for _, seeds in calls) == [1, 1, 2, 2] + for k, seeds in calls: + assert k in (5, 7) + assert 1 <= len(seeds) <= 2 + + run_params = load_df_from_npz(mock_cnmf.paths["nmf_replicate_parameters"]) + for _, row in run_params.iterrows(): + path = mock_cnmf.paths["iter_spectra"] % (int(row["n_components"]), int(row["iter"])) + assert os.path.exists(path), f"missing iter_spectra for k={row['n_components']} iter={row['iter']}" + + handed = sorted(s for _, seeds in calls for s in seeds) + assert handed == sorted(int(s) for s in run_params["nmf_seed"]) + + +def test_configure_nmf_engine_installs_gpu_factorize_at_default_batch_1(mock_cnmf, monkeypatch, tmp_path): + """Default GPU factorize uses batch=1: one seed per `_nmf_gpu_mu` launch.""" + counts_fn = generate_counts_file(tmp_path, "npz", np.float64) + mock_cnmf.prepare(counts_fn, components=[6], n_iter=2, densify=True) + + calls = [] + + def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + calls.append([int(s) for s in seeds]) + k = int(nmf_kwargs["n_components"]) + return [(np.zeros((k, X.shape[1])), np.zeros((X.shape[0], k))) for _ in seeds] + monkeypatch.setattr(nmf_gpu, "_nmf_gpu_mu", fake_mu) + + nmf_gpu.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu"}) + mock_cnmf.factorize(worker_i=0, total_workers=1) + + assert len(calls) == 2 + assert all(len(seeds) == 1 for seeds in calls) + run_params = load_df_from_npz(mock_cnmf.paths["nmf_replicate_parameters"]) + for _, row in run_params.iterrows(): + path = mock_cnmf.paths["iter_spectra"] % (int(row["n_components"]), int(row["iter"])) + assert os.path.exists(path), f"missing iter_spectra for k={row['n_components']} iter={row['iter']}" + + +# --------------------------------------------------------------------- +# cNMF <-> GPU engine integration: factorize / refit / consensus routing +# --------------------------------------------------------------------- def generate_positive_counts_file(tmp_path, cells=24, genes=12, dtype=np.float64): """Generate a small dense positive count table for consensus smoke tests.""" rng = np.random.default_rng(SEED) @@ -99,9 +198,6 @@ def fake_gpu_nmf_output(X, nmf_kwargs): return fixed_h.copy(), usages -# --------------------------------------------------------------------- -# GPU factorize wiring integration -# --------------------------------------------------------------------- def test_get_nmf_iter_params_default_cpu_engine_does_not_change_sklearn_kwargs(mock_cnmf): """Default CPU runs should not add engine/gpu keys to sklearn NMF kwargs.""" _replicate_params, run_params = mock_cnmf.get_nmf_iter_params(ks=[5, 7], n_iter=3, random_state_seed=14) @@ -114,7 +210,8 @@ def test_get_nmf_iter_params_default_cpu_engine_does_not_change_sklearn_kwargs(m def test_factorize_gpu_engine_passes_seed_components_run_params_and_gpu_kwargs(mock_cnmf, monkeypatch, tmp_path): - """cNMF factorize should pass n_components, random_state, run params, and GPU kwargs to `_nmf_gpu`.""" + """factorize_gpu should hand each replicate's seed, n_components, forwarded run params, and the + resolved GPU kwargs to the batch kernel _nmf_gpu_mu (default batch=1 -> one seed per launch).""" import cnmf.nmf_gpu as gpu_mod counts_fn = generate_counts_file(tmp_path, "txt", np.int64) @@ -122,35 +219,32 @@ def test_factorize_gpu_engine_passes_seed_components_run_params_and_gpu_kwargs(m captured = [] - def fake_nmf_gpu(self, X, nmf_kwargs): - captured.append(dict(nmf_kwargs)) - k = nmf_kwargs["n_components"] - return np.zeros((k, X.shape[1])), np.zeros((X.shape[0], k)) # (spectra, usages) + def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + captured.append((list(seeds), dict(nmf_kwargs), gpu_kwargs)) + k = int(nmf_kwargs["n_components"]) + return [(np.zeros((k, X.shape[1])), np.zeros((X.shape[0], k))) for _ in seeds] - monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) + monkeypatch.setattr(gpu_mod, "_nmf_gpu_mu", fake_mu) gpu_kwargs = {"device": "cpu", "dtype": "fp64"} gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs=gpu_kwargs) mock_cnmf.factorize(worker_i=0, total_workers=1) - assert len(captured) == 2 # n_iter=2 replicates for k=5 + assert len(captured) == 2 # n_iter=2 replicates, batch=1 -> 2 launches replicate_params = load_df_from_npz(mock_cnmf.paths["nmf_replicate_parameters"]) expected_seeds = set(replicate_params["nmf_seed"]) observed_seeds = set() - for kw in captured: - assert kw["n_components"] == 5 # set per replicate by factorize - assert kw["engine"] == "gpu" # adapter embedded the engine - assert kw["gpu"] == gpu_kwargs # ...and the resolved GPU kwargs - assert "beta_loss" in kw and "init" in kw # original run params forwarded - observed_seeds.add(kw["random_state"]) - assert observed_seeds == expected_seeds # exact seeds from prepared replicate params + for seeds, kw, gk in captured: + assert len(seeds) == 1 # default batch=1 -> one seed per launch + assert kw["n_components"] == 5 # set per k by factorize + assert "beta_loss" in kw and "init" in kw # original run params forwarded + assert gk == gpu_kwargs # resolved GPU kwargs passed through + observed_seeds.update(seeds) + assert observed_seeds == expected_seeds # exact seeds from prepared replicate params for iter_i in replicate_params["iter"]: assert os.path.exists(mock_cnmf.paths["iter_spectra"] % (5, iter_i)) -# --------------------------------------------------------------------- -# GPU consensus/refit wiring -# --------------------------------------------------------------------- def test_refit_usage_gpu_engine_passes_fixed_h_update_h_false_and_gpu_kwargs(mock_cnmf, monkeypatch, tmp_path): """cNMF refit_usage should route fixed-H consensus refits through the GPU adapter.""" import cnmf.nmf_gpu as gpu_mod @@ -172,7 +266,7 @@ def fake_nmf_gpu(self, X_arg, nmf_kwargs): captured.append((X_arg, dict(nmf_kwargs))) return fake_gpu_nmf_output(X_arg, nmf_kwargs) - monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) + monkeypatch.setattr(gpu_mod, "nmf_gpu", fake_nmf_gpu) gpu_kwargs = {"device": "cpu", "dtype": "fp64"} gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs=gpu_kwargs) @@ -213,7 +307,7 @@ def fake_nmf_gpu(self, X_arg, nmf_kwargs): captured.append((X_arg, dict(nmf_kwargs))) return fake_gpu_nmf_output(X_arg, nmf_kwargs) - monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) + monkeypatch.setattr(gpu_mod, "nmf_gpu", fake_nmf_gpu) gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu", "dtype": "fp64"}) spectra = mock_cnmf.refit_spectra(X, usage) @@ -258,7 +352,7 @@ def test_consensus_gpu_engine_smoke_writes_expected_outputs(mock_cnmf, monkeypat def fake_nmf_gpu(self, X_arg, nmf_kwargs): return fake_gpu_nmf_output(X_arg, nmf_kwargs) - monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) + monkeypatch.setattr(gpu_mod, "nmf_gpu", fake_nmf_gpu) gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu", "dtype": "fp64"}) mock_cnmf.consensus(k=2, density_threshold=2.0, local_neighborhood_size=0.5, @@ -309,41 +403,3 @@ def fake_cpu_nmf(X_arg, nmf_kwargs): assert kw["update_H"] is False assert np.allclose(kw["H"], spectra.values) assert usages.shape == (4, 2) - - -@pytest.mark.parametrize("file_format", ["txt", "npz", "h5ad"]) -@pytest.mark.parametrize("dtype", [np.int64, np.float32, np.float64]) -@pytest.mark.parametrize("densify", [True, False]) -def test_prepare(mock_cnmf, file_format, dtype, densify, tmp_path): - counts_fn = generate_counts_file(tmp_path, file_format, dtype) - - output_dir = tmp_path / "output" - os.makedirs(output_dir, exist_ok=True) - - mock_cnmf.prepare(counts_fn, components=[5, 10], n_iter=10, densify=densify) - - # Check if output files were created - expected_files = [ - mock_cnmf.paths['normalized_counts'], - mock_cnmf.paths['nmf_replicate_parameters'], - mock_cnmf.paths['nmf_run_parameters'], - mock_cnmf.paths['nmf_genes_list'], - mock_cnmf.paths['tpm'], - mock_cnmf.paths['tpm_stats'] - ] - - for file in expected_files: - assert os.path.exists(file), f"Expected output file {file} not found." - - # Clean up after test - for file in expected_files: - os.remove(file) - -@pytest.mark.parametrize("file_format", ["txt", "npz", "h5ad"]) -@pytest.mark.parametrize("dtype", [np.int64, np.float32, np.float64]) -@pytest.mark.parametrize("densify", [True, False]) -def test_prepare_raises_on_zero_count_cells(mock_cnmf, file_format, dtype, densify, tmp_path): - counts_fn = generate_counts_file(tmp_path, file_format, dtype, zero_count=True) - - with pytest.raises(Exception, match="Error: .* cells have zero counts of overdispersed genes.*"): - mock_cnmf.prepare(counts_fn, components=[5, 10], n_iter=10, densify=densify)