Wire GPU NMF engine into CLI

src/cnmf/cnmf.py

@@ -1234,6 +1234,11 @@ 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
+
     parser = argparse.ArgumentParser()
 
     parser.add_argument('command', type=str, choices=['prepare', 'factorize', 'combine', 'consensus', 'k_selection_plot'])
@@ -1257,11 +1262,18 @@ 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)
 
 
     args = parser.parse_args()
+    try:
+        engine_commands = ('factorize',)
+        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))
 
     if args.command == 'prepare':
         cnmf_obj.prepare(args.counts, components=args.components, n_iter=args.n_iter, densify=args.densify,

src/cnmf/nmf_gpu.py

@@ -69,6 +69,76 @@
 }
 
 
+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] NMF engine to use (default cpu)")
+    group.add_argument("--gpu-device", type=str, help="[factorize,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,gpu] Storage and matmul dtype for GPU NMF (default auto)")
+    group.add_argument("--gpu-allow-tf32", action="store_const", const=True, help="[factorize,gpu] Allow TF32 for CUDA fp32 matrix multiplication")
+    group.add_argument("--gpu-compile", action="store_const", const=True, help="[factorize,gpu] Enable torch.compile for the GPU NMF update step")
+    group.add_argument("--gpu-eps", type=float, help="[factorize,gpu] Multiplicative-update denominator guard")
+    group.add_argument("--gpu-check-every", type=int, help="[factorize,gpu] Eager-mode convergence check interval")
+    group.add_argument("--gpu-compile-block", type=int, help="[factorize,gpu] Number of MU iterations per compiled block")
+    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,
+    }
+    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",
+    ]
+    if any(getattr(args, name) is not None for name in gpu_arg_names):
+        commands = ", ".join(available_commands)
+        raise ValueError(f"NMF engine/GPU options are only valid with: {commands}")
+
+
+def configure_nmf_engine(cnmf_obj, engine="cpu", gpu_kwargs=None):
+    """Configure a cNMF instance with an optional GPU NMF adapter."""
+    if engine not in ("cpu", "gpu"):
+        raise ValueError("engine must be 'cpu' or 'gpu'")
+    if engine == "cpu":
+        return cnmf_obj
+
+    def _gpu_nmf(X, nmf_kwargs):
+        nmf_kwargs = dict(nmf_kwargs)
+        nmf_kwargs["engine"] = "gpu"
+        nmf_kwargs["gpu"] = gpu_kwargs or {}
+        return _nmf_gpu(cnmf_obj, X, nmf_kwargs)
+
+    cnmf_obj._nmf = _gpu_nmf
+    return cnmf_obj
+
+
 def _resolve_gpu_opts(gpu_kwargs):
     """Merge Nextflow-provided gpu_kwargs over defaults into a typed opts dict."""
     raw = dict(gpu_kwargs or {})
@@ -359,6 +429,9 @@ def factorize_nmf_gpu(X, nmf_kwargs, gpu_kwargs=None):
     return _to_nmf_output(H, W)
 
 
-def _nmf_gpu(self, X, nmf_kwargs, gpu_kwargs=None):
+def _nmf_gpu(self, X, nmf_kwargs):
     """cNMF adapter: same core NMF kernel, with `self` ignored for monkeypatch compatibility."""
+    nmf_kwargs = dict(nmf_kwargs)
+    gpu_kwargs = nmf_kwargs.pop("gpu", None)
+    nmf_kwargs.pop("engine", None)
     return factorize_nmf_gpu(X, nmf_kwargs, gpu_kwargs)

tests/test_nmf_gpu.py

@@ -133,7 +133,7 @@ def test_factorize_nmf_gpu_fp32_compute_still_returns_float64_numpy_outputs(kern
 
 
 def test_nmf_gpu_adapter_ignores_self_and_delegates_to_factorize_nmf_gpu(kernel, monkeypatch):
-    """Verify the cNMF monkeypatch adapter forwards args and ignores its bound `self`."""
+    """Verify the cNMF adapter extracts embedded GPU args and ignores its bound `self`."""
     calls = []
     sentinel = (object(), object())
 
@@ -143,13 +143,14 @@ def fake_factorize(X, nmf_kwargs, gpu_kwargs=None):
 
     monkeypatch.setattr(kernel, "factorize_nmf_gpu", fake_factorize)
     X = np.ones((3, 2))
-    nmf_kwargs = {"n_components": 1}
     gpu_kwargs = {"device": "cpu"}
+    nmf_kwargs = {"engine": "gpu", "gpu": gpu_kwargs, "n_components": 1}
 
-    result = kernel._nmf_gpu(object(), X, nmf_kwargs, gpu_kwargs)
+    result = kernel._nmf_gpu(object(), X, nmf_kwargs)
 
     assert result is sentinel
-    assert calls == [(X, nmf_kwargs, gpu_kwargs)]
+    assert calls == [(X, {"n_components": 1}, gpu_kwargs)]
+    assert nmf_kwargs == {"engine": "gpu", "gpu": gpu_kwargs, "n_components": 1}
 
 
 # ---------------------------------------------------------------------