From ff8ff06a3d23a354de1e2a10db673056d70f140f Mon Sep 17 00:00:00 2001 From: Adam Freedman Date: Mon, 24 Aug 2026 10:22:23 -0400 Subject: [PATCH 1/3] Add preflight cell-count gate to skip low-cell samples before the run A parse-time check counts called cells in each sample's CellRanger filtered_feature_bc_matrix (barcodes.tsv line count, no matrix load) and, per preflight_mode, skips (default) / errors / warns on samples below preflight_min_cells (default 100) before any job launches. Only the CellRanger filtered count is knowable preflight (emptydrops/cellbender_fromraw call cells at runtime), so it gates on that as the whole-sample viability signal; require_min_cells_for_pca + the low_quality_samples quarantine remain the runtime backstop for the other arms. 'skip' is dynamic (re-checked each run, so a re-sequenced sample returns automatically). common.smk: count_filtered_cells, preflight_min_cells_check, config validation; Snakefile: gate after excluded_samples; config.yaml: documented defaults; tests: test_preflight_checks.py (skip/error/warn/off/pass) + 2 config-validation cases. Co-Authored-By: Claude Opus 4.8 --- config/config.yaml | 10 ++++ tests/test_config_validation.py | 6 ++ tests/test_preflight_checks.py | 103 ++++++++++++++++++++++++++++++++ workflow/Snakefile | 33 ++++++++++ workflow/rules/common.smk | 43 +++++++++++++ 5 files changed, 195 insertions(+) create mode 100644 tests/test_preflight_checks.py diff --git a/config/config.yaml b/config/config.yaml index 6df43fb..4e18ffa 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -12,6 +12,16 @@ workflow_seed: 12345 # ID in samplesheet.tsv, remove that ID from this list so the new data is processed. excluded_samples: [] +# Preflight quality gate (runs before any job): counts the called cells in each sample's +# CellRanger filtered_feature_bc_matrix and handles samples with too few to cluster. +# preflight_mode: "skip" drops them from the DAG and reports them - re-checked every run, so a +# re-sequenced sample with more cells is picked up automatically; "error" hard-stops the run; +# "warn" only reports; "off" disables the check. +# Only the CellRanger filtered count is known preflight; the runtime min-cells guard and the +# results/low_quality_samples/ quarantine remain the backstop for the emptydrops/cellbender arms. +preflight_min_cells: 100 +preflight_mode: "skip" + emptydrop_removal_methods: ["tenx","emptydrops"] ambient_decon_methods: ["soupx","cellbender_fromraw"] doublet_removal_methods: ["doubletfinder", "scdblfinder"] diff --git a/tests/test_config_validation.py b/tests/test_config_validation.py index dcf25e6..86b8ca1 100644 --- a/tests/test_config_validation.py +++ b/tests/test_config_validation.py @@ -114,6 +114,10 @@ def test_default_config_has_required_keys_and_valid_values(): assert isinstance(config["min_ncount"], int) and config["min_ncount"] > 0 assert isinstance(config["max_mtdna"], (int, float)) and 0 <= config["max_mtdna"] <= 100 + assert config.get("preflight_mode", "skip") in {"off", "warn", "skip", "error"} + if "preflight_min_cells" in config: + assert isinstance(config["preflight_min_cells"], int) and config["preflight_min_cells"] > 0 + @pytest.mark.parametrize( "mutate, expected_message, override_results_dir", @@ -128,6 +132,8 @@ def test_default_config_has_required_keys_and_valid_values(): (lambda cfg: cfg.update({"resultsDir": ""}), "resultsDir must be a non-empty string", False), (lambda cfg: cfg.update({"workflow_mode": "bad_mode"}), "workflow_mode must be one of", True), (lambda cfg: cfg.update({"workflow_mode": "downsample_only", "downsampleRate": 1.5}), "downsampleRate must be > 0 and <= 1", True), + (lambda cfg: cfg.update({"preflight_min_cells": 0}), "preflight_min_cells must be a positive integer", True), + (lambda cfg: cfg.update({"preflight_mode": "bogus"}), "preflight_mode must be one of", True), ], ) def test_invalid_config_fails_early_with_clear_message(tmp_path, mutate, expected_message, override_results_dir): diff --git a/tests/test_preflight_checks.py b/tests/test_preflight_checks.py new file mode 100644 index 0000000..ac260aa --- /dev/null +++ b/tests/test_preflight_checks.py @@ -0,0 +1,103 @@ +"""Parse-time preflight cell-count gate. + +The `test` sample's CellRanger filtered matrix has 300 cells, so a threshold above 300 +exercises the "below threshold" behavior and a threshold at/below 300 (the default 100) +passes it. All checks run as `snakemake -np` dry runs in workflow_mode=preprocess so that +skipping the only sample simply yields an empty DAG (in preprocess_and_downsample mode an +emptied sample set trips the separate "no inputs for downsampling" guard instead).""" + +import shutil +import subprocess +from pathlib import Path + + +def run_dry_run(repo_root, results_dir, *config_overrides): + snakemake = shutil.which("snakemake") + assert snakemake is not None, "snakemake is not available on PATH" + cmd = [ + snakemake, + "-np", + "--profile", "none", + "--workflow-profile", "none", + "--snakefile", "workflow/Snakefile", + "--configfile", "config/config.yaml", + "--config", + "sampleTable=testdata/samplesheet_test.tsv", + f"resultsDir={results_dir}", + "workflow_mode=preprocess", + *config_overrides, + ] + return subprocess.run( + cmd, cwd=repo_root, text=True, capture_output=True, check=False, timeout=120 + ) + + +def combined_output(result): + return result.stdout + result.stderr + + +def test_preflight_passes_sample_above_threshold(tmp_path): + # 300-cell test sample vs the default preflight_min_cells=100 -> runs normally. + repo_root = Path(__file__).resolve().parents[1] + result = run_dry_run(repo_root, tmp_path / "results") + output = combined_output(result) + + assert result.returncode == 0, output + assert "tenx2seuratrds" in output + assert "below preflight_min_cells" not in output + + +def test_preflight_skip_drops_low_cell_sample(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + result = run_dry_run( + repo_root, tmp_path / "results", + "preflight_min_cells=500", "preflight_mode=skip", + ) + output = combined_output(result) + + assert result.returncode == 0, output + assert "below preflight_min_cells=500" in output + assert "test (300 cells)" in output + assert "skipping 1 low-cell sample(s)" in output + assert "tenx2seuratrds" not in output # the only sample was dropped from the DAG + + +def test_preflight_error_mode_hard_stops(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + result = run_dry_run( + repo_root, tmp_path / "results", + "preflight_min_cells=500", "preflight_mode=error", + ) + output = combined_output(result) + + assert result.returncode != 0 + assert "fewer than preflight_min_cells=500" in output + assert "test (300 cells)" in output + + +def test_preflight_warn_mode_reports_but_keeps_sample(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + result = run_dry_run( + repo_root, tmp_path / "results", + "preflight_min_cells=500", "preflight_mode=warn", + ) + output = combined_output(result) + + assert result.returncode == 0, output + assert "below preflight_min_cells=500" in output + assert "test (300 cells)" in output + assert "skipping" not in output + assert "tenx2seuratrds" in output # kept despite being below threshold + + +def test_preflight_off_mode_disables_check(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + result = run_dry_run( + repo_root, tmp_path / "results", + "preflight_min_cells=500", "preflight_mode=off", + ) + output = combined_output(result) + + assert result.returncode == 0, output + assert "[preflight]" not in output + assert "tenx2seuratrds" in output diff --git a/workflow/Snakefile b/workflow/Snakefile index 30126d4..3906566 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -18,6 +18,8 @@ config.setdefault("min_ncount", 1) config.setdefault("max_mtdna", 100) EXCLUDED_SAMPLES = validated_config["excluded_samples"] +PREFLIGHT_MIN_CELLS = validated_config["preflight_min_cells"] +PREFLIGHT_MODE = validated_config["preflight_mode"] if RUN_PREPROCESS: sampleinfo = pd.read_table(config["sampleTable"], dtype={"sampleid": str}) sampleinfo = validate_sample_sheet(sampleinfo, config["sampleTable"]) @@ -40,6 +42,37 @@ if RUN_PREPROCESS: ) SAMPLES = [s for s in SAMPLES if s not in EXCLUDED_SAMPLES] sampleinfo = sampleinfo[~sampleinfo["sampleid"].isin(EXCLUDED_SAMPLES)].reset_index(drop=True) + + # Preflight quality gate: before any job launches, count the called cells in each remaining + # sample's CellRanger filtered matrix and handle those too few to cluster. Only that count + # is knowable now (emptydrops/cellbender_fromraw call cells at runtime), so this gates on it + # as the whole-sample viability signal; require_min_cells_for_pca + the low_quality_samples + # quarantine remain the runtime backstop for the other arms. + if PREFLIGHT_MODE != "off" and SAMPLES: + preflight_counts, preflight_below = preflight_min_cells_check(sampleinfo, PREFLIGHT_MIN_CELLS) + if preflight_below: + detail = ", ".join(f"{s} ({preflight_counts[s]} cells)" for s in preflight_below) + if PREFLIGHT_MODE == "error": + raise ValueError( + f"[preflight] {len(preflight_below)} sample(s) have fewer than " + f"preflight_min_cells={PREFLIGHT_MIN_CELLS} called cells in their CellRanger " + f"filtered matrix: {detail}. Add them to excluded_samples, lower " + "preflight_min_cells, or set preflight_mode to 'skip' or 'warn'." + ) + print( + f"[preflight] {len(preflight_below)} sample(s) below " + f"preflight_min_cells={PREFLIGHT_MIN_CELLS} (CellRanger filtered cells): {detail}", + file=sys.stderr, + ) + if PREFLIGHT_MODE == "skip": + print( + f"[preflight] skipping {len(preflight_below)} low-cell sample(s) from the DAG; " + "re-checked every run, so a re-sequenced sample with more cells is picked up " + "automatically.", + file=sys.stderr, + ) + SAMPLES = [s for s in SAMPLES if s not in preflight_below] + sampleinfo = sampleinfo[~sampleinfo["sampleid"].isin(preflight_below)].reset_index(drop=True) else: sampleinfo = pd.DataFrame({"sampleid": [], "tenx_datadir": []}) SAMPLES = [] diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 7eee7b1..de031f5 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -1,3 +1,4 @@ +import gzip import os from pathlib import Path @@ -21,6 +22,9 @@ ALLOWED_EMPTYDROP_METHODS = {"tenx", "emptydrops"} ALLOWED_DECON_METHODS = {"soupx", "cellbender_fromraw"} ALLOWED_DOUBLET_METHODS = {"doubletfinder", "scdblfinder"} ALLOWED_POSTHOC_METHODS = {"threshold", "mad"} +ALLOWED_PREFLIGHT_MODES = {"off", "warn", "skip", "error"} +DEFAULT_PREFLIGHT_MIN_CELLS = 100 +DEFAULT_PREFLIGHT_MODE = "skip" def require_non_empty_string(config_values, key, errors): @@ -152,6 +156,12 @@ def validate_workflow_config(config_values): if isinstance(max_mtdna, bool) or not isinstance(max_mtdna, (int, float)) or not 0 <= max_mtdna <= 100: errors.append("max_mtdna must be a number between 0 and 100") + require_optional_positive_int(config_values, "preflight_min_cells", errors) + if config_values.get("preflight_mode", DEFAULT_PREFLIGHT_MODE) not in ALLOWED_PREFLIGHT_MODES: + errors.append( + "preflight_mode must be one of: " + ", ".join(sorted(ALLOWED_PREFLIGHT_MODES)) + ) + if workflow_mode in {"preprocess_and_downsample", "downsample_only"}: if "downsampleSeuratObjectDir" in config_values: require_non_empty_string(config_values, "downsampleSeuratObjectDir", errors) @@ -183,6 +193,8 @@ def validate_workflow_config(config_values): "doublet_methods": doublet_methods, "posthoc_methods": posthoc_methods, "excluded_samples": excluded_samples, + "preflight_min_cells": config_values.get("preflight_min_cells", DEFAULT_PREFLIGHT_MIN_CELLS), + "preflight_mode": config_values.get("preflight_mode", DEFAULT_PREFLIGHT_MODE), } @@ -249,6 +261,37 @@ def validate_sample_sheet(sampleinfo, sample_table): return validated +def count_filtered_cells(tenx_datadir): + """Number of called cells in a sample's CellRanger filtered matrix. + + Reads only the barcode list (barcodes.tsv[.gz]) line count -- no matrix is loaded -- so it + is cheap enough to run for every sample at parse time (the preflight check below).""" + matrix_dir = Path(tenx_datadir) / "filtered_feature_bc_matrix" + for name in ("barcodes.tsv.gz", "barcodes.tsv"): + barcodes = matrix_dir / name + if barcodes.exists(): + opener = gzip.open if name.endswith(".gz") else open + with opener(barcodes, "rt") as handle: + return sum(1 for line in handle if line.strip()) + raise FileNotFoundError(f"preflight cell count: no barcodes.tsv[.gz] under {matrix_dir}") + + +def preflight_min_cells_check(sampleinfo, min_cells): + """Per-sample CellRanger-filtered cell counts vs the preflight minimum. + + Returns (counts, below): counts maps each sampleid to its number of called cells in + filtered_feature_bc_matrix; below is the sorted list of sampleids with fewer than + min_cells. Only the CellRanger filtered count is knowable before the run (emptydrops and + cellbender_fromraw call cells at runtime), so this gates on that count as the whole-sample + viability signal; require_min_cells_for_pca remains the runtime backstop for the arms.""" + counts = { + str(sample_id): count_filtered_cells(data_dir) + for sample_id, data_dir in zip(sampleinfo["sampleid"], sampleinfo["tenx_datadir"]) + } + below = sorted(sample_id for sample_id, n in counts.items() if n < min_cells) + return counts, below + + def sample_tenx_dir(wildcards): return sampleinfo.loc[sampleinfo["sampleid"] == wildcards.sample, "tenx_datadir"].values[0] From 23fac5593e0aed24f277473de21a8fdceb5cee69 Mon Sep 17 00:00:00 2001 From: Adam Freedman Date: Mon, 24 Aug 2026 11:17:26 -0400 Subject: [PATCH 2/3] Add CellBender raw-matrix preflight (integrity error, droplet-sufficiency warning) For the cellbender_fromraw arm, before any job launches: hard-error on a raw matrix.mtx that is unreadable, empty/degenerate (zero features/barcodes/nonzeros), or has fewer barcodes than the sample's filtered matrix (raw must be a superset) -- a broken/wrong input to fix or exclude; and warn (non-blocking) when raw droplets are < 2x the CellRanger-filtered cell count (too few empties for CellBender's ambient estimate, often a filtered matrix on the raw path). These are input-integrity checks, so unlike the cell-count gate they are NOT governed by preflight_mode. Runs only when cellbender_fromraw is configured; reads only the MTX header. common.smk: read_mtx_dims, cellbender_preflight_checks; Snakefile: gate after the cell-count preflight; tests: 4 cases + a minimal-datadir fixture. Co-Authored-By: Claude Opus 4.8 --- tests/test_preflight_checks.py | 90 +++++++++++++++++++++++++++++++++- workflow/Snakefile | 16 ++++++ workflow/rules/common.smk | 64 ++++++++++++++++++++++++ 3 files changed, 168 insertions(+), 2 deletions(-) diff --git a/tests/test_preflight_checks.py b/tests/test_preflight_checks.py index ac260aa..df16cd6 100644 --- a/tests/test_preflight_checks.py +++ b/tests/test_preflight_checks.py @@ -6,12 +6,13 @@ skipping the only sample simply yields an empty DAG (in preprocess_and_downsample mode an emptied sample set trips the separate "no inputs for downsampling" guard instead).""" +import gzip import shutil import subprocess from pathlib import Path -def run_dry_run(repo_root, results_dir, *config_overrides): +def run_dry_run(repo_root, results_dir, *config_overrides, sample_table="testdata/samplesheet_test.tsv"): snakemake = shutil.which("snakemake") assert snakemake is not None, "snakemake is not available on PATH" cmd = [ @@ -22,7 +23,7 @@ def run_dry_run(repo_root, results_dir, *config_overrides): "--snakefile", "workflow/Snakefile", "--configfile", "config/config.yaml", "--config", - "sampleTable=testdata/samplesheet_test.tsv", + f"sampleTable={sample_table}", f"resultsDir={results_dir}", "workflow_mode=preprocess", *config_overrides, @@ -101,3 +102,88 @@ def test_preflight_off_mode_disables_check(tmp_path): assert result.returncode == 0, output assert "[preflight]" not in output assert "tenx2seuratrds" in output + + +# --- CellBender raw-matrix preflight (integrity -> error, droplet sufficiency -> warn) ------- + +def _write_gz(path, text): + with gzip.open(path, "wt") as handle: + handle.write(text) + + +def make_datadir(base, filtered_cells, raw_features, raw_barcodes, raw_nnz): + """A minimal 10x datadir: `filtered_cells` filtered barcodes and a raw matrix.mtx.gz whose + header advertises the given (features, barcodes, nnz). Enough to pass validate_sample_sheet + and exercise the CellBender raw-matrix checks. filtered_cells >= 100 so the cell-count gate + doesn't skip the sample before the CellBender check runs.""" + filtered = base / "filtered_feature_bc_matrix" + filtered.mkdir(parents=True) + _write_gz(filtered / "barcodes.tsv.gz", "".join(f"CELL{i}-1\n" for i in range(filtered_cells))) + + raw = base / "raw_feature_bc_matrix" + raw.mkdir(parents=True) + body = "1 1 1\n" if raw_nnz > 0 else "" + _write_gz( + raw / "matrix.mtx.gz", + "%%MatrixMarket matrix coordinate integer general\n" + f"{raw_features} {raw_barcodes} {raw_nnz}\n" + body, + ) + (base / "raw_feature_bc_matrix.h5").write_text("") # existence only; not read by the checks + return base + + +def write_sample_sheet(path, sampleid, datadir): + path.write_text(f"sampleid\ttenx_datadir\n{sampleid}\t{datadir}\n") + return path + + +def test_cellbender_preflight_warns_on_too_few_raw_droplets(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + datadir = make_datadir(tmp_path / "lowraw", filtered_cells=300, + raw_features=2000, raw_barcodes=400, raw_nnz=1000) # 400 < 2*300 + sheet = write_sample_sheet(tmp_path / "sheet.tsv", "lowraw", datadir) + result = run_dry_run(repo_root, tmp_path / "results", sample_table=sheet) + output = combined_output(result) + + assert result.returncode == 0, output # a warning is non-blocking + assert "[preflight:cellbender] WARNING" in output + assert "400 raw droplets vs 300 called cells" in output + assert "tenx2seuratrds" in output # the sample still runs + + +def test_cellbender_preflight_errors_on_empty_raw_matrix(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + datadir = make_datadir(tmp_path / "emptyraw", filtered_cells=300, + raw_features=0, raw_barcodes=0, raw_nnz=0) + sheet = write_sample_sheet(tmp_path / "sheet.tsv", "emptyraw", datadir) + result = run_dry_run(repo_root, tmp_path / "results", sample_table=sheet) + output = combined_output(result) + + assert result.returncode != 0 + assert "[preflight:cellbender]" in output + assert "empty/degenerate" in output + + +def test_cellbender_preflight_errors_when_raw_smaller_than_filtered(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + datadir = make_datadir(tmp_path / "mismatch", filtered_cells=300, + raw_features=2000, raw_barcodes=200, raw_nnz=1000) # raw < filtered + sheet = write_sample_sheet(tmp_path / "sheet.tsv", "mismatch", datadir) + result = run_dry_run(repo_root, tmp_path / "results", sample_table=sheet) + output = combined_output(result) + + assert result.returncode != 0 + assert "fewer barcodes (200) than the filtered matrix (300)" in output + + +def test_cellbender_preflight_passes_a_true_raw_matrix(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + datadir = make_datadir(tmp_path / "goodraw", filtered_cells=300, + raw_features=2000, raw_barcodes=20000, raw_nnz=5000) + sheet = write_sample_sheet(tmp_path / "sheet.tsv", "goodraw", datadir) + result = run_dry_run(repo_root, tmp_path / "results", sample_table=sheet) + output = combined_output(result) + + assert result.returncode == 0, output + assert "[preflight:cellbender]" not in output + assert "tenx2seuratrds" in output diff --git a/workflow/Snakefile b/workflow/Snakefile index 3906566..52590c9 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -73,6 +73,22 @@ if RUN_PREPROCESS: ) SAMPLES = [s for s in SAMPLES if s not in preflight_below] sampleinfo = sampleinfo[~sampleinfo["sampleid"].isin(preflight_below)].reset_index(drop=True) + + # CellBender raw-matrix sanity for the cellbender_fromraw arm. These are input-integrity + # checks, so (unlike the cell-count gate) they are NOT governed by preflight_mode: a + # malformed / empty / mismatched raw matrix hard-errors (fix the input or exclude the + # sample), while too few raw droplets for ambient estimation is a non-blocking warning. + # Reads only MTX headers, so it stays cheap. + if SAMPLES and "cellbender_fromraw" in validated_config["decon_methods"]: + cb_errors, cb_warnings = cellbender_preflight_checks(sampleinfo) + for message in cb_warnings: + print(f"[preflight:cellbender] WARNING: {message}", file=sys.stderr) + if cb_errors: + raise ValueError( + "[preflight:cellbender] raw-matrix problem(s) for the cellbender_fromraw arm:\n - " + + "\n - ".join(cb_errors) + + "\nFix the raw inputs or add these samples to excluded_samples." + ) else: sampleinfo = pd.DataFrame({"sampleid": [], "tenx_datadir": []}) SAMPLES = [] diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index de031f5..b045090 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -25,6 +25,7 @@ ALLOWED_POSTHOC_METHODS = {"threshold", "mad"} ALLOWED_PREFLIGHT_MODES = {"off", "warn", "skip", "error"} DEFAULT_PREFLIGHT_MIN_CELLS = 100 DEFAULT_PREFLIGHT_MODE = "skip" +DEFAULT_MIN_RAW_TO_CELL_RATIO = 2 def require_non_empty_string(config_values, key, errors): @@ -292,6 +293,69 @@ def preflight_min_cells_check(sampleinfo, min_cells): return counts, below +def read_mtx_dims(matrix_dir): + """(n_features, n_barcodes, nnz) from a 10x MatrixMarket matrix.mtx[.gz] header. + + Reads only the header (comment lines plus the single dims line), so it never loads the + matrix -- cheap enough for every sample at parse time.""" + for name in ("matrix.mtx.gz", "matrix.mtx"): + mtx = Path(matrix_dir) / name + if mtx.exists(): + opener = gzip.open if name.endswith(".gz") else open + with opener(mtx, "rt") as handle: + for line in handle: + if line.startswith("%"): + continue + parts = line.split() + if len(parts) < 3: + raise ValueError(f"malformed MatrixMarket dims line: {line.strip()!r}") + return int(parts[0]), int(parts[1]), int(parts[2]) + raise ValueError(f"no dimension line in {mtx}") + raise FileNotFoundError(f"no matrix.mtx[.gz] under {matrix_dir}") + + +def cellbender_preflight_checks(sampleinfo, min_raw_to_cell_ratio=DEFAULT_MIN_RAW_TO_CELL_RATIO): + """Raw-matrix sanity for the cellbender_fromraw arm, at parse time (reads only MTX headers). + + Returns (errors, warnings). These are input-integrity checks, deliberately NOT governed by + preflight_mode: + errors -- a broken or wrong raw input to fix: the raw matrix.mtx is unreadable, + empty/degenerate (zero features/barcodes/nonzeros), or has fewer barcodes than + the sample's filtered matrix (raw must be a superset of filtered). Hard-stops. + warnings -- the raw matrix has fewer than min_raw_to_cell_ratio x the called cells in + droplets: too few empty droplets for CellBender's ambient estimate, often a + sign the raw path actually points at a filtered matrix. Non-blocking. + Only meaningful when cellbender_fromraw is a configured decon method; the caller gates on that.""" + errors = [] + warnings = [] + for sample_id, data_dir in zip(sampleinfo["sampleid"], sampleinfo["tenx_datadir"]): + raw_dir = Path(data_dir) / "raw_feature_bc_matrix" + try: + n_features, n_barcodes, nnz = read_mtx_dims(raw_dir) + except Exception as exc: + errors.append(f"{sample_id}: raw matrix unreadable ({raw_dir}): {exc}") + continue + if n_features <= 0 or n_barcodes <= 0 or nnz <= 0: + errors.append( + f"{sample_id}: raw matrix is empty/degenerate " + f"(features={n_features}, barcodes={n_barcodes}, nonzeros={nnz})" + ) + continue + filtered_cells = count_filtered_cells(data_dir) + if n_barcodes < filtered_cells: + errors.append( + f"{sample_id}: raw matrix has fewer barcodes ({n_barcodes}) than the filtered " + f"matrix ({filtered_cells}); the raw and filtered inputs look mismatched" + ) + elif n_barcodes < min_raw_to_cell_ratio * filtered_cells: + warnings.append( + f"{sample_id}: {n_barcodes} raw droplets vs {filtered_cells} called cells " + f"(< {min_raw_to_cell_ratio}x); few empty droplets for CellBender's ambient " + "estimate - check the raw path is a true unfiltered matrix" + ) + return errors, warnings + + def sample_tenx_dir(wildcards): return sampleinfo.loc[sampleinfo["sampleid"] == wildcards.sample, "tenx_datadir"].values[0] From 276f4b6a66c57ebf6b65a298b4e9ea8fde879468 Mon Sep 17 00:00:00 2001 From: Adam Freedman Date: Mon, 24 Aug 2026 15:30:42 -0400 Subject: [PATCH 3/3] Add adaptive CellBender re-run based on the report's learning-curve assessment After the initial CellBender run, parse the HTML report's automated assessment: - "learning curve looks normal" -> keep the initial run. - "re-run with half the current learning rate" -> if cellbender_adaptive_rerun is on, re-run at half cellbender_learning_rate and keep the re-run only if ITS curve is normal, otherwise fall back to the initial run. - Any unrecognized assessment is a hard error so report-parsing edge cases surface. Both runs' reports are archived (_report_initial.html / _report_rerun.html) and a per-sample status file records the outcome. A new cellbender_adaptive_summary rule aggregates every sample's status into results/cellbender/cellbender_adaptive_summary.tsv, reporting which samples needed a re-run and whether it resolved the learning curve. New config keys (validated): cellbender_learning_rate (default 0.0001, CellBender's default) and cellbender_adaptive_rerun (default true). The wrapper's assess_report is a pure function; tests cover it plus every orchestration branch (via a fake cellbender), the run-level summary, and the two new config-validation cases. Co-Authored-By: Claude Opus 4.8 --- config/config.yaml | 9 + tests/test_cellbender_adaptive.py | 225 ++++++++++++++++++ tests/test_cellbender_rule.py | 10 + tests/test_conda_container_validation.py | 9 +- tests/test_config_validation.py | 3 + tests/test_sample_rule_output_files.txt | 2 + workflow/Snakefile | 17 ++ workflow/rules/cellbender.smk | 92 +++---- workflow/rules/common.smk | 25 ++ workflow/scripts/cellbender_adaptive_run.py | 196 +++++++++++++++ .../scripts/cellbender_adaptive_summary.py | 60 +++++ 11 files changed, 604 insertions(+), 44 deletions(-) create mode 100644 tests/test_cellbender_adaptive.py create mode 100644 workflow/scripts/cellbender_adaptive_run.py create mode 100644 workflow/scripts/cellbender_adaptive_summary.py diff --git a/config/config.yaml b/config/config.yaml index 4e18ffa..2e45cb1 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -26,6 +26,15 @@ emptydrop_removal_methods: ["tenx","emptydrops"] ambient_decon_methods: ["soupx","cellbender_fromraw"] doublet_removal_methods: ["doubletfinder", "scdblfinder"] posthoc_methods: ["threshold", "mad"] + +# CellBender (cellbender_fromraw arm) adaptive re-run. After the initial run, the HTML report's +# automated assessment is parsed: if the learning curve looks normal, that run is kept; if the +# report recommends halving the learning rate and cellbender_adaptive_rerun is true, CellBender is +# re-run at half cellbender_learning_rate and the re-run is kept only if ITS curve looks normal +# (otherwise the initial run is kept). results/cellbender/cellbender_adaptive_summary.tsv reports +# which samples needed a re-run and whether it resolved the issue. +cellbender_learning_rate: 0.0001 # CellBender's default; the re-run uses exactly half of this +cellbender_adaptive_rerun: true # set false to always keep the single initial run min_nfeature: 200 min_ncount: 500 max_mtdna: 5 diff --git a/tests/test_cellbender_adaptive.py b/tests/test_cellbender_adaptive.py new file mode 100644 index 0000000..3310fac --- /dev/null +++ b/tests/test_cellbender_adaptive.py @@ -0,0 +1,225 @@ +"""Adaptive CellBender re-run wrapper (workflow/scripts/cellbender_adaptive_run.py) and the +run-level summary (cellbender_adaptive_summary.py). + +`assess_report` is unit-tested directly. The orchestration is tested by driving `main()` against a +fake `cellbender` on PATH whose report verdict is scripted per run (initial vs re-run) through +environment variables, so every branch -- keep initial, re-run and resolve, re-run and fail to +resolve, adaptive off, and the hard error on an unrecognized report -- is exercised without a GPU.""" + +import importlib.util +import os +from pathlib import Path + +import pytest + +SCRIPTS = Path(__file__).resolve().parents[1] / "workflow" / "scripts" + + +def _load(module_name): + path = SCRIPTS / f"{module_name}.py" + spec = importlib.util.spec_from_file_location(module_name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +run_mod = _load("cellbender_adaptive_run") +summary_mod = _load("cellbender_adaptive_summary") + + +# A fake `cellbender`: advertises --seed, and on each real run writes the .h5/_filtered.h5 outputs +# plus a _report.html whose verdict is chosen by a per-invocation counter -- the first real run +# reads $FAKE_INITIAL_PHRASE, the second reads $FAKE_RERUN_PHRASE (values: normal|rerun|unknown). +FAKE_CELLBENDER = r"""#!/usr/bin/env bash +set -euo pipefail + +if [[ "${1:-}" == "remove-background" && "${2:-}" == "--help" ]]; then + printf '%s\n' "cellbender remove-background" " --seed INTEGER" + exit 0 +fi +if [[ "${1:-}" != "remove-background" ]]; then + echo "unexpected cellbender invocation: $*" >&2 + exit 2 +fi +shift + +output="" +while [[ $# -gt 0 ]]; do + case "$1" in + --output) output="$2"; shift 2 ;; + --input|--learning-rate|--seed) shift 2 ;; + *) shift ;; + esac +done + +n=1 +if [[ -f "${FAKE_COUNTER}" ]]; then n=$(( $(cat "${FAKE_COUNTER}") + 1 )); fi +echo "${n}" > "${FAKE_COUNTER}" +if [[ "${n}" == "1" ]]; then phrase="${FAKE_INITIAL_PHRASE:-normal}"; else phrase="${FAKE_RERUN_PHRASE:-normal}"; fi + +case "${phrase}" in + normal) summary="This learning curve looks normal." ;; + rerun) summary="Consider re-running with half the current learning rate to compare the results." ;; + *) summary="The assessment is inconclusive here." ;; +esac + +mkdir -p "$(dirname "${output}")" +printf 'fake raw for run %s\n' "${n}" > "${output}" +printf 'fake filtered for run %s\n' "${n}" > "${output%.h5}_filtered.h5" +printf '

