Adds test coverage for the GPU NMF CLI wiring

tests/test_nmf_gpu.py

@@ -42,6 +42,7 @@
 """
 
 from types import SimpleNamespace
+import argparse
 import builtins
 
 import numpy as np
@@ -57,7 +58,6 @@
     small_nonnegative_matrix,
 )
 
-
 # ---------------------------------------------------------------------
 # Test harness
 # ---------------------------------------------------------------------
@@ -507,6 +507,138 @@ def test_resolve_gpu_opts_floors_check_every_and_compile_block_to_at_least_one(k
     assert opts["compile_block"] == 1
 
 
+# ---------------------------------------------------------------------
+# GPU CLI and cNMF engine wiring
+# ---------------------------------------------------------------------
+def test_parse_gpu_args_defaults_to_none_until_user_sets_engine(kernel):
+    """CLI flags should default to None so absent options are distinguishable from explicit values."""
+    parser = argparse.ArgumentParser()
+    kernel.parse_gpu_args(parser)
+
+    args = parser.parse_args([])
+
+    assert args.engine is None
+    for name in ("gpu_device", "gpu_dtype", "gpu_allow_tf32", "gpu_compile",
+                 "gpu_eps", "gpu_check_every", "gpu_compile_block"):
+        assert getattr(args, name) is None, f"{name} should default to None"
+
+
+def test_gpu_kwargs_from_args_rejects_gpu_options_without_gpu_engine(kernel):
+    """GPU-specific CLI options should raise unless `--engine gpu` was explicitly selected."""
+    parser = argparse.ArgumentParser()
+    kernel.parse_gpu_args(parser)
+
+    for argv in (["--gpu-device", "cuda"], ["--engine", "cpu", "--gpu-device", "cuda"]):
+        with pytest.raises(ValueError, match="require --engine gpu"):
+            kernel.gpu_kwargs_from_args(parser.parse_args(argv))
+
+
+def test_gpu_kwargs_from_args_fills_defaults_when_gpu_engine_selected(kernel):
+    """`--engine gpu` alone should resolve missing GPU options from DEFAULT_GPU."""
+    parser = argparse.ArgumentParser()
+    kernel.parse_gpu_args(parser)
+
+    args = parser.parse_args(["--engine", "gpu"])
+
+    assert kernel.gpu_kwargs_from_args(args) == kernel.DEFAULT_GPU
+
+
+def test_gpu_kwargs_from_args_normalizes_cli_overrides(kernel):
+    """GPU CLI override values should normalize through the same resolver as config dict values."""
+    parser = argparse.ArgumentParser()
+    kernel.parse_gpu_args(parser)
+
+    args = parser.parse_args([
+        "--engine", "gpu",
+        "--gpu-device", "CUDA:0",        # device is lower-cased by _resolve_gpu_opts
+        "--gpu-dtype", "FP32",
+        "--gpu-allow-tf32",              # store_const flags -> True
+        "--gpu-compile",
+        "--gpu-eps", "1e-8",
+        "--gpu-check-every", "5",
+        "--gpu-compile-block", "100",
+    ])
+
+    assert kernel.gpu_kwargs_from_args(args) == {
+        "device": "cuda:0",
+        "dtype": "fp32",
+        "allow_tf32": True,
+        "compile": True,
+        "eps": 1e-8,
+        "check_every": 5,
+        "compile_block": 100,
+    }
+
+
+def test_validate_engine_args_for_command_rejects_non_factorize_gpu_options(kernel):
+    """Engine/GPU options should be accepted only for commands that support GPU factorization."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("command")
+    kernel.parse_gpu_args(parser)
+    supported = ("factorize",)
+
+    # factorize accepts engine/GPU options (no raise)
+    kernel.validate_engine_args_for_command(parser.parse_args(["factorize", "--engine", "gpu"]), supported)
+
+    # a non-factorize command carrying engine/GPU options is rejected
+    for argv in (["prepare", "--engine", "gpu"], ["consensus", "--gpu-device", "cuda"]):
+        with pytest.raises(ValueError, match="only valid with"):
+            kernel.validate_engine_args_for_command(parser.parse_args(argv), supported)
+
+    # a non-factorize command without engine/GPU options is fine
+    kernel.validate_engine_args_for_command(parser.parse_args(["prepare"]), supported)
+
+
+def test_configure_nmf_engine_cpu_leaves_cnmf_instance_unchanged(kernel):
+    """The default CPU engine should be a no-op so existing sklearn behavior is preserved."""
+    class DummyCNMF:
+        def _nmf(self, X, nmf_kwargs):
+            return "sklearn-path"
+
+    obj = DummyCNMF()
+    result = kernel.configure_nmf_engine(obj, engine="cpu", gpu_kwargs={"device": "cuda"})
+
+    assert result is obj
+    assert "_nmf" not in vars(obj)                  # no instance override added
+    assert obj._nmf("X", {}) == "sklearn-path"      # original class method intact
+
+
+def test_configure_nmf_engine_rejects_unknown_engine(kernel):
+    """Unknown engine names should fail loudly instead of silently using CPU."""
+    with pytest.raises(ValueError, match="engine must be 'cpu' or 'gpu'"):
+        kernel.configure_nmf_engine(object(), engine="tpu")
+
+
+def test_configure_nmf_engine_gpu_patches_instance_nmf_with_adapter(kernel, monkeypatch):
+    """The GPU engine should replace the instance `_nmf` callable with the GPU adapter path."""
+    captured = {}
+
+    def fake_nmf_gpu(self, X, nmf_kwargs):
+        captured["self"], captured["X"], captured["nmf_kwargs"] = self, X, dict(nmf_kwargs)
+        return ("spectra", "usages")
+
+    monkeypatch.setattr(kernel, "_nmf_gpu", fake_nmf_gpu)
+
+    class DummyCNMF:
+        def _nmf(self, X, nmf_kwargs):
+            return "sklearn-path"
+
+    obj = DummyCNMF()
+    gpu_kwargs = {"device": "cuda", "dtype": "fp32"}
+    result = kernel.configure_nmf_engine(obj, engine="gpu", gpu_kwargs=gpu_kwargs)
+
+    assert result is obj
+    assert "_nmf" in vars(obj)                                  # instance _nmf is now overridden
+
+    out = obj._nmf("Xdata", {"n_components": 5})
+
+    assert out == ("spectra", "usages")                        # dispatched through the GPU adapter
+    assert captured["self"] is obj and captured["X"] == "Xdata"
+    assert captured["nmf_kwargs"]["engine"] == "gpu"           # engine + gpu kwargs embedded first
+    assert captured["nmf_kwargs"]["gpu"] == gpu_kwargs
+    assert captured["nmf_kwargs"]["n_components"] == 5         # original kwargs preserved
+
+
 # ---------------------------------------------------------------------
 # Device selection
 # ---------------------------------------------------------------------

