Add batched GPU NMF support with sklearn-compatible multiplicative-update (MU) and Fast-HALS coordinate-descent (CD) solvers

.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:

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'],

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,

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),
+                )