Automated assessment

Summary

%s

\n' \ + "${summary}" > "${output%.h5}_report.html" +""" + + +def _install_fake(tmp_path, monkeypatch, initial_phrase, rerun_phrase="normal"): + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake = fake_bin / "cellbender" + fake.write_text(FAKE_CELLBENDER) + fake.chmod(0o755) + # Prepend so the fake wins as `cellbender` while bash/env/coreutils still resolve. + monkeypatch.setenv("PATH", f"{fake_bin}:{os.environ['PATH']}") + monkeypatch.setenv("FAKE_COUNTER", str(tmp_path / "counter")) + monkeypatch.setenv("FAKE_INITIAL_PHRASE", initial_phrase) + monkeypatch.setenv("FAKE_RERUN_PHRASE", rerun_phrase) + + +def _run(tmp_path, adaptive="true"): + """Drive main() and return (status_dict, results_dir Path).""" + input_h5 = tmp_path / "raw.h5" + input_h5.write_text("raw input") + results = tmp_path / "results" / "cellbender" + run_mod.main([ + "--input", str(input_h5), + "--output-base", str(results / "cellbender_test.h5"), + "--output-filtered", str(results / "cellbender_test_filtered.h5"), + "--report", str(results / "cellbender_test_report.html"), + "--status", str(results / "cellbender_test_adaptive_status.txt"), + "--sample", "test", + "--learning-rate", "0.0001", + "--adaptive", adaptive, + "--seed", "12345", + "--workdir", str(tmp_path / "scratch"), + ]) + status = {} + for line in (results / "cellbender_test_adaptive_status.txt").read_text().splitlines(): + key, value = line.split("\t", 1) + status[key] = value + return status, results + + +# --- assess_report unit tests -------------------------------------------------------------------- + +def test_assess_report_normal_ignores_tags_and_case(): + html = "

