From 499ff974af874cb774a283f4fc061ac7db4f3951 Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:40:12 +0000 Subject: [PATCH] added cd solver and refactored gpu kernal for solver routing --- .github/workflows/python-tests.yml | 2 + setup.py | 2 +- src/cnmf/cnmf.py | 34 +- src/cnmf/gpunmf/__init__.py | 165 ++++ src/cnmf/gpunmf/solver_cd.py | 358 +++++++++ src/cnmf/gpunmf/solver_cd_triton.py | 154 ++++ src/cnmf/gpunmf/solver_mu.py | 148 ++++ src/cnmf/gpunmf/utils.py | 366 +++++++++ src/cnmf/nmf_gpu.py | 608 --------------- tests/test_nmf_gpu.py | 1079 +++++++++++++++++++++------ tests/test_prepare.py | 248 +++++- tests/utils.py | 12 +- 12 files changed, 2283 insertions(+), 893 deletions(-) create mode 100644 src/cnmf/gpunmf/__init__.py create mode 100644 src/cnmf/gpunmf/solver_cd.py create mode 100644 src/cnmf/gpunmf/solver_cd_triton.py create mode 100644 src/cnmf/gpunmf/solver_mu.py create mode 100644 src/cnmf/gpunmf/utils.py delete mode 100644 src/cnmf/nmf_gpu.py diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 97b455c..fabc927 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -1,6 +1,8 @@ name: Python tests on: + # Includes direct pushes and merge commits pushed to the target branch. + push: pull_request: jobs: diff --git a/setup.py b/setup.py index f199766..19fb00e 100644 --- a/setup.py +++ b/setup.py @@ -49,7 +49,7 @@ def get_version(): 'pyyaml' ], extras_require={ - # GPU NMF engine (src/cnmf/nmf_gpu.py): PyTorch multiplicative-update kernel. + # GPU NMF engine: PyTorch MU and sklearn-compatible Fast-HALS kernels. # Install with: pip install -e ".[gpu]" 'gpu': ['torch>=2.0'], 'test': ['pytest'], diff --git a/src/cnmf/cnmf.py b/src/cnmf/cnmf.py index c659662..58f01bd 100755 --- a/src/cnmf/cnmf.py +++ b/src/cnmf/cnmf.py @@ -378,8 +378,8 @@ def prepare(self, counts_fn, components, n_iter = 100, densify=False, tpm_fn=Non max_NMF_iter : int, optional (default=1000) Maximum number of iterations per individual NMF run """ - - + + if counts_fn.endswith('.h5ad'): input_counts = sc.read(counts_fn) elif counts_fn.endswith('.mtx') or counts_fn.endswith('.mtx.gz'): @@ -626,7 +626,7 @@ def get_nmf_iter_params(self, ks, n_iter = 100, init=init ) - ## Coordinate descent is faster than multiplicative update but only works for frobenius + # Coordinate descent is faster than multiplicative update but only works for frobenius if beta_loss == 'frobenius': _nmf_kwargs['solver'] = 'cd' @@ -1234,10 +1234,7 @@ def main(): """ import sys, argparse - try: - from cnmf.nmf_gpu import configure_nmf_engine, gpu_kwargs_from_args, parse_gpu_args, validate_engine_args_for_command - except ImportError: - from nmf_gpu import configure_nmf_engine, gpu_kwargs_from_args, parse_gpu_args, validate_engine_args_for_command + from cnmf.gpunmf import configure_nmf_engine parser = argparse.ArgumentParser() @@ -1253,6 +1250,7 @@ def main(): parser.add_argument('--numgenes', type=int, help='[prepare] Number of high variance genes to use for matrix factorization.', default=2000) parser.add_argument('--tpm', type=str, help='[prepare] Pre-computed (cell x gene) TPM values as df.npz or tab separated txt file. If not provided TPM will be calculated automatically', default=None) parser.add_argument('--max-nmf-iter', type=int, help='[prepare] Max number of iterations per individual NMF run (default 1000)', default=1000) + parser.add_argument('--solver', type=str.lower, choices=['mu', 'cd'], help='[prepare] NMF solver to persist for factorization; cd requires Frobenius loss (default cd)', default='cd') parser.add_argument('--beta-loss', type=str, choices=['frobenius', 'kullback-leibler', 'itakura-saito'], help='[prepare] Loss function for NMF (default frobenius)', default='frobenius') parser.add_argument('--init', type=str, choices=['random', 'nndsvd'], help='[prepare] Initialization algorithm for NMF (default random)', default='random') parser.add_argument('--densify', dest='densify', help='[prepare] Treat the input data as non-sparse (default False)', action='store_true', default=False) @@ -1262,18 +1260,24 @@ def main(): parser.add_argument('--local-neighborhood-size', type=float, help='[consensus] Fraction of the number of replicates to use as nearest neighbors for local density filtering', default=0.30) parser.add_argument('--show-clustering', dest='show_clustering', help='[consensus] Produce a clustergram figure summarizing the spectra clustering', action='store_true') parser.add_argument('--build-reference', dest='build_reference', help='[consensus] Generates a reference spectra for use in starCAT', action='store_true', default=True) - parse_gpu_args(parser) + parser.add_argument("--engine", type=str.lower, choices=["cpu", "gpu"], help="[factorize,consensus] NMF engine to use (default cpu)", default="cpu") + parser.add_argument("--gpu-device", type=str, help="[factorize,consensus,gpu] Device for GPU NMF: auto, cpu, cuda, cuda:N, or mps") + parser.add_argument("--gpu-dtype", type=str.lower, choices=["auto", "fp32", "fp64", "bf16"], help="[factorize,consensus,gpu] Storage and matmul dtype for GPU NMF (default auto)") + parser.add_argument("--gpu-allow-tf32", action="store_const", const=True, help="[factorize,consensus,gpu] Allow TF32 for CUDA fp32 matrix multiplication") + parser.add_argument("--gpu-compile", action="store_const", const=True, help="[factorize,consensus,gpu] Enable torch.compile for the MU update step") + parser.add_argument("--gpu-eps", type=float, help="[factorize,consensus,gpu] Replacement for exactly-zero MU denominators") + parser.add_argument("--gpu-check-every", type=int, help="[factorize,consensus,gpu] Eager-mode convergence check interval") + parser.add_argument("--gpu-compile-block", type=int, help="[factorize,consensus,gpu] Number of MU iterations per compiled block") + parser.add_argument("--gpu-batch", type=int, help="[factorize] Replicates run per GPU solver launch; 1 = single-replicate") + args = parser.parse_args() - try: - engine_commands = ('factorize', 'consensus') - validate_engine_args_for_command(args, engine_commands) - except ValueError as e: - parser.error(str(e)) - cnmf_obj = cNMF(output_dir=args.output_dir, name=args.name) - cnmf_obj = configure_nmf_engine(cnmf_obj, engine=args.engine or 'cpu', gpu_kwargs=gpu_kwargs_from_args(args)) + try: + cnmf_obj = configure_nmf_engine(cNMF, args) + except ValueError as exc: + parser.error(str(exc)) if args.command == 'prepare': cnmf_obj.prepare(args.counts, components=args.components, n_iter=args.n_iter, densify=args.densify, diff --git a/src/cnmf/gpunmf/__init__.py b/src/cnmf/gpunmf/__init__.py new file mode 100644 index 0000000..48e8260 --- /dev/null +++ b/src/cnmf/gpunmf/__init__.py @@ -0,0 +1,165 @@ +"""Optional PyTorch NMF engine for cNMF. + +This package module owns solver routing and cNMF integration. Shared runtime +helpers and the MU/CD implementations remain in focused sibling modules. +""" + +import functools +from collections import OrderedDict + +import numpy as np + +from . import solver_cd, solver_mu, utils + + +__all__ = [ + "configure_nmf_engine", + "factorize_gpu", + "prepare_gpu", + "solver_cd", + "solver_mu", + "utils", +] + + +_GPU_SOLVERS = { + "mu": solver_mu._nmf_gpu_mu, + "cd": solver_cd._nmf_gpu_cd, +} + + +def _nmf_gpu_batch(X, seeds, nmf_kwargs, gpu_kwargs=None): + """Dispatch one same-k replicate batch to the selected GPU solver.""" + solver_name = str( + nmf_kwargs.get("solver", utils.DEFAULT_NMF["solver"]) + ).strip().lower() + try: + solver = _GPU_SOLVERS[solver_name] + except KeyError as exc: + available = ", ".join(sorted(_GPU_SOLVERS)) + raise ValueError( + f"GPU NMF solver {solver_name!r} is not available; " + f"available solvers: {available}" + ) from exc + return solver(X, seeds, nmf_kwargs, gpu_kwargs) + + +def _nmf_gpu(args, X, nmf_kwargs, gpu_kwargs=None): + """Single-replicate NMF adapter over the batch gateway.""" + gpu_kwargs = utils.gpu_kwargs_from_args(args) + (result,) = _nmf_gpu_batch( + X, [nmf_kwargs.get("random_state")], nmf_kwargs, gpu_kwargs + ) + return result + + +def configure_nmf_engine(cnmf_constructor, args): + """Construct and configure a cNMF instance for the selected execution engine.""" + engine = getattr(args, "engine", "cpu") + if engine not in ("cpu", "gpu"): + raise ValueError("engine must be 'cpu' or 'gpu'") + + if engine == "gpu": + utils._validate_engine_args(args, _GPU_SOLVERS) + + cnmf_obj = cnmf_constructor(output_dir=args.output_dir, name=args.name) + if engine == "cpu": + return cnmf_obj + + # patch cNMF to use GPU NMF engine + cnmf_obj._nmf = functools.partial(_nmf_gpu, args) + original_prepare = cnmf_obj.prepare + cnmf_obj.prepare = functools.partial(prepare_gpu, cnmf_obj, args, original_prepare) + cnmf_obj.factorize = functools.partial(factorize_gpu, cnmf_obj, args) + + return cnmf_obj + + +def prepare_gpu(cnmf_obj, args, original_prepare, *prepare_args, **prepare_kwargs): + """Prepare an unchanged cNMF run, then persist its explicit solver choice. + + Upstream ``cNMF.prepare`` derives the solver from ``beta_loss``. This + adapter keeps all matrix preparation in upstream cNMF and changes only the + saved factorization configuration used by factorize, resume, and consensus. + """ + import yaml + + result = original_prepare(*prepare_args, **prepare_kwargs) + + config_path = cnmf_obj.paths["nmf_run_parameters"] + with open(config_path, encoding="utf-8") as stream: + run_parameters = yaml.safe_load(stream) + if not isinstance(run_parameters, dict): + raise ValueError(f"invalid NMF run configuration: {config_path}") + run_parameters["solver"] = args.solver + with open(config_path, "w", encoding="utf-8") as stream: + yaml.safe_dump(run_parameters, stream, sort_keys=False) + + return result + + +def factorize_gpu( + cnmf_obj, + args, + worker_i=0, + total_workers=1, + skip_completed_runs=False, +): + """GPU ``factorize`` drop-in that batches same-k replicate seeds.""" + import pandas as pd + import scanpy as sc + import yaml + + try: + from ..cnmf import load_df_from_npz, save_df_to_npz, worker_filter + except ImportError: + from cnmf import load_df_from_npz, save_df_to_npz, worker_filter + + gpu_kwargs = utils.gpu_kwargs_from_args(args) + batch = utils._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"]) + with open(cnmf_obj.paths["nmf_run_parameters"], encoding="utf-8") as stream: + base_kwargs = yaml.load(stream, Loader=yaml.FullLoader) + + if skip_completed_runs: + pending = run_params.index[run_params["completed"] == False] + job_idx = worker_filter(pending, worker_i, total_workers) + else: + job_idx = worker_filter(range(len(run_params)), worker_i, total_workers) + + genes = norm_counts.var.index + X_dense = ( + norm_counts.X.toarray() + if hasattr(norm_counts.X, "toarray") + else np.asarray(norm_counts.X) + ) + by_k = OrderedDict() + for idx in job_idx: + params = run_params.iloc[idx, :] + by_k.setdefault(int(params["n_components"]), []).append( + (int(params["iter"]), int(params["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 = [iteration for iteration, _seed in chunk] + seeds = [seed for _iteration, seed in chunk] + print( + "[Worker %d]. k=%d: launching %d replicate(s), iters=%s." + % (worker_i, k, len(chunk), iters) + ) + results = _nmf_gpu_batch(X_dense, seeds, run_kwargs, gpu_kwargs) + for (spectra, _usages), iteration 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, iteration), + ) diff --git a/src/cnmf/gpunmf/solver_cd.py b/src/cnmf/gpunmf/solver_cd.py new file mode 100644 index 0000000..7fd234e --- /dev/null +++ b/src/cnmf/gpunmf/solver_cd.py @@ -0,0 +1,358 @@ +"""Batched sklearn-compatible Fast-HALS coordinate-descent NMF solver.""" + +import numpy as np + +from . import utils + +# --------------------------------------------------------------------- +# CD solver (sklearn-compatible Fast-HALS; batch-aware) +# --------------------------------------------------------------------- + + +def _numpy_staging_dtype(torch, dtype): + """Return the host dtype used for sklearn initialization before transfer.""" + if dtype is torch.float64: + return np.float64 + if dtype is torch.float32: + return np.float32 + raise TypeError("solver='cd' supports gpu dtype fp32 or fp64") + + +def _to_checked_custom_factor(value, shape, name, dtype): + """Validate and cast a custom CD factor.""" + if value is None: + raise ValueError(f"init='custom' requires {name}") + array = value.toarray() if hasattr(value, "toarray") else np.asarray(value) + if array.ndim != 2 or array.shape != shape: + raise ValueError(f"custom {name} shape must be {shape}; got {array.shape}") + if not np.isfinite(array).all(): + raise ValueError(f"custom {name} contains NaN/inf") + if array.size and array.min() < 0: + raise ValueError(f"custom {name} must be non-negative") + return np.ascontiguousarray(array, dtype=dtype) + + +def _cd_regularization(nmf_kwargs, n_samples, n_features): + """Return sklearn's sample/feature-scaled CD regularization terms.""" + alpha_w = float(nmf_kwargs.get("alpha_W", 0.0)) + alpha_h_raw = nmf_kwargs.get("alpha_H", "same") + alpha_h = alpha_w if alpha_h_raw == "same" else float(alpha_h_raw) + l1_ratio = float(nmf_kwargs.get("l1_ratio", 0.0)) + if alpha_w < 0 or alpha_h < 0: + raise ValueError("alpha_W and alpha_H must be non-negative") + if not 0.0 <= l1_ratio <= 1.0: + raise ValueError("l1_ratio must be in the range [0, 1]") + return ( + n_features * alpha_w * l1_ratio, + n_features * alpha_w * (1.0 - l1_ratio), + n_samples * alpha_h * l1_ratio, + n_samples * alpha_h * (1.0 - l1_ratio), + ) + + +def _validate_cd_runtime(torch, rc, nmf_kwargs): + """Validate the sklearn CD contract before allocating factor tensors.""" + legacy_regularization = {"alpha", "regularization"}.intersection(nmf_kwargs) + if legacy_regularization: + raise ValueError( + "solver='cd' does not accept deprecated alpha/regularization; " + "use alpha_W, alpha_H, and l1_ratio" + ) + beta_loss = nmf_kwargs.get("beta_loss", "frobenius") + if not (beta_loss == 2 or str(beta_loss).lower() == "frobenius"): + raise ValueError("solver='cd' supports only beta_loss='frobenius'") + if rc.dtype is torch.bfloat16: + raise ValueError("solver='cd' supports gpu dtype fp32 or fp64, not bf16") + if rc.max_iter < 1: + raise ValueError("max_iter must be at least 1 for solver='cd'") + if rc.tol < 0: + raise ValueError("tol must be non-negative for solver='cd'") + if not isinstance(nmf_kwargs.get("shuffle", False), (bool, np.bool_)): + raise ValueError("shuffle must be a boolean for solver='cd'") + + +def _hals_sweep_torch(factor, gram, cross, permutation, active): + """Apply one literal torch port of sklearn's serial CD coordinate sweep.""" + replicates, components, rows = factor.shape + violation = factor.new_zeros(replicates) + zero = factor.new_zeros(()) + cyclic = permutation is None + + for coordinate in range(components): + if cyclic: + component = coordinate + gram_row = gram[:, component, :] + cross_row = cross[:, component, :] + old_value = factor[:, component, :].clone() + else: + component = permutation[:, coordinate] + gram_row = gram.gather( + 1, component[:, None, None].expand(-1, 1, components) + ).squeeze(1) + factor_index = component[:, None, None].expand(-1, 1, rows) + cross_row = cross.gather(1, factor_index).squeeze(1) + old_value = factor.gather(1, factor_index).squeeze(1) + + # Match sklearn's _cdnmf_fast.pyx summation and Gauss-Seidel order. + gradient = -cross_row.clone() + for other_component in range(components): + gradient = ( + gradient + + gram_row[:, other_component, None] + * factor[:, other_component, :] + ) + + projected_gradient = gradient.where( + old_value != 0, zero.minimum(gradient) + ) + violation = violation + projected_gradient.abs().sum(dim=1) * active + + if cyclic: + hessian = gram[:, component, component, None] + else: + hessian = gram_row.gather(1, component[:, None]) + nonzero_hessian = hessian != 0 + safe_hessian = hessian.where(nonzero_hessian, hessian.new_ones(())) + candidate = (old_value - gradient / safe_hessian).clamp_min(0) + new_value = candidate.where(nonzero_hessian, old_value) + new_value = new_value.where(active[:, None], old_value) + + if cyclic: + factor[:, component, :] = new_value + else: + factor.scatter_(1, factor_index, new_value[:, None, :]) + + return violation + + +_HALS_CUDA_BACKEND_UNSET = object() +_HALS_CUDA_BACKEND = _HALS_CUDA_BACKEND_UNSET + + +def _get_hals_cuda_backend(): + """Resolve and cache the optional fused CUDA sweep.""" + global _HALS_CUDA_BACKEND + if _HALS_CUDA_BACKEND is _HALS_CUDA_BACKEND_UNSET: + try: + from .solver_cd_triton import hals_sweep_cuda + except (ImportError, ModuleNotFoundError): + hals_sweep_cuda = None + _HALS_CUDA_BACKEND = hals_sweep_cuda + return _HALS_CUDA_BACKEND + + +def _hals_sweep(factor, gram, cross, permutation, active): + """Use the fused CUDA sweep when available, otherwise use torch.""" + if factor.is_cuda: + hals_sweep_cuda = _get_hals_cuda_backend() + if hals_sweep_cuda is not None: + if permutation is None: + permutation = factor.new_tensor( + np.tile( + np.arange(factor.shape[1], dtype=np.int64), + (factor.shape[0], 1), + ) + ).long() + return hals_sweep_cuda(factor, gram, cross, permutation, active) + return _hals_sweep_torch(factor, gram, cross, permutation, active) + + +def _regularize_cd_products(gram, cross, l1_reg, l2_reg): + """Apply sklearn's L2 diagonal addition and L1 cross-product shift.""" + if l2_reg != 0.0: + gram.diagonal(dim1=-2, dim2=-1).add_(l2_reg) + if l1_reg != 0.0: + cross.sub_(l1_reg) + return gram, cross + + +def _batch_invariant_cd_products(factor, data): + """Build CD products with one identical 2D GEMM path per replicate. + + A single batched matmul may select a different fp32 CUDA reduction path as + the replicate count changes. Fast-HALS amplifies those small differences, + which can change projected-gradient stopping decisions. Keep the factor + state and HALS sweep batched, but compute each replicate's products through + the same two-dimensional matmul shape used by a batch of one. + """ + replicates, components, _ = factor.shape + gram = factor.new_empty((replicates, components, components)) + cross = factor.new_empty((replicates, components, data.shape[-1])) + for replicate in range(replicates): + replicate_factor = factor[replicate] + gram[replicate].copy_( + replicate_factor @ replicate_factor.transpose(-2, -1) + ) + cross[replicate].copy_(replicate_factor @ data) + return gram, cross + + +def _fit_cd( + torch, + Xg, + Wt, + H, + max_iter, + tol, + update_h, + regularization, + shuffle, + seeds, + tf32, + device, +): + """Run batched sklearn-compatible Fast-HALS to projected-gradient convergence.""" + if Xg.dtype != Wt.dtype or Xg.dtype != H.dtype: + raise RuntimeError( + "NMF runtime tensors must share dtype; got " + f"{sorted(map(str, {Xg.dtype, Wt.dtype, H.dtype}))}" + ) + + replicates, components, _ = Wt.shape + active = torch.ones(replicates, dtype=torch.bool, device=device) + n_iter = torch.zeros(replicates, dtype=torch.int64, device=device) + violation_init = Wt.new_zeros(replicates) + l1_w, l2_w, l1_h, l2_h = regularization + + if shuffle: + try: + from sklearn.utils import check_random_state + except ModuleNotFoundError as exc: + raise RuntimeError("scikit-learn is required for shuffled CD") from exc + rngs = [check_random_state(seed) for seed in seeds] + cyclic_permutation = None + else: + rngs = None + cyclic_permutation = ( + torch.arange(components, dtype=torch.int64, device=device) + .expand(replicates, -1) + .contiguous() + if Xg.is_cuda and _get_hals_cuda_backend() is not None + else None + ) + + def next_permutation(): + if rngs is None: + return cyclic_permutation + values = np.stack( + [rng.permutation(components) for rng in rngs], axis=0 + ).astype(np.int64, copy=False) + return torch.as_tensor(values, dtype=torch.int64, device=device) + + with torch.no_grad(), utils._cuda_tf32(torch, tf32, device): + for iteration in range(1, max_iter + 1): + gram, cross = _batch_invariant_cd_products( + H, Xg.transpose(-2, -1) + ) + gram, cross = _regularize_cd_products(gram, cross, l1_w, l2_w) + violation = _hals_sweep( + Wt, gram, cross, next_permutation(), active + ) + + if update_h: + gram, cross = _batch_invariant_cd_products(Wt, Xg) + gram, cross = _regularize_cd_products(gram, cross, l1_h, l2_h) + violation = violation + _hals_sweep( + H, gram, cross, next_permutation(), active + ) + + n_iter = n_iter.new_full((), iteration).where(active, n_iter) + if iteration == 1: + violation_init = violation.clone() + + zero_init = violation_init == 0 + denominator = violation_init.where( + ~zero_init, violation_init.new_ones(()) + ) + converged = active & ( + zero_init | ((violation / denominator) <= tol) + ) + active = active & ~converged + if not bool(active.any()): + break + + return Wt, H, n_iter + + +def _nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): + """Run full or fixed-H sklearn-compatible CD replicates.""" + torch = utils._loud_import_torch() + seeds = utils._normalize_seeds(seeds) + rc = utils._gpu_setup(torch, X, nmf_kwargs, gpu_kwargs) + _validate_cd_runtime(torch, rc, nmf_kwargs) + + host_dtype = _numpy_staging_dtype(torch, rc.dtype) + Xcompute = np.ascontiguousarray(rc.Xnp, dtype=host_dtype) + Xg = torch.as_tensor(Xcompute, dtype=rc.dtype, device=rc.device) + replicates = len(seeds) + update_h = nmf_kwargs.get("update_H", True) is not False + + if update_h: + init = nmf_kwargs.get("init") + if init == "custom": + W0 = _to_checked_custom_factor( + nmf_kwargs.get("W"), + (Xcompute.shape[0], rc.k), + "W", + host_dtype, + ) + H0 = _to_checked_custom_factor( + nmf_kwargs.get("H"), + (rc.k, Xcompute.shape[1]), + "H", + host_dtype, + ) + Wt0 = np.repeat(W0.T[None, :, :], replicates, axis=0) + Hs0 = np.repeat(H0[None, :, :], replicates, axis=0) + else: + Wt0 = np.empty( + (replicates, rc.k, Xcompute.shape[0]), dtype=host_dtype + ) + Hs0 = np.empty( + (replicates, rc.k, Xcompute.shape[1]), dtype=host_dtype + ) + for replicate, seed in enumerate(seeds): + W0, H0 = utils._init_wh(Xcompute, rc.k, seed, init) + Wt0[replicate] = W0.T + Hs0[replicate] = H0 + Wt = torch.as_tensor(Wt0, dtype=rc.dtype, device=rc.device) + H = torch.as_tensor(Hs0, dtype=rc.dtype, device=rc.device) + else: + H0 = utils._to_checked_fixed_h( + nmf_kwargs.get("H"), rc.k, Xcompute.shape[1] + ) + H0 = np.ascontiguousarray(H0, dtype=host_dtype) + # sklearn CD ignores supplied W and initializes fixed-H usages to zero. + Wt = torch.zeros( + (replicates, rc.k, Xcompute.shape[0]), + dtype=rc.dtype, + device=rc.device, + ) + H = ( + torch.as_tensor(H0, dtype=rc.dtype, device=rc.device) + .unsqueeze(0) + .expand(replicates, -1, -1) + .contiguous() + ) + + regularization = _cd_regularization( + nmf_kwargs, Xcompute.shape[0], Xcompute.shape[1] + ) + Wt, H, _ = _fit_cd( + torch, + Xg, + Wt, + H, + rc.max_iter, + rc.tol, + update_h, + regularization, + bool(nmf_kwargs.get("shuffle", False)), + seeds, + utils._want_tf32(torch, rc), + rc.device, + ) + + Hc = H.cpu().double().numpy() + Wc = Wt.transpose(-2, -1).cpu().double().numpy() + return [(Hc[r], Wc[r]) for r in range(replicates)] diff --git a/src/cnmf/gpunmf/solver_cd_triton.py b/src/cnmf/gpunmf/solver_cd_triton.py new file mode 100644 index 0000000..d946113 --- /dev/null +++ b/src/cnmf/gpunmf/solver_cd_triton.py @@ -0,0 +1,154 @@ +"""Fused CUDA sweep for sklearn-compatible Fast-HALS. + +This module is imported lazily by :mod:`cnmf.gpunmf.solver_cd`, so importing +cNMF does not require Triton. Each kernel program owns one replicate and one row block. +Rows and replicates are parallel; component updates retain sklearn's serial +Gauss-Seidel order. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _hals_sweep_kernel( + factor, + gram, + cross, + permutation, + active, + block_violation, + M: tl.constexpr, + K: tl.constexpr, + N_BLOCKS: tl.constexpr, + BLOCK: tl.constexpr, + IS_FP64: tl.constexpr, +): + program = tl.program_id(0) + replicate = program // N_BLOCKS + row_block = program - replicate * N_BLOCKS + rows = row_block * BLOCK + tl.arange(0, BLOCK) + replicate_offset = replicate.to(tl.int64) + rows_offset = rows.to(tl.int64) + row_mask = rows < M + is_active = tl.load(active + replicate) + + accumulator_dtype = tl.float64 if IS_FP64 else tl.float32 + violation = tl.zeros((BLOCK,), dtype=accumulator_dtype) + + for coordinate in range(0, K): + component = tl.load( + permutation + replicate_offset * K + coordinate + ) + cross_offset = ( + (replicate_offset * K + component) * M + rows_offset + ) + gradient = -tl.load(cross + cross_offset, mask=row_mask, other=0.0) + + # This order matches sklearn.decomposition._cdnmf_fast exactly. + for other_component in range(0, K): + gram_value = tl.load( + gram + + (replicate_offset * K + component) * K + + other_component + ) + factor_value = tl.load( + factor + + (replicate_offset * K + other_component) * M + + rows_offset, + mask=row_mask, + other=0.0, + ) + gradient += gram_value * factor_value + + factor_offset = ( + (replicate_offset * K + component) * M + rows_offset + ) + old_value = tl.load(factor + factor_offset, mask=row_mask, other=0.0) + projected_gradient = tl.where( + old_value == 0.0, tl.minimum(0.0, gradient), gradient + ) + violation += tl.where( + row_mask & (is_active != 0), tl.abs(projected_gradient), 0.0 + ) + + hessian = tl.load( + gram + (replicate_offset * K + component) * K + component + ) + # tl.where evaluates both branches, so guard the division separately. + safe_hessian = tl.where(hessian != 0.0, hessian, 1.0) + candidate = tl.maximum(old_value - gradient / safe_hessian, 0.0) + new_value = tl.where(hessian != 0.0, candidate, old_value) + tl.store( + factor + factor_offset, + new_value, + mask=row_mask & (is_active != 0), + ) + + tl.store( + block_violation + + replicate_offset * N_BLOCKS + + row_block.to(tl.int64), + tl.sum(violation, axis=0), + ) + + +def hals_sweep_cuda(factor, gram, cross, permutation, active): + """Update one factor in place and return violation per replicate.""" + if not factor.is_cuda: + raise ValueError("hals_sweep_cuda requires CUDA tensors") + if not factor.is_floating_point() or factor.element_size() not in (4, 8): + raise TypeError("CUDA Fast-HALS supports fp32 and fp64 tensors") + if factor.ndim != 3: + raise ValueError( + "factor must have shape [replicate, component, row]" + ) + + replicates, components, rows = factor.shape + if cross.shape != (replicates, components, rows): + raise ValueError( + "factor and cross must share [replicate, component, row] shape" + ) + if gram.shape != (replicates, components, components): + raise ValueError( + "gram must have shape [replicate, component, component]" + ) + if permutation.shape != (replicates, components): + raise ValueError( + "permutation must have shape [replicate, component]" + ) + if active.shape != (replicates,): + raise ValueError("active must have one entry per replicate") + if permutation.dtype not in (torch.int32, torch.int64): + raise TypeError("permutation must use an integer dtype") + if active.dtype != torch.bool: + raise TypeError("active must use boolean dtype") + + tensors = (factor, gram, cross, permutation, active) + if any(tensor.device != factor.device for tensor in tensors): + raise ValueError("all Fast-HALS tensors must be on the same CUDA device") + if gram.dtype != factor.dtype or cross.dtype != factor.dtype: + raise TypeError("factor, gram, and cross must share dtype") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("Fast-HALS CUDA tensors must be contiguous") + + block = 128 + n_blocks = triton.cdiv(rows, block) + block_violation = factor.new_empty((replicates, n_blocks)) + grid = (replicates * n_blocks,) + _hals_sweep_kernel[grid]( + factor, + gram, + cross, + permutation, + active, + block_violation, + M=rows, + K=components, + N_BLOCKS=n_blocks, + BLOCK=block, + IS_FP64=factor.element_size() == 8, + num_warps=4, + ) + return block_violation.sum(dim=1) diff --git a/src/cnmf/gpunmf/solver_mu.py b/src/cnmf/gpunmf/solver_mu.py new file mode 100644 index 0000000..02aafc6 --- /dev/null +++ b/src/cnmf/gpunmf/solver_mu.py @@ -0,0 +1,148 @@ +"""Batched PyTorch multiplicative-update NMF solver.""" + +import numpy as np + +from . import utils + +# --------------------------------------------------------------------- +# MU solver (batch-aware; R>=1 replicates per launch) +# --------------------------------------------------------------------- + + +def _validate_mu_runtime(nmf_kwargs): + """Reject sklearn options that this Frobenius-only MU kernel cannot honor.""" + legacy_regularization = {"alpha", "regularization"}.intersection(nmf_kwargs) + if legacy_regularization: + raise ValueError( + "GPU solver='mu' does not accept deprecated alpha/regularization; " + "use alpha_W and alpha_H" + ) + + beta_loss = nmf_kwargs.get("beta_loss", "frobenius") + if not (beta_loss == 2 or str(beta_loss).lower() == "frobenius"): + raise ValueError("GPU solver='mu' supports only beta_loss='frobenius'") + + alpha_w = float(nmf_kwargs.get("alpha_W", 0.0)) + alpha_h_raw = nmf_kwargs.get("alpha_H", "same") + alpha_h = alpha_w if alpha_h_raw == "same" else float(alpha_h_raw) + if alpha_w != 0.0 or alpha_h != 0.0: + raise ValueError( + "GPU solver='mu' does not yet support alpha_W/alpha_H regularization" + ) + + +def _mu_step(W, H, Xg, eps): + """One MU update, sklearn order: update W from old H, then H from new W. + + Accepts either 2D single-replicate tensors or stacked `[R,...]` replicate + tensors with shared `Xg[1,n,g]`. Operations are out-of-place for compile. + """ + Ht = H.transpose(-2, -1) # [g,k] or [R,g,k] + denominator = W @ (H @ Ht) + denominator = denominator.where(denominator != 0, eps) + W = W * ((Xg @ Ht) / denominator) # W *= XHᵀ / (W·HHᵀ) (uses old H) + Wt = W.transpose(-2, -1) # [k,n] or [R,k,n] + denominator = (Wt @ W) @ H + denominator = denominator.where(denominator != 0, eps) + H = H * ((Wt @ Xg) / denominator) # H *= WᵀX / (WᵀW·H) (uses new W) + return W, H + + +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] + denominator = W @ (H @ Ht) + denominator = denominator.where(denominator != 0, eps) + return W * ((Xg @ Ht) / denominator) + + +def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): + """Run full MU until `max_iter` or all replicate slices meet `tol`.""" + utils._check_runtime_tensors(Xg, W, H, eps) + xnorm2 = utils._sq_norm(Xg) # ‖X‖² once; error check avoids [R,n,g] + err_init = prev_err = None + with torch.no_grad(), utils._cuda_tf32(torch, tf32, device): + it = 0 + while it < max_iter: + n = min(block, max_iter - it) + for _ in range(n): # MU updates run inside the (compiled) step + W, H = step(W, H, Xg, eps) + it += n + err = utils._recon_err(Xg, W, H, xnorm2) + if err_init is None: + 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 fixed-H MU until `max_iter` or all replicate slices meet `tol`.""" + utils._check_runtime_tensors(Xg, W, H, eps) + xnorm2 = utils._sq_norm(Xg) # ‖X‖² once; error check avoids [R,n,g] + err_init = prev_err = None + with torch.no_grad(), utils._cuda_tf32(torch, tf32, device): + it = 0 + while it < max_iter: + n = min(block, max_iter - it) + for _ in range(n): + W = step(W, H, Xg, eps) + it += n + err = utils._recon_err(Xg, W, H, xnorm2) + if err_init is None: + 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 + + +def _nmf_gpu_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + """Run full or fixed-H same-k MU replicates; return one `(H, W)` per seed.""" + torch = utils._loud_import_torch() + + seeds = utils._normalize_seeds(seeds) + + rc = utils._gpu_setup(torch, X, nmf_kwargs, gpu_kwargs) + _validate_mu_runtime(nmf_kwargs) + init = nmf_kwargs.get("init") + update_h = nmf_kwargs.get("update_H", True) is not False + fixed_h = None if update_h else utils._to_checked_fixed_h( + nmf_kwargs.get("H"), rc.k, rc.Xnp.shape[1] + ) + + # 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 + + Ws, Hs = [], [] + for s in seeds: + W0, H0 = utils._init_wh(rc.Xnp, rc.k, s, init) # (usages, spectra) + Ws.append(np.ascontiguousarray(W0)) + if update_h: + Hs.append(np.ascontiguousarray(H0)) + W = torch.as_tensor(np.stack(Ws, 0), dtype=rc.dtype, device=rc.device) # [R, n, k] + if update_h: + H = torch.as_tensor(np.stack(Hs, 0), dtype=rc.dtype, device=rc.device) # [R, k, g] + step, block = utils._execution_plan(torch, rc.opt, rc.device, _mu_step) + W, H = _fit_mu( + torch, Xb, W, H, rc.eps, rc.max_iter, rc.tol, + step, block, utils._want_tf32(torch, rc), rc.device, + ) + else: + H = torch.as_tensor( + np.ascontiguousarray(fixed_h), dtype=rc.dtype, device=rc.device + ).unsqueeze(0) # [1, k, g] shared + step, block = utils._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, utils._want_tf32(torch, rc), rc.device, + ) + + Wc = W.cpu().double().numpy() + Hc = H.cpu().double().numpy() + if update_h: + return [(Hc[r], Wc[r]) for r in range(len(seeds))] + return [(Hc[0], Wc[r]) for r in range(len(seeds))] diff --git a/src/cnmf/gpunmf/utils.py b/src/cnmf/gpunmf/utils.py new file mode 100644 index 0000000..319d336 --- /dev/null +++ b/src/cnmf/gpunmf/utils.py @@ -0,0 +1,366 @@ +"""Shared configuration and runtime helpers for the GPU NMF solvers.""" + +import contextlib +from collections import namedtuple + +import numpy as np + + +_TRUTHY = {"1", "true", "yes", "on"} + +DEFAULT_NMF = { + "max_iter": 1000, + "tol": 1e-4, + "init": "random", + "solver": "mu", +} + + +DEFAULT_GPU = { + "device": "auto", + "dtype": "auto", + "allow_tf32": False, + "compile": False, + "eps": float(np.finfo(np.float32).eps), + "check_every": 10, + "compile_block": 1, + "batch": 1, +} + + +GPU_ARG_NAMES = ( + "engine", + "gpu_device", + "gpu_dtype", + "gpu_allow_tf32", + "gpu_compile", + "gpu_eps", + "gpu_check_every", + "gpu_compile_block", + "gpu_batch", +) + + +# --------------------------------------------------------------------- +# Engine argument validation and solver contract enforcement +# --------------------------------------------------------------------- + +def _validate_engine_args(args, available_solvers): + # sequencial validation of engine args for the given command + _validate_engine_args_for_command(args) + _validate_nmf_solver(args.solver, args.beta_loss, available_solvers) + + +def _validate_engine_args_for_command(args, available_commands=("prepare", "factorize", "consensus")): + """Engine/GPU CLI options are only valid for commands that support the selected engine.""" + available_commands = tuple(available_commands) + if args.command in available_commands: + return + + 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 _validate_nmf_solver(solver, beta_loss, available_solvers): + """Normalize a solver and enforce its loss-function contract.""" + available_solvers = tuple(available_solvers) + solver = str(solver).strip().lower() + if solver not in available_solvers: + available = ", ".join(sorted(available_solvers)) + raise ValueError( + f"solver must be one of: {available}" + ) + if solver == "cd" and not ( + beta_loss == 2 or str(beta_loss).strip().lower() == "frobenius" + ): + raise ValueError("solver='cd' supports only beta_loss='frobenius'") + return solver + +# --------------------------------------------------------------------- +# CLI argument parsing and option resolution +# --------------------------------------------------------------------- + +def gpu_kwargs_from_args(args): + """Collect parsed cNMF CLI GPU flags into a kernel gpu_kwargs dict.""" + raw = { + "device": args.gpu_device, + "dtype": args.gpu_dtype, + "allow_tf32": args.gpu_allow_tf32, + "compile": args.gpu_compile, + "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()): + raise ValueError("GPU options require --engine gpu") + return None + return _resolve_gpu_opts(raw) + + + +def _resolve_gpu_opts(gpu_kwargs): + """Merge Nextflow-provided gpu_kwargs over defaults into a typed opts dict.""" + raw = dict(gpu_kwargs or {}) + + def parse_bool(value, default): + return default if value is None else str(value).strip().lower() in _TRUTHY + + def parse_typed(value, default, cast, normalize=None): + parsed = default if value is None else cast(value) + return normalize(parsed) if normalize is not None else parsed + + def parse_positive_int(value, default): + return max(1, parse_typed(value, default, int)) + + return dict( + device = parse_typed(raw.get("device"), DEFAULT_GPU["device"], str, str.lower), + dtype = parse_typed(raw.get("dtype"), DEFAULT_GPU["dtype"], str, str.lower), + allow_tf32 = parse_bool(raw.get("allow_tf32"), DEFAULT_GPU["allow_tf32"]), + compile = parse_bool(raw.get("compile"), DEFAULT_GPU["compile"]), + 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 (device / dtype / TF32) +# --------------------------------------------------------------------- + + +def _select_device(torch, requested): + """Resolve device, raising for explicit unavailable CUDA/MPS requests.""" + gpu_availability = { + "cuda": torch.cuda.is_available, + "mps": torch.backends.mps.is_available, + } + valid_bases = {"cpu", *gpu_availability} + + if requested == "auto": + for candidate, is_available in gpu_availability.items(): + if is_available(): + return candidate + return "cpu" + + base = requested.split(":")[0] + if base not in valid_bases: + raise ValueError(f"device={requested!r} not recognized; use auto|cpu|cuda|cuda:N|mps.") + + if base in gpu_availability and not gpu_availability[base](): + raise RuntimeError(f"device={requested!r} requested but {base.upper()} is unavailable " + "(use device='auto' or 'cpu').") + return requested + + +def _select_storage(torch, requested, device): + """Resolve storage/matmul dtype for the selected device.""" + base = device.split(":")[0] + dtype_map = { + "fp32": torch.float32, + "fp64": torch.float64, + "bf16": torch.bfloat16, + } + + if requested == "auto": + requested = "fp64" if base == "cpu" else "fp32" + + if requested not in dtype_map: + choices = "|".join(["auto", *dtype_map]) + raise ValueError(f"dtype={requested!r} not recognized; use {choices}.") + + if requested == "fp64" and base == "mps": + raise RuntimeError("dtype='fp64' requested but MPS has no fp64 " + "(use dtype='auto'/'fp32', or device='cpu'/'cuda').") + + if requested == "bf16": + if base != "cuda": + raise RuntimeError("dtype='bf16' is only supported on CUDA in this kernel " + "(use dtype='auto'/'fp32' or device='cuda').") + is_supported = getattr(torch.cuda, "is_bf16_supported", None) + if callable(is_supported) and not is_supported(): + raise RuntimeError("dtype='bf16' requested but this CUDA device does not support bf16.") + + return dtype_map[requested] + + +@contextlib.contextmanager +def _cuda_tf32(torch, enable, device): + """Temporarily set CUDA TF32 matmul flags; no-op off CUDA.""" + if not device.startswith("cuda") or not torch.cuda.is_available(): + yield + return + prev_allow = torch.backends.cuda.matmul.allow_tf32 + prev_prec = torch.get_float32_matmul_precision() + torch.backends.cuda.matmul.allow_tf32 = enable + torch.set_float32_matmul_precision("high" if enable else "highest") + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = prev_allow + torch.set_float32_matmul_precision(prev_prec) + + +# --------------------------------------------------------------------- +# 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 + # preflight and/or row-blocked sparse loading before materializing full X in host RAM. + Xnp = X.toarray() if hasattr(X, "toarray") else np.asarray(X) + if Xnp.ndim != 2: + raise ValueError("NMF input X must be a 2D matrix") + if 0 in Xnp.shape: + raise ValueError("NMF input X must have at least one row and one column") + if not np.isfinite(Xnp).all(): + raise ValueError("NMF input X contains NaN/inf") + if Xnp.size and Xnp.min() < 0: + raise ValueError(f"NMF requires non-negative input X; found min(X) = {float(Xnp.min()):.4g}") + 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: + import torch + except ModuleNotFoundError as e: + if e.name != "torch": + raise + raise RuntimeError( + "PyTorch is required for GPU NMF but is not installed in the active environment. " + "Install cNMF with GPU support in that environment using " + "`python -m pip install -e \".[gpu]\"`, or install a CUDA-compatible torch build " + "before installing the cNMF GPU extra." + ) from e + return torch + + +def _loud_import_initialize_nmf(): + """Import sklearn's NMF initializer with an actionable environment/version error.""" + try: + from sklearn.decomposition._nmf import _initialize_nmf + except ModuleNotFoundError as e: + if e.name != "sklearn": + raise + raise RuntimeError( + "scikit-learn is required for sklearn-compatible NMF initialization but is not " + "installed in the active environment. Install cNMF into that environment using " + "`python -m pip install -e .` or `python -m pip install -e \".[gpu]\"`." + ) from e + except ImportError as e: + raise RuntimeError( + "The installed scikit-learn does not expose sklearn.decomposition._nmf._initialize_nmf. " + "Use a supported scikit-learn version or update the GPU NMF initializer adapter for " + "this scikit-learn version." + ) from e + return _initialize_nmf + + +def _init_wh(Xnp, k, seed, init): + """Initialize W/H with sklearn parity; `init=None` uses DEFAULT_NMF.""" + if init == "custom": + raise NotImplementedError("GPU NMF does not support init='custom'") + _initialize_nmf = _loud_import_initialize_nmf() + return _initialize_nmf( + Xnp, + n_components=k, + init=(init or DEFAULT_NMF["init"]), + random_state=seed, + ) + + +# --------------------------------------------------------------------- +# Shared solver runtime and device-tensor helpers +# --------------------------------------------------------------------- + + +# 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_eps(torch, eps, dtype, device): + """Create the exact-zero denominator replacement on runtime dtype/device.""" + return torch.tensor(eps, dtype=dtype, device=device) + + +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))}") + + +def _normalize_seeds(seeds): + """Return a non-empty list of integer or None replicate seeds.""" + normalized = [None if seed is None else int(seed) for seed in seeds] + if not normalized: + raise ValueError("seeds must be a non-empty list of per-replicate random states") + return normalized + + +def _recon_err(Xg, W, H, xnorm2): + """Per-replicate ‖X − WH‖_F without materializing the full residual.""" + Wt = W.transpose(-2, -1) + cross = ((Wt @ Xg) * H).sum(dim=(-2, -1)) + whnorm = ((Wt @ W) * (H @ H.transpose(-2, -1))).sum(dim=(-2, -1)) + return (xnorm2 - 2.0 * cross + whnorm).clamp_min(0).sqrt() + + +def _sq_norm(X): + """Return ‖X‖² using row chunks to avoid BLAS vector-length limits.""" + rows = X.shape[0] + step = max(1, (1 << 28) // max(1, X.numel() // max(1, rows))) + total = X.new_zeros(()) + for i in range(0, rows, step): + total = total + X[i:i + step].square().sum() + return total + + +def _execution_plan(torch, opt, device, step_fn): + """Return the solver step and convergence-check block 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"] diff --git a/src/cnmf/nmf_gpu.py b/src/cnmf/nmf_gpu.py deleted file mode 100644 index 1b92cbe..0000000 --- a/src/cnmf/nmf_gpu.py +++ /dev/null @@ -1,608 +0,0 @@ -#!/usr/bin/env python -"""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 -# --------------------------------------------------------------------- - - -_TRUTHY = {"1", "true", "yes", "on"} - - -DEFAULT_NMF = { - "max_iter": 1000, - "tol": 1e-4, - "init": "random", -} - - -_SKLEARN_EPSILON = float(np.finfo(np.float32).eps) - - -DEFAULT_GPU = { - "device": "auto", - "dtype": "auto", - "allow_tf32": False, - "compile": False, - "eps": _SKLEARN_EPSILON, - "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") - group.add_argument("--engine", type=str.lower, choices=["cpu", "gpu"], help="[factorize,consensus] NMF engine to use (default cpu)") - group.add_argument("--gpu-device", type=str, help="[factorize,consensus,gpu] Device for GPU NMF: auto, cpu, cuda, cuda:N, or mps") - group.add_argument("--gpu-dtype", type=str.lower, choices=["auto", "fp32", "fp64", "bf16"], help="[factorize,consensus,gpu] Storage and matmul dtype for GPU NMF (default auto)") - group.add_argument("--gpu-allow-tf32", action="store_const", const=True, help="[factorize,consensus,gpu] Allow TF32 for CUDA fp32 matrix multiplication") - group.add_argument("--gpu-compile", action="store_const", const=True, help="[factorize,consensus,gpu] Enable torch.compile for the GPU NMF update step") - group.add_argument("--gpu-eps", type=float, help="[factorize,consensus,gpu] Replacement for exactly-zero MU denominators") - 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 - - -def gpu_kwargs_from_args(args): - """Collect parsed cNMF CLI GPU flags into a kernel gpu_kwargs dict.""" - raw = { - "device": args.gpu_device, - "dtype": args.gpu_dtype, - "allow_tf32": args.gpu_allow_tf32, - "compile": args.gpu_compile, - "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()): - raise ValueError("GPU options require --engine gpu") - return None - return _resolve_gpu_opts(raw) - - -def validate_engine_args_for_command(args, available_commands): - """Engine/GPU CLI options are only valid for commands that support the selected engine.""" - available_commands = tuple(available_commands) - if args.command in available_commands: - return - - gpu_arg_names = [ - "engine", - "gpu_device", - "gpu_dtype", - "gpu_allow_tf32", - "gpu_compile", - "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 _resolve_gpu_opts(gpu_kwargs): - """Merge Nextflow-provided gpu_kwargs over defaults into a typed opts dict.""" - raw = dict(gpu_kwargs or {}) - - def parse_bool(value, default): - return default if value is None else str(value).strip().lower() in _TRUTHY - - def parse_typed(value, default, cast, normalize=None): - parsed = default if value is None else cast(value) - return normalize(parsed) if normalize is not None else parsed - - def parse_positive_int(value, default): - return max(1, parse_typed(value, default, int)) - - return dict( - device = parse_typed(raw.get("device"), DEFAULT_GPU["device"], str, str.lower), - dtype = parse_typed(raw.get("dtype"), DEFAULT_GPU["dtype"], str, str.lower), - allow_tf32 = parse_bool(raw.get("allow_tf32"), DEFAULT_GPU["allow_tf32"]), - compile = parse_bool(raw.get("compile"), DEFAULT_GPU["compile"]), - 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 (device / dtype / TF32) -# --------------------------------------------------------------------- - - -def _select_device(torch, requested): - """Resolve device, raising for explicit unavailable CUDA/MPS requests.""" - gpu_availability = { - "cuda": torch.cuda.is_available, - "mps": torch.backends.mps.is_available, - } - valid_bases = {"cpu", *gpu_availability} - - if requested == "auto": - for candidate, is_available in gpu_availability.items(): - if is_available(): - return candidate - return "cpu" - - base = requested.split(":")[0] - if base not in valid_bases: - raise ValueError(f"device={requested!r} not recognized; use auto|cpu|cuda|cuda:N|mps.") - - if base in gpu_availability and not gpu_availability[base](): - raise RuntimeError(f"device={requested!r} requested but {base.upper()} is unavailable " - "(use device='auto' or 'cpu').") - return requested - - -def _select_storage(torch, requested, device): - """Resolve storage/matmul dtype for the selected device.""" - base = device.split(":")[0] - dtype_map = { - "fp32": torch.float32, - "fp64": torch.float64, - "bf16": torch.bfloat16, - } - - if requested == "auto": - requested = "fp64" if base == "cpu" else "fp32" - - if requested not in dtype_map: - choices = "|".join(["auto", *dtype_map]) - raise ValueError(f"dtype={requested!r} not recognized; use {choices}.") - - if requested == "fp64" and base == "mps": - raise RuntimeError("dtype='fp64' requested but MPS has no fp64 " - "(use dtype='auto'/'fp32', or device='cpu'/'cuda').") - - if requested == "bf16": - if base != "cuda": - raise RuntimeError("dtype='bf16' is only supported on CUDA in this kernel " - "(use dtype='auto'/'fp32' or device='cuda').") - is_supported = getattr(torch.cuda, "is_bf16_supported", None) - if callable(is_supported) and not is_supported(): - raise RuntimeError("dtype='bf16' requested but this CUDA device does not support bf16.") - - return dtype_map[requested] - - -@contextlib.contextmanager -def _cuda_tf32(torch, enable, device): - """Temporarily set CUDA TF32 matmul flags; no-op off CUDA.""" - if not device.startswith("cuda") or not torch.cuda.is_available(): - yield - return - prev_allow = torch.backends.cuda.matmul.allow_tf32 - prev_prec = torch.get_float32_matmul_precision() - torch.backends.cuda.matmul.allow_tf32 = enable - torch.set_float32_matmul_precision("high" if enable else "highest") - try: - yield - finally: - torch.backends.cuda.matmul.allow_tf32 = prev_allow - torch.set_float32_matmul_precision(prev_prec) - - -# --------------------------------------------------------------------- -# 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 - # preflight and/or row-blocked sparse loading before materializing full X in host RAM. - Xnp = X.toarray() if hasattr(X, "toarray") else np.asarray(X) - if Xnp.ndim != 2: - raise ValueError("NMF input X must be a 2D matrix") - if 0 in Xnp.shape: - raise ValueError("NMF input X must have at least one row and one column") - if not np.isfinite(Xnp).all(): - raise ValueError("NMF input X contains NaN/inf") - if Xnp.size and Xnp.min() < 0: - raise ValueError(f"NMF requires non-negative input X; found min(X) = {float(Xnp.min()):.4g}") - 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: - import torch - except ModuleNotFoundError as e: - if e.name != "torch": - raise - raise RuntimeError( - "PyTorch is required for GPU NMF but is not installed in the active environment. " - "Install cNMF with GPU support in that environment using " - "`python -m pip install -e \".[gpu]\"`, or install a CUDA-compatible torch build " - "before installing the cNMF GPU extra." - ) from e - return torch - - -def _loud_import_initialize_nmf(): - """Import sklearn's NMF initializer with an actionable environment/version error.""" - try: - from sklearn.decomposition._nmf import _initialize_nmf - except ModuleNotFoundError as e: - if e.name != "sklearn": - raise - raise RuntimeError( - "scikit-learn is required for sklearn-compatible NMF initialization but is not " - "installed in the active environment. Install cNMF into that environment using " - "`python -m pip install -e .` or `python -m pip install -e \".[gpu]\"`." - ) from e - except ImportError as e: - raise RuntimeError( - "The installed scikit-learn does not expose sklearn.decomposition._nmf._initialize_nmf. " - "Use a supported scikit-learn version or update the GPU NMF initializer adapter for " - "this scikit-learn version." - ) from e - return _initialize_nmf - - -def _init_wh(Xnp, k, seed, init): - """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() - return _initialize_nmf( - Xnp, - n_components=k, - init=(init or DEFAULT_NMF["init"]), - random_state=seed, - ) - - -# --------------------------------------------------------------------- -# Multiplicative-update steppers (one MU iteration; batch-aware) -# --------------------------------------------------------------------- - - -def _mu_step(W, H, Xg, eps): - """One MU update, sklearn order: update W from old H, then H from new W. - - Accepts either 2D single-replicate tensors or stacked `[R,...]` replicate - tensors with shared `Xg[1,n,g]`. Operations are out-of-place for compile. - """ - Ht = H.transpose(-2, -1) # [g,k] or [R,g,k] - denominator = W @ (H @ Ht) - denominator = denominator.where(denominator != 0, eps) - W = W * ((Xg @ Ht) / denominator) # W *= XHᵀ / (W·HHᵀ) (uses old H) - Wt = W.transpose(-2, -1) # [k,n] or [R,k,n] - denominator = (Wt @ W) @ H - denominator = denominator.where(denominator != 0, eps) - H = H * ((Wt @ Xg) / denominator) # H *= WᵀX / (WᵀW·H) (uses new W) - return W, H - - -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] - denominator = W @ (H @ Ht) - denominator = denominator.where(denominator != 0, eps) - return W * ((Xg @ Ht) / denominator) - - -# --------------------------------------------------------------------- -# MU fit loops (iterate to convergence; batch-aware) -# --------------------------------------------------------------------- - - -def _recon_err(Xg, W, H, xnorm2): - """Per-replicate ‖X − WH‖_F via the identity ‖X‖² − 2⟨X,WH⟩ + ‖WH‖², using only small - [R,k,g]/[R,k,k] matmuls (WᵀX, WᵀW, HHᵀ) — it NEVER materializes the [R,n,g] product, so batched - fits scale to millions of cells (a full [R,n,g] residual would be ~R·n·g floats, tens of GB). - Returns a 0-dim tensor (unbatched) or shape-[R] (batched); xnorm2 = ‖X‖² is precomputed once.""" - Wt = W.transpose(-2, -1) - cross = ((Wt @ Xg) * H).sum(dim=(-2, -1)) # ⟨X, WH⟩ per replicate - whnorm = ((Wt @ W) * (H @ H.transpose(-2, -1))).sum(dim=(-2, -1)) # ‖WH‖² per replicate - return (xnorm2 - 2.0 * cross + whnorm).clamp_min(0).sqrt() - - -def _sq_norm(X): - """‖X‖² (every element squared, then summed) in row chunks. torch.dot is BLAS-backed and caps - its vector length at 2³¹ elements, so a flat dot overflows past ~2.1B entries (e.g. 2.4M cells - × 2000 genes = 4.8B); this chunked reduction is index-safe at any size and also bounds the - squaring temporary. Returns a 0-dim tensor.""" - rows = X.shape[0] - step = max(1, (1 << 28) // max(1, X.numel() // max(1, rows))) # ~2²⁸ elems/chunk (~1GB fp32) - total = X.new_zeros(()) - for i in range(0, rows, step): - total = total + X[i:i + step].square().sum() - return total - - -def _fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): - """Run full MU until `max_iter` or all replicate slices meet `tol`.""" - _check_runtime_tensors(Xg, W, H, eps) - xnorm2 = _sq_norm(Xg) # ‖X‖² once; error check avoids [R,n,g] - err_init = prev_err = None - with torch.no_grad(), _cuda_tf32(torch, tf32, device): - it = 0 - while it < max_iter: - n = min(block, max_iter - it) - for _ in range(n): # MU updates run inside the (compiled) step - W, H = step(W, H, Xg, eps) - it += n - err = _recon_err(Xg, W, H, xnorm2) - if err_init is None: - 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 fixed-H MU until `max_iter` or all replicate slices meet `tol`.""" - _check_runtime_tensors(Xg, W, H, eps) - xnorm2 = _sq_norm(Xg) # ‖X‖² once; error check avoids [R,n,g] - err_init = prev_err = None - with torch.no_grad(), _cuda_tf32(torch, tf32, device): - it = 0 - while it < max_iter: - n = min(block, max_iter - it) - for _ in range(n): - W = step(W, H, Xg, eps) - it += n - err = _recon_err(Xg, W, H, xnorm2) - if err_init is None: - 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 - - -# --------------------------------------------------------------------- -# 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 exact-zero denominator replacement 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 _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))}") - - -# --------------------------------------------------------------------- -# NMF kernels (batch-aware; R>=1 replicates per launch, R=1 = single instance) -# --------------------------------------------------------------------- - - -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") - - 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 - - # 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] - - 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 _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() - - 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") - - 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 - - # 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 - - -# --------------------------------------------------------------------- -# 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 `_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 _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 - # Densify X once. _nmf_gpu_mu -> _gpu_setup calls X.toarray(), so passing the sparse matrix - # would re-densify the full n×g matrix on *every* batch (O(#batches) redundant conversions plus - # RAM churn at millions of cells). Convert once here; downstream then sees a dense ndarray. - X_dense = norm_counts.X.toarray() if hasattr(norm_counts.X, "toarray") else np.asarray(norm_counts.X) - 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(X_dense, 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 4c31b5d..72761c2 100644 --- a/tests/test_nmf_gpu.py +++ b/tests/test_nmf_gpu.py @@ -1,4 +1,4 @@ -"""Reliability tests for the NMF GPU kernel (`src/cnmf/nmf_gpu.py`). +"""Reliability tests for the GPU NMF package (`cnmf.gpunmf`). 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 instance `_nmf` hook contract. MU update order, convergence, and iteration bounds Pins the sklearn-style W-then-H multiplicative-update order, early-stop @@ -45,9 +45,13 @@ and fp32 with identical initialization, seed, and iteration count. Small cases run routinely on the torch CPU backend and CUDA when available. An opt-in, memory-gated CUDA stress case uses a 100,000 x 20,000 matrix. + +sklearn CD parity + Pins sklearn's serial W-then-H Fast-HALS updates, regularization scaling, + shuffled coordinate stream, fixed-H zero initialization, and per-replicate + projected-gradient stopping for fp64 and fp32 batches. """ -import argparse import builtins import gc import os @@ -73,10 +77,10 @@ # --------------------------------------------------------------------- def test_kernel_loader_fails_when_kernel_file_is_missing(tmp_path): """Fail the test harness clearly if the kernel module is absent.""" - missing_kernel = tmp_path / "missing_nmf_gpu.py" + missing_kernel = tmp_path / "missing_gpunmf.py" with pytest.raises(pytest.fail.Exception, match="Required NMF GPU kernel file is missing"): - load_kernel_module(module_name="missing_nmf_gpu_for_test", kernel_path=missing_kernel) + load_kernel_module(module_name="missing_gpunmf_for_test", kernel_path=missing_kernel) # --------------------------------------------------------------------- @@ -88,7 +92,7 @@ def test_nmf_gpu_reconstructs_known_low_rank_matrix_with_small_relative_error(ke k = 3 X = low_rank_matrix(rank=k) - H, W = kernel._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": k, "max_iter": 600, "tol": 0, "random_state": 0}, {"device": "cpu", "check_every": 600}, @@ -104,7 +108,7 @@ def test_nmf_gpu_returns_spectra_then_usages_with_cnmf_orientation(kernel): require_nmf_runtime() X = small_nonnegative_matrix(cells=7, genes=5) - H, W = kernel._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 2, "max_iter": 2, "random_state": 0}, {"device": "cpu"}, @@ -118,7 +122,7 @@ 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._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 3, "max_iter": 3, "random_state": 0}, {"device": "cpu"}, @@ -132,7 +136,7 @@ def test_nmf_gpu_fp32_compute_still_returns_float64_numpy_outputs(kernel): require_nmf_runtime() X = small_nonnegative_matrix() - H, W = kernel._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 2, "max_iter": 1, "random_state": 0}, {"device": "cpu", "dtype": "fp32"}, @@ -142,27 +146,6 @@ def test_nmf_gpu_fp32_compute_still_returns_float64_numpy_outputs(kernel): assert W.dtype == np.float64 -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()) - - def fake_factorize(X, nmf_kwargs, gpu_kwargs=None): - calls.append((X, nmf_kwargs, gpu_kwargs)) - return sentinel - - 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) - - assert result is sentinel - assert calls == [(X, {"n_components": 1}, gpu_kwargs)] - assert nmf_kwargs == {"engine": "gpu", "gpu": gpu_kwargs, "n_components": 1} - - # --------------------------------------------------------------------- # MU update order, convergence, and iteration bounds # --------------------------------------------------------------------- @@ -180,7 +163,7 @@ def test_mu_step_updates_w_first_using_old_h_then_h_using_new_w(kernel): denominator = (expected_W.T @ expected_W) @ H0 denominator = denominator.where(denominator != 0, eps) expected_H = H0 * ((expected_W.T @ Xg) / denominator) - W, H = kernel._mu_step(W0, H0, Xg, eps) + W, H = kernel.solver_mu._mu_step(W0, H0, Xg, eps) assert torch.allclose(W, expected_W) assert torch.allclose(H, expected_H) @@ -191,7 +174,7 @@ def test_mu_step_fixed_h_matches_sklearn_exact_zero_protection(kernel): torch = require_nmf_runtime() eps = torch.tensor(np.finfo(np.float32).eps, dtype=torch.float64) - zero_result = kernel._mu_step_fixed_h( + zero_result = kernel.solver_mu._mu_step_fixed_h( torch.ones((1, 1), dtype=torch.float64), torch.zeros((1, 1), dtype=torch.float64), torch.ones((1, 1), dtype=torch.float64), @@ -201,7 +184,7 @@ def test_mu_step_fixed_h_matches_sklearn_exact_zero_protection(kernel): assert torch.isfinite(zero_result).all() tiny = torch.tensor(1e-12, dtype=torch.float64) - tiny_result = kernel._mu_step_fixed_h( + tiny_result = kernel.solver_mu._mu_step_fixed_h( torch.ones((1, 1), dtype=torch.float64), tiny.reshape(1, 1), torch.ones((1, 1), dtype=torch.float64), @@ -223,7 +206,7 @@ def no_change_step(W, H, Xg, eps): calls["count"] += 1 return W, H - kernel._fit_mu(torch, Xg, W, H, eps, 10, 1e-4, no_change_step, 1, False, "cpu") + kernel.solver_mu._fit_mu(torch, Xg, W, H, eps, 10, 1e-4, no_change_step, 1, False, "cpu") assert calls["count"] == 2 @@ -241,7 +224,7 @@ def no_change_step(W, H, Xg, eps): calls["count"] += 1 return W, H - kernel._fit_mu(torch, Xg, W, H, eps, 6, -1.0, no_change_step, 4, False, "cpu") + kernel.solver_mu._fit_mu(torch, Xg, W, H, eps, 6, -1.0, no_change_step, 4, False, "cpu") assert calls["count"] == 6 @@ -255,7 +238,7 @@ def test_check_runtime_tensors_rejects_mixed_dtypes(kernel): eps = torch.tensor(1e-9, dtype=torch.float32) with pytest.raises(RuntimeError, match="share dtype"): - kernel._check_runtime_tensors(Xg, W, H, eps) + kernel.utils._check_runtime_tensors(Xg, W, H, eps) # --------------------------------------------------------------------- @@ -266,14 +249,20 @@ def test_compile_mode_matches_eager_output_for_same_seed_and_options(kernel, mon torch = require_nmf_runtime() monkeypatch.setattr(torch, "compile", lambda fn: fn) X = small_nonnegative_matrix(cells=8, genes=6) - nmf_kwargs = {"n_components": 2, "max_iter": 4, "tol": -1.0, "random_state": 0} + nmf_kwargs = { + "n_components": 2, + "max_iter": 4, + "tol": -1.0, + "random_state": 0, + "solver": "mu", + } - eager_H, eager_W = kernel._nmf_gpu( + eager_H, eager_W = run_nmf_gpu(kernel, X, nmf_kwargs, {"device": "cpu", "dtype": "fp64", "compile": False, "check_every": 1}, ) - compiled_H, compiled_W = kernel._nmf_gpu( + compiled_H, compiled_W = run_nmf_gpu(kernel, X, nmf_kwargs, {"device": "cpu", "dtype": "fp64", "compile": True, "compile_block": 2}, @@ -288,12 +277,12 @@ def test_compile_mode_uses_explicit_multi_iteration_compile_block_when_requested torch = require_nmf_runtime() calls = [] monkeypatch.setattr(torch, "compile", lambda fn: calls.append(fn) or fn) - opt = dict(kernel.DEFAULT_GPU, compile=True, check_every=1, compile_block=3) + opt = dict(kernel.utils.DEFAULT_GPU, compile=True, check_every=1, compile_block=3) - step, block = kernel._execution_plan(torch, opt, "cpu") + step, block = kernel.utils._execution_plan(torch, opt, "cpu", kernel.solver_mu._mu_step) - assert calls == [kernel._mu_step] - assert step is kernel._mu_step + assert calls == [kernel.solver_mu._mu_step] + assert step is kernel.solver_mu._mu_step assert block == 3 @@ -307,8 +296,8 @@ def test_nmf_gpu_random_state_is_reproducible(kernel): kwargs = {"n_components": 3, "max_iter": 3, "random_state": 13} gpu = {"device": "cpu", "check_every": 3} - H1, W1 = kernel._nmf_gpu(X, kwargs, gpu) - H2, W2 = kernel._nmf_gpu(X, kwargs, gpu) + H1, W1 = run_nmf_gpu(kernel, X, kwargs, gpu) + H2, W2 = run_nmf_gpu(kernel, X, kwargs, gpu) assert np.allclose(H1, H2) assert np.allclose(W1, W2) @@ -318,10 +307,15 @@ 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"} + kwargs = { + "n_components": 3, + "max_iter": 0, + "init": "random", + "solver": "mu", + } - 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"}) + H1, W1 = run_nmf_gpu(kernel, X, dict(kwargs, random_state=1), {"device": "cpu"}) + H2, W2 = run_nmf_gpu(kernel, X, dict(kwargs, random_state=2), {"device": "cpu"}) assert not np.allclose(H1, H2) assert not np.allclose(W1, W2) @@ -335,9 +329,11 @@ def fake_initialize(X, n_components, init, random_state): seen.append(init) return np.ones((X.shape[0], n_components)), np.ones((n_components, X.shape[1])) - monkeypatch.setattr(kernel, "_loud_import_initialize_nmf", lambda: fake_initialize) + monkeypatch.setattr( + kernel.utils, "_loud_import_initialize_nmf", lambda: fake_initialize + ) - kernel._init_wh(small_nonnegative_matrix(), 2, 0, None) + kernel.utils._init_wh(small_nonnegative_matrix(), 2, 0, None) assert seen == ["random"] @@ -349,7 +345,7 @@ def test_random_init_matches_sklearn_initializer_contract(kernel): X = small_nonnegative_matrix() expected_W, expected_H = _initialize_nmf(X, n_components=3, init="random", random_state=5) - W, H = kernel._init_wh(X, 3, 5, "random") + W, H = kernel.utils._init_wh(X, 3, 5, "random") assert np.allclose(W, expected_W) assert np.allclose(H, expected_H) @@ -363,10 +359,12 @@ def fake_initialize(X, n_components, init, random_state): seen.append(init) return np.ones((X.shape[0], n_components)), np.ones((n_components, X.shape[1])) - monkeypatch.setattr(kernel, "_loud_import_initialize_nmf", lambda: fake_initialize) + monkeypatch.setattr( + kernel.utils, "_loud_import_initialize_nmf", lambda: fake_initialize + ) for init in ("nndsvd", "nndsvda", "nndsvdar"): - kernel._init_wh(small_nonnegative_matrix(), 2, 0, init) + kernel.utils._init_wh(small_nonnegative_matrix(), 2, 0, init) assert seen == ["nndsvd", "nndsvda", "nndsvdar"] @@ -374,7 +372,7 @@ def fake_initialize(X, n_components, init, random_state): 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") + kernel.utils._init_wh(small_nonnegative_matrix(), 2, 0, "custom") # --------------------------------------------------------------------- @@ -422,7 +420,7 @@ def _sklearn_mu_reference(X, nmf_kwargs): def _gpu_parity_result(kernel, X, nmf_kwargs, dtype_name, device): """Run the PyTorch kernel without TF32, compilation, or early stopping.""" - return kernel._nmf_gpu( + return run_nmf_gpu(kernel, X, nmf_kwargs, { @@ -562,25 +560,425 @@ def test_sklearn_mu_matches_cuda_on_100k_by_20k_matrix( torch.cuda.empty_cache() +# --------------------------------------------------------------------- +# sklearn coordinate-descent / Fast-HALS parity +# --------------------------------------------------------------------- +CD_PARITY_CASES = [ + pytest.param("fp64", np.float64, 5e-10, 5e-11, id="fp64"), + pytest.param("fp32", np.float32, 8e-5, 8e-6, id="fp32"), +] + + +def _cd_nmf_kwargs(n_components, seed, max_iter, **overrides): + """Return one deterministic sklearn CD contract for parity tests.""" + kwargs = { + "n_components": n_components, + "init": "random", + "random_state": seed, + "solver": "cd", + "beta_loss": "frobenius", + "tol": 0.0, + "max_iter": max_iter, + "alpha_W": 0.03, + "alpha_H": 0.02, + "l1_ratio": 0.25, + "shuffle": False, + } + kwargs.update(overrides) + return kwargs + + +def _sklearn_cd_reference(X, nmf_kwargs): + """Run sklearn CD from the same W/H inputs and return cNMF's H/W order.""" + pytest.importorskip("sklearn") + from sklearn.decomposition import non_negative_factorization + from sklearn.exceptions import ConvergenceWarning + + kwargs = dict(nmf_kwargs) + W = kwargs.pop("W", None) + H = kwargs.pop("H", None) + W = None if W is None else np.array(W, copy=True) + H = None if H is None else np.array(H, copy=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", ConvergenceWarning) + expected_W, expected_H, n_iter = non_negative_factorization( + X, W=W, H=H, **kwargs + ) + return expected_H, expected_W, n_iter + + +@pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", CD_PARITY_CASES) +@pytest.mark.parametrize("seed,max_iter", [(0, 1), (19, 8)]) +@pytest.mark.parametrize("alpha_h", [0.02, "same"]) +def test_sklearn_cd_matches_gpu_kernel_on_torch_cpu( + kernel, dtype_name, np_dtype, rtol, atol, seed, max_iter, alpha_h +): + """Match sklearn's update order and regularization in fp32 and fp64.""" + require_nmf_runtime() + X = np.random.default_rng(42).random((17, 11), dtype=np_dtype) + X += np_dtype(0.1) + nmf_kwargs = _cd_nmf_kwargs( + 3, seed, max_iter, alpha_H=alpha_h + ) + + expected_H, expected_W, expected_n_iter = _sklearn_cd_reference( + X, nmf_kwargs + ) + actual_H, actual_W = run_nmf_gpu(kernel, + X, + nmf_kwargs, + { + "device": "cpu", + "dtype": dtype_name, + "allow_tf32": False, + "compile": False, + }, + ) + + assert expected_n_iter == max_iter + np.testing.assert_allclose(actual_H, expected_H, rtol=rtol, atol=atol) + np.testing.assert_allclose(actual_W, expected_W, rtol=rtol, atol=atol) + np.testing.assert_allclose( + _relative_reconstruction_error(X, actual_H, actual_W), + _relative_reconstruction_error(X, expected_H, expected_W), + rtol=rtol, + atol=atol, + ) + + +@pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", CD_PARITY_CASES) +def test_cd_batched_seeds_match_independent_runs( + kernel, dtype_name, np_dtype, rtol, atol +): + """Batching must not change any seed's independent coordinate path.""" + require_nmf_runtime() + X = np.random.default_rng(7).random((19, 13), dtype=np_dtype) + X += np_dtype(0.1) + seeds = [23, 2, 41] + nmf_kwargs = _cd_nmf_kwargs(4, seed=0, max_iter=6) + gpu_kwargs = { + "device": "cpu", + "dtype": dtype_name, + "allow_tf32": False, + "compile": False, + } + + batched = kernel.solver_cd._nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs) + + assert len(batched) == len(seeds) + for (batched_H, batched_W), seed in zip(batched, seeds): + (single_H, single_W), = kernel.solver_cd._nmf_gpu_cd( + X, [seed], nmf_kwargs, gpu_kwargs + ) + np.testing.assert_allclose( + batched_H, single_H, rtol=rtol, atol=atol + ) + np.testing.assert_allclose( + batched_W, single_W, rtol=rtol, atol=atol + ) + + +@pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", CD_PARITY_CASES) +def test_sklearn_cd_fixed_h_matches_batched_gpu_refit( + kernel, dtype_name, np_dtype, rtol, atol +): + """Fixed-H CD must keep H and use sklearn's exact-zero W start.""" + require_nmf_runtime() + rng = np.random.default_rng(31) + X = rng.random((15, 9), dtype=np_dtype) + np_dtype(0.1) + fixed_H = rng.random((3, 9), dtype=np_dtype) + np_dtype(0.1) + seeds = [5, 29] + nmf_kwargs = _cd_nmf_kwargs( + 3, + seed=0, + max_iter=7, + update_H=False, + H=fixed_H, + ) + gpu_kwargs = { + "device": "cpu", + "dtype": dtype_name, + "allow_tf32": False, + "compile": False, + } + + actual = kernel.solver_cd._nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs) + + for actual_H, actual_W in actual: + expected_H, expected_W, _ = _sklearn_cd_reference(X, nmf_kwargs) + np.testing.assert_array_equal( + actual_H, fixed_H.astype(np.float64) + ) + np.testing.assert_array_equal(expected_H, fixed_H) + np.testing.assert_allclose( + actual_W, expected_W, rtol=rtol, atol=atol + ) + + +def test_sklearn_cd_custom_init_matches_gpu_kernel(kernel): + """Custom W/H should be consumed exactly as sklearn consumes them.""" + require_nmf_runtime() + rng = np.random.default_rng(52) + X = rng.random((13, 10)) + 0.1 + W0 = rng.random((13, 3)) + 0.1 + H0 = rng.random((3, 10)) + 0.1 + nmf_kwargs = _cd_nmf_kwargs( + 3, + seed=11, + max_iter=4, + init="custom", + W=W0, + H=H0, + ) + + expected_H, expected_W, _ = _sklearn_cd_reference(X, nmf_kwargs) + actual_H, actual_W = run_nmf_gpu(kernel, + X, nmf_kwargs, {"device": "cpu", "dtype": "fp64"} + ) + + np.testing.assert_allclose( + actual_H, expected_H, rtol=5e-10, atol=5e-11 + ) + np.testing.assert_allclose( + actual_W, expected_W, rtol=5e-10, atol=5e-11 + ) + + +def test_cd_batched_convergence_iterations_match_sklearn(kernel, monkeypatch): + """Each batch slice must stop at sklearn's projected-gradient iteration.""" + require_nmf_runtime() + X = np.random.default_rng(63).random((31, 17)) + 0.1 + seeds = [0, 7, 103] + nmf_kwargs = _cd_nmf_kwargs( + 4, + seed=0, + max_iter=200, + tol=1e-4, + alpha_W=0.0, + alpha_H=0.0, + l1_ratio=0.0, + ) + captured_n_iter = [] + real_fit_cd = kernel.solver_cd._fit_cd + + def capture_n_iter(*args, **kwargs): + result = real_fit_cd(*args, **kwargs) + captured_n_iter.extend(result[2].cpu().tolist()) + return result + + monkeypatch.setattr(kernel.solver_cd, "_fit_cd", capture_n_iter) + actual = kernel.solver_cd._nmf_gpu_cd( + X, + seeds, + nmf_kwargs, + {"device": "cpu", "dtype": "fp64", "allow_tf32": False}, + ) + + expected_n_iter = [] + for (actual_H, actual_W), seed in zip(actual, seeds): + expected_kwargs = dict(nmf_kwargs, random_state=seed) + expected_H, expected_W, n_iter = _sklearn_cd_reference( + X, expected_kwargs + ) + expected_n_iter.append(n_iter) + np.testing.assert_allclose( + actual_H, expected_H, rtol=5e-10, atol=5e-11 + ) + np.testing.assert_allclose( + actual_W, expected_W, rtol=5e-10, atol=5e-11 + ) + + assert captured_n_iter == expected_n_iter + + +def test_sklearn_cd_shuffled_batch_matches_per_seed_rng_stream(kernel): + """Shuffled CD must consume sklearn's W/H permutations per seed.""" + require_nmf_runtime() + X = np.random.default_rng(81).random((21, 12)) + 0.1 + seeds = [13, 47] + nmf_kwargs = _cd_nmf_kwargs( + 3, + seed=0, + max_iter=7, + shuffle=True, + alpha_W=0.0, + alpha_H=0.0, + l1_ratio=0.0, + ) + + actual = kernel.solver_cd._nmf_gpu_cd( + X, + seeds, + nmf_kwargs, + {"device": "cpu", "dtype": "fp64", "allow_tf32": False}, + ) + + for (actual_H, actual_W), seed in zip(actual, seeds): + expected_H, expected_W, _ = _sklearn_cd_reference( + X, dict(nmf_kwargs, random_state=seed) + ) + np.testing.assert_allclose( + actual_H, expected_H, rtol=5e-10, atol=5e-11 + ) + np.testing.assert_allclose( + actual_W, expected_W, rtol=5e-10, atol=5e-11 + ) + + +def test_nmf_gpu_batch_dispatches_explicit_cd(kernel, monkeypatch): + """An explicit CD solver must route through the registered CD kernel.""" + calls = [] + + def fake_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): + calls.append((X, seeds, nmf_kwargs, gpu_kwargs)) + return ["cd-result"] + + monkeypatch.setitem(kernel._GPU_SOLVERS, "cd", fake_cd) + X = np.ones((3, 2)) + seeds = [11] + nmf_kwargs = {"n_components": 1, "solver": "cd"} + gpu_kwargs = {"device": "cpu"} + + result = kernel._nmf_gpu_batch( + X, seeds, nmf_kwargs, gpu_kwargs + ) + + assert result == ["cd-result"] + assert len(calls) == 1 + actual_X, actual_seeds, actual_kwargs, actual_gpu_kwargs = calls[0] + assert actual_X is X + assert actual_seeds is seeds + assert actual_kwargs is nmf_kwargs + assert actual_gpu_kwargs is gpu_kwargs + + +@pytest.mark.parametrize( + "overrides,message", + [ + ({"beta_loss": "kullback-leibler"}, "beta_loss"), + ({"tol": -1.0}, "tol"), + ({"max_iter": 0}, "max_iter"), + ({"alpha": 0.1}, "alpha_W"), + ({"alpha_W": -0.1}, "alpha_W"), + ({"alpha_H": -0.1}, "alpha_H"), + ({"l1_ratio": 1.1}, "l1_ratio"), + ({"shuffle": "true"}, "shuffle"), + ], +) +def test_cd_rejects_options_outside_sklearn_contract( + kernel, overrides, message +): + """Invalid CD semantics should fail instead of silently changing solver behavior.""" + require_nmf_runtime() + nmf_kwargs = _cd_nmf_kwargs(2, seed=0, max_iter=1) + nmf_kwargs.update(overrides) + + with pytest.raises((TypeError, ValueError), match=message): + run_nmf_gpu(kernel, + small_nonnegative_matrix(cells=5, genes=4), + nmf_kwargs, + {"device": "cpu", "dtype": "fp64"}, + ) + + +@pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", CD_PARITY_CASES) +def test_sklearn_cd_matches_batched_cuda_when_available( + kernel, dtype_name, np_dtype, rtol, atol +): + """Match sklearn batches through the fused CUDA sweep when CUDA exists.""" + torch = require_nmf_runtime() + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + X = np.random.default_rng(101).random((23, 14), dtype=np_dtype) + X += np_dtype(0.1) + seeds = [3, 37] + nmf_kwargs = _cd_nmf_kwargs(4, seed=0, max_iter=5) + actual = kernel.solver_cd._nmf_gpu_cd( + X, + seeds, + nmf_kwargs, + { + "device": "cuda", + "dtype": dtype_name, + "allow_tf32": False, + "compile": False, + }, + ) + + cuda_rtol = max(rtol, 2e-4 if np_dtype is np.float32 else 2e-9) + cuda_atol = max(atol, 2e-5 if np_dtype is np.float32 else 2e-10) + for (actual_H, actual_W), seed in zip(actual, seeds): + expected_H, expected_W, _ = _sklearn_cd_reference( + X, dict(nmf_kwargs, random_state=seed) + ) + np.testing.assert_allclose( + actual_H, expected_H, rtol=cuda_rtol, atol=cuda_atol + ) + np.testing.assert_allclose( + actual_W, expected_W, rtol=cuda_rtol, atol=cuda_atol + ) + + +@pytest.mark.parametrize( + "dtype_name,np_dtype", + [("fp32", np.float32), ("fp64", np.float64)], +) +def test_cd_cuda_results_are_invariant_to_batch_width( + kernel, dtype_name, np_dtype +): + """Changing the replicate batch width must not change a CD trajectory.""" + torch = require_nmf_runtime() + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + X = np.random.default_rng(117).random((41, 23), dtype=np_dtype) + X += np_dtype(0.1) + seeds = [3, 37, 83] + nmf_kwargs = _cd_nmf_kwargs( + 5, + seed=0, + max_iter=12, + tol=0.0, + alpha_W=0.0, + alpha_H=0.0, + l1_ratio=0.0, + ) + gpu_kwargs = { + "device": "cuda", + "dtype": dtype_name, + "allow_tf32": False, + "compile": False, + } + + batched = kernel.solver_cd._nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs) + for (batched_H, batched_W), seed in zip(batched, seeds): + (single_H, single_W), = kernel.solver_cd._nmf_gpu_cd( + X, [seed], nmf_kwargs, gpu_kwargs + ) + np.testing.assert_array_equal(batched_H, single_H) + np.testing.assert_array_equal(batched_W, single_W) + + # --------------------------------------------------------------------- # Input validation and degenerate shapes # --------------------------------------------------------------------- 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]])) + kernel.utils._to_checked_array(np.array([[1.0, -0.1]])) 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]])) + kernel.utils._to_checked_array(np.array([[1.0, np.nan]])) 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]])) + kernel.utils._to_checked_array(np.array([[1.0, np.inf]])) def test_nmf_gpu_zero_matrix_does_not_crash_or_divide_by_zero(kernel): @@ -588,7 +986,7 @@ def test_nmf_gpu_zero_matrix_does_not_crash_or_divide_by_zero(kernel): require_nmf_runtime() X = np.zeros((5, 4)) - H, W = kernel._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 2, "max_iter": 5, "random_state": 0}, {"device": "cpu"}, @@ -601,12 +999,12 @@ 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._nmf_gpu( + H_row, W_row = run_nmf_gpu(kernel, np.array([[1.0, 2.0, 3.0]]), {"n_components": 1, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, ) - H_col, W_col = kernel._nmf_gpu( + H_col, W_col = run_nmf_gpu(kernel, np.array([[1.0], [2.0], [3.0]]), {"n_components": 1, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, @@ -622,17 +1020,17 @@ 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"): - kernel._to_checked_array(X) + kernel.utils._to_checked_array(X) with pytest.raises(ValueError, match="2D matrix"): - kernel._to_checked_array(np.array([1.0, 2.0])) + kernel.utils._to_checked_array(np.array([1.0, 2.0])) 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._nmf_gpu( + run_nmf_gpu(kernel, np.ones((3, 3)), {"n_components": 0, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, @@ -644,7 +1042,7 @@ def test_nmf_gpu_defines_behavior_when_k_exceeds_min_dimension(kernel): require_nmf_runtime() X = small_nonnegative_matrix(cells=3, genes=2) - H, W = kernel._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 4, "max_iter": 1, "random_state": 0}, {"device": "cpu"}, @@ -658,17 +1056,17 @@ def test_nmf_gpu_defines_behavior_when_k_exceeds_min_dimension(kernel): # --------------------------------------------------------------------- def test_resolve_gpu_opts_uses_defaults_when_gpu_kwargs_is_missing(kernel): """Missing gpu_kwargs should resolve exactly to the centralized DEFAULT_GPU values.""" - assert kernel._resolve_gpu_opts(None) == kernel.DEFAULT_GPU + assert kernel.utils._resolve_gpu_opts(None) == kernel.utils.DEFAULT_GPU def test_default_gpu_epsilon_matches_sklearn_exact_value(kernel): """Pin sklearn's float32 epsilon without importing its private EPSILON symbol.""" - assert kernel.DEFAULT_GPU["eps"] == float(np.finfo(np.float32).eps) + assert kernel.utils.DEFAULT_GPU["eps"] == float(np.finfo(np.float32).eps) def test_resolve_gpu_opts_dict_values_override_defaults(kernel): """Explicit gpu_kwargs values override defaults and are normalized to typed options.""" - opts = kernel._resolve_gpu_opts( + opts = kernel.utils._resolve_gpu_opts( { "device": "CUDA:1", "dtype": "FP32", @@ -697,16 +1095,16 @@ def test_resolve_gpu_opts_reads_only_gpu_kwargs_not_environment_variables(kernel monkeypatch.setenv("CNMF_GPU_DTYPE", "bf16") monkeypatch.setenv("CNMF_GPU_COMPILE", "true") - opts = kernel._resolve_gpu_opts({}) + opts = kernel.utils._resolve_gpu_opts({}) - assert opts["dtype"] == kernel.DEFAULT_GPU["dtype"] - assert opts["compile"] is kernel.DEFAULT_GPU["compile"] + assert opts["dtype"] == kernel.utils.DEFAULT_GPU["dtype"] + assert opts["compile"] is kernel.utils.DEFAULT_GPU["compile"] def test_resolve_gpu_opts_parses_truthy_boolean_strings(kernel): """Truthy strings accepted by Nextflow config should become real booleans.""" for value in ("1", "true", "TRUE", "yes", "on", True): - opts = kernel._resolve_gpu_opts({"allow_tf32": value, "compile": value}) + opts = kernel.utils._resolve_gpu_opts({"allow_tf32": value, "compile": value}) assert opts["allow_tf32"] is True assert opts["compile"] is True @@ -714,14 +1112,14 @@ def test_resolve_gpu_opts_parses_truthy_boolean_strings(kernel): def test_resolve_gpu_opts_parses_false_for_non_truthy_boolean_strings(kernel): """Non-truthy boolean strings should resolve to False.""" for value in ("0", "false", "off", "no", "", False): - opts = kernel._resolve_gpu_opts({"allow_tf32": value, "compile": value}) + opts = kernel.utils._resolve_gpu_opts({"allow_tf32": value, "compile": value}) assert opts["allow_tf32"] is False assert opts["compile"] is False def test_resolve_gpu_opts_coerces_numeric_strings_to_float_and_int(kernel): """Numeric config strings should be coerced to the expected float/int types.""" - opts = kernel._resolve_gpu_opts({"eps": "0.125", "check_every": "4", "compile_block": "5"}) + opts = kernel.utils._resolve_gpu_opts({"eps": "0.125", "check_every": "4", "compile_block": "5"}) assert opts["eps"] == 0.125 assert opts["check_every"] == 4 @@ -730,7 +1128,7 @@ def test_resolve_gpu_opts_coerces_numeric_strings_to_float_and_int(kernel): def test_resolve_gpu_opts_floors_check_every_and_compile_block_to_at_least_one(kernel): """Iteration cadence options should never resolve below one.""" - opts = kernel._resolve_gpu_opts({"check_every": 0, "compile_block": -3}) + opts = kernel.utils._resolve_gpu_opts({"check_every": 0, "compile_block": -3}) assert opts["check_every"] == 1 assert opts["compile_block"] == 1 @@ -741,15 +1139,15 @@ def test_resolve_gpu_opts_floors_check_every_and_compile_block_to_at_least_one(k # --------------------------------------------------------------------- def test_select_device_auto_prefers_cuda_then_mps_then_cpu(kernel): """Auto device selection should prefer CUDA, then MPS, then CPU.""" - assert kernel._select_device(fake_torch_backend(cuda_available=True, mps_available=True), "auto") == "cuda" - assert kernel._select_device(fake_torch_backend(cuda_available=False, mps_available=True), "auto") == "mps" - assert kernel._select_device(fake_torch_backend(cuda_available=False, mps_available=False), "auto") == "cpu" + assert kernel.utils._select_device(fake_torch_backend(cuda_available=True, mps_available=True), "auto") == "cuda" + assert kernel.utils._select_device(fake_torch_backend(cuda_available=False, mps_available=True), "auto") == "mps" + assert kernel.utils._select_device(fake_torch_backend(cuda_available=False, mps_available=False), "auto") == "cpu" def test_select_device_invalid_device_raises(kernel): """Unknown device names should fail loudly instead of falling back.""" with pytest.raises(ValueError, match="not recognized"): - kernel._select_device(fake_torch_backend(), "gpu") + kernel.utils._select_device(fake_torch_backend(), "gpu") def test_select_device_explicit_unavailable_cuda_or_mps_raises(kernel): @@ -757,9 +1155,9 @@ def test_select_device_explicit_unavailable_cuda_or_mps_raises(kernel): fake = fake_torch_backend(cuda_available=False, mps_available=False) with pytest.raises(RuntimeError, match="CUDA is unavailable"): - kernel._select_device(fake, "cuda") + kernel.utils._select_device(fake, "cuda") with pytest.raises(RuntimeError, match="MPS is unavailable"): - kernel._select_device(fake, "mps") + kernel.utils._select_device(fake, "mps") # --------------------------------------------------------------------- @@ -767,40 +1165,40 @@ def test_select_device_explicit_unavailable_cuda_or_mps_raises(kernel): # --------------------------------------------------------------------- def test_select_storage_auto_cpu_is_fp64(kernel): """Auto dtype on CPU should select fp64 for a stable reference path.""" - assert kernel._select_storage(fake_torch_backend(), "auto", "cpu") == "float64" + assert kernel.utils._select_storage(fake_torch_backend(), "auto", "cpu") == "float64" def test_select_storage_auto_gpu_is_fp32(kernel): """Auto dtype on GPU-class backends should select fp32.""" fake = fake_torch_backend() - assert kernel._select_storage(fake, "auto", "cuda:0") == "float32" - assert kernel._select_storage(fake, "auto", "mps") == "float32" + assert kernel.utils._select_storage(fake, "auto", "cuda:0") == "float32" + assert kernel.utils._select_storage(fake, "auto", "mps") == "float32" def test_select_storage_invalid_dtype_raises(kernel): """Unknown dtype names should fail with a clear configuration error.""" with pytest.raises(ValueError, match="not recognized"): - kernel._select_storage(fake_torch_backend(), "fp16", "cpu") + kernel.utils._select_storage(fake_torch_backend(), "fp16", "cpu") def test_select_storage_fp64_on_mps_raises(kernel): """MPS should reject explicit fp64 because this kernel treats MPS as fp32-only.""" with pytest.raises(RuntimeError, match="MPS has no fp64"): - kernel._select_storage(fake_torch_backend(), "fp64", "mps") + kernel.utils._select_storage(fake_torch_backend(), "fp64", "mps") def test_select_storage_bf16_is_cuda_only(kernel): """bf16 is accepted only for CUDA and means bf16 storage plus bf16 matmul operands.""" with pytest.raises(RuntimeError, match="only supported on CUDA"): - kernel._select_storage(fake_torch_backend(), "bf16", "cpu") + kernel.utils._select_storage(fake_torch_backend(), "bf16", "cpu") - assert kernel._select_storage(fake_torch_backend(bf16_supported=True), "bf16", "cuda") == "bfloat16" + assert kernel.utils._select_storage(fake_torch_backend(bf16_supported=True), "bf16", "cuda") == "bfloat16" def test_select_storage_bf16_checks_cuda_device_support(kernel): """CUDA bf16 requests should check the actual device capability.""" with pytest.raises(RuntimeError, match="does not support bf16"): - kernel._select_storage(fake_torch_backend(bf16_supported=False), "bf16", "cuda") + kernel.utils._select_storage(fake_torch_backend(bf16_supported=False), "bf16", "cuda") # --------------------------------------------------------------------- @@ -818,7 +1216,7 @@ def fake_import(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", fake_import) with pytest.raises(RuntimeError, match="PyTorch is required"): - kernel._loud_import_torch() + kernel.utils._loud_import_torch() def test_loud_import_sklearn_missing_has_actionable_error(kernel, monkeypatch): @@ -833,7 +1231,7 @@ def fake_import(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", fake_import) with pytest.raises(RuntimeError, match="scikit-learn is required"): - kernel._loud_import_initialize_nmf() + kernel.utils._loud_import_initialize_nmf() def test_loud_import_sklearn_incompatible_initializer_has_actionable_error(kernel, monkeypatch): @@ -848,7 +1246,7 @@ def fake_import(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", fake_import) with pytest.raises(RuntimeError, match="does not expose"): - kernel._loud_import_initialize_nmf() + kernel.utils._loud_import_initialize_nmf() # --------------------------------------------------------------------- @@ -860,7 +1258,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._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 2, "max_iter": 2, "random_state": 0}, {"device": "cpu"}, @@ -875,7 +1273,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._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 3, "max_iter": 1, "random_state": 0, "init": init}, {"device": "cpu"}, @@ -889,11 +1287,11 @@ def test_nndsvd_nndsvda_nndsvdar_initializers_return_valid_outputs(kernel): 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) + opt = dict(kernel.utils.DEFAULT_GPU, compile=True, check_every=4, compile_block=9) - step, block = kernel._execution_plan(fake_torch, opt, "mps") + step, block = kernel.utils._execution_plan(fake_torch, opt, "mps", kernel.solver_mu._mu_step) - assert step is kernel._mu_step + assert step is kernel.solver_mu._mu_step assert block == 4 @@ -906,12 +1304,17 @@ 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) + monkeypatch.setattr(kernel.solver_mu, "_fit_mu", fake_fit_mu) X = small_nonnegative_matrix(cells=4, genes=3) - kernel._nmf_gpu( + run_nmf_gpu(kernel, X, - {"n_components": 2, "max_iter": 1, "random_state": 0}, + { + "n_components": 2, + "max_iter": 1, + "random_state": 0, + "solver": "mu", + }, {"device": "cpu", "dtype": "fp32", "allow_tf32": True}, ) @@ -936,7 +1339,7 @@ def set_float32_matmul_precision(self, value): fake = FakeTorch() - with kernel._cuda_tf32(fake, True, "cpu"): + with kernel.utils._cuda_tf32(fake, True, "cpu"): assert fake.backends.cuda.matmul.allow_tf32 is False assert fake.precision == "highest" @@ -951,7 +1354,7 @@ def test_cuda_fp32_smoke_when_gpu_available(kernel): pytest.skip("CUDA is not available") X = small_nonnegative_matrix(cells=6, genes=5) - H, W = kernel._nmf_gpu( + H, W = run_nmf_gpu(kernel, X, {"n_components": 2, "max_iter": 2, "random_state": 0}, {"device": "cuda", "dtype": "fp32"}, @@ -973,12 +1376,17 @@ def fake_fit_mu(torch, Xg, W, H, eps, max_iter, tol, step, block, tf32, device): captured.append((Xg.dtype, W.dtype, H.dtype, eps.dtype, tf32, device)) return W, H - monkeypatch.setattr(kernel, "_fit_mu", fake_fit_mu) + monkeypatch.setattr(kernel.solver_mu, "_fit_mu", fake_fit_mu) X = small_nonnegative_matrix(cells=4, genes=3) - kernel._nmf_gpu( + run_nmf_gpu(kernel, X, - {"n_components": 2, "max_iter": 1, "random_state": 0}, + { + "n_components": 2, + "max_iter": 1, + "random_state": 0, + "solver": "mu", + }, {"device": "cuda", "dtype": "bf16"}, ) @@ -994,7 +1402,7 @@ def test_cuda_allow_tf32_scope_restores_previous_state_when_gpu_available(kernel prev_allow = torch.backends.cuda.matmul.allow_tf32 prev_precision = torch.get_float32_matmul_precision() - with kernel._cuda_tf32(torch, not prev_allow, "cuda"): + with kernel.utils._cuda_tf32(torch, not prev_allow, "cuda"): assert torch.backends.cuda.matmul.allow_tf32 is (not prev_allow) assert torch.backends.cuda.matmul.allow_tf32 is prev_allow @@ -1004,6 +1412,80 @@ def test_cuda_allow_tf32_scope_restores_previous_state_when_gpu_available(kernel # --------------------------------------------------------------------- # Batched-replicate factorize (--gpu-batch): batch-aware kernel parity # --------------------------------------------------------------------- +def test_nmf_gpu_batch_delegates_explicit_full_and_fixed_h_modes_to_mu(kernel, monkeypatch): + """The batch gate should pass explicit MU update_H modes to one solver.""" + calls = [] + + def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + mode = "fixed-h" if nmf_kwargs.get("update_H", True) is False else "full" + calls.append((mode, X, seeds, nmf_kwargs, gpu_kwargs)) + return [f"{mode}-result"] + + monkeypatch.setitem(kernel._GPU_SOLVERS, "mu", fake_mu) + + X = np.ones((3, 2)) + seeds = [7, 11] + gpu_kwargs = {"device": "cpu", "batch": 2} + + assert kernel._nmf_gpu_batch( + X, + seeds, + {"n_components": 1, "solver": "mu"}, + gpu_kwargs, + ) == ["full-result"] + assert kernel._nmf_gpu_batch( + X, + seeds, + {"n_components": 1, "solver": "mu", "update_H": False}, + gpu_kwargs, + ) == ["fixed-h-result"] + + assert [call[0] for call in calls] == ["full", "fixed-h"] + for _, actual_X, actual_seeds, _, actual_gpu_kwargs in calls: + assert actual_X is X + assert actual_seeds is seeds + assert actual_gpu_kwargs is gpu_kwargs + + +def test_nmf_gpu_batch_rejects_unknown_solver(kernel): + """The gateway should reject a solver that has not been registered.""" + with pytest.raises( + ValueError, match="solver 'als'.*available solvers: cd, mu" + ): + kernel._nmf_gpu_batch( + np.ones((3, 2)), + [7], + {"n_components": 1, "solver": "als"}, + {"device": "cpu"}, + ) + + +def test_nmf_gpu_batch_defaults_to_mu(kernel, monkeypatch): + """Missing solver configuration should route to MU.""" + calls = [] + + def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + calls.append((X, seeds, nmf_kwargs, gpu_kwargs)) + return ["mu-result"] + + monkeypatch.setitem(kernel._GPU_SOLVERS, "mu", fake_mu) + X = np.ones((3, 2)) + seeds = [7] + nmf_kwargs = {"n_components": 1} + gpu_kwargs = {"device": "cpu"} + + assert kernel.utils.DEFAULT_NMF["solver"] == "mu" + assert kernel._nmf_gpu_batch( + X, seeds, nmf_kwargs, gpu_kwargs + ) == ["mu-result"] + assert len(calls) == 1 + actual_X, actual_seeds, actual_kwargs, actual_gpu_kwargs = calls[0] + assert actual_X is X + assert actual_seeds is seeds + assert actual_kwargs is nmf_kwargs + assert actual_gpu_kwargs is gpu_kwargs + + 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) @@ -1017,14 +1499,19 @@ def test_nmf_gpu_mu_matches_single_kernel_for_each_seed(kernel): k = 3 X = low_rank_matrix(rank=k) seeds = [7, 3, 101] - nmf_kwargs = {"n_components": k, "max_iter": 300, "tol": 0} + nmf_kwargs = { + "n_components": k, + "max_iter": 300, + "tol": 0, + "solver": "mu", + } gpu_kwargs = {"device": "cpu"} - batched = kernel._nmf_gpu_mu(X, seeds, nmf_kwargs, gpu_kwargs) + batched = kernel.solver_mu._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) + Hs, Ws = run_nmf_gpu(kernel, 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) @@ -1036,10 +1523,15 @@ 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} + kw = { + "n_components": 2, + "max_iter": 80, + "tol": 0, + "solver": "mu", + } - (Hb, Wb), = kernel._nmf_gpu_mu(X, [5], kw, {"device": "cpu"}) - Hs, Ws = kernel._nmf_gpu(X, dict(kw, random_state=5), {"device": "cpu"}) + (Hb, Wb), = kernel.solver_mu._nmf_gpu_mu(X, [5], kw, {"device": "cpu"}) + Hs, Ws = run_nmf_gpu(kernel, 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 @@ -1050,7 +1542,7 @@ def test_nmf_gpu_mu_distinct_seeds_give_distinct_replicates(kernel): 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"}) + out = kernel.solver_mu._nmf_gpu_mu(X, [1, 2], {"n_components": 3, "max_iter": 50}, {"device": "cpu"}) assert not np.allclose(out[0][0], out[1][0]) @@ -1061,7 +1553,26 @@ def test_nmf_gpu_mu_rejects_empty_seeds(kernel): 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"}) + kernel.solver_mu._nmf_gpu_mu(X, [], {"n_components": 2, "max_iter": 1}, {"device": "cpu"}) + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"beta_loss": "kullback-leibler"}, "only beta_loss='frobenius'"), + ({"alpha_W": 0.1}, "does not yet support alpha_W/alpha_H"), + ({"alpha_H": 0.1}, "does not yet support alpha_W/alpha_H"), + ({"alpha": 0.1}, "does not accept deprecated alpha/regularization"), + ], +) +def test_nmf_gpu_mu_rejects_options_it_cannot_honor(kernel, override, message): + """GPU MU must fail loudly instead of silently solving a different objective.""" + require_nmf_runtime() + X = small_nonnegative_matrix(cells=6, genes=5) + kwargs = {"n_components": 2, "max_iter": 1, **override} + + with pytest.raises(ValueError, match=message): + kernel.solver_mu._nmf_gpu_mu(X, [1], kwargs, {"device": "cpu"}) def test_nmf_gpu_mu_results_are_invariant_to_batch_grouping(kernel): @@ -1072,11 +1583,11 @@ def test_nmf_gpu_mu_results_are_invariant_to_batch_grouping(kernel): 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 + batched = kernel.solver_mu._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 + (Hs, Ws), = kernel.solver_mu._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 @@ -1096,7 +1607,7 @@ 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") + Wout, Hout = kernel.solver_mu._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) @@ -1110,43 +1621,43 @@ def _fixed_spectra(k, genes, seed=0): 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): +def test_nmf_gpu_mu_fixed_h_mode_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"}) + (H, W), = kernel.solver_mu._nmf_gpu_mu(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): +def test_nmf_gpu_mu_fixed_h_mode_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"}) + batched = kernel.solver_mu._nmf_gpu_mu(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"}) + (Hs, Ws), = kernel.solver_mu._nmf_gpu_mu(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): +def test_nmf_gpu_mu_fixed_h_mode_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"}) + out = kernel.solver_mu._nmf_gpu_mu(X, [1, 2], kw, {"device": "cpu"}) assert not np.allclose(out[0][1], out[1][1]) @@ -1157,7 +1668,7 @@ def test_nmf_gpu_update_H_false_dispatches_to_fixed_h_refit(kernel): 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"}) + H, W = run_nmf_gpu(kernel, X, {"n_components": k, "max_iter": 20, "H": Hfix, "update_H": False}, {"device": "cpu"}) assert np.allclose(H, Hfix) and W.shape == (8, k) @@ -1165,86 +1676,114 @@ def test_nmf_gpu_update_H_false_dispatches_to_fixed_h_refit(kernel): # --------------------------------------------------------------------- # --gpu-batch config plumbing # --------------------------------------------------------------------- +def _engine_args(**overrides): + """Build the parsed CLI namespace consumed by the engine adapter.""" + values = { + "command": "factorize", + "name": "cNMF", + "output_dir": ".", + "engine": None, + "solver": "mu", + "beta_loss": "frobenius", + "gpu_device": None, + "gpu_dtype": None, + "gpu_allow_tf32": None, + "gpu_compile": None, + "gpu_eps": None, + "gpu_check_every": None, + "gpu_compile_block": None, + "gpu_batch": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def run_nmf_gpu(kernel, X, nmf_kwargs, gpu_kwargs=None): + """Exercise the single-replicate adapter through its parsed-argument API.""" + gpu_kwargs = dict(gpu_kwargs or {}) + args = _engine_args( + engine="gpu", + solver=nmf_kwargs.get("solver", "mu"), + beta_loss=nmf_kwargs.get("beta_loss", "frobenius"), + gpu_device=gpu_kwargs.get("device"), + gpu_dtype=gpu_kwargs.get("dtype"), + gpu_allow_tf32=gpu_kwargs.get("allow_tf32"), + gpu_compile=gpu_kwargs.get("compile"), + gpu_eps=gpu_kwargs.get("eps"), + gpu_check_every=gpu_kwargs.get("check_every"), + gpu_compile_block=gpu_kwargs.get("compile_block"), + gpu_batch=gpu_kwargs.get("batch"), + ) + return kernel._nmf_gpu(args, X, nmf_kwargs) + + 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 + assert kernel.utils.DEFAULT_GPU["batch"] == 1 + assert kernel.utils._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 + assert kernel.utils._resolve_gpu_opts({"batch": "4"})["batch"] == 4 + assert kernel.utils._resolve_gpu_opts({"batch": 0})["batch"] == 1 + assert kernel.utils._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()) +def test_gpu_kwargs_from_args_carries_batch_and_fills_its_default(kernel): + """A parsed --gpu-batch value flows through the shared option resolver.""" + args = _engine_args(engine="gpu", gpu_batch=8) + assert kernel.utils.gpu_kwargs_from_args(args)["batch"] == 8 - 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"]) + default_args = _engine_args(engine="gpu") assert default_args.gpu_batch is None - assert kernel.gpu_kwargs_from_args(default_args)["batch"] == 1 + assert kernel.utils.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([]) +def test_engine_args_default_to_none_until_user_selects_an_engine(kernel): + """Absent CLI options remain distinguishable from explicit values.""" + args = _engine_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"): + "gpu_eps", "gpu_check_every", "gpu_compile_block", "gpu_batch"): 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"]): + for args in ( + _engine_args(gpu_device="cuda"), + _engine_args(engine="cpu", gpu_device="cuda"), + ): with pytest.raises(ValueError, match="require --engine gpu"): - kernel.gpu_kwargs_from_args(parser.parse_args(argv)) + kernel.utils.gpu_kwargs_from_args(args) 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 = _engine_args(engine="gpu") - args = parser.parse_args(["--engine", "gpu"]) - - assert kernel.gpu_kwargs_from_args(args) == kernel.DEFAULT_GPU + assert kernel.utils.gpu_kwargs_from_args(args) == kernel.utils.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) == { + args = _engine_args( + engine="gpu", + gpu_device="CUDA:0", + gpu_dtype="FP32", + gpu_allow_tf32=True, + gpu_compile=True, + gpu_eps=1e-8, + gpu_check_every=5, + gpu_compile_block=100, + ) + + assert kernel.utils.gpu_kwargs_from_args(args) == { "device": "cuda:0", "dtype": "fp32", "allow_tf32": True, @@ -1258,111 +1797,179 @@ def test_gpu_kwargs_from_args_normalizes_cli_overrides(kernel): 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"], + for args in ( + _engine_args(command="factorize", engine="gpu"), + _engine_args(command="consensus", engine="gpu"), + _engine_args(command="consensus", gpu_device="cuda"), ): - kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) + kernel.utils._validate_engine_args_for_command(args, 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"], + for args in ( + _engine_args(command="prepare", engine="gpu"), + _engine_args(command="combine", gpu_device="cuda"), + _engine_args(command="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) + kernel.utils._validate_engine_args_for_command(args, supported) # non-engine commands without engine/GPU options are fine - kernel.validate_engine_args_for_command(parser.parse_args(["prepare"]), supported) + kernel.utils._validate_engine_args_for_command( + _engine_args(command="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"], + for args in ( + _engine_args(command="consensus", engine="gpu"), + _engine_args( + command="consensus", + engine="gpu", + gpu_device="CUDA:0", + gpu_dtype="FP32", + ), + _engine_args( + command="consensus", + gpu_allow_tf32=True, + gpu_compile=True, + ), ): - kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) + kernel.utils._validate_engine_args_for_command(args, 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"], + for args in ( + _engine_args(command="prepare", engine="gpu"), + _engine_args(command="combine", gpu_device="cuda"), + _engine_args(command="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) + kernel.utils._validate_engine_args_for_command(args, 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.""" +def test_configure_nmf_engine_cpu_constructs_unmodified_cnmf_instance(kernel): + """The CPU engine should construct cNMF without overriding its sklearn hook.""" class DummyCNMF: + def __init__(self, output_dir, name): + self.output_dir = output_dir + self.name = name + def _nmf(self, X, nmf_kwargs): return "sklearn-path" - obj = DummyCNMF() - result = kernel.configure_nmf_engine(obj, engine="cpu", gpu_kwargs={"device": "cuda"}) + result = kernel.configure_nmf_engine( + DummyCNMF, + _engine_args(engine="cpu", output_dir="runs", name="example"), + ) + + assert isinstance(result, DummyCNMF) + assert result.output_dir == "runs" + assert result.name == "example" + assert "_nmf" not in vars(result) # no instance override added + assert result._nmf("X", {}) == "sklearn-path" - 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_constructs_and_configures_once(kernel): + """The adapter should consume one namespace and construct one cNMF object.""" + class DummyCNMF: + def __init__(self, output_dir, name): + self.output_dir = output_dir + self.name = name + + args = _engine_args( + command="factorize", + engine="cpu", + output_dir="runs", + name="example", + ) + result = kernel.configure_nmf_engine(DummyCNMF, args) + + assert isinstance(result, DummyCNMF) + assert result.output_dir == "runs" + assert result.name == "example" + assert args.command == "factorize" + assert args.engine == "cpu" + + +def test_configure_nmf_engine_validates_before_construction(kernel): + """Invalid engine options must not create cNMF output directories.""" + constructed = [] + + def factory(**kwargs): + constructed.append(kwargs) + return object() + + args = _engine_args(engine="gpu", solver="cd", beta_loss="kullback-leibler") + with pytest.raises(ValueError, match="supports only beta_loss"): + kernel.configure_nmf_engine(factory, args) + + assert constructed == [] def test_configure_nmf_engine_rejects_unknown_engine(kernel): """Unknown engine names should fail loudly instead of silently using CPU.""" + constructed = [] + + def factory(**kwargs): + constructed.append(kwargs) + return object() + with pytest.raises(ValueError, match="engine must be 'cpu' or 'gpu'"): - kernel.configure_nmf_engine(object(), engine="tpu") + kernel.configure_nmf_engine(factory, _engine_args(engine="tpu")) + + assert constructed == [] -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.""" +def test_configure_nmf_engine_gpu_installs_instance_nmf_hook(kernel, monkeypatch): + """The GPU engine should bind its options to the instance `_nmf` hook.""" captured = {} - def fake_nmf_gpu(self, X, nmf_kwargs): - captured["self"], captured["X"], captured["nmf_kwargs"] = self, X, dict(nmf_kwargs) + def fake_nmf_gpu(args, X, nmf_kwargs, gpu_kwargs=None): + captured["args"] = args + captured["X"] = X + captured["nmf_kwargs"] = dict(nmf_kwargs) + captured["gpu_kwargs"] = gpu_kwargs return ("spectra", "usages") - monkeypatch.setattr(kernel, "nmf_gpu", fake_nmf_gpu) + monkeypatch.setattr(kernel, "_nmf_gpu", fake_nmf_gpu) class DummyCNMF: + def __init__(self, output_dir, name): + self.output_dir = output_dir + self.name = name + 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) + def prepare(self, *args, **kwargs): + return None + + args = _engine_args( + engine="gpu", + gpu_device="cuda", + gpu_dtype="fp32", + ) + result = kernel.configure_nmf_engine(DummyCNMF, args) - assert result is obj - assert "_nmf" in vars(obj) # instance _nmf is now overridden + assert isinstance(result, DummyCNMF) + assert "_nmf" in vars(result) # instance _nmf is now overridden - out = obj._nmf("Xdata", {"n_components": 5}) + out = result._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 + assert out == ("spectra", "usages") # dispatched through the GPU hook + assert captured["X"] == "Xdata" + assert captured["nmf_kwargs"] == {"n_components": 5} + assert captured["args"] is args + assert captured["gpu_kwargs"] is None def test_nmf_gpu_update_h_false_reconstructs_and_keeps_fixed_h(kernel): @@ -1372,7 +1979,7 @@ def test_nmf_gpu_update_h_false_reconstructs_and_keeps_fixed_h(kernel): 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( + H, W = run_nmf_gpu(kernel, X, {"n_components": 2, "max_iter": 5, "random_state": 0, "update_H": False, "H": fixed_h}, {"device": "cpu", "dtype": "fp64", "check_every": 5}, @@ -1385,7 +1992,7 @@ def test_nmf_gpu_update_h_false_reconstructs_and_keeps_fixed_h(kernel): 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) + kernel.utils._to_checked_fixed_h(None, 2, 3) invalid_cases = [ (np.array([1.0, 2.0, 3.0]), "2D"), @@ -1395,7 +2002,7 @@ def test_to_checked_fixed_h_rejects_missing_invalid_or_incompatible_h(kernel): ] for H, message in invalid_cases: with pytest.raises(ValueError, match=message): - kernel._to_checked_fixed_h(H, 2, 3) + kernel.utils._to_checked_fixed_h(H, 2, 3) def test_mu_step_fixed_h_matches_manual_w_only_update(kernel): @@ -1410,7 +2017,7 @@ def test_mu_step_fixed_h_matches_manual_w_only_update(kernel): denominator = W0 @ (H @ H.T) denominator = denominator.where(denominator != 0, eps) expected_W = W0 * ((Xg @ H.T) / denominator) - W = kernel._mu_step_fixed_h(W0, H, Xg, eps) + W = kernel.solver_mu._mu_step_fixed_h(W0, H, Xg, eps) assert torch.allclose(W, expected_W) assert torch.allclose(H, H_before) @@ -1430,7 +2037,7 @@ 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") + kernel.solver_mu._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} @@ -1439,7 +2046,7 @@ 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") + kernel.solver_mu._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 @@ -1447,10 +2054,10 @@ def test_execution_plan_for_fixed_h_compile_uses_fixed_h_step_and_compile_block( """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) + opt = dict(kernel.utils.DEFAULT_GPU, compile=True, check_every=1, compile_block=3) - step, block = kernel._execution_plan(fake_torch, opt, "cpu", kernel._mu_step_fixed_h) + step, block = kernel.utils._execution_plan(fake_torch, opt, "cpu", kernel.solver_mu._mu_step_fixed_h) - assert calls == [kernel._mu_step_fixed_h] - assert step is kernel._mu_step_fixed_h + assert calls == [kernel.solver_mu._mu_step_fixed_h] + assert step is kernel.solver_mu._mu_step_fixed_h assert block == 3 diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 83d1653..2d6eb32 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -3,9 +3,12 @@ import pandas as pd import scanpy as sc import os +import sys import scipy.sparse as sp +import yaml +from types import SimpleNamespace from cnmf import cNMF, save_df_to_npz, load_df_from_npz -from cnmf import nmf_gpu +import cnmf.gpunmf as gpunmf # Global parameters for data simulation NUM_CELLS = 100 @@ -18,6 +21,34 @@ def mock_cnmf(tmp_path): return cNMF(output_dir=str(tmp_path), name="test") + +def configure_gpu(cnmf_obj, *, command="factorize", solver="mu", + beta_loss="frobenius", gpu_kwargs=None): + """Configure an existing test instance through the CLI-shaped GPU adapter.""" + gpu_kwargs = dict(gpu_kwargs or {}) + args = SimpleNamespace( + command=command, + output_dir=cnmf_obj.output_dir, + name=cnmf_obj.name, + engine="gpu", + solver=solver, + beta_loss=beta_loss, + gpu_device=gpu_kwargs.get("device"), + gpu_dtype=gpu_kwargs.get("dtype"), + gpu_allow_tf32=gpu_kwargs.get("allow_tf32"), + gpu_compile=gpu_kwargs.get("compile"), + gpu_eps=gpu_kwargs.get("eps"), + gpu_check_every=gpu_kwargs.get("check_every"), + gpu_compile_block=gpu_kwargs.get("compile_block"), + gpu_batch=gpu_kwargs.get("batch"), + ) + configured = gpunmf.configure_nmf_engine( + lambda output_dir, name: cnmf_obj, + args, + ) + return configured, args + + def generate_counts_file(tmp_path, file_format, dtype=np.int64, zero_count=False): """ Generates a synthetic single-cell RNA-seq counts file in various formats. @@ -97,9 +128,120 @@ def test_prepare_raises_on_zero_count_cells(mock_cnmf, file_format, dtype, densi mock_cnmf.prepare(counts_fn, components=[5, 10], n_iter=10, densify=densify) +@pytest.mark.parametrize("solver", ["mu", "cd"]) +def test_main_prepare_persists_cli_solver(tmp_path, monkeypatch, solver): + """The CLI solver should flow through prepare into the saved run parameters.""" + from cnmf.cnmf import main + + counts_fn = generate_counts_file(tmp_path, "npz", np.float64) + output_dir = tmp_path / "cli-output" + run_name = f"solver-test-{solver}" + monkeypatch.setattr( + sys, + "argv", + [ + "cnmf", + "prepare", + "--output-dir", + str(output_dir), + "--name", + run_name, + "--counts", + counts_fn, + "--components", + "5", + "--n-iter", + "2", + "--numgenes", + "50", + "--densify", + "--engine", + "gpu", + "--solver", + solver, + ], + ) + + main() + + params_path = output_dir / run_name / "cnmf_tmp" / f"{run_name}.nmf_idvrun_params.yaml" + with open(params_path, encoding="utf-8") as stream: + assert yaml.safe_load(stream)["solver"] == solver + + +def test_gpunmf_prepare_solver_changes_only_saved_run_config(tmp_path): + """Selecting MU instead of CD must not change the matrix prepared by cNMF.""" + counts_fn = generate_counts_file(tmp_path, "npz", np.float64) + prepared = {} + + for solver in ("mu", "cd"): + cnmf_obj = cNMF(output_dir=str(tmp_path), name=f"prepared-{solver}") + cnmf_obj, _args = configure_gpu( + cnmf_obj, + command="prepare", + solver=solver, + ) + cnmf_obj.prepare( + counts_fn, + components=[5], + n_iter=2, + densify=True, + seed=14, + num_highvar_genes=50, + ) + normalized = sc.read(cnmf_obj.paths["normalized_counts"]) + with open(cnmf_obj.paths["nmf_run_parameters"], encoding="utf-8") as stream: + run_parameters = yaml.safe_load(stream) + prepared[solver] = (normalized, run_parameters) + + mu_counts, mu_parameters = prepared["mu"] + cd_counts, cd_parameters = prepared["cd"] + np.testing.assert_array_equal(mu_counts.X, cd_counts.X) + assert mu_counts.obs_names.equals(cd_counts.obs_names) + assert mu_counts.var_names.equals(cd_counts.var_names) + assert mu_parameters.pop("solver") == "mu" + assert cd_parameters.pop("solver") == "cd" + assert mu_parameters == cd_parameters + + +def test_main_rejects_cd_with_non_frobenius_before_creating_run(tmp_path, monkeypatch): + """Invalid CLI solver/loss combinations should have no filesystem side effects.""" + from cnmf.cnmf import main + + output_dir = tmp_path / "cli-output" + monkeypatch.setattr( + sys, + "argv", + [ + "cnmf", + "prepare", + "--output-dir", + str(output_dir), + "--name", + "must-not-exist", + "--engine", + "gpu", + "--solver", + "cd", + "--beta-loss", + "kullback-leibler", + ], + ) + + with pytest.raises(SystemExit): + main() + + assert not (output_dir / "must-not-exist").exists() + + 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, _args = configure_gpu( + mock_cnmf, + solver="mu", + gpu_kwargs={"device": "cpu", "batch": 2}, + ) mock_cnmf.prepare(counts_fn, components=[5, 7], n_iter=3, densify=True) calls = [] @@ -109,9 +251,8 @@ def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): 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) + monkeypatch.setitem(gpunmf._GPU_SOLVERS, "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. @@ -133,6 +274,11 @@ def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): 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, _args = configure_gpu( + mock_cnmf, + solver="mu", + gpu_kwargs={"device": "cpu"}, + ) mock_cnmf.prepare(counts_fn, components=[6], n_iter=2, densify=True) calls = [] @@ -142,9 +288,8 @@ def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): 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) + monkeypatch.setitem(gpunmf._GPU_SOLVERS, "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 @@ -209,13 +354,53 @@ def test_get_nmf_iter_params_default_cpu_engine_does_not_change_sklearn_kwargs(m assert set(run_params) == {"alpha_W", "alpha_H", "l1_ratio", "beta_loss", "solver", "tol", "max_iter", "init"} +def test_get_nmf_iter_params_retains_upstream_solver_inference(mock_cnmf): + """The unwrapped cNMF helper should retain its upstream loss-based behavior.""" + _replicate_params, frobenius_params = mock_cnmf.get_nmf_iter_params( + ks=[5], n_iter=2, beta_loss="frobenius" + ) + _replicate_params, kl_params = mock_cnmf.get_nmf_iter_params( + ks=[5], n_iter=2, beta_loss="kullback-leibler" + ) + + assert frobenius_params["solver"] == "cd" + assert kl_params["solver"] == "mu" + + +def test_gpunmf_prepare_validates_solver_before_reading_counts(mock_cnmf, tmp_path): + """The adapter should reject an invalid solver/loss before constructing cNMF.""" + + with pytest.raises( + ValueError, + match="solver='cd' supports only beta_loss='frobenius'", + ): + configure_gpu( + mock_cnmf, + command="prepare", + solver="cd", + beta_loss="kullback-leibler", + ) + + def test_factorize_gpu_engine_passes_seed_components_run_params_and_gpu_kwargs(mock_cnmf, monkeypatch, tmp_path): """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 + import cnmf.gpunmf as gpu_mod counts_fn = generate_counts_file(tmp_path, "txt", np.int64) - mock_cnmf.prepare(counts_fn, components=[5], n_iter=2, densify=True, seed=14) + gpu_kwargs = {"device": "cpu", "dtype": "fp64"} + mock_cnmf, _args = configure_gpu( + mock_cnmf, + solver="mu", + gpu_kwargs=gpu_kwargs, + ) + mock_cnmf.prepare( + counts_fn, + components=[5], + n_iter=2, + densify=True, + seed=14, + ) captured = [] @@ -224,9 +409,7 @@ def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): 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_mu", fake_mu) - gpu_kwargs = {"device": "cpu", "dtype": "fp64"} - gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs=gpu_kwargs) + monkeypatch.setitem(gpu_mod._GPU_SOLVERS, "mu", fake_mu) mock_cnmf.factorize(worker_i=0, total_workers=1) @@ -238,7 +421,7 @@ def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): 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 + assert gk == gpu_mod.utils._resolve_gpu_opts(gpu_kwargs) observed_seeds.update(seeds) assert observed_seeds == expected_seeds # exact seeds from prepared replicate params for iter_i in replicate_params["iter"]: @@ -247,7 +430,7 @@ def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): 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 + import cnmf.gpunmf as gpu_mod write_minimal_nmf_run_params(mock_cnmf) X = pd.DataFrame( @@ -262,24 +445,24 @@ def test_refit_usage_gpu_engine_passes_fixed_h_update_h_false_and_gpu_kwargs(moc ) captured = [] - def fake_nmf_gpu(self, X_arg, nmf_kwargs): - captured.append((X_arg, dict(nmf_kwargs))) + def fake_nmf_gpu(args, X_arg, nmf_kwargs, gpu_kwargs=None): + captured.append((X_arg, dict(nmf_kwargs), gpu_mod.utils.gpu_kwargs_from_args(args))) 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) + mock_cnmf, _args = configure_gpu(mock_cnmf, gpu_kwargs=gpu_kwargs) usages = mock_cnmf.refit_usage(X, spectra) assert len(captured) == 1 - X_arg, kw = captured[0] + X_arg, kw, actual_gpu_kwargs = captured[0] assert X_arg is X assert kw["n_components"] == 2 assert np.allclose(kw["H"], spectra.values) assert kw["update_H"] is False - assert kw["engine"] == "gpu" - assert kw["gpu"] == gpu_kwargs + assert "engine" not in kw and "gpu" not in kw + assert actual_gpu_kwargs == gpu_mod.utils._resolve_gpu_opts(gpu_kwargs) assert "beta_loss" in kw and "init" in kw assert list(usages.index) == list(X.index) assert list(usages.columns) == list(spectra.index) @@ -288,7 +471,7 @@ def fake_nmf_gpu(self, X_arg, nmf_kwargs): def test_refit_spectra_gpu_engine_routes_through_transposed_refit_usage(mock_cnmf, monkeypatch, tmp_path): """cNMF refit_spectra should use the same GPU fixed-H path through transposed refit_usage.""" - import cnmf.nmf_gpu as gpu_mod + import cnmf.gpunmf as gpu_mod write_minimal_nmf_run_params(mock_cnmf) X = pd.DataFrame( @@ -303,17 +486,20 @@ def test_refit_spectra_gpu_engine_routes_through_transposed_refit_usage(mock_cnm ) captured = [] - def fake_nmf_gpu(self, X_arg, nmf_kwargs): - captured.append((X_arg, dict(nmf_kwargs))) + def fake_nmf_gpu(args, X_arg, nmf_kwargs, gpu_kwargs=None): + captured.append((X_arg, dict(nmf_kwargs), gpu_mod.utils.gpu_kwargs_from_args(args))) return fake_gpu_nmf_output(X_arg, nmf_kwargs) - monkeypatch.setattr(gpu_mod, "nmf_gpu", fake_nmf_gpu) - gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu", "dtype": "fp64"}) + monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) + mock_cnmf, _args = configure_gpu( + mock_cnmf, + gpu_kwargs={"device": "cpu", "dtype": "fp64"}, + ) spectra = mock_cnmf.refit_spectra(X, usage) assert len(captured) == 1 - X_arg, kw = captured[0] + X_arg, kw, _ = captured[0] assert X_arg.shape == (3, 4) # genes x cells after transpose assert np.allclose(kw["H"], usage.T.values) # programs x cells fixed H assert kw["update_H"] is False @@ -325,7 +511,7 @@ def fake_nmf_gpu(self, X_arg, nmf_kwargs): def test_consensus_gpu_engine_smoke_writes_expected_outputs(mock_cnmf, monkeypatch, tmp_path): """A tiny CPU-backed GPU-engine consensus run should write the expected consensus outputs.""" - import cnmf.nmf_gpu as gpu_mod + import cnmf.gpunmf as gpu_mod counts_fn = generate_positive_counts_file(tmp_path) mock_cnmf.prepare(counts_fn, components=[2], n_iter=3, densify=True, @@ -349,11 +535,15 @@ def test_consensus_gpu_engine_smoke_writes_expected_outputs(mock_cnmf, monkeypat ) save_df_to_npz(merged, mock_cnmf.paths["merged_spectra"] % 2) - def fake_nmf_gpu(self, X_arg, nmf_kwargs): + def fake_nmf_gpu(args, X_arg, nmf_kwargs, gpu_kwargs=None): return fake_gpu_nmf_output(X_arg, nmf_kwargs) - monkeypatch.setattr(gpu_mod, "nmf_gpu", fake_nmf_gpu) - gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu", "dtype": "fp64"}) + monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) + mock_cnmf, _args = configure_gpu( + mock_cnmf, + command="consensus", + gpu_kwargs={"device": "cpu", "dtype": "fp64"}, + ) mock_cnmf.consensus(k=2, density_threshold=2.0, local_neighborhood_size=0.5, show_clustering=False, refit_usage=False) diff --git a/tests/utils.py b/tests/utils.py index cb7b2fa..8db84d4 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2,6 +2,7 @@ from pathlib import Path from types import SimpleNamespace +import importlib import importlib.util import json import os @@ -18,19 +19,22 @@ os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") REPO_ROOT = Path(__file__).resolve().parents[1] -KERNEL_PATH = REPO_ROOT / "src" / "cnmf" / "nmf_gpu.py" +KERNEL_PATH = REPO_ROOT / "src" / "cnmf" / "gpunmf" / "__init__.py" DOWNLOAD_PYTEST_DATA_PATH = REPO_ROOT / "download_pytest_data.py" # --------------------------------------------------------------------- # Standalone NMF GPU kernel helpers # --------------------------------------------------------------------- -def load_kernel_module(module_name="nmf_gpu", kernel_path=KERNEL_PATH): - """Load the standalone kernel script as an importable module for tests.""" +def load_kernel_module(module_name="cnmf.gpunmf", kernel_path=KERNEL_PATH): + """Load the GPU engine through its real package import path.""" kernel_path = Path(kernel_path) if not kernel_path.exists(): pytest.fail(f"Required NMF GPU kernel file is missing: {kernel_path}", pytrace=False) + if kernel_path.resolve() == KERNEL_PATH.resolve(): + return importlib.import_module("cnmf.gpunmf") + if module_name in sys.modules: return sys.modules[module_name] @@ -46,7 +50,7 @@ def load_kernel_module(module_name="nmf_gpu", kernel_path=KERNEL_PATH): @pytest.fixture(scope="session") def kernel(): - """Loaded `src/cnmf/nmf_gpu.py` module under test.""" + """Loaded `cnmf.gpunmf` package module under test.""" return load_kernel_module()