Add cNMF GPU consensus test coverage

tests/test_nmf_gpu.py

@@ -570,25 +570,68 @@ def test_gpu_kwargs_from_args_normalizes_cli_overrides(kernel):
     }
 
 
-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."""
+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",)
-
-    # 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"]):
+    supported = ("factorize", "consensus")
+
+    # factorize and consensus accept engine/GPU options (no raise)
+    for argv in (
+        ["factorize", "--engine", "gpu"],
+        ["consensus", "--engine", "gpu"],
+        ["consensus", "--gpu-device", "cuda"],
+    ):
+        kernel.validate_engine_args_for_command(parser.parse_args(argv), supported)
+
+    # non-engine commands carrying engine/GPU options are rejected
+    for argv in (
+        ["prepare", "--engine", "gpu"],
+        ["combine", "--gpu-device", "cuda"],
+        ["k_selection_plot", "--gpu-dtype", "fp32"],
+    ):
         with pytest.raises(ValueError, match="only valid with"):
             kernel.validate_engine_args_for_command(parser.parse_args(argv), supported)
 
-    # a non-factorize command without engine/GPU options is fine
+    # non-engine commands without engine/GPU options are fine
     kernel.validate_engine_args_for_command(parser.parse_args(["prepare"]), supported)
 
 
+# ---------------------------------------------------------------------
+# GPU CLI and cNMF consensus wiring
+# ---------------------------------------------------------------------
+def test_validate_engine_args_accepts_consensus_gpu_options(kernel):
+    """`consensus --engine gpu` and consensus GPU flags should be valid CLI input."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("command")
+    kernel.parse_gpu_args(parser)
+    supported = ("factorize", "consensus")
+
+    for argv in (
+        ["consensus", "--engine", "gpu"],
+        ["consensus", "--engine", "gpu", "--gpu-device", "CUDA:0", "--gpu-dtype", "FP32"],
+        ["consensus", "--gpu-allow-tf32", "--gpu-compile"],
+    ):
+        kernel.validate_engine_args_for_command(parser.parse_args(argv), supported)
+
+
+def test_validate_engine_args_rejects_gpu_options_for_non_engine_commands(kernel):
+    """GPU flags should still be rejected for prepare/combine/k_selection_plot."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("command")
+    kernel.parse_gpu_args(parser)
+    supported = ("factorize", "consensus")
+
+    for argv in (
+        ["prepare", "--engine", "gpu"],
+        ["combine", "--gpu-device", "cuda"],
+        ["k_selection_plot", "--gpu-check-every", "2"],
+    ):
+        with pytest.raises(ValueError, match="only valid with"):
+            kernel.validate_engine_args_for_command(parser.parse_args(argv), supported)
+
+
 def test_configure_nmf_engine_cpu_leaves_cnmf_instance_unchanged(kernel):
     """The default CPU engine should be a no-op so existing sklearn behavior is preserved."""
     class DummyCNMF:
@@ -786,6 +829,98 @@ def test_nndsvd_nndsvda_nndsvdar_initializers_return_valid_outputs(kernel):
         assert_valid_nmf_output(X, H, W, 3)
 
 
+# ---------------------------------------------------------------------
+# Consensus fixed-H kernel behavior
+# ---------------------------------------------------------------------
+def test_factorize_nmf_gpu_update_h_false_keeps_fixed_h_and_updates_only_w(kernel):
+    """Consensus refit should preserve supplied H and optimize usages W only."""
+    require_nmf_runtime()
+    fixed_h = np.array([[1.0, 0.3, 0.6], [0.2, 1.1, 0.4]], dtype=np.float64)
+    true_w = np.array([[1.0, 0.5], [0.4, 1.2], [1.5, 0.3], [0.7, 0.9]], dtype=np.float64)
+    X = true_w @ fixed_h
+
+    H, W = kernel.factorize_nmf_gpu(
+        X,
+        {"n_components": 2, "max_iter": 5, "random_state": 0, "update_H": False, "H": fixed_h},
+        {"device": "cpu", "dtype": "fp64", "check_every": 5},
+    )
+
+    assert_valid_nmf_output(X, H, W, 2)
+    assert np.allclose(H, fixed_h)
+
+
+def test_to_checked_fixed_h_rejects_missing_invalid_or_incompatible_h(kernel):
+    """Fixed-H consensus refit should fail clearly for invalid supplied spectra."""
+    with pytest.raises(ValueError, match="requires a fixed H"):
+        kernel._to_checked_fixed_h(None, 2, 3)
+
+    invalid_cases = [
+        (np.array([1.0, 2.0, 3.0]), "2D"),
+        (np.ones((3, 3)), "shape"),
+        (np.array([[1.0, np.nan, 0.2], [0.4, 0.5, 0.6]]), "NaN/inf"),
+        (np.array([[1.0, -0.1, 0.2], [0.4, 0.5, 0.6]]), "non-negative"),
+    ]
+    for H, message in invalid_cases:
+        with pytest.raises(ValueError, match=message):
+            kernel._to_checked_fixed_h(H, 2, 3)
+
+
+def test_mu_step_fixed_h_matches_manual_w_only_update(kernel):
+    """One fixed-H MU step should match the manual W-only Frobenius update."""
+    torch = require_nmf_runtime()
+    Xg = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64)
+    W0 = torch.tensor([[0.5, 0.7], [0.9, 1.1]], dtype=torch.float64)
+    H = torch.tensor([[0.6, 0.8], [1.0, 1.2]], dtype=torch.float64)
+    eps = torch.tensor(1e-9, dtype=torch.float64)
+    H_before = H.clone()
+
+    expected_W = W0 * ((Xg @ H.T) / (W0 @ (H @ H.T) + eps))
+    W = kernel._mu_step_fixed_h(W0, H, Xg, eps)
+
+    assert torch.allclose(W, expected_W)
+    assert torch.allclose(H, H_before)
+
+
+def test_fit_mu_fixed_h_respects_early_stop_and_max_iter_block_bounds(kernel):
+    """Fixed-H fit loop should share early-stop and no-overrun guarantees with full MU."""
+    torch = require_nmf_runtime()
+    Xg = torch.full((3, 2), 2.0, dtype=torch.float64)
+    W = torch.ones((3, 1), dtype=torch.float64)
+    H = torch.ones((1, 2), dtype=torch.float64)
+    eps = torch.tensor(1e-9, dtype=torch.float64)
+
+    early_calls = {"count": 0}
+
+    def no_change_step_early(W, H, Xg, eps):
+        early_calls["count"] += 1
+        return W
+
+    kernel._fit_mu_fixed_h(torch, Xg, W, H, eps, 10, 1e-4, no_change_step_early, 1, False, "cpu")
+    assert early_calls["count"] == 2
+
+    max_iter_calls = {"count": 0}
+
+    def no_change_step_max_iter(W, H, Xg, eps):
+        max_iter_calls["count"] += 1
+        return W
+
+    kernel._fit_mu_fixed_h(torch, Xg, W, H, eps, 6, -1.0, no_change_step_max_iter, 4, False, "cpu")
+    assert max_iter_calls["count"] == 6
+
+
+def test_execution_plan_for_fixed_h_compile_uses_fixed_h_step_and_compile_block(kernel):
+    """Compiled consensus refit should compile _mu_step_fixed_h and use compile_block."""
+    calls = []
+    fake_torch = SimpleNamespace(compile=lambda fn: calls.append(fn) or fn)
+    opt = dict(kernel.DEFAULT_GPU, compile=True, check_every=1, compile_block=3)
+
+    step, block = kernel._execution_plan(fake_torch, opt, "cpu", kernel._mu_step_fixed_h)
+
+    assert calls == [kernel._mu_step_fixed_h]
+    assert step is kernel._mu_step_fixed_h
+    assert block == 3
+
+
 # ---------------------------------------------------------------------
 # Backend-specific execution behavior
 # ---------------------------------------------------------------------

tests/test_prepare.py

@@ -59,6 +59,46 @@ def generate_counts_file(tmp_path, file_format, dtype=np.int64, zero_count=False
     return str(counts_fn)
 
 
+def generate_positive_counts_file(tmp_path, cells=24, genes=12, dtype=np.float64):
+    """Generate a small dense positive count table for consensus smoke tests."""
+    rng = np.random.default_rng(SEED)
+    data = rng.poisson(lam=4.0, size=(cells, genes)).astype(dtype) + 1
+    df = pd.DataFrame(
+        data,
+        columns=[f"gene{i}" for i in range(genes)],
+        index=[f"cell{i}" for i in range(cells)],
+    )
+    counts_fn = tmp_path / "positive_counts.txt"
+    df.to_csv(counts_fn, sep="\t")
+    return str(counts_fn)
+
+
+def write_minimal_nmf_run_params(cnmf_obj):
+    """Write the NMF kwargs file required by refit_usage/refit_spectra."""
+    run_params = {
+        "alpha_W": 0.0,
+        "alpha_H": 0.0,
+        "l1_ratio": 0.0,
+        "beta_loss": "frobenius",
+        "solver": "mu",
+        "tol": 1e-4,
+        "max_iter": 5,
+        "init": "random",
+    }
+    cnmf_obj.save_nmf_iter_params(pd.DataFrame(), run_params)
+    return run_params
+
+
+def fake_gpu_nmf_output(X, nmf_kwargs):
+    """Return deterministic positive factors matching the cNMF `(spectra, usages)` contract."""
+    k = int(nmf_kwargs["n_components"])
+    fixed_h = np.asarray(nmf_kwargs["H"], dtype=np.float64)
+    row_scale = np.linspace(1.0, 2.0, X.shape[0], dtype=np.float64).reshape(-1, 1)
+    col_scale = np.arange(1, k + 1, dtype=np.float64).reshape(1, -1)
+    usages = row_scale * col_scale
+    return fixed_h.copy(), usages
+
+
 # ---------------------------------------------------------------------
 # GPU factorize wiring integration
 # ---------------------------------------------------------------------
@@ -108,6 +148,169 @@ def fake_nmf_gpu(self, X, nmf_kwargs):
         assert os.path.exists(mock_cnmf.paths["iter_spectra"] % (5, iter_i))
 
 
+# ---------------------------------------------------------------------
+# GPU consensus/refit wiring
+# ---------------------------------------------------------------------
+def test_refit_usage_gpu_engine_passes_fixed_h_update_h_false_and_gpu_kwargs(mock_cnmf, monkeypatch, tmp_path):
+    """cNMF refit_usage should route fixed-H consensus refits through the GPU adapter."""
+    import cnmf.nmf_gpu as gpu_mod
+
+    write_minimal_nmf_run_params(mock_cnmf)
+    X = pd.DataFrame(
+        np.arange(12, dtype=np.float64).reshape(4, 3) + 1,
+        index=[f"cell{i}" for i in range(4)],
+        columns=[f"gene{i}" for i in range(3)],
+    )
+    spectra = pd.DataFrame(
+        [[1.0, 0.3, 0.6], [0.2, 1.1, 0.4]],
+        index=["program_a", "program_b"],
+        columns=X.columns,
+    )
+    captured = []
+
+    def fake_nmf_gpu(self, X_arg, nmf_kwargs):
+        captured.append((X_arg, dict(nmf_kwargs)))
+        return fake_gpu_nmf_output(X_arg, nmf_kwargs)
+
+    monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu)
+    gpu_kwargs = {"device": "cpu", "dtype": "fp64"}
+    gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs=gpu_kwargs)
+
+    usages = mock_cnmf.refit_usage(X, spectra)
+
+    assert len(captured) == 1
+    X_arg, kw = 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 "beta_loss" in kw and "init" in kw
+    assert list(usages.index) == list(X.index)
+    assert list(usages.columns) == list(spectra.index)
+    assert usages.shape == (4, 2)
+
+
+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
+
+    write_minimal_nmf_run_params(mock_cnmf)
+    X = pd.DataFrame(
+        np.arange(12, dtype=np.float64).reshape(4, 3) + 1,
+        index=[f"cell{i}" for i in range(4)],
+        columns=[f"gene{i}" for i in range(3)],
+    )
+    usage = pd.DataFrame(
+        [[1.0, 0.2], [0.8, 0.4], [0.6, 0.7], [0.4, 1.0]],
+        index=X.index,
+        columns=["program_a", "program_b"],
+    )
+    captured = []
+
+    def fake_nmf_gpu(self, X_arg, nmf_kwargs):
+        captured.append((X_arg, dict(nmf_kwargs)))
+        return fake_gpu_nmf_output(X_arg, nmf_kwargs)
+
+    monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu)
+    gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu", "dtype": "fp64"})
+
+    spectra = mock_cnmf.refit_spectra(X, usage)
+
+    assert len(captured) == 1
+    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
+    assert kw["n_components"] == usage.shape[1]
+    assert list(spectra.index) == list(usage.columns)
+    assert list(spectra.columns) == list(X.columns)
+    assert spectra.shape == (2, 3)
+
+
+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
+
+    counts_fn = generate_positive_counts_file(tmp_path)
+    mock_cnmf.prepare(counts_fn, components=[2], n_iter=3, densify=True,
+                      seed=14, num_highvar_genes=6, max_NMF_iter=5)
+
+    norm_counts = sc.read(mock_cnmf.paths["normalized_counts"])
+    genes = list(norm_counts.var.index)
+    rng = np.random.default_rng(SEED)
+    prototypes = rng.random((2, len(genes))) + 0.2
+    merged = pd.DataFrame(
+        np.vstack([
+            prototypes[0] * 0.98,
+            prototypes[0],
+            prototypes[0] * 1.02,
+            prototypes[1] * 0.98,
+            prototypes[1],
+            prototypes[1] * 1.02,
+        ]),
+        index=[f"iter{i}_topic{j}" for i in range(3) for j in range(1, 3)],
+        columns=genes,
+    )
+    save_df_to_npz(merged, mock_cnmf.paths["merged_spectra"] % 2)
+
+    def fake_nmf_gpu(self, X_arg, nmf_kwargs):
+        return fake_gpu_nmf_output(X_arg, nmf_kwargs)
+
+    monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu)
+    gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs={"device": "cpu", "dtype": "fp64"})
+
+    mock_cnmf.consensus(k=2, density_threshold=2.0, local_neighborhood_size=0.5,
+                        show_clustering=False, refit_usage=False)
+
+    density = "2_0"
+    expected_files = [
+        mock_cnmf.paths["consensus_spectra"] % (2, density),
+        mock_cnmf.paths["consensus_usages"] % (2, density),
+        mock_cnmf.paths["gene_spectra_tpm"] % (2, density),
+        mock_cnmf.paths["gene_spectra_score"] % (2, density),
+        mock_cnmf.paths["starcat_spectra"] % (2, density),
+    ]
+    for path in expected_files:
+        assert os.path.exists(path), f"Expected consensus output {path} not found."
+
+    assert load_df_from_npz(mock_cnmf.paths["consensus_spectra"] % (2, density)).shape == (2, len(genes))
+    assert load_df_from_npz(mock_cnmf.paths["consensus_usages"] % (2, density)).shape == (norm_counts.n_obs, 2)
+
+
+def test_consensus_default_cpu_engine_keeps_original_sklearn_nmf_path(mock_cnmf, monkeypatch, tmp_path):
+    """Consensus without GPU engine should not inject engine/gpu kwargs into sklearn NMF calls."""
+    write_minimal_nmf_run_params(mock_cnmf)
+    X = pd.DataFrame(
+        np.arange(12, dtype=np.float64).reshape(4, 3) + 1,
+        index=[f"cell{i}" for i in range(4)],
+        columns=[f"gene{i}" for i in range(3)],
+    )
+    spectra = pd.DataFrame(
+        [[1.0, 0.3, 0.6], [0.2, 1.1, 0.4]],
+        index=["program_a", "program_b"],
+        columns=X.columns,
+    )
+    captured = []
+
+    def fake_cpu_nmf(X_arg, nmf_kwargs):
+        captured.append(dict(nmf_kwargs))
+        return fake_gpu_nmf_output(X_arg, nmf_kwargs)
+
+    monkeypatch.setattr(mock_cnmf, "_nmf", fake_cpu_nmf)
+
+    usages = mock_cnmf.refit_usage(X, spectra)
+
+    assert len(captured) == 1
+    kw = captured[0]
+    assert "engine" not in kw
+    assert "gpu" not in kw
+    assert kw["update_H"] is False
+    assert np.allclose(kw["H"], spectra.values)
+    assert usages.shape == (4, 2)
+
+
 @pytest.mark.parametrize("file_format", ["txt", "npz", "h5ad"])
 @pytest.mark.parametrize("dtype", [np.int64, np.float32, np.float64])
 @pytest.mark.parametrize("densify", [True, False])