Summary

This learning curve looks NORMAL.

" + assert run_mod.assess_report(html) == "normal" + + +def test_assess_report_detects_rerun_suggestion(): + html = "

Consider re-running with half the current learning rate to compare.

" + assert run_mod.assess_report(html) == "rerun" + + +def test_assess_report_unknown_when_neither_phrase_present(): + assert run_mod.assess_report("

Training finished. Elbo converged.

") == "unknown" + + +# --- orchestration ------------------------------------------------------------------------------- + +def test_normal_first_try_keeps_initial(tmp_path, monkeypatch): + _install_fake(tmp_path, monkeypatch, initial_phrase="normal") + status, results = _run(tmp_path) + + assert status["outcome"] == "NORMAL_FIRST_TRY" + assert status["reran"] == "false" + assert status["kept_run"] == "initial" + assert status["rerun_learning_rate"] == "NA" + assert (results / "cellbender_test.h5").read_text() == "fake raw for run 1\n" + assert (results / "cellbender_test_report_initial.html").exists() + assert not (results / "cellbender_test_report_rerun.html").exists() + + +def test_rerun_resolves_keeps_rerun(tmp_path, monkeypatch): + _install_fake(tmp_path, monkeypatch, initial_phrase="rerun", rerun_phrase="normal") + status, results = _run(tmp_path) + + assert status["outcome"] == "RERUN_RESOLVED" + assert status["reran"] == "true" + assert status["kept_run"] == "rerun" + assert status["initial_learning_rate"] == "0.0001" + assert status["rerun_learning_rate"] == "5e-05" # exactly half + # The kept output is the second (re-run) result, and both reports are archived. + assert (results / "cellbender_test.h5").read_text() == "fake raw for run 2\n" + assert (results / "cellbender_test_report_initial.html").exists() + assert (results / "cellbender_test_report_rerun.html").exists() + + +def test_rerun_does_not_resolve_keeps_initial(tmp_path, monkeypatch): + _install_fake(tmp_path, monkeypatch, initial_phrase="rerun", rerun_phrase="rerun") + status, results = _run(tmp_path) + + assert status["outcome"] == "RERUN_DID_NOT_RESOLVE" + assert status["reran"] == "true" + assert status["kept_run"] == "initial" + assert status["rerun_verdict"] == "rerun" + # Kept output falls back to the initial run even though a re-run happened. + assert (results / "cellbender_test.h5").read_text() == "fake raw for run 1\n" + assert (results / "cellbender_test_report_rerun.html").exists() + + +def test_adaptive_off_does_not_rerun(tmp_path, monkeypatch): + _install_fake(tmp_path, monkeypatch, initial_phrase="rerun", rerun_phrase="normal") + status, results = _run(tmp_path, adaptive="false") + + assert status["outcome"] == "RERUN_SUGGESTED_BUT_ADAPTIVE_OFF" + assert status["reran"] == "false" + assert status["kept_run"] == "initial" + assert (results / "cellbender_test.h5").read_text() == "fake raw for run 1\n" + assert not (results / "cellbender_test_report_rerun.html").exists() + + +def test_unknown_initial_report_hard_errors(tmp_path, monkeypatch): + _install_fake(tmp_path, monkeypatch, initial_phrase="unknown") + with pytest.raises(SystemExit) as excinfo: + _run(tmp_path) + assert "could not classify" in str(excinfo.value) + + +def test_unknown_rerun_report_hard_errors(tmp_path, monkeypatch): + _install_fake(tmp_path, monkeypatch, initial_phrase="rerun", rerun_phrase="unknown") + with pytest.raises(SystemExit) as excinfo: + _run(tmp_path) + assert "re-run report" in str(excinfo.value) + + +# --- run-level summary --------------------------------------------------------------------------- + +def _write_status(path, **fields): + path.write_text("".join(f"{k}\t{v}\n" for k, v in fields.items())) + return path + + +def test_summary_aggregates_status_files(tmp_path, capsys): + a = _write_status( + tmp_path / "a.txt", sample="alpha", outcome="NORMAL_FIRST_TRY", + initial_learning_rate="0.0001", initial_verdict="normal", reran="false", + rerun_learning_rate="NA", rerun_verdict="NA", kept_run="initial", + ) + b = _write_status( + tmp_path / "b.txt", sample="beta", outcome="RERUN_DID_NOT_RESOLVE", + initial_learning_rate="0.0001", initial_verdict="rerun", reran="true", + rerun_learning_rate="5e-05", rerun_verdict="rerun", kept_run="initial", + ) + out = tmp_path / "summary.tsv" + # Pass b before a to confirm the summary sorts rows by sample id. + summary_mod.main(["--output", str(out), str(b), str(a)]) + + lines = out.read_text().splitlines() + assert lines[0].split("\t") == list(summary_mod.COLUMNS) + assert lines[1].startswith("alpha\t") + assert lines[2].startswith("beta\t") + + printed = capsys.readouterr().out + assert "2 sample(s)" in printed + assert "RERUN_DID_NOT_RESOLVE: 1" in printed + assert "beta" in printed # named as re-ran and as unresolved diff --git a/tests/test_cellbender_rule.py b/tests/test_cellbender_rule.py index 9aa2e4d..b43e90e 100644 --- a/tests/test_cellbender_rule.py +++ b/tests/test_cellbender_rule.py @@ -62,6 +62,8 @@ printf 'fake raw cellbender output for %s\\n' "${input}" > "${output}" filtered="${output%.h5}_filtered.h5" printf 'fake filtered cellbender output for %s\\n' "${input}" > "${filtered}" +report="${output%.h5}_report.html" +printf '

