From 0c810862adbae7e12b53d776159fda596414036b Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Tue, 25 Aug 2026 07:00:27 +0000 Subject: [PATCH 1/2] add reconstruction error calc with filtered spectra --- src/cnmf/cnmf.py | 4 ++++ src/cnmf/qc.py | 13 +++++++++++++ 2 files changed, 17 insertions(+) diff --git a/src/cnmf/cnmf.py b/src/cnmf/cnmf.py index 4bac813..b400b32 100755 --- a/src/cnmf/cnmf.py +++ b/src/cnmf/cnmf.py @@ -935,6 +935,8 @@ def consensus(self, k, density_threshold=0.5, local_neighborhood_size=0.30, show columns = ['stats']) return(consensus_stats) + + density_filtered_refit_usages = rf_usages # Re-order usage by total contribution norm_usages = rf_usages.div(rf_usages.sum(axis=1), axis=0) @@ -1095,6 +1097,8 @@ def consensus(self, k, density_threshold=0.5, local_neighborhood_size=0.30, show l2_spectra=l2_spectra, local_density=local_density, kmeans_cluster_labels=kmeans_cluster_labels, + normalized_counts=norm_counts.X, + refit_usages=density_filtered_refit_usages.values, ) diff --git a/src/cnmf/qc.py b/src/cnmf/qc.py index cba5d40..31ed4e5 100644 --- a/src/cnmf/qc.py +++ b/src/cnmf/qc.py @@ -138,6 +138,7 @@ def get_r2(X, U, H, log2p1=False, scale=False): return r2, sse, tss + def plot_heatmap(df, outfile, scale_rows=False): # df: rows=sources, cols=conditions if scale_rows: @@ -210,11 +211,20 @@ def run_qc( l2_spectra, local_density, kmeans_cluster_labels, + normalized_counts, + refit_usages, ): """Calculate and save QC using state already produced by ``consensus``.""" # local_density contains all spectra, while l2_spectra has already been density-filtered. nruns = len(local_density) nruns_per_gep = nruns/k + if sp.issparse(normalized_counts): + normalized_counts = np.asarray(normalized_counts.todense()) + run_r2, run_sse, run_tss = get_r2( + normalized_counts, + refit_usages, + median_spectra.values, + ) #----------------------------------------------------------------------- # Calculate how many iterations feed into each GEP @@ -258,6 +268,9 @@ def run_qc( annotation['run_davies_bouldin'] = davies_bouldin_score(l2_spectra.values, kmeans_cluster_labels) annotation['run_median_density'] = np.median(local_density.iloc[:, 0]) annotation['run_mean_density'] = np.mean(local_density.iloc[:, 0]) + annotation['run_r2'] = run_r2 + annotation['run_sse'] = run_sse + annotation['run_tss'] = run_tss annotation.index = range(1, len(annotation)+1) annotation = annotation.loc[reorder.index,:] From 22eac9bb0a8b4dcfb1929377e095437b4b4e004d Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Tue, 25 Aug 2026 08:11:45 +0000 Subject: [PATCH 2/2] add regress test on a sythentic dataset --- .github/workflows/python-tests.yml | 25 ++ .../generate_original_results.py | 319 ++++++++++++++++ .../generate_synthetic_dataset.py | 147 ++++++++ .../reference_scripts/manifest.yaml | 6 + .../sc_blipper_fba3745/run_qc_cnmf.py | 349 ++++++++++++++++++ .../fixtures/synthetic_cnmf_k3/run_spec.json | 37 ++ tests/test_consensus_qc.py | 14 + tests/test_synthetic_regression.py | 301 +++++++++++++++ 8 files changed, 1198 insertions(+) create mode 100644 tests/fixtures/synthetic_cnmf_k3/generate_original_results.py create mode 100644 tests/fixtures/synthetic_cnmf_k3/generate_synthetic_dataset.py create mode 100644 tests/fixtures/synthetic_cnmf_k3/reference_scripts/manifest.yaml create mode 100755 tests/fixtures/synthetic_cnmf_k3/reference_scripts/sc_blipper_fba3745/run_qc_cnmf.py create mode 100644 tests/fixtures/synthetic_cnmf_k3/run_spec.json create mode 100644 tests/test_synthetic_regression.py diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index fabc927..125ca2a 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -19,10 +19,35 @@ jobs: python-version: "3.10" cache: pip + - name: Read synthetic regression source + id: synthetic-regression + run: | + python - <<'PY' + import json + import os + from pathlib import Path + + spec = json.loads( + Path("tests/fixtures/synthetic_cnmf_k3/run_spec.json").read_text() + ) + with open(os.environ["GITHUB_OUTPUT"], "a", encoding="utf-8") as output: + print(f"repository={spec['upstream']['repository']}", file=output) + print(f"commit={spec['upstream']['commit']}", file=output) + PY + + - name: Check out pinned upstream cNMF + uses: actions/checkout@v4 + with: + repository: ${{ steps.synthetic-regression.outputs.repository }} + ref: ${{ steps.synthetic-regression.outputs.commit }} + path: .upstream/dylkot-cnmf + - name: Install dependencies run: | python -m pip install --upgrade pip python -m pip install -e ".[gpu,test]" - name: Run tests + env: + CNMF_UPSTREAM_SOURCE: ${{ github.workspace }}/.upstream/dylkot-cnmf run: bash run_tests.sh diff --git a/tests/fixtures/synthetic_cnmf_k3/generate_original_results.py b/tests/fixtures/synthetic_cnmf_k3/generate_original_results.py new file mode 100644 index 0000000..58fe2cf --- /dev/null +++ b/tests/fixtures/synthetic_cnmf_k3/generate_original_results.py @@ -0,0 +1,319 @@ +#!/usr/bin/env python3 +"""Generate stage-separated reference outputs with a pinned original cNMF source.""" + +from argparse import ArgumentParser +import hashlib +import importlib +from importlib.metadata import version +import json +import os +from pathlib import Path +import shutil +from statistics import median +import subprocess +import sys +from time import perf_counter + +import pandas as pd + + +FIXTURE_DIR = Path(__file__).resolve().parent +DEFAULT_SPEC = FIXTURE_DIR / "run_spec.json" +DEFAULT_QC_SCRIPT = ( + FIXTURE_DIR + / "reference_scripts" + / "sc_blipper_fba3745" + / "run_qc_cnmf.py" +) + + +def load_spec(path=DEFAULT_SPEC): + with open(path, encoding="utf-8") as stream: + return json.load(stream) + + +def load_upstream_cnmf(source_dir, expected_sha, expected_version): + """Import cNMF only after verifying the checked-out upstream commit.""" + source_dir = Path(source_dir).resolve() + actual_sha = subprocess.check_output( + ["git", "-C", str(source_dir), "rev-parse", "HEAD"], + text=True, + ).strip() + if actual_sha != expected_sha: + raise RuntimeError( + f"upstream cNMF checkout is {actual_sha}; expected {expected_sha}" + ) + + sys.path.insert(0, str(source_dir / "src")) + from cnmf import __version__ as cnmf_version + from cnmf import cNMF + + if cnmf_version != expected_version: + raise RuntimeError( + f"upstream cNMF version is {cnmf_version}; expected {expected_version}" + ) + return cNMF, cnmf_version + + +def files_under(path): + return {candidate for candidate in Path(path).rglob("*") if candidate.is_file()} + + +def snapshot_new_files(run_dir, before, stage_dir, name): + """Copy files created by one cNMF stage while preserving run-relative paths.""" + created = sorted(files_under(run_dir) - before) + destination_root = Path(stage_dir) / name + for source in created: + destination = destination_root / source.relative_to(run_dir) + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + return created + + +def timed_stage(name, function, run_dir, result_dir, fixture_name): + before = files_under(run_dir) + started = perf_counter() + function() + elapsed = perf_counter() - started + created = snapshot_new_files( + run_dir, + before, + Path(result_dir) / name, + fixture_name, + ) + return elapsed, created + + +def factorize_with_iteration_tracking(cnmf_obj): + """Run unmodified cNMF factorization while observing sklearn's n_iter result.""" + cnmf_module = importlib.import_module("cnmf.cnmf") + original_factorization = cnmf_module.non_negative_factorization + records = [] + + def recording_factorization(*args, **kwargs): + usages, spectra, n_iter = original_factorization(*args, **kwargs) + records.append( + { + "replicate": len(records), + "nmf_seed": int(kwargs["random_state"]), + "n_components": int(kwargs["n_components"]), + "n_iter": int(n_iter), + } + ) + return usages, spectra, n_iter + + cnmf_module.non_negative_factorization = recording_factorization + try: + cnmf_obj.factorize(worker_i=0, total_workers=1) + finally: + cnmf_module.non_negative_factorization = original_factorization + return records + + +def generate(input_dir, work_dir, result_dir, source_dir, spec_path, standalone_qc_script): + spec = load_spec(spec_path) + fixture = spec["fixture"] + nmf = spec["nmf"] + upstream = spec["upstream"] + name = fixture["name"] + k = nmf["k"] + n_replicates = nmf["replicates"] + seed = fixture["seed"] + max_nmf_iter = nmf["max_iter"] + density_threshold = nmf["density_threshold"] + local_neighborhood_size = nmf["local_neighborhood_size"] + source_dir = Path(source_dir).resolve() + cnmf_class, cnmf_version = load_upstream_cnmf( + source_dir, + upstream["commit"], + upstream["version"], + ) + + input_dir = Path(input_dir).resolve() + work_dir = Path(work_dir).resolve() + result_dir = Path(result_dir).resolve() + standalone_qc_script = Path(standalone_qc_script).resolve() + qc_script_sha = hashlib.sha256(standalone_qc_script.read_bytes()).hexdigest() + expected_qc_script_sha = spec["qc_reference"]["sha256"] + if qc_script_sha != expected_qc_script_sha: + raise RuntimeError( + f"standalone QC script SHA256 is {qc_script_sha}; " + f"expected {expected_qc_script_sha}" + ) + work_dir.mkdir(parents=True, exist_ok=True) + result_dir.mkdir(parents=True, exist_ok=True) + + cnmf_obj = cnmf_class(output_dir=str(work_dir), name=name) + run_dir = work_dir / name + timings = {} + stage_files = {} + + timings["prepare"], stage_files["prepare"] = timed_stage( + "prepare", + lambda: cnmf_obj.prepare( + counts_fn=str(input_dir / "synthetic_counts.h5ad"), + components=[k], + n_iter=n_replicates, + seed=seed, + genes_file=str(input_dir / "genes.txt"), + beta_loss=nmf["beta_loss"], + init=nmf["init"], + max_NMF_iter=max_nmf_iter, + ), + run_dir, + result_dir, + name, + ) + factorization_iterations = [] + + def factorize(): + factorization_iterations.extend(factorize_with_iteration_tracking(cnmf_obj)) + + timings["factorize"], stage_files["factorize"] = timed_stage( + "factorize", + factorize, + run_dir, + result_dir, + name, + ) + factorization_iterations_path = ( + result_dir / "factorize" / name / "factorization_iterations.tsv" + ) + pd.DataFrame(factorization_iterations).to_csv( + factorization_iterations_path, + sep="\t", + index=False, + ) + stage_files["factorize"].append(factorization_iterations_path) + timings["combine"], stage_files["combine"] = timed_stage( + "combine", + lambda: cnmf_obj.combine(components=[k]), + run_dir, + result_dir, + name, + ) + timings["consensus"], stage_files["consensus"] = timed_stage( + "consensus", + lambda: cnmf_obj.consensus( + k=k, + density_threshold=density_threshold, + local_neighborhood_size=local_neighborhood_size, + show_clustering=True, + close_clustergram_fig=True, + ), + run_dir, + result_dir, + name, + ) + + standalone_qc_dir = result_dir / "standalone_qc" / name + standalone_qc_dir.mkdir(parents=True, exist_ok=True) + standalone_qc_output = standalone_qc_dir / name + started = perf_counter() + qc_environment = dict(os.environ) + qc_environment["PYTHONPATH"] = os.pathsep.join( + filter( + None, + (str(source_dir / "src"), qc_environment.get("PYTHONPATH")), + ) + ) + subprocess.run( + [ + sys.executable, + str(standalone_qc_script), + "--spectra", + cnmf_obj.paths["merged_spectra"] % k, + "--density", + cnmf_obj.paths["local_density_cache"] % k, + "--norm_counts", + cnmf_obj.paths["normalized_counts"], + "--params", + cnmf_obj.paths["nmf_run_parameters"], + "--output", + str(standalone_qc_output), + "--density_threshold", + str(density_threshold), + "--k", + str(k), + ], + check=True, + env=qc_environment, + ) + timings["standalone_qc"] = perf_counter() - started + stage_files["standalone_qc"] = sorted(files_under(standalone_qc_dir)) + + manifest = { + "source_repository": upstream["repository"], + "source_sha": upstream["commit"], + "standalone_qc_script": { + "file": standalone_qc_script.name, + "sha256": qc_script_sha, + }, + "cnmf_version": cnmf_version, + "dependencies": { + package: version(package) + for package in ("anndata", "numpy", "pandas", "scanpy", "scikit-learn", "scipy") + }, + "configuration": { + "name": name, + "k": k, + "replicates": n_replicates, + "seed": seed, + "solver": nmf["solver"], + "beta_loss": nmf["beta_loss"], + "init": nmf["init"], + "max_nmf_iter": max_nmf_iter, + "density_threshold": density_threshold, + "local_neighborhood_size": local_neighborhood_size, + }, + "factorization_iterations": { + "minimum": min(record["n_iter"] for record in factorization_iterations), + "median": median(record["n_iter"] for record in factorization_iterations), + "maximum": max(record["n_iter"] for record in factorization_iterations), + "reached_cap": sum( + record["n_iter"] >= max_nmf_iter + for record in factorization_iterations + ), + "all_converged_before_cap": all( + record["n_iter"] < max_nmf_iter + for record in factorization_iterations + ), + }, + "stages": { + stage: { + "elapsed_seconds": timings[stage], + "file_count": len(stage_files[stage]), + } + for stage in ("prepare", "factorize", "combine", "consensus", "standalone_qc") + }, + } + with open(result_dir / "run_manifest.json", "w", encoding="utf-8") as stream: + json.dump(manifest, stream, indent=2, sort_keys=True) + stream.write("\n") + + +def main(): + parser = ArgumentParser() + parser.add_argument("--input-dir", type=Path, required=True) + parser.add_argument("--work-dir", type=Path, required=True) + parser.add_argument("--result-dir", type=Path, required=True) + parser.add_argument("--source-dir", type=Path, required=True) + parser.add_argument("--spec", type=Path, default=DEFAULT_SPEC) + parser.add_argument( + "--standalone-qc-script", + type=Path, + default=DEFAULT_QC_SCRIPT, + ) + args = parser.parse_args() + generate( + input_dir=args.input_dir, + work_dir=args.work_dir, + result_dir=args.result_dir, + source_dir=args.source_dir, + spec_path=args.spec, + standalone_qc_script=args.standalone_qc_script, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/fixtures/synthetic_cnmf_k3/generate_synthetic_dataset.py b/tests/fixtures/synthetic_cnmf_k3/generate_synthetic_dataset.py new file mode 100644 index 0000000..15606cb --- /dev/null +++ b/tests/fixtures/synthetic_cnmf_k3/generate_synthetic_dataset.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +"""Generate the deterministic synthetic three-program cNMF fixture.""" + +from argparse import ArgumentParser +import json +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd +from scipy import sparse + + +DEFAULT_SPEC = Path(__file__).resolve().parent / "run_spec.json" + + +def load_spec(path=DEFAULT_SPEC): + with open(path, encoding="utf-8") as stream: + return json.load(stream) + + +def generate_dataset(spec): + """Return counts and their known three-program generative factors.""" + fixture = spec["fixture"] + seed = fixture["seed"] + n_programs = fixture["programs"] + n_genes = fixture["genes"] + if (n_programs, n_genes) != (3, 12): + raise ValueError("the synthetic generator currently requires 3 programs and 12 genes") + rng = np.random.default_rng(seed) + + spectra = np.full((n_programs, n_genes), 0.2, dtype=np.float64) + spectra[0, 0:4] = [8.0, 6.0, 4.0, 2.0] + spectra[1, 4:8] = [7.0, 5.0, 4.0, 2.0] + spectra[2, 8:12] = [6.0, 5.0, 3.0, 2.0] + + group_sizes = fixture["cell_groups"] + labels = [ + label + for label in ("program_1", "program_2", "program_3", "mixed") + for _ in range(group_sizes[label]) + ] + if len(labels) != fixture["cells"]: + raise ValueError("cell_groups must sum to fixture.cells") + usages = np.empty((len(labels), n_programs), dtype=np.float64) + dominant_usage = { + "program_1": np.array([2.8, 0.15, 0.10]), + "program_2": np.array([0.15, 2.2, 0.10]), + "program_3": np.array([0.10, 0.15, 1.6]), + } + mixed_usages = np.array( + [ + [1.4, 1.1, 0.6], + [1.2, 0.8, 1.0], + [0.8, 1.4, 0.7], + [1.0, 0.7, 1.3], + [1.5, 0.9, 0.5], + [0.7, 1.2, 1.1], + ] + ) + + mixed_index = 0 + for cell_index, label in enumerate(labels): + if label == "mixed": + usages[cell_index] = mixed_usages[mixed_index] + mixed_index += 1 + else: + scale = 0.85 + 0.05 * (cell_index % 7) + usages[cell_index] = dominant_usage[label] * scale + + expected_rate = usages @ spectra + 0.5 + counts = rng.poisson(expected_rate).astype(np.int32) + return counts, usages, spectra, labels + + +def write_dataset(output_dir, spec): + """Write the synthetic AnnData input and inspectable truth tables.""" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + counts, usages, spectra, labels = generate_dataset(spec) + cell_ids = [f"cell_{index:03d}" for index in range(counts.shape[0])] + gene_ids = [f"gene_{index:02d}" for index in range(counts.shape[1])] + program_names = [ + f"program_{index}" for index in range(1, spec["fixture"]["programs"] + 1) + ] + + obs = pd.DataFrame(index=pd.Index(cell_ids, name="cell_id")) + obs["true_program"] = pd.Categorical( + labels, + categories=program_names + ["mixed"], + ) + obs["library_size"] = counts.sum(axis=1) + + var = pd.DataFrame(index=pd.Index(gene_ids, name="gene_id")) + var["marker_program"] = pd.Categorical( + [program_names[index // 4] for index in range(spec["fixture"]["genes"])], + categories=program_names, + ) + + adata = ad.AnnData(X=sparse.csr_matrix(counts), obs=obs, var=var) + adata.obsm["true_usage"] = usages + adata.varm["true_spectra"] = spectra.T + adata.uns["fixture_seed"] = spec["fixture"]["seed"] + adata.uns["program_names"] = program_names + adata.write_h5ad(output_dir / "synthetic_counts.h5ad") + + pd.DataFrame(counts, index=cell_ids, columns=gene_ids).to_csv( + output_dir / "synthetic_counts.tsv", + sep="\t", + index_label="cell_id", + ) + pd.DataFrame(usages, index=cell_ids, columns=program_names).to_csv( + output_dir / "true_usages.tsv", + sep="\t", + index_label="cell_id", + ) + pd.DataFrame(spectra, index=program_names, columns=gene_ids).to_csv( + output_dir / "true_spectra.tsv", + sep="\t", + index_label="program", + ) + (output_dir / "genes.txt").write_text( + "\n".join(gene_ids) + "\n", + encoding="utf-8", + ) + + print( + f"wrote {counts.shape[0]} cells x {counts.shape[1]} genes " + f"with seed {spec['fixture']['seed']} to {output_dir}" + ) + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--output-dir", + type=Path, + required=True, + ) + parser.add_argument("--spec", type=Path, default=DEFAULT_SPEC) + args = parser.parse_args() + write_dataset(args.output_dir, load_spec(args.spec)) + + +if __name__ == "__main__": + main() diff --git a/tests/fixtures/synthetic_cnmf_k3/reference_scripts/manifest.yaml b/tests/fixtures/synthetic_cnmf_k3/reference_scripts/manifest.yaml new file mode 100644 index 0000000..08b086a --- /dev/null +++ b/tests/fixtures/synthetic_cnmf_k3/reference_scripts/manifest.yaml @@ -0,0 +1,6 @@ +scripts: + sc_blipper_fba3745/run_qc_cnmf.py: + repository: https://github.com/TrynkaLab/sc-blipper.git + commit: fba3745343567a3d68980c80dca347c84492e7b8 + source_path: bin/run_qc_cnmf.py + sha256: 7656754f55633e546a67757f2c0514103698c91b345220bfb6b9fac9a8151df8 diff --git a/tests/fixtures/synthetic_cnmf_k3/reference_scripts/sc_blipper_fba3745/run_qc_cnmf.py b/tests/fixtures/synthetic_cnmf_k3/reference_scripts/sc_blipper_fba3745/run_qc_cnmf.py new file mode 100755 index 0000000..a112df7 --- /dev/null +++ b/tests/fixtures/synthetic_cnmf_k3/reference_scripts/sc_blipper_fba3745/run_qc_cnmf.py @@ -0,0 +1,349 @@ +#!/usr/bin/env python + +import numpy as np +import pandas as pd +import argparse +from tqdm import tqdm +from sklearn.metrics import pairwise_distances +from itertools import combinations +from joblib import Parallel, delayed +from sklearn.cluster import KMeans +from sklearn.metrics import silhouette_score, calinski_harabasz_score, davies_bouldin_score +import scanpy as sc +from cnmf import cNMF +import yaml +from sklearn.decomposition import non_negative_factorization +from sklearn.preprocessing import scale +import scipy.sparse as sp +import matplotlib.pyplot as plt +import seaborn as sns + +def pairwise_dist(mat, grouping, dist='sqeuclidean', sample_correct=True, n_jobs=-1, verbose=True): + """Average of pairwise distances between cells of each group. + Adapted from: https://github.com/sanderlab/scPerturb/blob/master/package/src/scperturb/edistance.py + + Arguments + --------- + mat: :class:`~np.array` + Data matrix. + grouping: 'str' + Grouping variable. + dist: `str` for any distance in scipy.spatial.distance (default: `sqeuclidean`) + Distance metric to use in embedding space. + sample_correct: `bool` (default: `True`) + Whether make the estimator for sigma more unbiased (dividing by N-1 instead of N, similar to sample and population variance). + n_jobs: `int` (default: `-1`) + Number of jobs to use for parallelization. If `n_jobs=1`, no parallelization is used. The default uses all available threads. + verbose: `bool` (default: `True`) + Whether to show a progress bar iterating over all groups. + + Returns + ------- + pwd: pandas.DataFrame + DataFrame with average of pairwise PCA distances between all groups. + """ + + X = mat.copy() + y = grouping.copy() + groups = pd.unique(y) + df = pd.DataFrame(index=groups, columns=groups, dtype=float) + fct = tqdm if verbose else lambda x: x + combis = list(combinations(groups, 2)) + [(x,x) for x in groups] + def one_step(pair): + p1, p2 = pair + x1 = X[y==p1].copy() + N = len(x1) + x2 = X[y==p2].copy() + pwd = pairwise_distances(x1, x2, metric=dist) + M = len(x2)-1 if (p1==p2) & sample_correct else len(x2) + factor = N * M + mean_pwd = np.sum(pwd) / factor + return (p1, p2, mean_pwd) + res = Parallel(n_jobs=n_jobs)(delayed(one_step)(pair) for pair in fct(combis)) + for p1, p2, val in res: + df.loc[p1, p2] = val + df.loc[p2, p1] = val + #df.index.name = obs_key + #df.columns.name = obs_key + #df.name = 'pairwise PCA distances' + return df + +def edist(mat, grouping, pwd=None, dist='sqeuclidean', sample_correct=True, n_jobs=1, verbose=True): + """Computes the edistance to control. Accepts precomputed pwd. + Computes the pairwise E-distances between all groups of rows defined in + grouping. + + Adapted from https://github.com/sanderlab/scPerturb/blob/master/package/src/scperturb/edistance.py + Arguments + --------- + mat: :class:`~np.array` + Data matrix. + grouping: `str` + Array with grouping information. + dist: `str` for any distance in scipy.spatial.distance (default: `sqeuclidean`) + Distance metric to use in embedding space. + sample_correct: `bool` (default: `True`) + Whether make the estimator for sigma more unbiased (dividing by N-1 instead of N, similar to sample and population variance). + n_jobs: `int` (default: `-1`) + Number of jobs to use for parallelization. If `n_jobs=1`, no parallelization is used. The default uses all available threads. + verbose: `bool` (default: `True`) + Whether to show a progress bar iterating over all groups. + + Returns + ------- + estats: pandas.DataFrame + DataFrame with pairwise E-distances between all groups. + """ + pwd = pairwise_dist(mat, grouping, + dist=dist, sample_correct=sample_correct, + n_jobs=n_jobs, verbose=verbose) if pwd is None else pwd + # derive basic statistics + sigmas = np.diag(pwd) + deltas = pwd + estats = 2 * deltas - sigmas - sigmas[:, np.newaxis] + return estats + + +def load_df_from_npz(filename): + with np.load(filename, allow_pickle=True) as f: + obj = pd.DataFrame(**f) + return obj + +def get_r2(X, U, H, log2p1=False, scale=False): + if type(X) not in [np.ndarray]: + raise ValueError("X must be a numpy array") + if type(U) not in [np.ndarray]: + raise ValueError("U must be a numpy array") + if type(H) not in [np.ndarray]: + raise ValueError("H must be a numpy array") + + # Input matrix + if X is None: + if log2p1: + X = np.log2(X+1) + if scale: + X = scale(X, axis=0) + + # Reconstructed matrix + X_pred = U.dot(H) + if log2p1: + X_pred = np.log2(X_pred+1) + if scale: + X_pred = scale(X_pred, axis=0) + + X_centered = X - np.mean(X) + sse = np.sum((X - X_pred)**2) + tss = np.sum(X_centered**2) + r2 = 1 - (sse / tss) + + return r2, sse, tss + +def plot_heatmap(df, outfile, scale_rows=False): + # df: rows=sources, cols=conditions + if scale_rows: + df_plot = df.apply( + lambda r: (r - r.mean()) / (r.std() if r.std() != 0 else 1), + axis=1, + ) + else: + df_plot = df.copy() + + nrows, ncols = df_plot.shape + side = 2 + (max(nrows, ncols) * 0.5) # keep overall figure large enough + fig, ax = plt.subplots(figsize=(side, side)) + + sns.heatmap( + df_plot, + ax=ax, + cmap="RdBu_r", + center=0, + cbar_kws={"shrink": 0.5}, + ) + + ax.set_box_aspect(1) # force the heatmap panel to be square + fig.tight_layout() + fig.savefig(outfile) + plt.close(fig) + + + +def nmf_variance_explained(X, U, H): + """ + This is a crude way of estimating the variance explained by each program by leaving one out. + + :param X: Description + :param U: Description + :param H: Description + """ + if sp.issparse(X): + X=np.asarray(X.todense()) + + if type(X) not in [np.ndarray, np.matrix]: + raise ValueError("X must be a numpy array") + if type(U) not in [np.ndarray]: + raise ValueError("U must be a numpy array") + if type(H) not in [np.ndarray]: + raise ValueError("H must be a numpy array") + + k = H.shape[0] + selectors = [np.arange(k) != i for i in range(k)] + r2_base, sse_base, tss_base = get_r2(X, U, H) + + r2_wo = [] + sse_wo = [] + tss_wo = [] + for selector in selectors: + r2, error, tss = get_r2(X, U[:,selector], H[selector,:]) + r2_wo.append(r2) + sse_wo.append(error) + tss_wo.append(tss) + + return r2_base, sse_base, tss_base, r2_wo, sse_wo, tss_wo + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="Calculate QC metrics on an NMF run") + parser.add_argument('--spectra', required=True, help='path to merged spectra npz') + parser.add_argument('--density', required=True, help='path to local density npz') + parser.add_argument('--norm_counts', required=True, help='path to norm_counts h5ad') + parser.add_argument('--params', required=False, help='path to nmf params yaml (can be skipped if --skip_varexp is set)', default=None) + parser.add_argument('--density_threshold', required=True, type=float, help='density threshold (float)') + parser.add_argument('--output', required=True, help='output prefix') + parser.add_argument('--k', required=True, type=int, help='number of clusters') + parser.add_argument('--skip_varexp', dest='skip_varexp', action='store_true', help='Skip calculating variance explained (default: false)') + parser.set_defaults(skip_varexp=False) + params = parser.parse_args() + + # This script is partially reverse engineered from cNMF + # https://github.com/dylkot/cNMF/blob/main/src/cnmf/cnmf.py + + #--------------------------------------------- + print("Re-clustering spectra") + + # Re-construct the clustering cNMF does + merged_spectra = load_df_from_npz(params.spectra) + + # L2 normalize the raw spectra + l2_spectra = (merged_spectra.T/np.sqrt((merged_spectra**2).sum(axis=1))).T + + # Filter for density + local_density = load_df_from_npz(params.density) + + nruns = l2_spectra.shape[0] + nruns_per_gep = nruns/params.k + + density_filter = local_density.iloc[:, 0] < params.density_threshold + l2_spectra = l2_spectra.loc[density_filter, :] + if l2_spectra.shape[0] == 0: + raise RuntimeError("Zero components remain after density filtering. Consider increasing density threshold") + + # Cluster + kmeans_model = KMeans(n_clusters=params.k, n_init=10, random_state=1) + kmeans_model.fit(l2_spectra) + kmeans_cluster_labels = pd.Series(kmeans_model.labels_+1, index=l2_spectra.index) + kmeans_cluster_labels = kmeans_cluster_labels + + # Find median usage for each gene across cluster + # Row index is the cluster labels + median_spectra = l2_spectra.groupby(kmeans_cluster_labels).median() + + # Normalize median spectra to probability distributions. + median_spectra = (median_spectra.T/median_spectra.sum(1)).T + + #----------------------------------------------------------------------- + # Calculate how many iterations feed into each GEP + unique, counts = np.unique(kmeans_cluster_labels, return_counts=True) + annotation = pd.DataFrame({ + 'gep': None, + 'cluster': unique, + 'iter_count': counts, + 'iter_perc': (counts/nruns_per_gep)*100, + 'k': params.k + }).sort_values('cluster', ascending=True) + + #----------------------------------------------------------------------- + print("Calculating e-distance metrics") + + # Calculate the edistance + spectra_edist = edist(l2_spectra, kmeans_cluster_labels.to_numpy()) + + # Re-order to match the clusters + spectra_edist = spectra_edist.reindex(index=annotation['cluster']).reindex(columns=annotation['cluster']) + + # Calculate the minimal e-distance to other clusters for each cluster + annotation['min_edist'] = spectra_edist.mask(np.eye(len(spectra_edist), dtype=bool)).min(axis=1).to_numpy() + + #----------------------------------------------------------------------- + # Number of non-zero genes in spectra + annotation['nonzero_genes'] = (median_spectra != 0).sum(axis=1).to_numpy() + annotation['nonzero_perc'] = (annotation['nonzero_genes'] / median_spectra.shape[1]) *100 + + #----------------------------------------------------------------------- + # This is needed to be able to properly order the GEPs by total usage across all cells, + # which is important for consistency with the main cNMF pipeline. + # Secondly its also needed to estimate the error metrics. + print("Re-fitting NMF with fixed H to get usages") + + norm_counts = sc.read(params.norm_counts) + refit_nmf_kwargs = yaml.load(open(params.params), Loader=yaml.FullLoader) + + if type(median_spectra) is pd.DataFrame: + refit_nmf_kwargs.update(dict(n_components = median_spectra.shape[0], H = median_spectra.values, update_H = False)) + else: + refit_nmf_kwargs.update(dict(n_components = median_spectra.shape[0], H = median_spectra, update_H = False)) + + (usages, _, niter) = non_negative_factorization(norm_counts.X, **refit_nmf_kwargs) + rf_usages = pd.DataFrame(usages, index=norm_counts.obs.index, columns=median_spectra.index) + + #----------------------------------------------------------------------- + # Re-order based on normalized usages, this is important to keep consistency + # with the main cNMF pipeline where GEPs are ordered by total usage across all cells. + norm_usages = rf_usages.div(rf_usages.sum(axis=1), axis=0) + annotation['total_usage'] = norm_usages.values.sum(axis=0) + reorder = norm_usages.sum(axis=0).sort_values(ascending=False) + + #----------------------------------------------------------------------- + # Calculate the LOO variance explained + # Read the 'counts' data that the nmf was fit on + if not params.skip_varexp: + + print("Estimating variance explained per program by leaving one out") + # Baseline r2 + print(type(norm_counts.X)) + r2_base, sse_base, tss_base, r2_wo, sse_wo, tss_wo = nmf_variance_explained(norm_counts.X, rf_usages.values, median_spectra.values) + + annotation['r2_diff'] = r2_base - np.array(r2_wo) + annotation['loo_r2'] = r2_wo + annotation['loo_sse'] = sse_wo + annotation['loo_tss'] = tss_wo + annotation['run_r2'] = r2_base + annotation['run_sse'] = sse_base + annotation['run_tss'] = tss_base + + #----------------------------------------------------------------------- + # Calculate clustering metrics, this is a sanity check to see if the clusters are well separated in the spectra space overall + annotation['run_iter_count'] = l2_spectra.shape[0] + annotation['run_iter_count_perc'] = (l2_spectra.shape[0]/nruns)*100 + annotation['run_silhouette'] = silhouette_score(l2_spectra.values, kmeans_cluster_labels, metric='euclidean') + annotation['run_calinski_harabasz'] = calinski_harabasz_score(l2_spectra.values, kmeans_cluster_labels) + annotation['run_davies_bouldin'] = davies_bouldin_score(l2_spectra.values, kmeans_cluster_labels) + annotation['run_median_density'] = np.median(local_density.iloc[:, 0]) + annotation['run_mean_density'] = np.mean(local_density.iloc[:, 0]) + + annotation.index = range(1, len(annotation)+1) + annotation = annotation.loc[reorder.index,:] + annotation['gep'] = range(1, len(annotation)+1) + annotation.index = range(1, len(annotation)+1) + + # Re order edits + spectra_edist = spectra_edist.loc[annotation['cluster'], annotation['cluster']] + spectra_edist.index = range(1, len(annotation)+1) + spectra_edist.columns = range(1, len(annotation)+1) + + # Save the results + prefix = f"k_{params.k}.dt_{params.density_threshold}".replace(".", "_") + annotation.to_csv(f"{params.output}.{prefix}.annotation.tsv", sep="\t", index=False) + spectra_edist.to_csv(f"{params.output}.{prefix}.edist.tsv", sep="\t", index=True) + + plot_heatmap(spectra_edist, f"{params.output}.{prefix}.edist.pdf", scale_rows=False) \ No newline at end of file diff --git a/tests/fixtures/synthetic_cnmf_k3/run_spec.json b/tests/fixtures/synthetic_cnmf_k3/run_spec.json new file mode 100644 index 0000000..d6ab838 --- /dev/null +++ b/tests/fixtures/synthetic_cnmf_k3/run_spec.json @@ -0,0 +1,37 @@ +{ + "fixture": { + "name": "synthetic_cnmf_k3", + "seed": 20260825, + "cells": 36, + "genes": 12, + "programs": 3, + "cell_groups": { + "program_1": 14, + "program_2": 10, + "program_3": 6, + "mixed": 6 + } + }, + "upstream": { + "repository": "dylkot/cNMF", + "commit": "5dbc5baaa0b9079b55bce554d801caa235a50457", + "version": "1.7.1" + }, + "nmf": { + "engine": "cpu", + "solver": "cd", + "beta_loss": "frobenius", + "init": "random", + "k": 3, + "replicates": 30, + "max_iter": 50, + "density_threshold": 0.1, + "local_neighborhood_size": 0.3 + }, + "qc_reference": { + "repository": "TrynkaLab/sc-blipper", + "commit": "fba3745343567a3d68980c80dca347c84492e7b8", + "source_path": "bin/run_qc_cnmf.py", + "sha256": "7656754f55633e546a67757f2c0514103698c91b345220bfb6b9fac9a8151df8" + } +} diff --git a/tests/test_consensus_qc.py b/tests/test_consensus_qc.py index d13d72e..b4c2692 100644 --- a/tests/test_consensus_qc.py +++ b/tests/test_consensus_qc.py @@ -55,6 +55,15 @@ def test_run_qc_matches_run_qc_cnmf_reference_output(tmp_path): ) median_spectra = l2_spectra.groupby(kmeans_cluster_labels).median() median_spectra = median_spectra.div(median_spectra.sum(axis=1), axis=0) + refit_usages = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 2.0, 0.0], + [0.0, 0.0, 3.0], + [1.0, 1.0, 1.0], + ] + ) + normalized_counts = refit_usages @ median_spectra.values local_density = pd.DataFrame( {"local_density": [0.10, 0.20, 0.30, 0.40, 0.50, 0.60]}, index=l2_spectra.index, @@ -74,6 +83,8 @@ def test_run_qc_matches_run_qc_cnmf_reference_output(tmp_path): l2_spectra=l2_spectra, local_density=local_density, kmeans_cluster_labels=kmeans_cluster_labels, + normalized_counts=normalized_counts, + refit_usages=refit_usages, ) annotation_path = tmp_path / f"reference.{prefix}.annotation.tsv" @@ -99,6 +110,9 @@ def test_run_qc_matches_run_qc_cnmf_reference_output(tmp_path): "run_davies_bouldin": [0.31744325019893016] * 3, "run_median_density": [0.35] * 3, "run_mean_density": [0.35] * 3, + "run_r2": [1.0] * 3, + "run_sse": [0.0] * 3, + "run_tss": [3.3251911190193537] * 3, } ) pd.testing.assert_frame_equal( diff --git a/tests/test_synthetic_regression.py b/tests/test_synthetic_regression.py new file mode 100644 index 0000000..6d73700 --- /dev/null +++ b/tests/test_synthetic_regression.py @@ -0,0 +1,301 @@ +"""Runtime-generated regression tests against a pinned upstream cNMF release.""" + +import json +import os +from pathlib import Path +import shutil +from types import SimpleNamespace +import subprocess +import sys + +import numpy as np +import pandas as pd +import pytest +import scanpy as sc +import yaml + +from cnmf import cNMF, load_df_from_npz +from cnmf.gpunmf import configure_nmf_engine + + +FIXTURE_ROOT = Path(__file__).parent / "fixtures" / "synthetic_cnmf_k3" +SPEC_PATH = FIXTURE_ROOT / "run_spec.json" +GENERATOR = FIXTURE_ROOT / "generate_synthetic_dataset.py" +REFERENCE_GENERATOR = FIXTURE_ROOT / "generate_original_results.py" + +with open(SPEC_PATH, encoding="utf-8") as stream: + SPEC = json.load(stream) + +NAME = SPEC["fixture"]["name"] +K = SPEC["nmf"]["k"] +N_REPLICATES = SPEC["nmf"]["replicates"] +SEED = SPEC["fixture"]["seed"] +MAX_NMF_ITER = SPEC["nmf"]["max_iter"] +DENSITY_THRESHOLD = SPEC["nmf"]["density_threshold"] +LOCAL_NEIGHBORHOOD_SIZE = SPEC["nmf"]["local_neighborhood_size"] + + +@pytest.fixture(scope="module") +def generated_fixture(tmp_path_factory): + """Generate deterministic input and pinned upstream outputs for this run.""" + upstream_source = os.environ.get("CNMF_UPSTREAM_SOURCE") + if not upstream_source: + pytest.skip( + "set CNMF_UPSTREAM_SOURCE to the pinned dylkot/cNMF checkout " + "to run synthetic regression tests" + ) + + upstream_source = Path(upstream_source).resolve() + if not (upstream_source / ".git").exists(): + pytest.fail(f"CNMF_UPSTREAM_SOURCE is not a Git checkout: {upstream_source}") + + runtime_root = tmp_path_factory.mktemp("synthetic-cnmf-regression") + input_root = runtime_root / "input" + reference_root = runtime_root / "reference" + subprocess.run( + [ + sys.executable, + str(GENERATOR), + "--output-dir", + str(input_root), + "--spec", + str(SPEC_PATH), + ], + check=True, + ) + subprocess.run( + [ + sys.executable, + str(REFERENCE_GENERATOR), + "--input-dir", + str(input_root), + "--work-dir", + str(runtime_root / "upstream-work"), + "--result-dir", + str(reference_root), + "--source-dir", + str(upstream_source), + "--spec", + str(SPEC_PATH), + ], + check=True, + ) + + with open(reference_root / "run_manifest.json", encoding="utf-8") as stream: + manifest = json.load(stream) + assert manifest["source_sha"] == SPEC["upstream"]["commit"] + assert manifest["configuration"]["seed"] == SEED + assert manifest["configuration"]["max_nmf_iter"] == MAX_NMF_ITER + return SimpleNamespace( + input_root=input_root, + reference_root=reference_root, + manifest=manifest, + ) + + +def _reference_cnmf(reference_root, stage): + return cNMF(output_dir=str(reference_root / stage), name=NAME) + + +def _stage_reference(reference_root, stage, output_dir): + """Copy one generated reference stage into a temporary cNMF run.""" + shutil.copytree( + reference_root / stage / NAME, + Path(output_dir) / NAME, + dirs_exist_ok=True, + ) + + +def _dense(matrix): + return matrix.toarray() if hasattr(matrix, "toarray") else np.asarray(matrix) + + +def _assert_adata_equal(actual_path, expected_path): + actual = sc.read(actual_path) + expected = sc.read(expected_path) + pd.testing.assert_frame_equal(actual.obs, expected.obs) + pd.testing.assert_frame_equal(actual.var, expected.var) + np.testing.assert_allclose( + _dense(actual.X), + _dense(expected.X), + rtol=1e-12, + atol=1e-12, + ) + + +def _assert_npz_equal(actual_path, expected_path, rtol=1e-12, atol=1e-12): + pd.testing.assert_frame_equal( + load_df_from_npz(actual_path), + load_df_from_npz(expected_path), + check_dtype=False, + rtol=rtol, + atol=atol, + ) + + +def test_prepare_matches_dylkot_reference(generated_fixture, tmp_path): + """Current prepare must reproduce the canonical dylkot/cNMF stage.""" + actual = cNMF(output_dir=str(tmp_path), name=NAME) + expected = _reference_cnmf(generated_fixture.reference_root, "prepare") + + actual.prepare( + counts_fn=str(generated_fixture.input_root / "synthetic_counts.h5ad"), + components=[K], + n_iter=N_REPLICATES, + seed=SEED, + genes_file=str(generated_fixture.input_root / "genes.txt"), + beta_loss=SPEC["nmf"]["beta_loss"], + init=SPEC["nmf"]["init"], + max_NMF_iter=MAX_NMF_ITER, + ) + + for key in ("normalized_counts", "tpm"): + _assert_adata_equal(actual.paths[key], expected.paths[key]) + for key in ("nmf_replicate_parameters", "tpm_stats"): + _assert_npz_equal(actual.paths[key], expected.paths[key]) + assert Path(actual.paths["nmf_genes_list"]).read_text() == Path( + expected.paths["nmf_genes_list"] + ).read_text() + with open(actual.paths["nmf_run_parameters"], encoding="utf-8") as stream: + actual_parameters = yaml.safe_load(stream) + with open(expected.paths["nmf_run_parameters"], encoding="utf-8") as stream: + expected_parameters = yaml.safe_load(stream) + assert actual_parameters == expected_parameters + + +def test_gpu_cd_factorize_matches_dylkot_cpu_reference( + generated_fixture, + tmp_path, +): + """GPU-engine FP64 CD on CPU must reproduce canonical CPU-CD spectra.""" + pytest.importorskip("torch") + + _stage_reference(generated_fixture.reference_root, "prepare", tmp_path) + args = SimpleNamespace( + command="factorize", + engine="gpu", + solver="cd", + beta_loss="frobenius", + output_dir=str(tmp_path), + name=NAME, + gpu_device="cpu", + gpu_dtype="fp64", + gpu_allow_tf32=False, + gpu_compile=False, + gpu_eps=None, + gpu_check_every=1, + gpu_compile_block=None, + gpu_batch=N_REPLICATES, + ) + actual = configure_nmf_engine(cNMF, args) + expected = _reference_cnmf(generated_fixture.reference_root, "factorize") + + actual.factorize(worker_i=0, total_workers=1) + + for replicate in range(N_REPLICATES): + _assert_npz_equal( + actual.paths["iter_spectra"] % (K, replicate), + expected.paths["iter_spectra"] % (K, replicate), + rtol=1e-10, + atol=1e-12, + ) + + +def test_combine_matches_dylkot_reference(generated_fixture, tmp_path): + """Current combine must reproduce the canonical merged spectra exactly.""" + _stage_reference(generated_fixture.reference_root, "prepare", tmp_path) + _stage_reference(generated_fixture.reference_root, "factorize", tmp_path) + actual = cNMF(output_dir=str(tmp_path), name=NAME) + expected = _reference_cnmf(generated_fixture.reference_root, "combine") + + actual.combine(components=[K]) + + _assert_npz_equal( + actual.paths["merged_spectra"] % K, + expected.paths["merged_spectra"] % K, + ) + + +@pytest.fixture(scope="module") +def consensus_case(generated_fixture, tmp_path_factory): + """Run consensus once for the consensus and QC assertions.""" + output_dir = tmp_path_factory.mktemp("synthetic-regression-consensus") + _stage_reference(generated_fixture.reference_root, "prepare", output_dir) + _stage_reference(generated_fixture.reference_root, "combine", output_dir) + actual = cNMF(output_dir=str(output_dir), name=NAME) + actual.consensus( + k=K, + density_threshold=DENSITY_THRESHOLD, + local_neighborhood_size=LOCAL_NEIGHBORHOOD_SIZE, + show_clustering=False, + close_clustergram_fig=True, + ) + return actual, generated_fixture.reference_root + + +def test_consensus_matches_dylkot_reference(consensus_case): + """Current consensus numerics must match canonical cNMF outputs.""" + actual, reference_root = consensus_case + expected = _reference_cnmf(reference_root, "consensus") + threshold = str(DENSITY_THRESHOLD).replace(".", "_") + + _assert_npz_equal( + actual.paths["local_density_cache"] % K, + expected.paths["local_density_cache"] % K, + ) + for key in ( + "consensus_spectra", + "consensus_usages", + "gene_spectra_tpm", + "gene_spectra_score", + "starcat_spectra", + ): + _assert_npz_equal( + actual.paths[key] % (K, threshold), + expected.paths[key] % (K, threshold), + rtol=1e-9, + atol=1e-11, + ) + + +def test_consensus_qc_matches_run_qc_cnmf_reference(consensus_case): + """Integrated QC must match shared columns from the original script.""" + actual, reference_root = consensus_case + prefix = "k_3_dt_0_1" + actual_prefix = Path(actual.output_dir) / NAME / NAME + expected_prefix = reference_root / "standalone_qc" / NAME / NAME + + actual_annotation = pd.read_csv( + f"{actual_prefix}.{prefix}.annotation.tsv", + sep="\t", + ) + expected_annotation = pd.read_csv( + f"{expected_prefix}.{prefix}.annotation.tsv", + sep="\t", + ) + pd.testing.assert_frame_equal( + actual_annotation, + expected_annotation.loc[:, actual_annotation.columns], + check_dtype=False, + rtol=1e-9, + atol=1e-11, + ) + + actual_edist = pd.read_csv( + f"{actual_prefix}.{prefix}.edist.tsv", + sep="\t", + index_col=0, + ) + expected_edist = pd.read_csv( + f"{expected_prefix}.{prefix}.edist.tsv", + sep="\t", + index_col=0, + ) + pd.testing.assert_frame_equal( + actual_edist, + expected_edist, + check_dtype=False, + rtol=1e-9, + atol=1e-11, + ) + assert Path(f"{actual_prefix}.{prefix}.edist.pdf").stat().st_size > 0