From 458b85248fb85c0b90b7c43ced7227b18b2c429b Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Wed, 1 Jul 2026 23:50:08 +0900 Subject: [PATCH] Add cNMF GPU consensus coverage --- tests/test_nmf_gpu.py | 155 +++++++++++++++++++++++++++++--- tests/test_prepare.py | 203 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 348 insertions(+), 10 deletions(-) diff --git a/tests/test_nmf_gpu.py b/tests/test_nmf_gpu.py index 5aee6b1..0252c64 100644 --- a/tests/test_nmf_gpu.py +++ b/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 # --------------------------------------------------------------------- diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 2cbb020..eefec72 100644 --- a/tests/test_prepare.py +++ b/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])