Automated assessment

Summary

This learning curve looks normal.

\\n' > "${report}" """ @@ -86,6 +88,8 @@ def test_cellbender_rule_uses_cellbender_filtered_output_convention(tmp_path): results_dir = tmp_path / "results" base_output = results_dir / "cellbender" / "cellbender_test.h5" filtered_output = results_dir / "cellbender" / "cellbender_test_filtered.h5" + report_output = results_dir / "cellbender" / "cellbender_test_report.html" + status_output = results_dir / "cellbender" / "cellbender_test_adaptive_status.txt" env = os.environ.copy() env["PATH"] = f"{fake_bin}{os.pathsep}{env.get('PATH', '')}" @@ -125,3 +129,9 @@ def test_cellbender_rule_uses_cellbender_filtered_output_convention(tmp_path): assert result.returncode == 0, result.stdout + result.stderr assert base_output.read_text().startswith("fake raw cellbender output") assert filtered_output.read_text().startswith("fake filtered cellbender output") + # The adaptive wrapper also produces the report and per-sample status; a normal-looking + # learning curve means the single initial run is kept, with no re-run. + assert "This learning curve looks normal" in report_output.read_text() + status = status_output.read_text() + assert "outcome\tNORMAL_FIRST_TRY" in status + assert "reran\tfalse" in status diff --git a/tests/test_conda_container_validation.py b/tests/test_conda_container_validation.py index f494cb0..a5fefc1 100644 --- a/tests/test_conda_container_validation.py +++ b/tests/test_conda_container_validation.py @@ -136,7 +136,14 @@ def test_workflow_container_declarations_are_explicit_and_recognized(): containers.append((rule_path, uri)) assert containers, "no workflow container declarations found" - assert containers == [(root / "workflow/rules/cellbender.smk", CELLBENDER_CONTAINER_URI)] + # The `cellbender` rule and the `cellbender_adaptive_summary` rule (which reuses the same + # image for a stdlib aggregation) are the only container declarations, and both pin the + # identical CellBender digest. + cellbender_smk = root / "workflow/rules/cellbender.smk" + assert containers == [ + (cellbender_smk, CELLBENDER_CONTAINER_URI), + (cellbender_smk, CELLBENDER_CONTAINER_URI), + ] for _, uri in containers: assert uri.startswith("docker://") assert ":" in uri.removeprefix("docker://"), f"container URI is missing a tag: {uri}" diff --git a/tests/test_config_validation.py b/tests/test_config_validation.py index 86b8ca1..49f7f40 100644 --- a/tests/test_config_validation.py +++ b/tests/test_config_validation.py @@ -134,6 +134,9 @@ def test_default_config_has_required_keys_and_valid_values(): (lambda cfg: cfg.update({"workflow_mode": "downsample_only", "downsampleRate": 1.5}), "downsampleRate must be > 0 and <= 1", True), (lambda cfg: cfg.update({"preflight_min_cells": 0}), "preflight_min_cells must be a positive integer", True), (lambda cfg: cfg.update({"preflight_mode": "bogus"}), "preflight_mode must be one of", True), + (lambda cfg: cfg.update({"cellbender_learning_rate": 0}), "cellbender_learning_rate must be a positive number", True), + (lambda cfg: cfg.update({"cellbender_learning_rate": "fast"}), "cellbender_learning_rate must be a positive number", True), + (lambda cfg: cfg.update({"cellbender_adaptive_rerun": "yes"}), "cellbender_adaptive_rerun must be a boolean", True), ], ) def test_invalid_config_fails_early_with_clear_message(tmp_path, mutate, expected_message, override_results_dir): diff --git a/tests/test_sample_rule_output_files.txt b/tests/test_sample_rule_output_files.txt index 3a85591..e2378b3 100644 --- a/tests/test_sample_rule_output_files.txt +++ b/tests/test_sample_rule_output_files.txt @@ -1,5 +1,7 @@ testdata/results/cellbender/cellbender_test.h5 +testdata/results/cellbender/cellbender_test_adaptive_status.txt testdata/results/cellbender/cellbender_test_filtered.h5 +testdata/results/cellbender/cellbender_test_report.html testdata/results/cellbender_fromraw/seurat_cellbender_fromraw_test.rds testdata/results/cellbender_fromraw/seurat_cellbender_fromraw_test_markergenes.csv testdata/results/doubletfinder/seurat_doubletfinder_cellbender_fromraw_test.rds diff --git a/workflow/Snakefile b/workflow/Snakefile index 52590c9..4d5b56f 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -100,6 +100,8 @@ DECON_METHODS = validated_config["decon_methods"] DOUBLET_METHODS = validated_config["doublet_methods"] EMPTYDROP_METHODS = validated_config["emptydrop_methods"] POSTHOC_METHODS = validated_config["posthoc_methods"] +CELLBENDER_LEARNING_RATE = validated_config["cellbender_learning_rate"] +CELLBENDER_ADAPTIVE = validated_config["cellbender_adaptive_rerun"] DOWNSAMPLE_RESULTS_DIR = config.get("downsampleResultsDir", f"{RESULTS_DIR}/downsampling").rstrip("/") DOWNSAMPLE_SEURAT_OBJECT_DIR = config.get( "downsampleSeuratObjectDir", @@ -183,6 +185,20 @@ PREPROCESS_PREFIXES = ( PREPROCESS_MARKER_TARGETS = marker_targets(PREPROCESS_PREFIXES) PREPROCESS_SEURAT_TARGETS = rds_targets(PREPROCESS_PREFIXES) +# Per-sample adaptive-run status files (one per cellbender_fromraw sample) and the run-level +# summary that aggregates them. Requested as targets so the summary is always built for the +# cellbender_fromraw arm, giving a single place to see which samples needed a learning-rate re-run. +CELLBENDER_STATUS = ( + expand(f"{RESULTS_DIR}/cellbender/cellbender_{{sample}}_adaptive_status.txt", sample=SAMPLES) + if RUN_PREPROCESS and "cellbender_fromraw" in DECON_METHODS + else [] +) +CELLBENDER_ADAPTIVE_SUMMARY = ( + [f"{RESULTS_DIR}/cellbender/cellbender_adaptive_summary.tsv"] + if CELLBENDER_STATUS + else [] +) + if RUN_DOWNSAMPLE: if WORKFLOW_MODE == "downsample_only": @@ -223,6 +239,7 @@ DOWNSAMPLE_TARGET_REGEX = ( ALL_TARGETS = [] if RUN_PREPROCESS: ALL_TARGETS.extend(PREPROCESS_MARKER_TARGETS) + ALL_TARGETS.extend(CELLBENDER_ADAPTIVE_SUMMARY) if RUN_DOWNSAMPLE: ALL_TARGETS.extend(DOWNSAMPLE_TARGETS) diff --git a/workflow/rules/cellbender.smk b/workflow/rules/cellbender.smk index ca848d2..4ff8fdc 100644 --- a/workflow/rules/cellbender.smk +++ b/workflow/rules/cellbender.smk @@ -1,15 +1,20 @@ rule cellbender: input: - tenx_raw_h5_input + raw=tenx_raw_h5_input, + script="workflow/scripts/cellbender_adaptive_run.py" output: base=f"{RESULTS_DIR}/cellbender/cellbender_{{sample}}.h5", - filtered=f"{RESULTS_DIR}/cellbender/cellbender_{{sample}}_filtered.h5" + filtered=f"{RESULTS_DIR}/cellbender/cellbender_{{sample}}_filtered.h5", + report=f"{RESULTS_DIR}/cellbender/cellbender_{{sample}}_report.html", + status=f"{RESULTS_DIR}/cellbender/cellbender_{{sample}}_adaptive_status.txt" log: f"{RESULTS_DIR}/logs/cellbender/cellbender_{{sample}}.log" params: - # per-sample directory ONLY for checkpoints / temp + # per-sample directory ONLY for checkpoints / temp (holds the initial/ and rerun/ runs) workdir="scratch/cellbender/{sample}", - seed=WORKFLOW_SEED + seed=WORKFLOW_SEED, + learning_rate=CELLBENDER_LEARNING_RATE, + adaptive="true" if CELLBENDER_ADAPTIVE else "false" container: "docker://us.gcr.io/broad-dsde-methods/cellbender@sha256:093f2caf1ce4acae4541ea45e52ab7b220ca131ec73b4d1f664b85fe12850bae" resources: @@ -18,52 +23,53 @@ rule cellbender: gres = "gpu:1", runtime = 2880 shell: + # The adaptive wrapper runs CellBender (once, or twice with a halved learning rate when the + # report's automated assessment recommends it), copies the chosen run's outputs to the + # declared paths, keeps both runs' reports, and writes the per-sample status file. It + # resolves relative paths against the launch dir, so we do NOT cd before invoking it. r""" set -euo pipefail - - # base_dir is the directory from which Snakemake was launched exec > {log} 2>&1 - base_dir=$(pwd) export PYTHONHASHSEED={params.seed} - seed_args=() - if cellbender remove-background --help 2>&1 | grep -q -- "--seed"; then - seed_args=(--seed {params.seed}) - fi - - input_h5="{input}" - case "${{input_h5}}" in - /*) ;; - *) input_h5="${{base_dir}}/${{input_h5}}" ;; - esac - - base_output="{output.base}" - case "${{base_output}}" in - /*) ;; - *) base_output="${{base_dir}}/${{base_output}}" ;; - esac - - filtered_output="{output.filtered}" - case "${{filtered_output}}" in - /*) ;; - *) filtered_output="${{base_dir}}/${{filtered_output}}" ;; - esac - - expected_filtered_output="${{base_output%.h5}}_filtered.h5" - if [ "${{expected_filtered_output}}" != "${{filtered_output}}" ]; then - echo "CellBender derives filtered output from --output as ${{expected_filtered_output}}, but the rule declares ${{filtered_output}}" >&2 - exit 1 - fi - - # Make sure output and scratch dirs exist on the host - mkdir -p "$(dirname "${{base_output}}")" + mkdir -p "$(dirname "{output.base}")" mkdir -p "{params.workdir}" - # Run CellBender in the per-sample scratch dir so ckpt.tar.gz is unique. - cd "{params.workdir}" + python "{input.script}" \ + --input "{input.raw}" \ + --output-base "{output.base}" \ + --output-filtered "{output.filtered}" \ + --report "{output.report}" \ + --status "{output.status}" \ + --sample "{wildcards.sample}" \ + --learning-rate "{params.learning_rate}" \ + --adaptive "{params.adaptive}" \ + --seed "{params.seed}" \ + --workdir "{params.workdir}" + + test -s "{output.base}" + test -s "{output.filtered}" + test -s "{output.report}" + test -s "{output.status}" + """ - cellbender remove-background --cuda "${{seed_args[@]}}" --input "${{input_h5}}" --output "${{base_output}}" - test -s "${{base_output}}" - test -s "${{filtered_output}}" +rule cellbender_adaptive_summary: + # Aggregate every sample's adaptive-run status into one run-level TSV, reporting which samples + # needed a re-run and whether halving the learning rate resolved the learning curve. + input: + status=CELLBENDER_STATUS, + script="workflow/scripts/cellbender_adaptive_summary.py" + output: + f"{RESULTS_DIR}/cellbender/cellbender_adaptive_summary.tsv" + log: + f"{RESULTS_DIR}/logs/cellbender/cellbender_adaptive_summary.log" + # Pure stdlib aggregation; reuse the CellBender image (already pulled for the `cellbender` + # rule) rather than maintain a separate conda env just to run Python. + container: + "docker://us.gcr.io/broad-dsde-methods/cellbender@sha256:093f2caf1ce4acae4541ea45e52ab7b220ca131ec73b4d1f664b85fe12850bae" + shell: + r""" + set -euo pipefail + python "{input.script}" --output "{output}" {input.status} > {log} 2>&1 """ diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index b045090..aef3ce4 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -26,6 +26,9 @@ ALLOWED_PREFLIGHT_MODES = {"off", "warn", "skip", "error"} DEFAULT_PREFLIGHT_MIN_CELLS = 100 DEFAULT_PREFLIGHT_MODE = "skip" DEFAULT_MIN_RAW_TO_CELL_RATIO = 2 +# CellBender's own default learning rate; the adaptive re-run uses exactly half of this. +DEFAULT_CELLBENDER_LEARNING_RATE = 0.0001 +DEFAULT_CELLBENDER_ADAPTIVE_RERUN = True def require_non_empty_string(config_values, key, errors): @@ -45,6 +48,19 @@ def require_optional_positive_int(config_values, key, errors): require_positive_int(config_values, key, errors) +def require_optional_positive_number(config_values, key, errors): + if key not in config_values: + return + value = config_values.get(key) + if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0: + errors.append(f"{key} must be a positive number") + + +def require_optional_bool(config_values, key, errors): + if key in config_values and not isinstance(config_values.get(key), bool): + errors.append(f"{key} must be a boolean (true or false)") + + def validate_method_list(config_values, key, allowed_values, errors): value = config_values.get(key) if not isinstance(value, list) or not value: @@ -163,6 +179,9 @@ def validate_workflow_config(config_values): "preflight_mode must be one of: " + ", ".join(sorted(ALLOWED_PREFLIGHT_MODES)) ) + require_optional_positive_number(config_values, "cellbender_learning_rate", errors) + require_optional_bool(config_values, "cellbender_adaptive_rerun", errors) + if workflow_mode in {"preprocess_and_downsample", "downsample_only"}: if "downsampleSeuratObjectDir" in config_values: require_non_empty_string(config_values, "downsampleSeuratObjectDir", errors) @@ -196,6 +215,12 @@ def validate_workflow_config(config_values): "excluded_samples": excluded_samples, "preflight_min_cells": config_values.get("preflight_min_cells", DEFAULT_PREFLIGHT_MIN_CELLS), "preflight_mode": config_values.get("preflight_mode", DEFAULT_PREFLIGHT_MODE), + "cellbender_learning_rate": config_values.get( + "cellbender_learning_rate", DEFAULT_CELLBENDER_LEARNING_RATE + ), + "cellbender_adaptive_rerun": config_values.get( + "cellbender_adaptive_rerun", DEFAULT_CELLBENDER_ADAPTIVE_RERUN + ), } diff --git a/workflow/scripts/cellbender_adaptive_run.py b/workflow/scripts/cellbender_adaptive_run.py new file mode 100644 index 0000000..49a0c1e --- /dev/null +++ b/workflow/scripts/cellbender_adaptive_run.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python +"""Adaptive CellBender run. + +Run `cellbender remove-background`, read the automated assessment in the HTML report, and if it +recommends halving the learning rate, re-run at half the learning rate and keep whichever run's +learning curve is "normal" -- falling back to the initial run if the re-run is still not normal. +A per-sample status file records what happened, and BOTH runs' reports are kept for audit. + +An unrecognized assessment (neither the "learning curve looks normal" nor the "re-run with half +the current learning rate" phrasing) is a HARD ERROR, so report-parsing edge cases surface +instead of being silently mishandled. + +Invoked by workflow/rules/cellbender.smk inside the CellBender container. `assess_report` is a +pure function so it can be unit-tested on its own. +""" + +import argparse +import os +import re +import shutil +import subprocess +import sys + +NORMAL_PHRASE = "this learning curve looks normal" +RERUN_PHRASE = "consider re-running with half the current learning rate" + + +def assess_report(html_text): + """Classify a CellBender report's automated assessment. + + Returns 'normal', 'rerun', or 'unknown'. HTML tags are stripped and whitespace collapsed + before matching, so minor formatting differences don't affect the result.""" + text = re.sub(r"<[^>]+>", " ", html_text) + text = re.sub(r"\s+", " ", text).lower() + if NORMAL_PHRASE in text: + return "normal" + if RERUN_PHRASE in text: + return "rerun" + return "unknown" + + +def assess_report_file(report_path): + if not os.path.exists(report_path): + raise FileNotFoundError(f"CellBender report not found: {report_path}") + with open(report_path, encoding="utf-8", errors="replace") as handle: + return assess_report(handle.read()) + + +def cellbender_supports_seed(): + try: + completed = subprocess.run( + ["cellbender", "remove-background", "--help"], + capture_output=True, text=True, check=False, + ) + except FileNotFoundError: + return False + return "--seed" in (completed.stdout + completed.stderr) + + +def run_outputs(run_dir, base_name): + """(base, filtered, report) paths CellBender writes for --output run_dir/base_name.""" + base = os.path.join(run_dir, base_name) + stem = base[:-3] if base.endswith(".h5") else base + return base, f"{stem}_filtered.h5", f"{stem}_report.html" + + +def run_cellbender(input_h5, output_h5, learning_rate, seed, seed_supported, run_dir): + """Run cellbender remove-background in its own run_dir (own cwd -> own ckpt.tar.gz).""" + os.makedirs(run_dir, exist_ok=True) + cmd = [ + "cellbender", "remove-background", "--cuda", + "--input", str(input_h5), + "--output", str(output_h5), + "--learning-rate", format(learning_rate, ".10g"), + ] + if seed_supported and seed is not None: + cmd += ["--seed", str(seed)] + print(f"[cellbender-adaptive] running: {' '.join(cmd)} (cwd={run_dir})", flush=True) + subprocess.run(cmd, cwd=run_dir, check=True) + + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--input", required=True) + ap.add_argument("--output-base", required=True) + ap.add_argument("--output-filtered", required=True) + ap.add_argument("--report", required=True) + ap.add_argument("--status", required=True) + ap.add_argument("--sample", required=True) + ap.add_argument("--learning-rate", type=float, required=True) + ap.add_argument("--adaptive", required=True, help='"true" or "false"') + ap.add_argument("--seed", default=None) + ap.add_argument("--workdir", required=True) + args = ap.parse_args(argv) + + adaptive = str(args.adaptive).strip().lower() in ("1", "true", "yes") + input_h5 = os.path.abspath(args.input) + output_base = os.path.abspath(args.output_base) + output_filtered = os.path.abspath(args.output_filtered) + output_report = os.path.abspath(args.report) + output_status = os.path.abspath(args.status) + base_name = os.path.basename(output_base) # cellbender_{sample}.h5 + seed_supported = cellbender_supports_seed() + + workdir = os.path.abspath(args.workdir) + initial_dir = os.path.join(workdir, "initial") + rerun_dir = os.path.join(workdir, "rerun") + initial_lr = args.learning_rate + + # --- initial run ------------------------------------------------------- + init_base, init_filtered, init_report = run_outputs(initial_dir, base_name) + run_cellbender(input_h5, init_base, initial_lr, args.seed, seed_supported, initial_dir) + initial_verdict = assess_report_file(init_report) + print(f"[cellbender-adaptive] {args.sample}: initial run " + f"(lr={format(initial_lr, '.10g')}) -> {initial_verdict}", flush=True) + if initial_verdict == "unknown": + sys.exit( + f"[cellbender-adaptive] {args.sample}: could not classify the initial report " + f"({init_report}) - found neither the 'learning curve looks normal' nor the " + "'re-run with half the current learning rate' phrasing. Inspect the report and " + "extend workflow/scripts/cellbender_adaptive_run.py to handle this case." + ) + + reran = False + rerun_lr = None + rerun_verdict = None + chosen_dir = initial_dir + + if initial_verdict == "normal": + outcome = "NORMAL_FIRST_TRY" + elif not adaptive: + # initial suggested a re-run, but adaptive mode is off -> keep the single run. + outcome = "RERUN_SUGGESTED_BUT_ADAPTIVE_OFF" + else: + reran = True + rerun_lr = initial_lr / 2.0 + rr_base, rr_filtered, rr_report = run_outputs(rerun_dir, base_name) + run_cellbender(input_h5, rr_base, rerun_lr, args.seed, seed_supported, rerun_dir) + rerun_verdict = assess_report_file(rr_report) + print(f"[cellbender-adaptive] {args.sample}: re-run " + f"(lr={format(rerun_lr, '.10g')}) -> {rerun_verdict}", flush=True) + if rerun_verdict == "unknown": + sys.exit( + f"[cellbender-adaptive] {args.sample}: could not classify the re-run report " + f"({rr_report}). Inspect the report and extend " + "workflow/scripts/cellbender_adaptive_run.py to handle this case." + ) + if rerun_verdict == "normal": + chosen_dir = rerun_dir + outcome = "RERUN_RESOLVED" + else: + chosen_dir = initial_dir + outcome = "RERUN_DID_NOT_RESOLVE" + + kept_run = "rerun" if chosen_dir == rerun_dir else "initial" + + # --- copy the chosen run's outputs to the declared paths --------------- + chosen_base, chosen_filtered, chosen_report = run_outputs(chosen_dir, base_name) + os.makedirs(os.path.dirname(output_base), exist_ok=True) + shutil.copy2(chosen_base, output_base) + shutil.copy2(chosen_filtered, output_filtered) + shutil.copy2(chosen_report, output_report) + + # --- keep BOTH runs' reports for audit --------------------------------- + stem = output_report[:-len("_report.html")] if output_report.endswith("_report.html") \ + else os.path.splitext(output_report)[0] + initial_archive = f"{stem}_report_initial.html" + shutil.copy2(init_report, initial_archive) + rerun_archive = "NA" + if reran: + rerun_archive = f"{stem}_report_rerun.html" + shutil.copy2(rr_report, rerun_archive) + + # --- per-sample status ------------------------------------------------- + fields = [ + ("sample", args.sample), + ("outcome", outcome), + ("adaptive", str(adaptive).lower()), + ("initial_learning_rate", format(initial_lr, ".10g")), + ("initial_verdict", initial_verdict), + ("reran", str(reran).lower()), + ("rerun_learning_rate", format(rerun_lr, ".10g") if rerun_lr is not None else "NA"), + ("rerun_verdict", rerun_verdict if rerun_verdict is not None else "NA"), + ("kept_run", kept_run), + ("report", output_report), + ("report_initial", initial_archive), + ("report_rerun", rerun_archive), + ] + with open(output_status, "w") as handle: + for key, value in fields: + handle.write(f"{key}\t{value}\n") + print(f"[cellbender-adaptive] {args.sample}: outcome={outcome}, kept={kept_run}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/workflow/scripts/cellbender_adaptive_summary.py b/workflow/scripts/cellbender_adaptive_summary.py new file mode 100644 index 0000000..f372095 --- /dev/null +++ b/workflow/scripts/cellbender_adaptive_summary.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python +"""Aggregate the per-sample CellBender adaptive-run status files into one run-level summary. + +Writes a TSV (one row per sample) and prints a short report: how many samples needed a re-run, +which were resolved by halving the learning rate, and which were not (kept the initial run).""" + +import argparse +from collections import Counter +from pathlib import Path + +COLUMNS = [ + "sample", "outcome", "initial_learning_rate", "initial_verdict", + "reran", "rerun_learning_rate", "rerun_verdict", "kept_run", +] + + +def read_status(path): + record = {} + for line in Path(path).read_text().splitlines(): + if "\t" in line: + key, value = line.split("\t", 1) + record[key] = value + return record + + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--output", required=True) + ap.add_argument("status_files", nargs="*") + args = ap.parse_args(argv) + + rows = sorted((read_status(p) for p in args.status_files), + key=lambda r: r.get("sample", "")) + + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + with open(args.output, "w") as handle: + handle.write("\t".join(COLUMNS) + "\n") + for row in rows: + handle.write("\t".join(row.get(col, "NA") for col in COLUMNS) + "\n") + + counts = Counter(row.get("outcome", "NA") for row in rows) + reran = [row["sample"] for row in rows if row.get("reran") == "true"] + unresolved = [row["sample"] for row in rows + if row.get("outcome") == "RERUN_DID_NOT_RESOLVE"] + + print(f"[cellbender-adaptive-summary] {len(rows)} sample(s):") + for outcome in ("NORMAL_FIRST_TRY", "RERUN_RESOLVED", "RERUN_DID_NOT_RESOLVE", + "RERUN_SUGGESTED_BUT_ADAPTIVE_OFF"): + if counts.get(outcome): + print(f" {outcome}: {counts[outcome]}") + if reran: + print(f" re-ran with halved learning rate: {', '.join(reran)}") + if unresolved: + print(f" re-run did NOT resolve the learning curve (kept initial): " + f"{', '.join(unresolved)}") + print(f"[cellbender-adaptive-summary] wrote {args.output}") + + +if __name__ == "__main__": + main()