tests/test_prepare.py

@@ -58,6 +58,56 @@ def generate_counts_file(tmp_path, file_format, dtype=np.int64, zero_count=False
 
     return str(counts_fn)
 
+
+# ---------------------------------------------------------------------
+# GPU factorize wiring integration
+# ---------------------------------------------------------------------
+def test_get_nmf_iter_params_default_cpu_engine_does_not_change_sklearn_kwargs(mock_cnmf):
+    """Default CPU runs should not add engine/gpu keys to sklearn NMF kwargs."""
+    _replicate_params, run_params = mock_cnmf.get_nmf_iter_params(ks=[5, 7], n_iter=3, random_state_seed=14)
+
+    # GPU wiring must not leak into the sklearn factorization kwargs.
+    assert "engine" not in run_params
+    assert "gpu" not in run_params
+    # Only the existing cNMF/sklearn factorization keys are present.
+    assert set(run_params) == {"alpha_W", "alpha_H", "l1_ratio", "beta_loss", "solver", "tol", "max_iter", "init"}
+
+
+def test_factorize_gpu_engine_passes_seed_components_run_params_and_gpu_kwargs(mock_cnmf, monkeypatch, tmp_path):
+    """cNMF factorize should pass n_components, random_state, run params, and GPU kwargs to `_nmf_gpu`."""
+    import cnmf.nmf_gpu 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)
+
+    captured = []
+
+    def fake_nmf_gpu(self, X, nmf_kwargs):
+        captured.append(dict(nmf_kwargs))
+        k = nmf_kwargs["n_components"]
+        return np.zeros((k, X.shape[1])), np.zeros((X.shape[0], k))    # (spectra, usages)
+
+    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.factorize(worker_i=0, total_workers=1)
+
+    assert len(captured) == 2                       # n_iter=2 replicates for k=5
+    replicate_params = load_df_from_npz(mock_cnmf.paths["nmf_replicate_parameters"])
+    expected_seeds = set(replicate_params["nmf_seed"])
+    observed_seeds = set()
+    for kw in captured:
+        assert kw["n_components"] == 5              # set per replicate by factorize
+        assert kw["engine"] == "gpu"               # adapter embedded the engine
+        assert kw["gpu"] == gpu_kwargs             # ...and the resolved GPU kwargs
+        assert "beta_loss" in kw and "init" in kw  # original run params forwarded
+        observed_seeds.add(kw["random_state"])
+    assert observed_seeds == expected_seeds          # exact seeds from prepared replicate params
+    for iter_i in replicate_params["iter"]:
+        assert os.path.exists(mock_cnmf.paths["iter_spectra"] % (5, iter_i))
+
+
 @pytest.mark.parametrize("file_format", ["txt", "npz", "h5ad"])
 @pytest.mark.parametrize("dtype", [np.int64, np.float32, np.float64])
 @pytest.mark.parametrize("densify", [True, False])