diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4130f65 --- /dev/null +++ b/.gitignore @@ -0,0 +1,38 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +dist/ +build/ +.eggs/ +*.egg + +# Editable install +src/*.egg-info/ + +# Testing +.pytest_cache/ +.coverage +htmlcov/ + +# Type checking +.mypy_cache/ + +# Data and databases — large files not tracked (scripts are tracked) +data/**/*.tsv +data/**/*.csv +data/**/*.pdf +data/**/*.png +data/**/*.txt +data/**/gtdb_gather/ +data/**/gtdb_subdb/ +*.zip +*.sbt.zip +*.sig + +# macOS +.DS_Store + +# Editors +.vscode/ +.idea/ diff --git a/CLAUDE.md b/CLAUDE.md index b271e1f..05fdfe7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -39,7 +39,7 @@ Each stage can be run independently (step-by-step mode) or the full pipeline run ## Selector strategies -Six selection algorithms are implemented. All share the same interface (see `src/refrover/selectors/base.py`). +Seven selection algorithms are implemented. All share the same interface (see `src/refrover/selectors/base.py`): `select(matrix, query_id) -> list[str]`, returning prototype sample_ids. | Strategy ID | Class | Description | Mathematical basis | |-------------|-------|-------------|-------------------| @@ -48,7 +48,10 @@ Six selection algorithms are implemented. All share the same interface (see `src | `kmedoids` | `KMedoidsSelector` | k-medoids clustering in Jaccard space; medoids as prototypes | Minimizes within-cluster sum of distances | | `archetype` | `ArchetypeSelector` | Archetype analysis: finds assemblies on the convex hull of the similarity space | Cutler & Breiman (1994); extremes that span the space | | `greedy_var` | `GreedyVarSelector` | Greedy coverage-variance maximization: selects prototypes predicted to maximize cross-sample variance | Estimated via similarity proxy | -| `feedback` | `FeedbackSelector` | Coverage-feedback: map → measure actual variance → iterate | Empirical; most accurate, most expensive | +| `containment` | `ContainmentSelector` | MaxMin on cross-sample read **containment** profiles (not Jaccard); diversity in read-mapping space | Greedy distance maximization on 1 − Pearson(containment columns) | +| `feedback` | `FeedbackSelector` | Coverage-feedback: map → measure actual variance → iterate | Empirical; most accurate, most expensive (not yet implemented) | + +**Containment vs. Jaccard selectors**: the Jaccard selectors operate on an assembly-vs-assembly similarity matrix; `containment` operates on a reads-vs-assembly containment matrix (`containment[i,j]` = fraction of sample i's read k-mers in assembly j), which directly measures how well reads will map. It is a `BaseSelector` subclass but is fed a different matrix — pass `--containment-matrix` (CLI) or `containment_matrix=` (`RefRoverPipeline`). On the rewilded-mouse validation it was the recommended strategy at small k. The matrix is produced by `data/mouse_rewilding/compute_containment.py` (slated to move into the package as `refrover.containment`). **Archetype vs. k-medoids distinction**: k-medoids finds central representatives; archetypes find extreme points that together span the full diversity space. For differential coverage, span is more valuable than centrality — archetypes are the theoretically preferred approach. Benchmark both. @@ -93,7 +96,21 @@ refrover align --assignments assignments.tsv --manifest samples.tsv --outdir refrover coverage --bams bams/ --outdir raw_coverage/ refrover format --coverage raw_coverage/ --binners metabat2,semibin2 --outdir formatted/ -# Benchmark selectors against each other on a dataset with known ground truth +# Containment-based selection (computes the reads-vs-assemblies matrix, then selects) +refrover containment --read-sketches read_sigs/ --assembly-sketches sketches/ --outdir cont/ +refrover select --containment-matrix cont/containment_matrix.tsv \ + --manifest samples.tsv --selector containment --k 5 --min-jaccard 0.05 \ + --outdir assignments/ + +# Rank selectors cheaply (Tier 1): alignment-free variance-explained proxy +refrover rank-selectors \ + --containment-matrix cont/containment_matrix.tsv \ + --selectors random,maxmin,kmedoids,greedy_var,containment \ + --k-range 3,5,8,10 \ + --outdir ranking/ +# (or --sketches sketches/ to rank on the Jaccard matrix instead) + +# Benchmark selectors by MAG yield (Tier 2): full pipeline + CheckM2 (not yet implemented) refrover benchmark \ --manifest samples.tsv \ --truth community_truth.tsv \ @@ -103,6 +120,8 @@ refrover benchmark \ --outdir benchmark_results/ ``` +**Two-tier selector evaluation.** Tier 1 (`refrover rank-selectors`, `benchmark.rank_selectors`) ranks selectors in seconds with no alignment, by the fraction of total cross-sample variance the selected prototypes explain via orthogonal projection (correlated/redundant selections don't double-count). Tier 2 (`refrover benchmark`, not yet implemented) is the gold standard: CheckM2-passing MAGs per CPU-hour. The open research question is whether the Tier-1 ranking predicts the Tier-2 ranking — if so, selector choice can be made cheaply from Tier 1 alone. + ### Input manifest format Tab-separated, one row per sample. `long_reads` and `r2` are optional. @@ -175,7 +194,8 @@ refrover/ │ ├── pipeline.py # RefRoverPipeline orchestrator │ ├── io.py # manifest read/write, validation │ ├── sketch.py # sourmash wrapper (sketch + pairwise compare) -│ ├── similarity.py # similarity matrix operations, Jaccard filtering +│ ├── similarity.py # similarity matrix operations, Jaccard filtering, containment loader +│ ├── containment.py # reads-vs-assemblies containment matrix (feeds ContainmentSelector) │ ├── selectors/ │ │ ├── __init__.py # exports all selectors + SELECTOR_REGISTRY dict │ │ ├── base.py # BaseSelector ABC: select(similarity_matrix, query_id) -> list[str] diff --git a/data/mouse_rewilding/build_community_subdb.py b/data/mouse_rewilding/build_community_subdb.py new file mode 100644 index 0000000..c8d9b06 --- /dev/null +++ b/data/mouse_rewilding/build_community_subdb.py @@ -0,0 +1,203 @@ +""" +Build a community-specific sub-database for fast gather. + +Step 1: Run full gather on a handful of representative samples against GTDB RS226 + to identify all species present in this dataset. +Step 2: Extract those signatures from the full database into a small sub-database. +Step 3: All 183 samples can then be gathered against the sub-database in minutes. + +At ~26 min per sample for the full 143k-sig GTDB scan, 3 representative samples +takes ~80 min. The resulting sub-database (~2000-3000 sigs) reduces per-sample +gather time to ~seconds. Total time for 183 samples: ~1 hour vs. ~80 hours. +""" + +import subprocess +import sys +import time +from pathlib import Path +import pandas as pd + +DB_PATH = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/gtdb-reps-rs226-k31.dna.zip") +READS_DIR = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/sourmash_sketch_reads") +MANIFEST = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/sample_manifest.tsv") +OUT_DIR = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/gtdb_gather") +SUBDB_DIR = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/gtdb_subdb") +SUBDB_PATH = SUBDB_DIR / "community_species.zip" +THRESHOLD_BP = 5_000 + + +def select_representative_samples(manifest: pd.DataFrame, n: int = 5) -> list[str]: + """ + Pick representative samples that together span the dataset's diversity. + Strategy: one per trial, favouring POST timepoints (more rewilded signal). + """ + chosen = [] + for trial in manifest["Trial_ID"].unique(): + sub = manifest[manifest["Trial_ID"] == trial] + # prefer POST; fall back to any + post = sub[sub["Time_Point"] == "POST"] + pool = post if len(post) > 0 else sub + sid = pool.iloc[0]["Sample_ID"] + chosen.append(sid) + if len(chosen) >= n: + break + return chosen + + +def run_gather_full_db(sig_path: Path, out_csv: Path) -> int: + """Run gather against the full GTDB database. Returns match count.""" + if out_csv.exists(): + return sum(1 for _ in open(out_csv)) - 1 + cmd = [ + "sourmash", "gather", + str(sig_path), str(DB_PATH), + "-k", "31", + "--threshold-bp", str(THRESHOLD_BP), + "-o", str(out_csv), + "--quiet", + ] + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0 or not out_csv.exists(): + print(f" FAILED: {result.stderr[:200]}", file=sys.stderr) + return 0 + return sum(1 for _ in open(out_csv)) - 1 + + +def build_subdb(gather_csvs: list[Path]) -> Path: + """ + Extract all matched species from gather results and build a sub-database. + Returns path to the sub-database zip. + """ + SUBDB_DIR.mkdir(parents=True, exist_ok=True) + + # Collect all unique md5 hashes across gather results + all_md5s = set() + for csv_path in gather_csvs: + df = pd.read_csv(csv_path, usecols=["md5"]) + all_md5s.update(df["md5"].str[:8].tolist()) # sourmash uses 8-char prefix + + print(f" {len(all_md5s)} unique species signatures across all representative samples") + + # Write a picklist file for sourmash extract + picklist_path = SUBDB_DIR / "community_picklist.csv" + with open(picklist_path, "w") as f: + f.write("md5short\n") + for md5 in sorted(all_md5s): + f.write(md5 + "\n") + + # Extract matching signatures from the full database + cmd = [ + "sourmash", "sig", "extract", + str(DB_PATH), + "--picklist", f"{picklist_path}:md5short:md5short", + "-o", str(SUBDB_PATH), + "--quiet", + ] + print(f" Extracting {len(all_md5s)} signatures from GTDB...", flush=True) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"sourmash sig extract failed:\n{result.stderr}") + + size_mb = SUBDB_PATH.stat().st_size / 1e6 + print(f" Sub-database: {SUBDB_PATH} ({len(all_md5s)} signatures, {size_mb:.0f} MB)") + return SUBDB_PATH + + +def main(): + OUT_DIR.mkdir(parents=True, exist_ok=True) + manifest = pd.read_csv(MANIFEST, sep="\t") + + # ── Step 1: Full gather on representative samples ──────────────────────── + if SUBDB_PATH.exists(): + print(f"Sub-database already exists: {SUBDB_PATH}", flush=True) + else: + reps = select_representative_samples(manifest, n=5) + print(f"Representative samples for sub-database construction: {reps}", flush=True) + + gather_csvs = [] + t_total = time.time() + for i, sid in enumerate(reps, 1): + t0 = time.time() + sig_path = READS_DIR / f"{sid}.sig" + out_csv = OUT_DIR / f"{sid}_gather.csv" + print(f"[{i}/{len(reps)}] Gathering {sid} against full GTDB...", flush=True) + if not sig_path.exists(): + print(f" WARNING: sig not found for {sid}, skipping") + continue + n = run_gather_full_db(sig_path, out_csv) + elapsed = time.time() - t0 + print(f" {n} species in {elapsed/60:.1f} min", flush=True) + gather_csvs.append(out_csv) + + print(f"\nBuilding community sub-database...", flush=True) + build_subdb([p for p in gather_csvs if p.exists()]) + + # ── Step 2: Gather all samples against sub-database ───────────────────── + print(f"\nGathering all {len(manifest)} samples against sub-database...", flush=True) + sample_ids = manifest["Sample_ID"].tolist() + todo = [ + sid for sid in sample_ids + if not (OUT_DIR / f"{sid}_gather.csv").exists() + and (READS_DIR / f"{sid}.sig").exists() + ] + already_done = len(sample_ids) - len(todo) + print(f"{already_done} done, {len(todo)} to go", flush=True) + + t_total = time.time() + for i, sid in enumerate(todo, 1): + t0 = time.time() + sig_path = READS_DIR / f"{sid}.sig" + out_csv = OUT_DIR / f"{sid}_gather.csv" + cmd = [ + "sourmash", "gather", + str(sig_path), str(SUBDB_PATH), + "-k", "31", + "--threshold-bp", str(THRESHOLD_BP), + "-o", str(out_csv), + "--quiet", + ] + result = subprocess.run(cmd, capture_output=True, text=True) + n = sum(1 for _ in open(out_csv)) - 1 if out_csv.exists() else 0 + elapsed = time.time() - t0 + eta_min = (time.time() - t_total) / i * (len(todo) - i) / 60 + print(f"[{i}/{len(todo)}] {sid}: {n} matches, {elapsed:.0f}s (ETA {eta_min:.0f} min)", flush=True) + + # ── Step 3: Aggregate ──────────────────────────────────────────────────── + print("\nAggregating gather results...", flush=True) + records = [] + for sid in sample_ids: + csv_path = OUT_DIR / f"{sid}_gather.csv" + if not csv_path.exists(): + continue + df = pd.read_csv(csv_path, usecols=["name", "f_unique_to_query"]) + if df.empty: + continue + df["sample_id"] = sid + records.append(df[["sample_id", "name", "f_unique_to_query"]]) + + if not records: + print("No gather results found.") + return + + long_df = pd.concat(records, ignore_index=True) + wide_df = long_df.pivot_table( + index="name", columns="sample_id", values="f_unique_to_query", fill_value=0 + ) + wide_df = wide_df.reindex(columns=sample_ids, fill_value=0) + + out_wide = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/gtdb_species_matrix.tsv") + wide_df.to_csv(out_wide, sep="\t") + print(f"Species matrix: {wide_df.shape[0]:,} species x {wide_df.shape[1]} samples -> {out_wide}", flush=True) + + out_long = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/gtdb_gather_long.tsv") + long_df.to_csv(out_long, sep="\t", index=False) + print(f"Long format: {len(long_df):,} rows -> {out_long}", flush=True) + + prevalence = (wide_df > 0).sum(axis=1) + print(f"\nTop 10 most prevalent species:") + for name, prev in prevalence.nlargest(10).items(): + print(f" {prev:3d}/{len(sample_ids)} samples {name[:70]}") + + +if __name__ == "__main__": + main() diff --git a/data/mouse_rewilding/compare_selectors.py b/data/mouse_rewilding/compare_selectors.py new file mode 100644 index 0000000..4273069 --- /dev/null +++ b/data/mouse_rewilding/compare_selectors.py @@ -0,0 +1,339 @@ +""" +Compare prototype selector strategies on the cross-sample containment matrix. + +Selectors evaluated +------------------- + random Baseline: k assemblies drawn uniformly at random from candidates + top_containment Sample-centric: k assemblies with highest containment(query, assembly) + top_var Variance-centric: k assemblies with highest cross-sample containment variance + maxmin Diversity-centric: greedy MaxMin on containment-profile Pearson distances + +Metric: fraction of total cross-sample containment variance "explained" by selected +prototypes, measured as sum(var(selected columns)) / sum(var(all columns)). +Note: this sums individual variances without orthogonalising, so correlated selections +overcount shared variance. It is a conservative upper bound; correlated selectors will +look better than they are in practice. + +Adaptive k +---------- +For each sample, runs the MaxMin greedy selector while tracking the marginal +*residual* variance gain at each step (variance of the new column after regressing +out variance already explained by selected columns). The elbow of this gain curve +is the per-sample optimal k — the point of diminishing returns. + +Outputs +------- + selector_comparison.tsv long-format (sample, selector, k, variance_explained) + selector_agreement_k5.tsv pairwise mean Jaccard overlap of selections at k=5 + adaptive_k_ground_truth.tsv per-sample optimal k + metadata + saturation_curves.tsv per-sample per-step cumulative and marginal variance +""" + +import sys +import numpy as np +import pandas as pd +from pathlib import Path + +DATA_DIR = Path(__file__).parent +MIN_CONTAINMENT = 0.02 +K_VALUES = [3, 5, 8, 10] +MAX_K = 20 + +# ── Helpers ─────────────────────────────────────────────────────────────────── + +def residual_var(y: np.ndarray, X_cols: list[np.ndarray]) -> float: + """Variance of y unexplained by the columns in X_cols (OLS residual).""" + if not X_cols: + return float(np.var(y, ddof=1)) + X = np.column_stack(X_cols) + coef, _, _, _ = np.linalg.lstsq(X, y, rcond=None) + return float(np.var(y - X @ coef, ddof=1)) + + +def maxmin_greedy(dist_sub: np.ndarray, seed: int, k: int) -> list[int]: + """ + Greedy MaxMin on a precomputed distance matrix (local indices). + Returns list of local indices of length min(k, n). + """ + n = dist_sub.shape[0] + k = min(k, n) + selected = [seed] + remaining = [i for i in range(n) if i != seed] + if not remaining: + return selected + min_dist = dist_sub[seed, remaining].copy() + for _ in range(k - 1): + if not remaining: + break + best_local = int(np.argmax(min_dist)) + best = remaining[best_local] + selected.append(best) + remaining.pop(best_local) + min_dist = np.delete(min_dist, best_local) + if remaining: + min_dist = np.minimum(min_dist, dist_sub[best, remaining]) + return selected + + +# ── Selectors ───────────────────────────────────────────────────────────────── + +def select_random(k, nc, rng): + return list(rng.choice(nc, size=min(k, nc), replace=False)) + +def select_top_containment(query_sub, k): + return list(np.argsort(query_sub)[::-1][:k]) + +def select_top_var(col_vars_sub, k): + return list(np.argsort(col_vars_sub)[::-1][:k]) + +def select_maxmin(dist_sub, query_sub, k): + seed = int(np.argmax(query_sub)) # highest containment = most likely to map + return maxmin_greedy(dist_sub, seed, k) + + +# ── Saturation curve ────────────────────────────────────────────────────────── + +def saturation_curve(cont_sub: np.ndarray, col_vars_sub: np.ndarray, + total_var: float, max_k: int = MAX_K): + """ + Greedy forward selection by residual variance (NOT MaxMin order). + + At each step, add the assembly that contributes the most unique cross-sample + variance not already explained by the selected set. This guarantees + monotonically decreasing marginal gains — a prerequisite for elbow detection. + + Uses incremental Gram-Schmidt so cost is O(n_samp × n_cand × max_k). + + Returns + ------- + gains : ndarray, length <= max_k + gains[i] = marginal residual variance of the (i+1)-th selected assembly, + normalised by total_var. + cumvars : ndarray, same length + Cumulative sum of gains. + selected_order : list of local column indices in selection order. + """ + n_samp, n_cand = cont_sub.shape + + # Mean-centre each column (needed for correct Pearson-style projection) + cols = cont_sub - cont_sub.mean(axis=0) + + # Seed: highest individual variance + first = int(np.argmax(col_vars_sub)) + selected = [first] + remaining = list(range(n_cand)) + remaining.remove(first) + + # Initialise orthonormal basis Q for the column space of selected assemblies + v0 = cols[:, first] + norm0 = np.linalg.norm(v0) + Q = (v0 / norm0).reshape(-1, 1) if norm0 > 1e-12 else np.zeros((n_samp, 1)) + + gains = [col_vars_sub[first] / total_var] + cumvars = [gains[0]] + + limit = min(max_k - 1, len(remaining)) + for _ in range(limit): + if not remaining: + break + # Compute residuals for all remaining columns at once (vectorised) + Y = cols[:, remaining] # n_samp × n_remaining + proj = Q @ (Q.T @ Y) # n_samp × n_remaining + resids = Y - proj # n_samp × n_remaining + resid_vars = np.var(resids, axis=0, ddof=1) + + best_local = int(np.argmax(resid_vars)) + best = remaining[best_local] + gain = resid_vars[best_local] / total_var + gains.append(gain) + cumvars.append(cumvars[-1] + gain) + + # Extend Q with the new column (Gram-Schmidt) + v = cols[:, best] - Q @ (Q.T @ cols[:, best]) + norm = np.linalg.norm(v) + if norm > 1e-12: + Q = np.column_stack([Q, v / norm]) + + selected.append(best) + remaining.pop(best_local) + + return np.array(gains), np.array(cumvars), selected + + +def elbow_from_gains(gains: np.ndarray, saturation_fraction: float = 0.1, + k_min: int = 1) -> int: + """ + Threshold-based elbow: k = first step where marginal gain < fraction × first_gain. + + Works correctly on a monotonically decreasing gains curve. The largest-drop + heuristic fails when MaxMin selection order mixes high- and low-variance steps. + """ + if len(gains) < 2 or gains[0] == 0: + return max(k_min, 1) + threshold = saturation_fraction * gains[0] + idxs = np.where(gains < threshold)[0] + if len(idxs) == 0: + return max(k_min, len(gains)) + return max(k_min, int(idxs[0])) + + +# ── Main ────────────────────────────────────────────────────────────────────── + +def main(): + print("Loading data...", flush=True) + cont = pd.read_csv(DATA_DIR / "containment_wide.tsv", sep="\t", index_col=0) + manifest = pd.read_csv(DATA_DIR / "sample_manifest.tsv", sep="\t") + samples = cont.index.tolist() + assemblies = cont.columns.tolist() + print(f" {len(samples)} samples × {len(assemblies)} assemblies") + + col_vars = cont.var(axis=0) + total_var = float(col_vars.sum()) + asm_idx = {a: i for i, a in enumerate(assemblies)} + + print("Pre-computing containment-profile distance matrix...", flush=True) + corr_mat = cont.corr().values.astype(float) # assemblies × assemblies + np.fill_diagonal(corr_mat, 1.0) + dist_mat = np.clip(1.0 - corr_mat, 0.0, 2.0) + np.fill_diagonal(dist_mat, 0.0) + + # ── Saturation curves + adaptive k ──────────────────────────────────────── + print("Computing saturation curves per sample...", flush=True) + ak_rows = [] + curve_rows = [] + + for i, sid in enumerate(samples, 1): + row = cont.loc[sid] + cands = row[row >= MIN_CONTAINMENT].index.tolist() + if len(cands) < 2: + continue + + cidx = np.array([asm_idx[c] for c in cands]) + cont_sub = cont.values[:, cidx] + col_vars_sub = col_vars.values[cidx] + + gains, cumvars, sel_order = saturation_curve(cont_sub, col_vars_sub, total_var) + opt_k = elbow_from_gains(gains) + + ak_rows.append({ + "sample_id": sid, + "optimal_k": opt_k, + "n_candidates": len(cands), + }) + for step, (g, cv) in enumerate(zip(gains, cumvars), start=1): + curve_rows.append({ + "sample_id": sid, + "k": step, + "marginal_var": g, + "cumulative_var": cv, + }) + + if i % 30 == 0: + print(f" {i}/{len(samples)}", flush=True) + + ak_df = pd.DataFrame(ak_rows).merge( + manifest[["Sample_ID", "Mouse_Strain", "Time_Point", "Trial_ID"]], + left_on="sample_id", right_on="Sample_ID", how="left" + ) + ak_df.to_csv(DATA_DIR / "adaptive_k_ground_truth.tsv", sep="\t", index=False) + + curve_df = pd.DataFrame(curve_rows) + curve_df.to_csv(DATA_DIR / "saturation_curves.tsv", sep="\t", index=False) + + print(f" Adaptive k: median={ak_df['optimal_k'].median():.0f} " + f"mean={ak_df['optimal_k'].mean():.1f} " + f"range=[{ak_df['optimal_k'].min()}, {ak_df['optimal_k'].max()}]") + + # ── Selector comparison ──────────────────────────────────────────────────── + print("\nComputing variance explained per selector × k × sample...", flush=True) + SELECTORS = ["random", "top_containment", "top_var", "maxmin"] + comp_rows = [] + agree_sets: dict[tuple, dict] = {(s, k): {} for s in SELECTORS for k in K_VALUES} + + for i, sid in enumerate(samples, 1): + row = cont.loc[sid] + cands = row[row >= MIN_CONTAINMENT].index.tolist() + if not cands: + continue + nc = len(cands) + cidx = np.array([asm_idx[c] for c in cands]) + query_sub = row.values[cidx] + cv_sub = col_vars.values[cidx] + dist_sub = dist_mat[np.ix_(cidx, cidx)] + + for k in K_VALUES: + for sel in SELECTORS: + if sel == "random": + rng_i = np.random.default_rng(42 + abs(hash(sid)) % 10**6) + local_idx = select_random(k, nc, rng_i) + elif sel == "top_containment": + local_idx = select_top_containment(query_sub, k) + elif sel == "top_var": + local_idx = select_top_var(cv_sub, k) + else: # maxmin + local_idx = select_maxmin(dist_sub, query_sub, k) + + sel_names = [cands[j] for j in local_idx] + + # Simple metric: sum of individual column variances (overcounts correlated cols) + ve = col_vars[sel_names].sum() / total_var if total_var > 0 else 0.0 + + # Unique metric: fraction of ALL 182 column variances explained via + # orthogonal projection onto the selected columns (no double-counting). + # Computation: Q from QR of selected columns; project full matrix onto Q; + # captured = (total_var - sum residual variances) / total_var + sel_global = [asm_idx[n] for n in sel_names] + X = cont.values - cont.values.mean(axis=0) # centre all cols + Xp = X[:, sel_global] # n_samp × k + if Xp.shape[1] > 0: + Q, _ = np.linalg.qr(Xp, mode="reduced") + proj = Q @ (Q.T @ X) # n_samp × n_assemblies + resid_all = X - proj + resid_vars = np.var(resid_all, axis=0, ddof=1).sum() + uve = (total_var - resid_vars) / total_var if total_var > 0 else 0.0 + else: + uve = 0.0 + + comp_rows.append({ + "sample_id": sid, "selector": sel, "k": k, + "variance_explained": ve, + "unique_var_explained": uve, + "n_selected": len(sel_names), + "n_candidates": nc, + }) + agree_sets[(sel, k)][sid] = set(sel_names) + + if i % 30 == 0: + print(f" {i}/{len(samples)}", flush=True) + + comp_df = pd.DataFrame(comp_rows) + comp_df.to_csv(DATA_DIR / "selector_comparison.tsv", sep="\t", index=False) + + summary = (comp_df.groupby(["selector", "k"])[["variance_explained", "unique_var_explained"]] + .mean().round(4)) + print("\nMean variance explained by selector and k:") + print(summary.to_string()) + + # ── Pairwise agreement at k=5 ────────────────────────────────────────────── + print("\nComputing pairwise selector agreement (Jaccard) at k=5...", flush=True) + agree_df = pd.DataFrame(index=SELECTORS, columns=SELECTORS, dtype=float) + for s1 in SELECTORS: + for s2 in SELECTORS: + jacs = [] + for sid in samples: + a = agree_sets.get((s1, 5), {}).get(sid) + b = agree_sets.get((s2, 5), {}).get(sid) + if a is None or b is None or not (a | b): + continue + jacs.append(len(a & b) / len(a | b)) + agree_df.loc[s1, s2] = round(np.mean(jacs), 3) if jacs else 0.0 + + agree_df.to_csv(DATA_DIR / "selector_agreement_k5.tsv", sep="\t") + print("Selector agreement (mean Jaccard at k=5):") + print(agree_df.to_string()) + + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/data/mouse_rewilding/compute_containment.py b/data/mouse_rewilding/compute_containment.py new file mode 100644 index 0000000..ced99b3 --- /dev/null +++ b/data/mouse_rewilding/compute_containment.py @@ -0,0 +1,114 @@ +""" +Compute cross-sample containment matrix from sourmash .sig files. + +For each query (read sketch), computes the fraction of its k-mer hashes +present in each reference (assembly sketch): containment(reads_i, asm_j). + +Outputs a long-format TSV: query_id, reference_id, containment +""" + +import json +import os +import glob +import time +import numpy as np +import pandas as pd +from pathlib import Path +from multiprocessing import Pool, cpu_count + +READ_DIR = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/sourmash_sketch_reads") +ASM_DIR = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/sourmash_sketch_assemblies") +MANIFEST = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/sample_manifest.tsv") +OUT_LONG = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/containment_long.tsv") +OUT_WIDE = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/containment_wide.tsv") +N_WORKERS = max(1, cpu_count() - 2) + + +def load_sig(path: Path) -> np.ndarray: + with open(path) as f: + data = json.load(f) + mins = data[0]["signatures"][0]["mins"] + arr = np.array(mins, dtype=np.uint64) + arr.sort() + return arr + + +def load_all_assemblies(sample_ids: list[str]) -> dict[str, np.ndarray]: + asms = {} + for sid in sample_ids: + p = ASM_DIR / f"{sid}.sig" + if p.exists(): + asms[sid] = load_sig(p) + print(f"Loaded {len(asms)} assembly sketches", flush=True) + return asms + + +# Module-level global so forked workers share without pickling +_ASM_HASHES: dict[str, np.ndarray] = {} +_ASM_IDS: list[str] = [] + + +def _init_worker(asm_hashes, asm_ids): + global _ASM_HASHES, _ASM_IDS + _ASM_HASHES = asm_hashes + _ASM_IDS = asm_ids + + +def _compute_row(query_id: str) -> list[tuple[str, str, float]] | None: + read_sig = READ_DIR / f"{query_id}.sig" + if not read_sig.exists(): + print(f" no read sig: {query_id}", flush=True) + return None + + read_h = load_sig(read_sig) + n_read = len(read_h) + if n_read == 0: + return None + + rows = [] + for ref_id in _ASM_IDS: + asm_h = _ASM_HASHES[ref_id] + n_intersect = len(np.intersect1d(read_h, asm_h, assume_unique=True)) + rows.append((query_id, ref_id, n_intersect / n_read)) + return rows + + +def main(): + manifest = pd.read_csv(MANIFEST, sep="\t") + sample_ids = manifest["Sample_ID"].tolist() + + print(f"Loading {len(sample_ids)} assembly sketches...", flush=True) + t0 = time.time() + asm_hashes = load_all_assemblies(sample_ids) + asm_ids = sorted(asm_hashes.keys()) + print(f" done in {time.time()-t0:.1f}s — {len(asm_ids)} assemblies", flush=True) + + print(f"Computing containment matrix ({len(sample_ids)} queries × {len(asm_ids)} refs) " + f"using {N_WORKERS} workers...", flush=True) + t0 = time.time() + + with Pool( + processes=N_WORKERS, + initializer=_init_worker, + initargs=(asm_hashes, asm_ids), + ) as pool: + results = pool.map(_compute_row, sample_ids) + + elapsed = time.time() - t0 + print(f" done in {elapsed/60:.1f} min", flush=True) + + # Flatten to long format + records = [rec for rows in results if rows for rec in rows] + long_df = pd.DataFrame(records, columns=["query_id", "reference_id", "containment"]) + long_df.to_csv(OUT_LONG, sep="\t", index=False) + print(f"Long format: {len(long_df):,} rows → {OUT_LONG}", flush=True) + + # Pivot to wide matrix (rows=queries, cols=refs) + wide_df = long_df.pivot(index="query_id", columns="reference_id", values="containment") + wide_df = wide_df.reindex(index=sample_ids, columns=asm_ids) + wide_df.to_csv(OUT_WIDE, sep="\t") + print(f"Wide matrix: {wide_df.shape} → {OUT_WIDE}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/data/mouse_rewilding/compute_species_weights.py b/data/mouse_rewilding/compute_species_weights.py new file mode 100644 index 0000000..008fcd4 --- /dev/null +++ b/data/mouse_rewilding/compute_species_weights.py @@ -0,0 +1,186 @@ +""" +Compute per-assembly species-level weights for ContainmentSelector. + +After running explore_species.R (which produces gtdb_species_stats.tsv), +this script estimates how much differential coverage signal each assembly +would provide across samples. + +Weight derivation +----------------- +We don't have GTDB gather results for assemblies (only for reads). Instead +we use the cross-sample containment matrix as a proxy: + + weight_j = CV of containment(reads_i, assembly_j) across all samples i, + restricted to assemblies above a mean-containment floor. + +This captures: assemblies whose reads-mapping rate varies strongly across +samples are the most informative for differential coverage binning. + +The three-regime filter (from the R analysis) can optionally restrict to +"sweet-spot" species — mid-abundance, high-CV — and use the fraction of each +assembly's content coming from those species as an additional weight factor. + +Outputs +------- + assembly_weights.tsv : sample_id, weight, cv_containment, mean_containment +""" + +import argparse +import sys +from pathlib import Path +import numpy as np +import pandas as pd + +DATA_DIR = Path(__file__).parent +CONTAINMENT_WIDE = DATA_DIR / "containment_wide.tsv" +SPECIES_STATS = DATA_DIR / "gtdb_species_stats.tsv" +SPECIES_MATRIX = DATA_DIR / "gtdb_species_matrix.tsv" +OUT_WEIGHTS = DATA_DIR / "assembly_weights.tsv" + + +def containment_based_weights( + containment_wide: pd.DataFrame, + min_mean_containment: float = 0.005, +) -> pd.Series: + """ + Per-assembly weight = CV of containment across all samples. + + Assemblies that are consistently high (stable core) or consistently + low (rare / absent everywhere) get low weight. Assemblies with high + variance in how well reads map to them are the most informative. + """ + mean_cont = containment_wide.mean(axis=0) + std_cont = containment_wide.std(axis=0) + cv_cont = std_cont / (mean_cont + 1e-9) + + # Exclude assemblies too rarely detected to carry differential signal + eligible = mean_cont >= min_mean_containment + weights = cv_cont.where(eligible, other=0.0) + return weights + + +def species_regime_weights( + species_stats: pd.DataFrame, + species_matrix: pd.DataFrame, + containment_wide: pd.DataFrame, + low_pct: float = 33.0, + high_pct: float = 67.0, +) -> pd.Series: + """ + Weight assemblies by the fraction of their content from 'sweet-spot' species: + mid-abundance, high-CV taxa that vary most across samples. + + Because we don't have assembly-level GTDB gather, we approximate the + species composition of each assembly by: for each sample (assembly), + use that sample's own GTDB gather result to identify its dominant species. + + weight_j = sum over species s detected in sample j of: + CV(s) × f_unique_to_query(s, j) + ...restricted to mid-abundance species + """ + if species_stats is None or species_matrix is None: + print(" Species stats not available; using containment-only weights", file=sys.stderr) + return pd.Series(dtype=float) + + # Identify mid-abundance species (the "sweet spot") + log_mean = np.log10(species_stats["mean_abund"] + 1e-9) + lo = np.percentile(log_mean, low_pct) + hi = np.percentile(log_mean, high_pct) + sweetspot = species_stats[ + (log_mean >= lo) & (log_mean <= hi) + ].set_index("name") + + if sweetspot.empty: + return pd.Series(dtype=float) + + # For each assembly (column in species_matrix), compute weighted CV score + assembly_ids = containment_wide.columns.tolist() + sample_ids = species_matrix.columns.tolist() + + weights = {} + for asm_id in assembly_ids: + if asm_id not in sample_ids: + weights[asm_id] = 0.0 + continue + + # Species abundances in this sample + abund = species_matrix.loc[ + species_matrix.index.intersection(sweetspot.index), asm_id + ] + if abund.empty or abund.sum() == 0: + weights[asm_id] = 0.0 + continue + + # Weight = sum of (abundance × CV) for sweet-spot species + cv_vals = sweetspot.loc[abund.index, "cv"] + weights[asm_id] = float((abund * cv_vals).sum()) + + return pd.Series(weights, dtype=float) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--use-species", action="store_true", + help="Incorporate species-level sweet-spot weighting (requires gtdb outputs)") + parser.add_argument("--min-mean-containment", type=float, default=0.005) + parser.add_argument("--low-pct", type=float, default=33.0, + help="Lower percentile for mid-abundance species window") + parser.add_argument("--high-pct", type=float, default=67.0, + help="Upper percentile for mid-abundance species window") + args = parser.parse_args() + + # ── Load data ───────────────────────────────────────────────────────────── + if not CONTAINMENT_WIDE.exists(): + sys.exit(f"ERROR: {CONTAINMENT_WIDE} not found. Run compute_containment.py first.") + + print(f"Loading containment matrix...", flush=True) + cont = pd.read_csv(CONTAINMENT_WIDE, sep="\t", index_col=0) + print(f" {cont.shape[0]} samples x {cont.shape[1]} assemblies") + + # ── Compute weights ─────────────────────────────────────────────────────── + print("Computing containment-based weights...", flush=True) + w_cont = containment_based_weights(cont, args.min_mean_containment) + + if args.use_species and SPECIES_STATS.exists() and SPECIES_MATRIX.exists(): + print("Computing species-regime weights...", flush=True) + sp_stats = pd.read_csv(SPECIES_STATS, sep="\t") + sp_matrix = pd.read_csv(SPECIES_MATRIX, sep="\t", index_col=0) + w_species = species_regime_weights( + sp_stats, sp_matrix, cont, + low_pct=args.low_pct, high_pct=args.high_pct, + ) + # Normalise each weight vector to [0, 1] then take geometric mean + def _norm(s): + r = s - s.min() + return r / (r.max() + 1e-9) + w_combined = (_norm(w_cont) * _norm(w_species.reindex(w_cont.index, fill_value=0))) ** 0.5 + label = "combined (containment × species)" + else: + w_combined = w_cont + label = "containment CV only" + + # ── Summary and output ──────────────────────────────────────────────────── + mean_cont = cont.mean(axis=0) + std_cont = cont.std(axis=0) + cv_cont = std_cont / (mean_cont + 1e-9) + + result = pd.DataFrame({ + "sample_id": w_combined.index, + "weight": w_combined.values, + "cv_containment": cv_cont.reindex(w_combined.index).values, + "mean_containment": mean_cont.reindex(w_combined.index).values, + }).sort_values("weight", ascending=False) + + result.to_csv(OUT_WEIGHTS, sep="\t", index=False) + print(f"\nWeights ({label}) -> {OUT_WEIGHTS}") + + print(f"\nTop 10 highest-weight assemblies:") + print(result.head(10).to_string(index=False)) + + print(f"\nBottom 10 (lowest signal):") + print(result.tail(10).to_string(index=False)) + + +if __name__ == "__main__": + main() diff --git a/data/mouse_rewilding/explore_containment.R b/data/mouse_rewilding/explore_containment.R new file mode 100644 index 0000000..e83aa08 --- /dev/null +++ b/data/mouse_rewilding/explore_containment.R @@ -0,0 +1,160 @@ +library(tidyverse) +library(pheatmap) +library(RColorBrewer) + +DATA_DIR <- "~/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding" + +long <- read_tsv(file.path(DATA_DIR, "containment_long.tsv")) +wide <- read_tsv(file.path(DATA_DIR, "containment_wide.tsv")) +manifest <- read_tsv(file.path(DATA_DIR, "sample_manifest.tsv")) + +# ── 1. Heatmap ──────────────────────────────────────────────────────────────── + +# Order samples: Trial → Timepoint → Subject +sample_order <- manifest |> + arrange(Trial_ID, Time_Point, Subject_ID) |> + pull(Sample_ID) + +# Wide matrix: rows = queries (reads), cols = references (assemblies) +mat <- wide |> + column_to_rownames("query_id") |> + as.matrix() + +# Reorder to match sample_order (some cols may be missing — m1082_POST assembly absent) +row_order <- intersect(sample_order, rownames(mat)) +col_order <- intersect(sample_order, colnames(mat)) +mat <- mat[row_order, col_order] + +# Annotation bars for heatmap +ann_row <- manifest |> + filter(Sample_ID %in% row_order) |> + arrange(match(Sample_ID, row_order)) |> + select(Sample_ID, Trial_ID, Time_Point, Mouse_Strain) |> + column_to_rownames("Sample_ID") + +ann_col <- manifest |> + filter(Sample_ID %in% col_order) |> + arrange(match(Sample_ID, col_order)) |> + select(Sample_ID, Trial_ID, Time_Point) |> + column_to_rownames("Sample_ID") + +ann_colors <- list( + Trial_ID = setNames(brewer.pal(5, "Set1"), paste0("T00", 1:5)), + Time_Point = c(PRE = "#4393C3", POST = "#D6604D"), + Mouse_Strain = c(C57 = "#984EA3", NYOB = "#FF7F00") +) + +pdf(file.path(DATA_DIR, "heatmap_containment.pdf"), width = 14, height = 12) +pheatmap( + mat, + cluster_rows = FALSE, + cluster_cols = FALSE, + annotation_row = ann_row, + annotation_col = ann_col, + annotation_colors = ann_colors, + color = colorRampPalette(c("white", "#2166AC"))(100), + breaks = seq(0, 0.40, length.out = 101), + show_rownames = FALSE, + show_colnames = FALSE, + main = "Cross-sample containment: reads_i ∩ assembly_j / reads_i", + fontsize = 8 +) +dev.off() +message("Saved: heatmap_containment.pdf") + +# ── 2. μ and CV per reference assembly (the w_j signal) ────────────────────── + +col_stats <- long |> + group_by(reference_id) |> + summarise( + mean_c = mean(containment), + sd_c = sd(containment), + cv = sd_c / mean_c, + n = n(), + .groups = "drop" + ) |> + left_join(manifest |> select(Sample_ID, Trial_ID, Time_Point, Mouse_Strain), + by = c("reference_id" = "Sample_ID")) + +# Three-regime scatter +p_regimes <- ggplot(col_stats, aes(x = mean_c, y = cv, color = Time_Point, shape = Mouse_Strain)) + + geom_point(size = 2.5, alpha = 0.8) + + geom_vline(xintercept = c(0.05, 0.20), linetype = "dashed", color = "grey60") + + annotate("text", x = 0.025, y = max(col_stats$cv, na.rm=TRUE)*0.95, + label = "sparse\n(hard to recover)", hjust = 0.5, size = 3, color = "grey40") + + annotate("text", x = 0.125, y = max(col_stats$cv, na.rm=TRUE)*0.95, + label = "mid-abundance\nhigh variance\n← target", hjust = 0.5, size = 3, color = "#E41A1C") + + annotate("text", x = 0.30, y = max(col_stats$cv, na.rm=TRUE)*0.95, + label = "abundant / stable\n(recovers anyway)", hjust = 0.5, size = 3, color = "grey40") + + scale_color_manual(values = c(PRE = "#4393C3", POST = "#D6604D")) + + labs( + x = "Mean containment (μ_j)", + y = "CV of containment (CV_j = σ/μ)", + title = "Per-assembly selection value signal", + subtitle = "Mid-μ / high-CV assemblies are where prototype selection matters most" + ) + + theme_bw(base_size = 12) + +ggsave(file.path(DATA_DIR, "regimes_scatter.pdf"), p_regimes, width = 8, height = 6) +message("Saved: regimes_scatter.pdf") + +# Facet by trial to see if the regime distribution shifts across trials +p_by_trial <- p_regimes + + facet_wrap(~Trial_ID, nrow = 2) + + theme(legend.position = "bottom") +ggsave(file.path(DATA_DIR, "regimes_scatter_by_trial.pdf"), p_by_trial, width = 12, height = 8) +message("Saved: regimes_scatter_by_trial.pdf") + +# ── 3. Within-mouse vs. cross-mouse hierarchy ───────────────────────────────── + +meta_key <- manifest |> select(Sample_ID, Subject_ID, Time_Point, Trial_ID) + +pairwise <- long |> + left_join(meta_key, by = c("query_id" = "Sample_ID")) |> + rename(q_subject = Subject_ID, q_tp = Time_Point, q_trial = Trial_ID) |> + left_join(meta_key, by = c("reference_id" = "Sample_ID")) |> + rename(r_subject = Subject_ID, r_tp = Time_Point, r_trial = Trial_ID) |> + filter(query_id != reference_id) |> + mutate( + pair_type = case_when( + q_subject == r_subject ~ "Within-mouse", + q_trial == r_trial & q_subject != r_subject ~ "Cross-mouse,\nsame trial", + TRUE ~ "Cross-trial" + ) |> factor(levels = c("Within-mouse", "Cross-mouse,\nsame trial", "Cross-trial")), + query_label = paste0(q_tp, " reads"), + ref_label = paste0(r_tp, " assembly") + ) + +p_hierarchy <- pairwise |> + filter(q_tp == "PRE", r_tp == "POST") |> # most interpretable contrast + ggplot(aes(x = pair_type, y = containment, fill = pair_type)) + + geom_violin(alpha = 0.7, draw_quantiles = 0.5) + + geom_jitter(width = 0.15, size = 0.4, alpha = 0.3) + + scale_fill_brewer(palette = "Set2") + + labs( + x = NULL, + y = "Containment (PRE reads → POST assembly)", + title = "Containment hierarchy: within-mouse > same-trial > cross-trial", + subtitle = "Validates that the matrix captures biological structure" + ) + + theme_bw(base_size = 12) + + theme(legend.position = "none") + +ggsave(file.path(DATA_DIR, "hierarchy_violin.pdf"), p_hierarchy, width = 7, height = 5) +message("Saved: hierarchy_violin.pdf") + +# ── 4. Summary table ───────────────────────────────────────────────────────── + +summary_tbl <- col_stats |> + arrange(desc(cv)) |> + select(reference_id, Trial_ID, Time_Point, Mouse_Strain, mean_c, cv) |> + rename(mean_containment = mean_c) + +write_tsv(summary_tbl, file.path(DATA_DIR, "assembly_selection_values.tsv")) +message("Saved: assembly_selection_values.tsv") + +cat("\n--- Top 10 highest-CV assemblies (best prototype candidates) ---\n") +print(head(summary_tbl, 10), n = 10) + +cat("\n--- Bottom 10 lowest-CV assemblies (least discriminating) ---\n") +print(tail(summary_tbl, 10), n = 10) diff --git a/data/mouse_rewilding/explore_selectors.R b/data/mouse_rewilding/explore_selectors.R new file mode 100644 index 0000000..187a69f --- /dev/null +++ b/data/mouse_rewilding/explore_selectors.R @@ -0,0 +1,178 @@ +library(tidyverse) +library(ggplot2) + +SELECTOR_LABELS <- c( + random = "Random", + top_containment = "Top containment", + top_var = "Top variance", + maxmin = "MaxMin diversity" +) +SELECTOR_COLORS <- c( + random = "#999999", + top_containment = "#E69F00", + top_var = "#56B4E9", + maxmin = "#009E73" +) + +dir.create("plots", showWarnings = FALSE) + +# ── 1. Variance-explained curves ────────────────────────────────────────────── + +comp <- read_tsv("selector_comparison.tsv", show_col_types = FALSE) |> + mutate(selector = factor(selector, levels = names(SELECTOR_LABELS))) + +comp_long <- comp |> + pivot_longer(c(variance_explained, unique_var_explained), + names_to = "metric", values_to = "value") |> + mutate(metric = recode(metric, + variance_explained = "Sum of variances\n(overcounts correlated selections)", + unique_var_explained = "Unique variance\n(orthogonal projection onto selected set)" + )) + +comp_summary <- comp_long |> + group_by(selector, k, metric) |> + summarise( + mean_ve = mean(value), + se_ve = sd(value) / sqrt(n()), + .groups = "drop" + ) + +p_curves <- ggplot(comp_summary, aes(x = k, y = mean_ve, color = selector)) + + geom_ribbon(aes(ymin = mean_ve - se_ve, ymax = mean_ve + se_ve, fill = selector), + alpha = 0.12, color = NA) + + geom_line(linewidth = 1.1) + + geom_point(size = 2.5) + + facet_wrap(~metric, scales = "free_y") + + scale_color_manual(values = SELECTOR_COLORS, labels = SELECTOR_LABELS) + + scale_fill_manual(values = SELECTOR_COLORS, labels = SELECTOR_LABELS) + + scale_x_continuous(breaks = c(3, 5, 8, 10)) + + scale_y_continuous(labels = scales::percent_format(accuracy = 1)) + + labs( + title = "Variance explained by selector and k", + subtitle = "Left: naive metric (overcounts); Right: orthogonal projection (correct)", + x = "k (number of prototypes per sample)", + y = "Variance explained (mean +/- SE)", + color = "Selector", fill = "Selector" + ) + + theme_bw(base_size = 11) + + theme(legend.position = "right", + strip.text = element_text(size = 9)) + +ggsave("plots/selector_variance_curves.pdf", p_curves, width = 10, height = 5) +ggsave("plots/selector_variance_curves.png", p_curves, width = 10, height = 5, dpi = 150) +cat("Saved plots/selector_variance_curves.pdf\n") + +# ── 2. Per-sample distribution of variance explained ───────────────────────── + +p_dist <- comp |> + filter(k == 5) |> + mutate(selector = fct_reorder(selector, unique_var_explained, .fun = median)) |> + ggplot(aes(x = selector, y = unique_var_explained, fill = selector)) + + geom_violin(trim = TRUE, scale = "width", alpha = 0.8) + + geom_boxplot(width = 0.15, fill = "white", outlier.size = 0.8, outlier.alpha = 0.5) + + scale_fill_manual(values = SELECTOR_COLORS, labels = SELECTOR_LABELS) + + scale_x_discrete(labels = SELECTOR_LABELS) + + scale_y_continuous(labels = scales::percent_format(accuracy = 1)) + + labs( + title = "Per-sample unique variance explained at k=5", + subtitle = "Orthogonal projection metric (removes correlated redundancy)", + x = NULL, y = "Unique variance explained" + ) + + theme_bw(base_size = 12) + + theme(legend.position = "none", axis.text.x = element_text(angle = 20, hjust = 1)) + +ggsave("plots/selector_variance_dist.pdf", p_dist, width = 6, height = 5) +ggsave("plots/selector_variance_dist.png", p_dist, width = 6, height = 5, dpi = 150) +cat("Saved plots/selector_variance_dist.pdf\n") + +# ── 3. Selector agreement heatmap at k=5 ───────────────────────────────────── + +agree_raw <- read_tsv("selector_agreement_k5.tsv", show_col_types = FALSE) +agree <- agree_raw |> + rename(s1 = 1) |> + pivot_longer(-s1, names_to = "s2", values_to = "jaccard") |> + mutate( + s1 = factor(s1, levels = names(SELECTOR_LABELS), labels = SELECTOR_LABELS), + s2 = factor(s2, levels = names(SELECTOR_LABELS), labels = SELECTOR_LABELS), + ) + +p_agree <- ggplot(agree, aes(x = s2, y = s1, fill = jaccard)) + + geom_tile(color = "white") + + geom_text(aes(label = sprintf("%.2f", jaccard)), size = 4) + + scale_fill_gradient(low = "#f7fbff", high = "#08519c", name = "Mean\nJaccard", + limits = c(0, 1)) + + labs( + title = "Selector agreement at k=5", + subtitle = "Mean pairwise Jaccard overlap of selected prototype sets across all samples", + x = NULL, y = NULL + ) + + coord_fixed() + + theme_bw(base_size = 11) + + theme(axis.text.x = element_text(angle = 25, hjust = 1)) + +ggsave("plots/selector_agreement_heatmap.pdf", p_agree, width = 6, height = 5) +ggsave("plots/selector_agreement_heatmap.png", p_agree, width = 6, height = 5, dpi = 150) +cat("Saved plots/selector_agreement_heatmap.pdf\n") + +# ── 4. Adaptive k distribution ──────────────────────────────────────────────── + +ak <- read_tsv("adaptive_k_ground_truth.tsv", show_col_types = FALSE) + +cat("\nAdaptive k summary:\n") +print(ak |> count(optimal_k) |> arrange(optimal_k)) + +cat("\nAdaptive k by strain and timepoint:\n") +print(ak |> group_by(Mouse_Strain, Time_Point) |> + summarise(median_k = median(optimal_k), mean_k = mean(optimal_k), .groups = "drop")) + +p_ak <- ak |> + filter(!is.na(Mouse_Strain)) |> + ggplot(aes(x = factor(optimal_k), fill = Time_Point)) + + geom_bar(position = "dodge") + + facet_wrap(~Mouse_Strain) + + scale_fill_manual(values = c(PRE = "#6baed6", POST = "#2ca25f")) + + labs( + title = "Per-sample optimal k (from containment saturation curve)", + subtitle = "Elbow of MaxMin residual-variance gain curve", + x = "Optimal k", y = "Number of samples", + fill = "Timepoint" + ) + + theme_bw(base_size = 11) + +ggsave("plots/adaptive_k_distribution.pdf", p_ak, width = 7, height = 4) +ggsave("plots/adaptive_k_distribution.png", p_ak, width = 7, height = 4, dpi = 150) +cat("Saved plots/adaptive_k_distribution.pdf\n") + +# ── 5. Saturation curves for representative samples ─────────────────────────── + +curves <- read_tsv("saturation_curves.tsv", show_col_types = FALSE) +manifest <- read_tsv("sample_manifest.tsv", show_col_types = FALSE) + +# Pick 3 representative sample IDs (one per strain/timepoint combo) +rep_samples <- manifest |> + group_by(Mouse_Strain, Time_Point) |> + slice_head(n = 1) |> + pull(Sample_ID) |> + intersect(unique(curves$sample_id)) |> + head(6) + +p_sat <- curves |> + filter(sample_id %in% rep_samples) |> + left_join(select(manifest, sample_id = Sample_ID, Mouse_Strain, Time_Point), by = "sample_id") |> + mutate(label = paste0(Mouse_Strain, " ", Time_Point, "\n(", sample_id, ")")) |> + ggplot(aes(x = k)) + + geom_col(aes(y = marginal_var), fill = "#6baed6", alpha = 0.7) + + geom_line(aes(y = cumulative_var), color = "#08519c", linewidth = 1) + + facet_wrap(~label, ncol = 3) + + scale_y_continuous(labels = scales::percent_format(accuracy = 0.1)) + + labs( + title = "Saturation curves for representative samples", + subtitle = "Bars = marginal residual variance; line = cumulative variance explained", + x = "k (prototype number, MaxMin order)", + y = "Fraction of total cross-sample variance" + ) + + theme_bw(base_size = 10) + +ggsave("plots/saturation_curves.pdf", p_sat, width = 9, height = 5) +ggsave("plots/saturation_curves.png", p_sat, width = 9, height = 5, dpi = 150) +cat("Saved plots/saturation_curves.pdf\n") diff --git a/data/mouse_rewilding/explore_species.R b/data/mouse_rewilding/explore_species.R new file mode 100644 index 0000000..8964811 --- /dev/null +++ b/data/mouse_rewilding/explore_species.R @@ -0,0 +1,184 @@ +library(tidyverse) +library(ggplot2) + +# ── Load data ────────────────────────────────────────────────────────────────── + +species_mat <- read_tsv("gtdb_species_matrix.tsv") |> + column_to_rownames("name") + +manifest <- read_tsv("sample_manifest.tsv") + +cat("Species matrix:", nrow(species_mat), "species x", ncol(species_mat), "samples\n") +cat("Non-zero entries:", sum(species_mat > 0), "\n") +cat("Sparsity:", round(mean(species_mat == 0) * 100, 1), "%\n\n") + +# ── Per-species summary stats ────────────────────────────────────────────────── + +species_stats <- tibble( + name = rownames(species_mat), + mean_abund = rowMeans(species_mat), + sd_abund = apply(species_mat, 1, sd), + prevalence = rowSums(species_mat > 0), # number of samples detected in + max_abund = apply(species_mat, 1, max), +) |> + mutate( + cv = sd_abund / (mean_abund + 1e-9), # coefficient of variation + log10_mean = log10(mean_abund + 1e-9), + ) |> + arrange(desc(mean_abund)) + +cat("Top 20 species by mean abundance:\n") +print(head(select(species_stats, name, mean_abund, cv, prevalence), 20)) + +# ── CV vs mean scatter (the key diagnostic) ────────────────────────────────── +# Three-regime hypothesis: +# Low abundance / high CV -> rare / intermittent taxa (noisy, little differential signal) +# Mid abundance / high CV -> core taxa varying between treatments (HIGH VALUE for selection) +# High abundance / low CV -> stable core (always present, constant depth) + +# Restrict to species detected in at least 5 samples (filter extreme sparsity) +detected <- filter(species_stats, prevalence >= 5) +cat("\nSpecies in >=5 samples:", nrow(detected), "\n") + +p_cv_mean <- ggplot(detected, aes(x = log10_mean, y = cv)) + + geom_point(aes(color = prevalence, size = prevalence), alpha = 0.6) + + scale_color_viridis_c(name = "Prevalence\n(n samples)", option = "plasma") + + scale_size_continuous(name = "Prevalence", range = c(0.5, 3)) + + geom_smooth(method = "loess", se = TRUE, color = "black", linewidth = 0.8) + + labs( + title = "Species-level abundance: CV vs mean", + subtitle = paste0(nrow(detected), " species detected in >=5 samples"), + x = "log10(mean f_unique_to_query)", + y = "Coefficient of variation" + ) + + theme_bw(base_size = 11) + +ggsave("plots/species_cv_vs_mean.pdf", p_cv_mean, width = 8, height = 6) +ggsave("plots/species_cv_vs_mean.png", p_cv_mean, width = 8, height = 6, dpi = 150) +cat("Saved plots/species_cv_vs_mean.pdf\n") + +# ── SD vs mean (log-log) ────────────────────────────────────────────────────── +# At assembly level, SD scaled linearly with mean. Does this persist at species level? + +p_sd_mean <- ggplot(detected, aes(x = log10_mean, y = log10(sd_abund + 1e-9))) + + geom_point(aes(color = prevalence), alpha = 0.5, size = 0.8) + + scale_color_viridis_c(option = "plasma") + + geom_smooth(method = "lm", se = TRUE, color = "firebrick") + + labs( + title = "SD vs mean (log-log) — species level", + subtitle = "Linear relationship indicates proportional variance; deviation = differential signal", + x = "log10(mean)", y = "log10(SD)" + ) + + theme_bw(base_size = 11) + +ggsave("plots/species_sd_vs_mean.pdf", p_sd_mean, width = 7, height = 5) +cat("Saved plots/species_sd_vs_mean.pdf\n") + +# ── Strain comparison: C57 vs NYOB ─────────────────────────────────────────── + +c57_ids <- manifest |> filter(Mouse_Strain == "C57") |> pull(Sample_ID) +nyob_ids <- manifest |> filter(Mouse_Strain == "NYOB") |> pull(Sample_ID) + +c57_mat <- species_mat[, intersect(colnames(species_mat), c57_ids), drop = FALSE] +nyob_mat <- species_mat[, intersect(colnames(species_mat), nyob_ids), drop = FALSE] + +strain_stats <- tibble( + name = rownames(species_mat), + mean_c57 = rowMeans(c57_mat), + mean_nyob = rowMeans(nyob_mat), + cv_c57 = apply(c57_mat, 1, sd) / (rowMeans(c57_mat) + 1e-9), + cv_nyob = apply(nyob_mat, 1, sd) / (rowMeans(nyob_mat) + 1e-9), + prev_c57 = rowSums(c57_mat > 0), + prev_nyob = rowSums(nyob_mat > 0), +) |> + mutate( + log10_mean_c57 = log10(mean_c57 + 1e-9), + log10_mean_nyob = log10(mean_nyob + 1e-9), + ) + +# Scatter: C57 mean vs NYOB mean (species consistently present in both vs strain-specific) +p_strain <- strain_stats |> + filter(prev_c57 >= 3 | prev_nyob >= 3) |> + ggplot(aes(x = log10_mean_c57, y = log10_mean_nyob)) + + geom_point(alpha = 0.4, size = 0.8, color = "steelblue") + + geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "grey40") + + labs( + title = "Per-species mean abundance: C57 vs NYOB", + subtitle = "Points above diagonal = higher in NYOB; below = higher in C57", + x = "log10(mean) in C57 samples", + y = "log10(mean) in NYOB samples" + ) + + theme_bw(base_size = 11) + +ggsave("plots/species_c57_vs_nyob.pdf", p_strain, width = 6, height = 6) +cat("Saved plots/species_c57_vs_nyob.pdf\n") + +# ── PRE vs POST: within-mouse variance ──────────────────────────────────────── + +pre_ids <- manifest |> filter(Time_Point == "PRE") |> pull(Sample_ID) +post_ids <- manifest |> filter(Time_Point == "POST") |> pull(Sample_ID) + +pre_mat <- species_mat[, intersect(colnames(species_mat), pre_ids), drop = FALSE] +post_mat <- species_mat[, intersect(colnames(species_mat), post_ids), drop = FALSE] + +timepoint_stats <- tibble( + name = rownames(species_mat), + mean_pre = rowMeans(pre_mat), + mean_post = rowMeans(post_mat), + cv_pre = apply(pre_mat, 1, sd) / (rowMeans(pre_mat) + 1e-9), + cv_post = apply(post_mat, 1, sd) / (rowMeans(post_mat) + 1e-9), +) |> + mutate( + log2fc = log2((mean_post + 1e-9) / (mean_pre + 1e-9)) + ) |> + filter(rowMeans(pre_mat) > 1e-6 | rowMeans(post_mat) > 1e-6) + +# Volcano-style: log2FC (POST/PRE) vs mean abundance +p_volcano <- timepoint_stats |> + ggplot(aes(x = log2fc, y = log10((mean_pre + mean_post)/2 + 1e-9))) + + geom_point(aes(color = abs(log2fc) > 1), alpha = 0.5, size = 0.8) + + scale_color_manual(values = c("grey60", "firebrick"), guide = "none") + + geom_vline(xintercept = c(-1, 1), linetype = "dashed", color = "grey40") + + labs( + title = "Species-level PRE vs POST shift", + subtitle = "Red = |log2FC| > 1 (more than 2x change after rewilding)", + x = "log2(POST / PRE mean abundance)", + y = "log10(mean abundance)" + ) + + theme_bw(base_size = 11) + +ggsave("plots/species_pre_vs_post_volcano.pdf", p_volcano, width = 7, height = 5) +cat("Saved plots/species_pre_vs_post_volcano.pdf\n") + +# ── Summary: how many species fall into each regime? ───────────────────────── + +# Define regimes on detected species (prev >= 5) +det <- filter(species_stats, prevalence >= 5) + +# Use log10 mean quantiles to define abundance regimes +low_thresh <- quantile(det$log10_mean, 0.33) +high_thresh <- quantile(det$log10_mean, 0.67) + +regime_stats <- det |> + mutate(regime = case_when( + log10_mean < low_thresh ~ "low_abund", + log10_mean > high_thresh ~ "high_abund", + TRUE ~ "mid_abund" + )) |> + group_by(regime) |> + summarise( + n = n(), + mean_cv = mean(cv), + median_cv = median(cv), + .groups = "drop" + ) + +cat("\nSpecies by abundance regime (detected in >=5 samples):\n") +print(regime_stats) +cat("\nThree-regime hypothesis: mid-abundance species should have HIGHEST CV\n") +cat("(they vary most across samples -> most useful for differential coverage)\n") + +# ── Save summary table ──────────────────────────────────────────────────────── + +write_tsv(species_stats, "gtdb_species_stats.tsv") +cat("\nFull species stats -> gtdb_species_stats.tsv\n") diff --git a/data/mouse_rewilding/run_gtdb_gather.py b/data/mouse_rewilding/run_gtdb_gather.py new file mode 100644 index 0000000..3ed824a --- /dev/null +++ b/data/mouse_rewilding/run_gtdb_gather.py @@ -0,0 +1,138 @@ +""" +Run sourmash gather for all samples against the GTDB RS226 database. + +Uses the sourmash CLI (subprocess) per sample — the CLI's Rust backend +is ~3-5x faster than the Python API for the linear prefetch scan. + +At ~26 min/sample on an M-chip laptop, 183 samples = ~80 hours. +Practical options: + - Run on a subset with --n-samples N (e.g. 20 = ~9 hours) + - Run on a compute server where the linear scan is faster + - Resume-capable: already-done samples are skipped automatically + +Output: one CSV per sample in gtdb_gather/, plus an aggregated +species x sample matrix (gtdb_species_matrix.tsv). +""" + +import csv +import subprocess +import sys +import time +import argparse +from pathlib import Path +import pandas as pd + +DB_PATH = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/gtdb-reps-rs226-k31.dna.zip") +READS_DIR = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/sourmash_sketch_reads") +MANIFEST = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/sample_manifest.tsv") +OUT_DIR = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/gtdb_gather") +THRESHOLD_BP = 5_000 + + +def run_gather_cli(sig_path: Path, out_csv: Path) -> int: + """Run sourmash gather via CLI. Returns number of matches.""" + cmd = [ + "sourmash", "gather", + str(sig_path), + str(DB_PATH), + "-k", "31", + "--threshold-bp", str(THRESHOLD_BP), + "-o", str(out_csv), + "--quiet", + ] + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + print(f" FAILED: {result.stderr[:200]}", file=sys.stderr) + return 0 + if not out_csv.exists(): + out_csv.write_text("intersect_bp,f_orig_query,f_match,f_unique_to_query,f_unique_weighted,average_abund,median_abund,std_abund,filename,name,md5,f_match_orig,unique_intersect_bp,gather_result_rank,remaining_bp,query_filename,query_name,query_md5,query_bp,ksize,moltype,scaled,query_n_hashes,query_abundance,query_containment_ani,match_containment_ani,average_containment_ani,max_containment_ani,potential_false_negative,n_unique_weighted_found,sum_weighted_found,total_weighted_hashes\n") + return 0 + return sum(1 for _ in open(out_csv)) - 1 # lines minus header + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--n-samples", type=int, default=None, + help="Process only the first N samples (for testing; default: all)") + parser.add_argument("--aggregate-only", action="store_true", + help="Skip gather, just re-aggregate existing CSVs") + args = parser.parse_args() + + OUT_DIR.mkdir(parents=True, exist_ok=True) + manifest = pd.read_csv(MANIFEST, sep="\t") + sample_ids = manifest["Sample_ID"].tolist() + + if args.n_samples: + sample_ids = sample_ids[:args.n_samples] + print(f"Processing {args.n_samples} samples (subset mode)", flush=True) + + todo = [ + sid for sid in sample_ids + if not (OUT_DIR / f"{sid}_gather.csv").exists() + and (READS_DIR / f"{sid}.sig").exists() + ] + already_done = len([sid for sid in sample_ids if (OUT_DIR / f"{sid}_gather.csv").exists()]) + print(f"{already_done} already done, {len(todo)} to process", flush=True) + + if todo and not args.aggregate_only: + # Estimate total time + if already_done == 0: + print(f" Note: first sample took ~26 min on M-chip laptop. " + f"ETA for {len(todo)} samples: ~{len(todo)*26/60:.0f} hours", flush=True) + t_total = time.time() + for i, sid in enumerate(todo, 1): + t0 = time.time() + print(f"[{i}/{len(todo)}] {sid}...", end=" ", flush=True) + + sig_path = READS_DIR / f"{sid}.sig" + out_csv = OUT_DIR / f"{sid}_gather.csv" + n_matches = run_gather_cli(sig_path, out_csv) + + elapsed = time.time() - t0 + done_so_far = already_done + i + eta_min = (time.time() - t_total) / i * (len(todo) - i) / 60 + print(f"{n_matches} matches, {elapsed/60:.1f} min (ETA {eta_min:.0f} min)", flush=True) + + # ── Aggregate into species × sample matrix ─────────────────────────────── + print("\nAggregating gather results...", flush=True) + records = [] + for sid in sample_ids: + csv_path = OUT_DIR / f"{sid}_gather.csv" + if not csv_path.exists(): + continue + df = pd.read_csv(csv_path, usecols=["name", "f_unique_to_query"]) + if df.empty: + continue + df["sample_id"] = sid + records.append(df[["sample_id", "name", "f_unique_to_query"]]) + + if not records: + print("No gather results found.") + return + + long_df = pd.concat(records, ignore_index=True) + + # Pivot to wide: rows = species, columns = samples + wide_df = long_df.pivot_table( + index="name", columns="sample_id", values="f_unique_to_query", fill_value=0 + ) + wide_df = wide_df.reindex(columns=sample_ids, fill_value=0) + + out_wide = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/gtdb_species_matrix.tsv") + wide_df.to_csv(out_wide, sep="\t") + print(f"Species matrix: {wide_df.shape[0]:,} species x {wide_df.shape[1]} samples -> {out_wide}", flush=True) + + out_long = Path("/Users/danielsprockett/Documents/Projects/WF22_RefRover/RefRover/data/mouse_rewilding/gtdb_gather_long.tsv") + long_df.to_csv(out_long, sep="\t", index=False) + print(f"Long format: {len(long_df):,} rows -> {out_long}", flush=True) + + # Quick summary + n_species = wide_df.shape[0] + prevalence = (wide_df > 0).sum(axis=1) + print(f"\nTop 10 most prevalent species:") + for name, prev in prevalence.nlargest(10).items(): + print(f" {prev:3d}/{len(sample_ids)} samples {name[:70]}") + + +if __name__ == "__main__": + main() diff --git a/opus_suggestions.md b/opus_suggestions.md new file mode 100644 index 0000000..9e801e9 --- /dev/null +++ b/opus_suggestions.md @@ -0,0 +1,257 @@ +# RefRover — Code & Analysis Review + +Review date: 2026-06-10. Scope: `src/refrover/` package, `data/mouse_rewilding/` analysis +scripts, and the bridge between them. Test suite passes (69 passed) and the validation +analysis (183-sample GTDB gather + selector comparison) is complete. + +The headline finding: **the science has moved ahead of the package.** The validated approach +— cross-sample *containment* with MaxMin selection and a residual-variance saturation curve for +adaptive k — lives entirely in `data/mouse_rewilding/*.py` analysis scripts. The shippable +package still revolves around assembly-vs-assembly *Jaccard* selectors, and the one class that +implements the validated method (`ContainmentSelector`) is unreachable from the CLI or pipeline. +Closing that gap is the single highest-value change. + +Suggestions are ordered by priority within each section. File/line references are to the state +reviewed. + +--- + +## 1. Correctness issues (fix first) + +### 1.1 MetaBAT2 output has fabricated variance — wrong for differential binning +`coverage.py:38-44` runs `coverm contig --methods mean` only. `formatters/metabat2.py` then +writes the per-sample variance columns as **0** because the variance was never computed. But +MetaBAT2's `jgi_summarize_bam_contig_depths` format *uses* the variance column, and SemiBin2 can +too. Shipping zero-variance defeats the differential signal the whole tool exists to produce. + +Fix: request both statistics — `coverm contig --methods mean variance` — and thread the variance +columns through to the formatters. This is a real bug, not a cosmetic one: it silently degrades +every MetaBAT2 run. + +### 1.2 Hybrid coverage merging is documented but not implemented +`CLAUDE.md` states hybrid short+long depth merging is "handled in `coverage.py`." It is not — +`run_coverm` runs one CoverM invocation over whatever BAMs exist, with no per-contig +read-count-weighted merge (`coverage.py:38-49`). Either implement the merge or correct the doc. +Right now a hybrid sample silently gets whichever BAMs `align` happened to write, concatenated +without weighting. + +### 1.3 Inconsistent "query is always included" contract across selectors +`MaxMinSelector`, `GreedyVarSelector`, `ArchetypeSelector`, and `KMedoidsSelector` all force the +query's own assembly into the result. `RandomSelector` does **not** (`random.py:17-18` — +`rng.sample(candidates, k)` with no query guarantee). The query's own assembly is the one +reference reads are guaranteed to map to, so omitting it skews the baseline comparison and means +"random" isn't a clean control for the others. Either document that random deliberately omits the +query, or make the contract uniform (recommended: a `BaseSelector._ensure_query_first()` helper +that every selector calls). + +### 1.4 `ContainmentSelector` assumes assembly ID == sample ID +`containment.py:97` — `own_assembly = query_id # assumes assembly ID matches sample ID`. The +manifest treats `sample_id` and `assembly` as separate columns (`io.py:4`). For the mouse dataset +they happen to coincide; for any caller where they don't, the wrong column gets pinned first. Pass +the manifest (or an explicit `sample_id → assembly_id` map) into the selector instead of assuming. + +### 1.5 ID provenance is fragile end-to-end +Sketch names (`matrix_from_sigs` uses `sig.name or filename stem`, `similarity.py:80`), +sourmash-compare CSV header labels, manifest `sample_id`, and manifest `assembly` are assumed +equal at several boundaries but never reconciled. `load_similarity_matrix` only checks the matrix +is square (`similarity.py:26`), not that its labels match the manifest. A single mislabeled sketch +produces silently wrong assignments. Add one validation point that asserts the similarity/ +containment matrix index is a superset of `manifest.sample_id`, and normalize sketch names to +`sample_id` at sketch time. + +--- + +## 2. Performance issues + +### 2.1 `ContainmentSelector` recomputes correlations in a Python loop +`containment.py:134-160` calls `np.corrcoef` pairwise inside the greedy loop — O(k · n) corrcoef +calls, each O(samples). For 182 candidates this is the dominant cost and is entirely avoidable. +`compare_selectors.py` already shows the fix: compute the full correlation matrix **once** +(`cont.corr()`), convert to a distance matrix, and index into it. Port that pattern into the +class. Expect ~100× speedup at this scale. + +### 2.2 `GreedyVarSelector` uses scalar `.loc` in a triple loop +`greedy_var.py:33-38` — `min(1.0 - float(sim_matrix.loc[c, s]) for s in selected)` inside +`for c in remaining` inside the `while` loop. Pandas scalar `.loc` is ~microseconds each; at +183 samples × k × n_candidates this is seconds per query for no reason. Convert `candidates` to a +NumPy submatrix once and do the min-distance update vectorized (the same `min_dist` array trick +`MaxMinSelector` already uses, `maxmin.py:31-42`). Also: `var_scores` includes the self-term +(`sim_matrix.loc[c, all_ids]` contains `c,c = 1.0`), mildly inflating variance — drop the diagonal. + +### 2.3 GTDB gather is serial and memory-bound +`build_community_subdb.py` gathers samples one at a time; the log shows per-sample times swinging +from 13 s to 990 s under memory pressure. A process pool (or GNU parallel / `sourmash multigather` +with a manifest) bounded to a sane worker count would cut wall-clock and smooth the variance. Even +2–3 workers against the 55 MB sub-database would help, since the sub-DB fits in memory many times +over. + +--- + +## 3. Architecture & API consistency + +### 3.1 Two selector interfaces, no shared contract +`BaseSelector.select(square_jaccard, query_id) -> list[str]` returns **sample IDs**; +`ContainmentSelector.select(rectangular_containment, query_id) -> list[str]` returns **assembly +IDs** and isn't a `BaseSelector` subclass (`containment.py:44`). The pipeline can't treat them +polymorphically, which is why containment never reached the CLI. Recommended: define one protocol + +```python +class Selector(Protocol): + def select(self, query_id: str, *, matrix: pd.DataFrame, manifest: pd.DataFrame) -> list[str]: + ... # returns assembly IDs, always +``` + +Let both Jaccard and containment selectors implement it; resolve sample→assembly via the manifest +in one place. Returning assembly IDs uniformly also removes the implicit sample==assembly coupling +flagged in 1.4. + +### 3.2 Adaptive k works in `select` but not in `run`/pipeline +`cli.py select` supports `--k auto` (`cli.py:90-104`), but `RefRoverPipeline.run()` and `cli.py run` +accept only a fixed integer `k` (`pipeline.py:64-68`, `cli.py:189`). End-to-end runs therefore +can't use the feature the research is about. Move the per-query k estimation into the pipeline's +selection loop and have `cli.py run` accept `--k auto` too. + +### 3.3 Unimplemented strategies are exposed as if they work +`SELECTOR_REGISTRY` includes `feedback` (`selectors/__init__.py:15`), which raises +`NotImplementedError`; `benchmark.py:31` is a stub; `cli.py select` advertises `greedy_var` but not +`containment`. A user picking `feedback` gets a crash, not a helpful message. Either gate +unimplemented entries behind a clear "experimental/unimplemented" error at construction, or remove +them from the CLI `Choice` lists until they exist. + +### 3.4 `print`/`click.echo` instead of logging +Subprocess wrappers swallow stdout/stderr and re-raise on failure (`coverage.py:45-47`, +`sketch.py`, `align.py`). There's no `--verbose`, no run log, no timing. A `logging`-based setup +with a `--verbose/-v` flag and a per-run `refrover.log` would make the multi-hour pipeline +debuggable without code edits. + +--- + +## 4. Streamline the analysis → converge it with the package + +The analysis scripts and the package have forked. Three concrete merges: + +### 4.1 Move containment computation into the package +`compute_containment.py` lives only in `data/mouse_rewilding/` but computes the matrix that the +*best-performing* selector depends on. Promote it to `refrover/containment.py` with a +`refrover containment` subcommand (sourmash containment / prefetch under the hood). Then containment +is a first-class signal, not an out-of-tree preprocessing step. + +### 4.2 Promote the variance-explained proxy into `benchmark.py` +`compare_selectors.py` computes the orthogonal-projection "unique variance explained" metric — a +fast, **alignment-free** ranking of selectors. This is exactly what `benchmark.py` should do as its +cheap first tier, *before* spending CPU-days on CheckM2. Make `run_benchmark` report both: +(1) the proxy ranking (seconds), and (2) optionally the CheckM2 gold standard (hours). The proxy is +also the only way to benchmark at the 183-sample scale without an aligner. + +### 4.3 Port the validated adaptive-k method into `adaptive_k.py` +The package's three estimators (`scree_elbow`, `similarity_gap`, `saturation_curve`) all run on the +**Jaccard** matrix as a proxy. The analysis showed the trustworthy method is greedy **residual +variance** on the **containment** matrix with a 10%-of-first-gain elbow (`compare_selectors.py`: +`saturation_curve` + `elbow_from_gains`), which gave a stable k=4 across all 183 samples. Add it as +a `containment_saturation` method and make it the default when a containment matrix is available. +Keep the Jaccard estimators as fallbacks for the no-containment case. + +### 4.4 Have the analysis import the package, not reimplement it +`compare_selectors.py` reimplements `random` and `maxmin` selection inline. If those drift from +`selectors/maxmin.py` the benchmark stops measuring what ships. Import the real classes (once 3.1 +gives them a common interface) so the experiment exercises production code. + +### 4.5 Minor analysis cleanups +- Consolidate `run_gtdb_gather.py` and `build_community_subdb.py` — they overlap substantially. +- Factor the repeated tidyverse/theme boilerplate in the four `explore_*.R` scripts into a sourced + `_setup.R`. +- The sub-DB build prints `Sub-database: ... (0 signatures)` (`bd8ykcrxj` log line 15) even on + success — the count is read from the wrong place. Cosmetic but misleading; report + `len(all_md5s)` or the actual sig count. + +--- + +## 5. How to select an optimal Selector + +This is the central open question (CLAUDE.md research Q1), and the validation work already gives a +defensible answer and a reusable procedure. + +**What the data showed (183 samples):** +- Selectors disagree strongly — mean Jaccard overlap 0.01–0.31 at k=5. The choice is *not* + cosmetic. +- On the **unique-variance** metric (orthogonal projection, the unbiased one): + MaxMin wins at small k (k=3: 76% vs top-var 62%); top-var edges ahead by k=5 (87% vs 84%); all + converge to ~90% by k=10. "Top containment" — picking assemblies most similar to the query — + loses at every k because it selects correlated, redundant references. +- The naive "sum of variances" metric inverts this ranking, which is exactly why an unbiased metric + matters. + +**Recommended selection procedure (a three-tier funnel):** + +1. **Tier 0 — agreement screen (seconds).** Compute pairwise selection overlap. If all selectors + agree (high Jaccard), pick the cheapest (random/maxmin) and stop — the choice doesn't matter for + this dataset. The mouse data failed this screen, so proceed. + +2. **Tier 1 — variance-explained proxy (seconds–minutes, no alignment).** Rank selectors by unique + cross-sample variance explained at the adaptive k. This is the orthogonal-projection metric from + `compare_selectors.py`. It needs only the containment (or similarity) matrix and discriminates + selectors cleanly. Use it as the routine, scalable decision rule. + +3. **Tier 2 — MAG yield per CPU-hour (hours, gold standard).** For the 2–3 finalists from Tier 1, + run the full pipeline on one or two datasets and count CheckM2-passing MAGs (≥50% complete, + ≤10% contaminated) per CPU-hour. This is the metric that actually matters and the one that + *validates the proxy*. Do it once to establish that Tier-1 ranking predicts Tier-2 ranking; if + it does, you rarely need Tier 2 again. + +**The key research deliverable** is the Tier-1↔Tier-2 correlation. If unique-variance-explained +predicts CheckM2 MAG yield, you have a principled, cheap selector-choice rule and a publishable +result. If it doesn't, that itself tells you the proxy is missing something (e.g., mapping rate, +assembly fragmentation) and points to the next metric. + +**Practical default to ship now:** MaxMin on containment profiles at adaptive k (≈4–5 for cohort +data). It's the robust choice across k, it's cheap, and it dominates at the small k that adaptive-k +actually selects. Make it the documented default; expose top-var as the alternative for users who +fix a larger k. + +--- + +## 6. Additional functionality worth adding + +- **`refrover rank-selectors` subcommand.** Operationalize §5 Tiers 0–1: given sketches or a + containment matrix, run every selector at a k-range and print the agreement screen + variance + proxy. A pre-flight that tells the user which selector to use before committing to alignment. +- **`refrover containment` subcommand** (per §4.1) — first-class containment computation. +- **Min-threshold calibration helper.** Research Q3 (what `min_jaccard`/`min_containment` to use) + is answerable from the data: at `min_containment=0.1` some samples drop to 0 candidates, at 0.02 + the median is ~172. A small utility that plots candidate-count vs threshold and picks the knee + would stop users from silently starving samples of references. +- **Run provenance file.** Write `params.json` per run (selector, k, thresholds, tool versions, + input hashes). Essential for reproducibility and for the benchmark bookkeeping. +- **Resumability beyond file-existence.** A small state file (which samples are sketched/aligned/ + covered) makes the multi-hour pipeline restartable after interruption without re-globbing. +- **Smoke integration test.** Tests currently mock all external tools. One opt-in test + (`@pytest.mark.integration`) that runs tiny real FASTA/FASTQ through bwa-mem2 + coverm would catch + interface drift (e.g., the BAM glob in `coverage.py:34` vs. what `align` actually names files). + +--- + +## 7. Testing & infrastructure + +- **No CI.** Add a GitHub Actions workflow running `ruff`, `mypy`, and `pytest` on push. The + `[dev]` extras already declare all three (`pyproject.toml`). +- **Selector property tests.** Every selector should satisfy invariants worth asserting once in a + shared parametrized test: returns ≤ k IDs, returns only candidates above threshold, returns + unique IDs, and (per §1.3) includes the query. Right now each selector has its own ad-hoc tests + and the contract isn't enforced uniformly. +- **`coverage.py` / `align.py` have no unit tests** beyond the mocked pipeline test. The + CoverM/BWA argument construction is exactly the kind of thing that breaks silently; a test that + asserts the constructed command list would be cheap insurance. + +--- + +## 8. Suggested order of work + +1. Fix MetaBAT2 variance (§1.1) and the query-inclusion contract (§1.3) — correctness, small. +2. Unify the selector interface and wire `ContainmentSelector` into CLI + pipeline (§3.1, 3.2) — + unlocks the validated method. +3. Port containment computation and the adaptive-k saturation method into the package (§4.1, 4.3). +4. Implement `benchmark.py` Tier-1 proxy + `rank-selectors` subcommand (§4.2, §6) — makes selector + choice routine and scalable. +5. Performance passes on the two slow selectors (§2.1, 2.2). +6. CI + property tests + provenance (§7, §6). +7. Run the Tier-1↔Tier-2 validation once (§5) — the publishable result. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e4a82b1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "refrover" +version = "0.1.0" +requires-python = ">=3.11" +dependencies = [ + "click>=8.0", + "pandas>=2.0", + "numpy>=1.24", + "scipy>=1.10", + "scikit-learn>=1.3", + "sourmash>=4.8", + "archetypes>=0.12", +] + +[project.optional-dependencies] +dev = ["pytest>=7.0", "ruff", "mypy"] + +[project.scripts] +refrover = "refrover.cli:main" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.ruff] +line-length = 100 +target-version = "py311" diff --git a/src/refrover/__init__.py b/src/refrover/__init__.py new file mode 100644 index 0000000..3ba0853 --- /dev/null +++ b/src/refrover/__init__.py @@ -0,0 +1,6 @@ +from importlib.metadata import version, PackageNotFoundError + +try: + __version__ = version("refrover") +except PackageNotFoundError: + __version__ = "0.1.0-dev" diff --git a/src/refrover/adaptive_k.py b/src/refrover/adaptive_k.py new file mode 100644 index 0000000..0e1f0af --- /dev/null +++ b/src/refrover/adaptive_k.py @@ -0,0 +1,274 @@ +""" +Adaptive k selection: determine the number of prototypes per sample +from the dataset's similarity structure rather than using a fixed global k. + +The core idea: as k increases, each additional prototype adds less new +differential signal (diminishing returns). The optimal k is at the elbow +of the signal-saturation curve. + +Three estimators are implemented: + +scree_elbow + PCoA on the Jaccard distance matrix; fit an exponential decay to the + eigenvalue spectrum; elbow = point where marginal explained variance + drops below a threshold. Analogous to scree plot in PCA. + +similarity_gap + Sort candidate similarities to the query in descending order. + Find the largest gap in the sorted similarity values — the natural + break between "similar enough to add signal" and "too distant to matter". + Simple, fast, interpretable. + +saturation_curve + For increasing k, estimate the marginal diversity gain from adding + the k-th prototype (using MaxMin distance). Stop when the gain drops + below a fraction of the initial gain. + +containment_saturation + The method validated on the rewilded-mouse dataset. Greedy forward + selection by *residual* cross-sample variance (incremental Gram-Schmidt), + which guarantees a monotonically decreasing marginal-gain curve, then the + elbow is the first k whose gain falls below a fraction of the first gain. + Intended for a containment matrix but works on any feature matrix. The + other three estimators use a Jaccard matrix as a proxy; this one measures + the quantity that actually matters (variance the prototypes span). + +All return an integer k_hat. The CLI exposes --k auto to trigger this. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +from typing import Literal + + +AdaptiveMethod = Literal[ + "scree_elbow", "similarity_gap", "saturation_curve", "containment_saturation" +] + + +def estimate_k( + sim_matrix: pd.DataFrame, + query_id: str, + min_jaccard: float = 0.1, + method: AdaptiveMethod = "similarity_gap", + k_min: int = 3, + k_max: int = 20, +) -> int: + """ + Estimate the optimal number of prototypes for query_id. + + Parameters + ---------- + sim_matrix : square Jaccard similarity DataFrame + query_id : row/column label for the query sample + min_jaccard : minimum similarity threshold (same as selector) + method : which estimator to use + k_min, k_max : bounds on returned k + + Returns + ------- + int : estimated k, clipped to [k_min, k_max] + """ + if query_id not in sim_matrix.index: + raise KeyError(f"query_id '{query_id}' not found in similarity matrix") + + # Restrict to valid candidates + row = sim_matrix.loc[query_id] + candidates = row[row >= min_jaccard].index.tolist() + n_candidates = len(candidates) + + if n_candidates <= k_min: + return n_candidates + + if method == "scree_elbow": + k_hat = _scree_elbow(sim_matrix, candidates, k_min, k_max) + elif method == "similarity_gap": + k_hat = _similarity_gap(row, candidates, k_min, k_max) + elif method == "saturation_curve": + k_hat = _saturation_curve(sim_matrix, query_id, candidates, k_min, k_max) + elif method == "containment_saturation": + k_hat = _containment_saturation(sim_matrix, candidates, k_min, k_max) + else: + raise ValueError(f"Unknown adaptive k method: {method!r}") + + return int(np.clip(k_hat, k_min, min(k_max, n_candidates))) + + +# ── Estimators ──────────────────────────────────────────────────────────────── + +def _scree_elbow( + sim_matrix: pd.DataFrame, + candidates: list[str], + k_min: int, + k_max: int, +) -> int: + """ + PCoA scree elbow: eigenvalue-based estimate of intrinsic dimensionality. + + Converts Jaccard similarity to distance, double-centres, takes eigenvalues. + The elbow in the sorted positive eigenvalue spectrum indicates the number + of meaningful axes of variation — a natural estimate for k. + """ + sub = sim_matrix.loc[candidates, candidates] + dist = 1.0 - sub.values.astype(float) + np.fill_diagonal(dist, 0.0) + + # Double-centring (classical MDS) + n = len(candidates) + H = np.eye(n) - np.ones((n, n)) / n + B = -0.5 * H @ (dist ** 2) @ H + + eigvals = np.linalg.eigvalsh(B) + eigvals = np.sort(eigvals)[::-1] + pos = eigvals[eigvals > 0] + + if len(pos) < 2: + return k_min + + # Kneedle-style: find index with maximum curvature in normalised scree plot + x = np.arange(len(pos), dtype=float) + y = pos / pos.sum() + # Straight line from first to last point + line_y = y[0] + (y[-1] - y[0]) * x / x[-1] + deviations = y - line_y + elbow_idx = int(np.argmax(deviations)) + 1 # 1-indexed k + + return max(k_min, elbow_idx) + + +def _similarity_gap( + query_row: pd.Series, + candidates: list[str], + k_min: int, + k_max: int, +) -> int: + """ + Largest-gap heuristic: sort candidate similarities descending, find the + biggest drop. Prototypes above the gap add meaningful signal; below it + are barely different from the background. + """ + sims = query_row[candidates].sort_values(ascending=False).values + if len(sims) < 2: + return k_min + + gaps = np.diff(sims) # all negative (descending) + gap_idx = int(np.argmin(gaps)) # index of the biggest drop + k_hat = gap_idx + 1 # number of prototypes above the gap + + return max(k_min, k_hat) + + +def _saturation_curve( + sim_matrix: pd.DataFrame, + query_id: str, + candidates: list[str], + k_min: int, + k_max: int, +) -> int: + """ + Greedy diversity saturation: incrementally add the MaxMin-greedy prototype + and track the marginal gain in minimum pairwise distance. The elbow is the + largest consecutive drop in the gains sequence — the point where adding + another prototype stops providing meaningfully new diversity. + """ + if len(candidates) == 0: + return k_min + + sub = sim_matrix.loc[candidates, candidates] + dist = 1.0 - sub.values.astype(float) + np.fill_diagonal(dist, 0.0) + idx = {c: i for i, c in enumerate(candidates)} + + first = query_id if query_id in idx else candidates[0] + selected = [idx[first]] + remaining = [i for i in range(len(candidates)) if i != idx[first]] + + min_dist = dist[selected[0], remaining].copy() + gains = [] + + while remaining and len(selected) < k_max: + best_local_idx = int(np.argmax(min_dist)) + gain = float(min_dist[best_local_idx]) + gains.append(gain) + + best = remaining[best_local_idx] + selected.append(best) + remaining.pop(best_local_idx) + min_dist = np.delete(min_dist, best_local_idx) + + if remaining: + min_dist = np.minimum(min_dist, dist[best, remaining]) + + if len(gains) < 2: + return max(k_min, len(gains) + 1) + + # Elbow = largest consecutive drop in gains + drops = np.diff(gains) # all non-positive (gains are non-increasing) + elbow_idx = int(np.argmin(drops)) # index of biggest drop + # +1 for seed, +1 because we want k prototypes *including* the seed + k_hat = elbow_idx + 2 + + return max(k_min, k_hat) + + +def _containment_saturation( + sim_matrix: pd.DataFrame, + candidates: list[str], + k_min: int, + k_max: int, + saturation_fraction: float = 0.1, +) -> int: + """ + Greedy residual-variance saturation (the validated method). + + Forward-selects candidate columns by the cross-sample variance each adds + that is *not* already explained by the selected set (incremental + Gram-Schmidt orthogonalisation), giving a monotonically decreasing + marginal-gain curve. The elbow is the first step whose gain falls below + `saturation_fraction` of the first (largest) gain. + + Unlike _saturation_curve, this measures variance directly rather than + MaxMin distance, and the monotone curve makes a threshold elbow robust. + """ + sub = sim_matrix[candidates].to_numpy(dtype=float) + cols = sub - sub.mean(axis=0) + col_vars = np.var(cols, axis=0, ddof=1) + n_cand = cols.shape[1] + + first = int(np.argmax(col_vars)) + remaining = [i for i in range(n_cand) if i != first] + + v0 = cols[:, first] + norm0 = np.linalg.norm(v0) + Q = (v0 / norm0).reshape(-1, 1) if norm0 > 1e-12 else np.zeros((cols.shape[0], 1)) + + first_gain = float(col_vars[first]) + if first_gain <= 0.0: + return k_min + + gains = [first_gain] + limit = min(k_max - 1, len(remaining)) + for _ in range(limit): + if not remaining: + break + Y = cols[:, remaining] + resid = Y - Q @ (Q.T @ Y) + resid_vars = np.var(resid, axis=0, ddof=1) + + best_local = int(np.argmax(resid_vars)) + best = remaining[best_local] + gains.append(float(resid_vars[best_local])) + + v = cols[:, best] - Q @ (Q.T @ cols[:, best]) + norm = np.linalg.norm(v) + if norm > 1e-12: + Q = np.column_stack([Q, v / norm]) + remaining.pop(best_local) + + gains_arr = np.asarray(gains) + threshold = saturation_fraction * gains_arr[0] + below = np.where(gains_arr < threshold)[0] + k_hat = int(below[0]) if len(below) else len(gains_arr) + return max(k_min, k_hat) diff --git a/src/refrover/align.py b/src/refrover/align.py new file mode 100644 index 0000000..71c2ecb --- /dev/null +++ b/src/refrover/align.py @@ -0,0 +1,275 @@ +""" +Alignment wrappers: BWA-MEM2, BWA, minimap2. + +For each sample, concatenates its assigned prototype assemblies into a +combined reference, builds an index, aligns reads, and writes a sorted BAM. +""" + +import logging +import subprocess +from pathlib import Path + +import pandas as pd + +log = logging.getLogger(__name__) + + +def run_alignment( + manifest: pd.DataFrame, + assignments: pd.DataFrame, + *, + outdir: Path | str, + aligner: str = "bwa-mem2", + threads: int = 8, + force: bool = False, + bwa_path: str = "bwa-mem2", + minimap2_path: str = "minimap2", + samtools_path: str = "samtools", +) -> list[Path]: + """ + For each sample, align reads against its assigned prototype assemblies. + + Steps per sample: + 1. Concatenate prototype FASTAs into a combined reference. + 2. Build aligner index (bwa-mem2 index or minimap2 index). + 3. Align reads (paired/single/long). + 4. Sort and index the BAM. + + Returns list of sorted BAM paths, in manifest order. + Skips existing BAMs unless force=True. + """ + outdir = Path(outdir) + outdir.mkdir(parents=True, exist_ok=True) + + asgn_map = dict(zip(assignments["sample_id"], assignments["prototype_ids"])) + bam_paths: list[Path] = [] + + for _, row in manifest.iterrows(): + sid = row["sample_id"] + if sid not in asgn_map: + log.warning("No assignment for sample %s — skipping", sid) + continue + + bam_path = outdir / f"{sid}.sorted.bam" + bam_paths.append(bam_path) + + if bam_path.exists() and not force: + log.debug("Skipping %s (BAM exists)", sid) + continue + + prototype_ids = [p.strip() for p in asgn_map[sid].split(",")] + sample_dir = outdir / sid + sample_dir.mkdir(exist_ok=True) + + # Locate prototype FASTA paths from manifest + proto_fastas = _resolve_fastas(manifest, prototype_ids, row["assembly"]) + + # Build combined reference + ref_fa = sample_dir / "reference.fasta" + _concat_fastas(proto_fastas, ref_fa) + + # Align + if aligner in ("bwa-mem2", "bwa"): + _align_bwa( + row, ref_fa, bam_path, sample_dir, + aligner_bin=bwa_path if aligner == "bwa-mem2" else "bwa", + samtools_bin=samtools_path, + threads=threads, + ) + elif aligner == "minimap2": + _align_minimap2( + row, ref_fa, bam_path, sample_dir, + minimap2_bin=minimap2_path, + samtools_bin=samtools_path, + threads=threads, + ) + else: + raise ValueError(f"Unknown aligner: {aligner!r}") + + log.info("Aligned %s → %s", sid, bam_path) + + return bam_paths + + +# ── Internal helpers ────────────────────────────────────────────────────────── + +def _resolve_fastas( + manifest: pd.DataFrame, + prototype_ids: list[str], + own_assembly: str, +) -> list[Path]: + """Return FASTA paths for the given prototype IDs, using the manifest for lookup.""" + asm_map = dict(zip(manifest["sample_id"], manifest["assembly"])) + paths = [] + for pid in prototype_ids: + if pid in asm_map: + paths.append(Path(asm_map[pid])) + else: + # Fallback: assume the prototype ID is already a path + paths.append(Path(pid)) + return paths + + +def _concat_fastas(fastas: list[Path], out_fa: Path) -> None: + """Concatenate FASTA files into a single reference.""" + with open(out_fa, "wb") as fout: + for fa in fastas: + with open(fa, "rb") as f: + fout.write(f.read()) + + +def _run(cmd: list[str], description: str) -> None: + """Run a subprocess command; raise RuntimeError on failure.""" + log.debug("Running: %s", " ".join(cmd)) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"{description} failed (exit {result.returncode}):\n" + f" cmd: {' '.join(cmd)}\n" + f" stderr: {result.stderr[:500]}" + ) + + +def _align_bwa( + sample_row: pd.Series, + ref_fa: Path, + out_bam: Path, + work_dir: Path, + aligner_bin: str, + samtools_bin: str, + threads: int, +) -> None: + """Index with bwa-mem2/bwa, align, sort into out_bam.""" + # Index + _run([aligner_bin, "index", str(ref_fa)], f"bwa index {ref_fa.name}") + + # Build align command + r1 = sample_row.get("r1") + r2 = sample_row.get("r2") if "r2" in sample_row.index else None + long_reads = sample_row.get("long_reads") if "long_reads" in sample_row.index else None + + sam_path = work_dir / "aln.sam" + rg = f"@RG\\tID:{sample_row['sample_id']}\\tSM:{sample_row['sample_id']}" + + if pd.notna(r1): + mem_cmd = [aligner_bin, "mem", "-t", str(threads), "-R", rg, str(ref_fa), str(r1)] + if pd.notna(r2): + mem_cmd.append(str(r2)) + with open(sam_path, "w") as fout: + result = subprocess.run(mem_cmd, stdout=fout, stderr=subprocess.PIPE, text=True) + if result.returncode != 0: + raise RuntimeError(f"bwa mem failed:\n{result.stderr[:500]}") + + if pd.notna(long_reads): + _merge_long_reads_bwa( + sample_row, ref_fa, sam_path, work_dir, aligner_bin, samtools_bin, threads + ) + + elif pd.notna(long_reads): + # Long-reads only: fall through to minimap2 + raise ValueError( + f"Sample {sample_row['sample_id']}: long-reads-only with bwa aligner. " + "Use --aligner minimap2 for long-read or hybrid samples." + ) + + _sort_index(sam_path, out_bam, samtools_bin, threads) + + +def _merge_long_reads_bwa( + sample_row: pd.Series, + ref_fa: Path, + short_sam: Path, + work_dir: Path, + aligner_bin: str, + samtools_bin: str, + threads: int, +) -> None: + """ + For hybrid samples: align long reads with minimap2, merge with short-read SAM. + Modifies short_sam in place (appends long-read alignments). + """ + long_sam = work_dir / "long.sam" + _run( + ["minimap2", "-ax", "map-ont", "-t", str(threads), + str(ref_fa), str(sample_row["long_reads"])], + "minimap2 long reads", + ) + # Append long-read alignments (skip header lines that start with @) + with open(short_sam, "a") as fout, open(long_sam) as fin: + for line in fin: + if not line.startswith("@"): + fout.write(line) + + +def _align_minimap2( + sample_row: pd.Series, + ref_fa: Path, + out_bam: Path, + work_dir: Path, + minimap2_bin: str, + samtools_bin: str, + threads: int, +) -> None: + """Align with minimap2 (long reads or hybrid), sort into out_bam.""" + r1 = sample_row.get("r1") if "r1" in sample_row.index else None + long_reads = sample_row.get("long_reads") if "long_reads" in sample_row.index else None + sam_path = work_dir / "aln.sam" + + if pd.notna(long_reads): + preset = "map-ont" # or map-pb for PacBio + reads = [str(long_reads)] + if pd.notna(r1): + # Hybrid: align short reads with sr preset, long with map-ont, then merge + _align_minimap2_preset(minimap2_bin, ref_fa, [str(r1)], work_dir / "short.sam", + preset="sr", threads=threads) + _align_minimap2_preset(minimap2_bin, ref_fa, reads, work_dir / "long.sam", + preset=preset, threads=threads) + _merge_sams([work_dir / "short.sam", work_dir / "long.sam"], sam_path) + else: + _align_minimap2_preset(minimap2_bin, ref_fa, reads, sam_path, + preset=preset, threads=threads) + elif pd.notna(r1): + reads = [str(r1)] + r2 = sample_row.get("r2") if "r2" in sample_row.index else None + if pd.notna(r2): + reads.append(str(r2)) + _align_minimap2_preset(minimap2_bin, ref_fa, reads, sam_path, + preset="sr", threads=threads) + else: + raise ValueError(f"Sample {sample_row['sample_id']}: no reads to align") + + _sort_index(sam_path, out_bam, samtools_bin, threads) + + +def _align_minimap2_preset( + minimap2_bin: str, + ref_fa: Path, + reads: list[str], + out_sam: Path, + preset: str, + threads: int, +) -> None: + cmd = [minimap2_bin, "-ax", preset, "-t", str(threads), str(ref_fa)] + reads + with open(out_sam, "w") as fout: + result = subprocess.run(cmd, stdout=fout, stderr=subprocess.PIPE, text=True) + if result.returncode != 0: + raise RuntimeError(f"minimap2 ({preset}) failed:\n{result.stderr[:500]}") + + +def _merge_sams(sam_paths: list[Path], out_sam: Path) -> None: + """Merge SAM files, keeping header only from the first.""" + with open(out_sam, "w") as fout: + for i, p in enumerate(sam_paths): + with open(p) as fin: + for line in fin: + if i == 0 or not line.startswith("@"): + fout.write(line) + + +def _sort_index(sam_path: Path, out_bam: Path, samtools_bin: str, threads: int) -> None: + """Sort SAM → BAM and index.""" + _run( + [samtools_bin, "sort", "-@", str(threads), "-o", str(out_bam), str(sam_path)], + f"samtools sort → {out_bam.name}", + ) + _run([samtools_bin, "index", str(out_bam)], f"samtools index {out_bam.name}") diff --git a/src/refrover/benchmark.py b/src/refrover/benchmark.py new file mode 100644 index 0000000..abfffac --- /dev/null +++ b/src/refrover/benchmark.py @@ -0,0 +1,157 @@ +""" +Benchmark selector strategies against each other. + +Two tiers: + +Tier 1 — variance-explained proxy (cheap, alignment-free). + For each (selector, k), run the selector for every sample and measure the + fraction of total cross-sample variance the selected prototypes explain, + via orthogonal projection of the full matrix onto the selected columns + (so correlated/redundant selections are not double-counted). This ranks + selectors in seconds from a similarity or containment matrix alone, with no + alignment. See rank_selectors(). + +Tier 2 — MAG yield per CPU-hour (gold standard, expensive). + Run the full pipeline + a binner + CheckM2 and count passing MAGs per + compute. Not yet implemented (needs CheckM2 and a labelled community); + run_benchmark() is the entry point. + +The research question is whether the Tier-1 ranking predicts the Tier-2 ranking. +If it does, selector choice can be made cheaply from Tier 1 alone. +""" + +import warnings +from pathlib import Path + +import numpy as np +import pandas as pd + +from refrover.selectors import SELECTOR_REGISTRY + +# Selectors that can't run unattended in a benchmark (unimplemented). +_SKIP_SELECTORS = {"feedback"} +DEFAULT_SELECTORS = ["random", "maxmin", "kmedoids", "archetype", "greedy_var", "containment"] + + +def unique_variance_explained(matrix: pd.DataFrame, selected_ids: list[str]) -> float: + """ + Fraction of total cross-sample column variance explained by selected columns. + + The full matrix (samples × features) is mean-centred; the selected columns + span a subspace; we project all columns onto it and report + (total_var − residual_var) / total_var. Orthogonal projection means columns + correlated with the selected set count once, not multiple times — so this + rewards prototypes that span independent axes of variation rather than + piling onto the same one. + + Returns a value in [0, 1] (1.0 when the selected columns span the column + space, e.g. all columns selected). + """ + cols = list(matrix.columns) + pos = {c: i for i, c in enumerate(cols)} + idx = [pos[s] for s in selected_ids if s in pos] + + X = matrix.to_numpy(dtype=float) + X = X - X.mean(axis=0) + total_var = float(np.var(X, axis=0, ddof=1).sum()) + if total_var == 0.0 or not idx: + return 0.0 + + Xp = X[:, idx] + Q, _ = np.linalg.qr(Xp) # orthonormal basis for the selected subspace + resid = X - Q @ (Q.T @ X) + resid_var = float(np.var(resid, axis=0, ddof=1).sum()) + return (total_var - resid_var) / total_var + + +def rank_selectors( + matrix: pd.DataFrame, + *, + selectors: list[str] | None = None, + k_values: list[int], + min_similarity: float = 0.1, + query_ids: list[str] | None = None, +) -> pd.DataFrame: + """ + Rank selectors by the Tier-1 variance-explained proxy on a single matrix. + + The same matrix (Jaccard similarity or containment) is fed to every selector, + so the comparison is apples-to-apples. Selectors that can't run in the + environment (e.g. archetype without the `archetypes` package) are skipped + with a warning rather than aborting the whole run. + + Returns a DataFrame: selector, k, mean_var_explained, median_var_explained, + n_samples — sorted by k then descending mean_var_explained. + """ + selectors = selectors or DEFAULT_SELECTORS + query_ids = query_ids if query_ids is not None else list(matrix.index) + + records = [] + for name in selectors: + if name in _SKIP_SELECTORS or name not in SELECTOR_REGISTRY: + warnings.warn(f"Skipping selector '{name}' (unavailable).", stacklevel=2) + continue + cls = SELECTOR_REGISTRY[name] + selector_failed = False + + for k in k_values: + sel = cls(k=k, min_jaccard=min_similarity) + # Make the random baseline reproducible across the benchmark. + if hasattr(sel, "seed") and getattr(sel, "seed") is None: + sel.seed = 0 + + scores = [] + for q in query_ids: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # fewer-than-k candidate notices + try: + chosen = sel.select(matrix, q) + except (ImportError, NotImplementedError) as exc: + warnings.warn( + f"Skipping selector '{name}': {exc}", stacklevel=2 + ) + selector_failed = True + break + except (KeyError, ValueError): + continue + scores.append(unique_variance_explained(matrix, chosen)) + + if selector_failed: + break + if scores: + records.append({ + "selector": name, + "k": k, + "mean_var_explained": float(np.mean(scores)), + "median_var_explained": float(np.median(scores)), + "n_samples": len(scores), + }) + + df = pd.DataFrame.from_records( + records, + columns=["selector", "k", "mean_var_explained", "median_var_explained", "n_samples"], + ) + if not df.empty: + df = df.sort_values(["k", "mean_var_explained"], ascending=[True, False]) + return df.reset_index(drop=True) + + +def run_benchmark( + manifest: pd.DataFrame, + truth_path: Path | str, + selectors: list[str], + k_values: list[int], + outdir: Path | str, + threads: int = 8, + checkm2_db: Path | str | None = None, +) -> pd.DataFrame: + """ + Tier 2 — MAG yield per CPU-hour against a labelled ground-truth community. + + For each (selector, k): run RefRoverPipeline, bin with MetaBAT2, evaluate with + CheckM2, and report passing MAGs per compute. Not yet implemented. + """ + raise NotImplementedError( + "Tier-2 benchmark (CheckM2 MAG yield) not yet implemented. " + "Use rank_selectors() for the alignment-free Tier-1 variance proxy." + ) diff --git a/src/refrover/cli.py b/src/refrover/cli.py new file mode 100644 index 0000000..2a6f168 --- /dev/null +++ b/src/refrover/cli.py @@ -0,0 +1,372 @@ +""" +RefRover CLI entry point. + +Subcommands: + sketch Sketch assemblies with sourmash + select Prototype selection (assigns prototypes per sample) + align Align reads to selected prototypes + coverage Compute per-contig depth with CoverM + format Format coverage tables for downstream binners + run Full pipeline end-to-end + benchmark Compare selectors on a dataset with known ground truth +""" + +import click +from pathlib import Path + + +@click.group() +@click.version_option() +def main(): + """RefRover: differential coverage tables for metagenomic binners.""" + + +# ── sketch ──────────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--manifest", required=True, type=click.Path(exists=True), help="Sample manifest TSV") +@click.option("--outdir", required=True, type=click.Path(), help="Output directory for .sig files") +@click.option("--ksize", default=31, show_default=True) +@click.option("--scaled", default=1000, show_default=True) +@click.option("--threads", default=1, show_default=True) +@click.option("--force", is_flag=True, help="Re-sketch even if output exists") +def sketch(manifest, outdir, ksize, scaled, threads, force): + """Sketch assemblies with sourmash sketch dna.""" + from refrover.io import read_manifest + from refrover.sketch import sketch_assemblies + + df = read_manifest(manifest) + sig_paths = sketch_assemblies( + df["assembly"].tolist(), + outdir=outdir, + ksize=ksize, + scaled=scaled, + threads=threads, + force=force, + ) + click.echo(f"Sketched {len(sig_paths)} assemblies → {outdir}") + + +# ── containment ─────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--read-sketches", required=True, type=click.Path(exists=True), + help="Directory of read .sig files (one per sample, stem = sample_id)") +@click.option("--assembly-sketches", required=True, type=click.Path(exists=True), + help="Directory of assembly .sig files (one per sample, stem = sample_id)") +@click.option("--ksize", default=31, show_default=True) +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--force", is_flag=True) +def containment(read_sketches, assembly_sketches, ksize, outdir, force): + """Compute the cross-sample reads-vs-assemblies containment matrix.""" + from refrover.containment import ( + compute_containment_matrix, sigs_by_sample, write_containment_matrix, + ) + + outdir = Path(outdir) + outdir.mkdir(parents=True, exist_ok=True) + out_tsv = outdir / "containment_matrix.tsv" + + if out_tsv.exists() and not force: + click.echo(f"Containment matrix already exists at {out_tsv} (use --force to redo)") + return + + read_sigs = sigs_by_sample(read_sketches) + asm_sigs = sigs_by_sample(assembly_sketches) + if not read_sigs: + raise click.ClickException(f"No .sig files found in {read_sketches}") + if not asm_sigs: + raise click.ClickException(f"No .sig files found in {assembly_sketches}") + + click.echo(f"Computing containment: {len(read_sigs)} read × {len(asm_sigs)} " + f"assembly sketches...") + matrix = compute_containment_matrix(read_sigs, asm_sigs, ksize=ksize) + write_containment_matrix(matrix, out_tsv) + click.echo(f"Containment matrix ({matrix.shape[0]}×{matrix.shape[1]}) → {out_tsv}") + + +# ── select ──────────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--sketches", type=click.Path(exists=True), + help="Directory of assembly .sig files (Jaccard selectors)") +@click.option("--containment-matrix", type=click.Path(exists=True), + help="Cross-sample containment matrix TSV (required for --selector containment)") +@click.option("--manifest", required=True, type=click.Path(exists=True)) +@click.option("--selector", default="archetype", show_default=True, + type=click.Choice(["random", "maxmin", "kmedoids", "archetype", + "greedy_var", "containment"])) +@click.option("--k", default="5", show_default=True, + help="Prototypes per sample, or 'auto' to infer from similarity structure") +@click.option("--adaptive-k-method", default="similarity_gap", show_default=True, + type=click.Choice(["scree_elbow", "similarity_gap", "saturation_curve", + "containment_saturation"])) +@click.option("--min-jaccard", default=0.1, show_default=True, + help="Minimum similarity/containment floor for candidate assemblies") +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--force", is_flag=True) +def select(sketches, containment_matrix, manifest, selector, k, adaptive_k_method, + min_jaccard, outdir, force): + """Select prototype assemblies for each sample.""" + import pandas as pd + from refrover.io import read_manifest, write_assignments + from refrover.selectors import SELECTOR_REGISTRY, CONTAINMENT_SELECTORS + from refrover.similarity import matrix_from_sigs, load_containment_matrix + + outdir = Path(outdir) + outdir.mkdir(parents=True, exist_ok=True) + assignments_path = outdir / "assignments.tsv" + + if assignments_path.exists() and not force: + click.echo(f"Assignments already exist at {assignments_path} (use --force to redo)") + return + + df = read_manifest(manifest) + + # Containment selectors run on a reads-vs-assembly containment matrix; + # Jaccard selectors run on an assembly-vs-assembly matrix from the sketches. + if selector in CONTAINMENT_SELECTORS: + if not containment_matrix: + raise click.ClickException( + f"--selector {selector} requires --containment-matrix" + ) + click.echo(f"Loading containment matrix from {containment_matrix}...") + sim_matrix = load_containment_matrix(containment_matrix) + else: + if not sketches: + raise click.ClickException( + f"--selector {selector} requires --sketches" + ) + sig_dir = Path(sketches) + sig_paths = sorted(sig_dir.glob("*.sig")) + if not sig_paths: + raise click.ClickException(f"No .sig files found in {sketches}") + click.echo(f"Computing similarity matrix from {len(sig_paths)} sketches...") + sim_matrix = matrix_from_sigs(sig_paths) + + use_adaptive_k = (str(k).lower() == "auto") + fixed_k = None if use_adaptive_k else int(k) + + sel_cls = SELECTOR_REGISTRY[selector] + + rows = [] + for _, row in df.iterrows(): + sid = row["sample_id"] + + if use_adaptive_k: + from refrover.adaptive_k import estimate_k + k_i = estimate_k(sim_matrix, sid, min_jaccard=min_jaccard, + method=adaptive_k_method) + else: + k_i = fixed_k + + sel = sel_cls(k=k_i, min_jaccard=min_jaccard) + prototypes = sel.select(sim_matrix, sid) + rows.append({ + "sample_id": sid, + "prototype_ids": ",".join(prototypes), + "n_prototypes": len(prototypes), + "k_estimated": k_i if use_adaptive_k else None, + }) + + assignments = pd.DataFrame(rows) + write_assignments(assignments, assignments_path) + click.echo(f"Assignments → {assignments_path}") + + +# ── align ───────────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--assignments", required=True, type=click.Path(exists=True)) +@click.option("--manifest", required=True, type=click.Path(exists=True)) +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--aligner", default="bwa-mem2", show_default=True, + type=click.Choice(["bwa-mem2", "bwa", "minimap2"])) +@click.option("--threads", default=8, show_default=True) +@click.option("--force", is_flag=True) +def align(assignments, manifest, outdir, aligner, threads, force): + """Align reads to selected prototype assemblies.""" + from refrover.align import run_alignment + from refrover.io import read_manifest + import pandas as pd + + df = read_manifest(manifest) + asgn = pd.read_csv(assignments, sep="\t") + bam_paths = run_alignment(df, asgn, outdir=outdir, aligner=aligner, + threads=threads, force=force) + click.echo(f"Aligned {len(bam_paths)} samples → {outdir}") + + +# ── coverage ────────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--bams", required=True, type=click.Path(exists=True), + help="Directory of sorted BAM files") +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--threads", default=8, show_default=True) +@click.option("--force", is_flag=True) +def coverage(bams, outdir, threads, force): + """Compute per-contig depth with CoverM.""" + from refrover.coverage import run_coverm + run_coverm(bams_dir=bams, outdir=outdir, threads=threads, force=force) + click.echo(f"Coverage → {outdir}") + + +# ── format ──────────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--coverage", "coverage_dir", required=True, type=click.Path(exists=True)) +@click.option("--binners", default="generic", show_default=True, + help="Comma-separated list: metabat2,semibin2,maxbin2,concoct,generic") +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--force", is_flag=True) +def format(coverage_dir, binners, outdir, force): + """Format coverage tables for downstream binners.""" + from refrover.formatters import format_for_binner + from refrover.coverage import load_coverage + + outdir = Path(outdir) + outdir.mkdir(parents=True, exist_ok=True) + + binner_list = [b.strip() for b in binners.split(",")] + cov_tsv = Path(coverage_dir) / "coverage.tsv" + df = load_coverage(cov_tsv) + + for binner in binner_list: + out = format_for_binner(df, binner=binner, outdir=outdir, force=force) + click.echo(f"{binner} → {out}") + + +# ── run ─────────────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--manifest", required=True, type=click.Path(exists=True)) +@click.option("--selector", default="archetype", show_default=True, + type=click.Choice(["random", "maxmin", "kmedoids", "archetype", + "greedy_var", "containment"])) +@click.option("--k", default=5, show_default=True) +@click.option("--min-jaccard", default=0.1, show_default=True, + help="Minimum similarity/containment floor for candidate assemblies") +@click.option("--containment-matrix", type=click.Path(exists=True), + help="Cross-sample containment matrix TSV (required for --selector containment)") +@click.option("--aligner", default="bwa-mem2", show_default=True, + type=click.Choice(["bwa-mem2", "bwa", "minimap2"])) +@click.option("--binners", default="generic", show_default=True) +@click.option("--threads", default=8, show_default=True) +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--force", is_flag=True) +def run(manifest, selector, k, min_jaccard, containment_matrix, aligner, binners, + threads, outdir, force): + """Run the full RefRover pipeline end-to-end.""" + from refrover.pipeline import RefRoverPipeline + from refrover.io import read_manifest + from refrover.selectors import SELECTOR_REGISTRY, CONTAINMENT_SELECTORS + from refrover.similarity import load_containment_matrix + + df = read_manifest(manifest) + sel_cls = SELECTOR_REGISTRY[selector] + sel = sel_cls(k=k, min_jaccard=min_jaccard) + binner_list = [b.strip() for b in binners.split(",")] + + cont_df = None + if selector in CONTAINMENT_SELECTORS: + if not containment_matrix: + raise click.ClickException( + f"--selector {selector} requires --containment-matrix" + ) + cont_df = load_containment_matrix(containment_matrix) + + pipeline = RefRoverPipeline( + manifest=df, + selector=sel, + binners=binner_list, + threads=threads, + outdir=Path(outdir), + force=force, + containment_matrix=cont_df, + ) + results = pipeline.run() + click.echo(f"Done. Coverage tables: {list(results.coverage_tables.keys())}") + + +# ── rank-selectors ──────────────────────────────────────────────────────────── + +@main.command(name="rank-selectors") +@click.option("--sketches", type=click.Path(exists=True), + help="Directory of assembly .sig files (Jaccard matrix)") +@click.option("--containment-matrix", type=click.Path(exists=True), + help="Cross-sample containment matrix TSV (alternative to --sketches)") +@click.option("--selectors", default="random,maxmin,kmedoids,greedy_var,containment", + show_default=True, help="Comma-separated selector IDs to rank") +@click.option("--k-range", default="3,5,8,10", show_default=True, + help="Comma-separated k values to test") +@click.option("--min-jaccard", default=0.1, show_default=True, + help="Minimum similarity/containment floor for candidate assemblies") +@click.option("--outdir", required=True, type=click.Path()) +def rank_selectors_cmd(sketches, containment_matrix, selectors, k_range, min_jaccard, outdir): + """Rank selectors by the alignment-free variance-explained proxy (Tier 1).""" + from refrover.benchmark import rank_selectors + from refrover.similarity import matrix_from_sigs, load_containment_matrix + + if bool(sketches) == bool(containment_matrix): + raise click.ClickException( + "Provide exactly one of --sketches or --containment-matrix." + ) + + if containment_matrix: + matrix = load_containment_matrix(containment_matrix) + else: + sig_paths = sorted(Path(sketches).glob("*.sig")) + if not sig_paths: + raise click.ClickException(f"No .sig files found in {sketches}") + matrix = matrix_from_sigs(sig_paths) + + selector_list = [s.strip() for s in selectors.split(",")] + k_list = [int(k.strip()) for k in k_range.split(",")] + + ranking = rank_selectors( + matrix, selectors=selector_list, k_values=k_list, min_similarity=min_jaccard, + ) + + outdir = Path(outdir) + outdir.mkdir(parents=True, exist_ok=True) + out_tsv = outdir / "selector_ranking.tsv" + ranking.to_csv(out_tsv, sep="\t", index=False) + + click.echo(ranking.to_string(index=False)) + click.echo(f"\nRanking → {out_tsv}") + + +# ── benchmark ───────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--manifest", required=True, type=click.Path(exists=True)) +@click.option("--truth", required=True, type=click.Path(exists=True), + help="Ground truth community TSV") +@click.option("--selectors", default="random,maxmin,kmedoids,archetype", + show_default=True, help="Comma-separated selector IDs to benchmark") +@click.option("--k-range", default="3,5,8,10", show_default=True, + help="Comma-separated k values to test") +@click.option("--checkm2-db", type=click.Path(exists=True), default=None, + help="Path to CheckM2 database (required for MAG quality evaluation)") +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--threads", default=8, show_default=True) +def benchmark(manifest, truth, selectors, k_range, checkm2_db, outdir, threads): + """Benchmark selectors against each other on a labelled dataset.""" + from refrover.benchmark import run_benchmark + from refrover.io import read_manifest + + df = read_manifest(manifest) + selector_list = [s.strip() for s in selectors.split(",")] + k_list = [int(k.strip()) for k in k_range.split(",")] + + run_benchmark( + manifest=df, + truth_path=truth, + selectors=selector_list, + k_values=k_list, + checkm2_db=checkm2_db, + outdir=Path(outdir), + threads=threads, + ) + click.echo(f"Benchmark results → {outdir}") diff --git a/src/refrover/containment.py b/src/refrover/containment.py new file mode 100644 index 0000000..f3d191d --- /dev/null +++ b/src/refrover/containment.py @@ -0,0 +1,78 @@ +""" +Cross-sample containment matrix computation. + +For each query (read sketch) and reference (assembly sketch), computes +containment(reads_i, assembly_j) = |hashes(reads_i) ∩ hashes(assembly_j)| / +|hashes(reads_i)| — the fraction of the query's read k-mers found in the +assembly. This is the signal ContainmentSelector selects on: it directly +predicts how well sample i's reads will map to assembly j. + +The matrix is square-ish and labelled by sample_id on both axes (rows = query +read samples, columns = candidate assemblies). It need not be symmetric: +containment is directional. +""" + +from pathlib import Path + +import pandas as pd + + +def _load_minhash(sig_path: Path | str, ksize: int): + """Load a single MinHash from a sourmash .sig file at the given ksize.""" + import sourmash + + sig = next(sourmash.load_file_as_signatures(str(sig_path), ksize=ksize)) + return sig.minhash + + +def sigs_by_sample(sketch_dir: Path | str) -> dict[str, Path]: + """Map sample_id -> .sig path for every sketch in a directory (stem = id).""" + sketch_dir = Path(sketch_dir) + return {p.stem: p for p in sorted(sketch_dir.glob("*.sig"))} + + +def compute_containment_matrix( + read_sigs: dict[str, Path | str], + assembly_sigs: dict[str, Path | str], + *, + ksize: int = 31, +) -> pd.DataFrame: + """ + Compute the reads-vs-assemblies containment matrix. + + Parameters + ---------- + read_sigs : dict + sample_id -> read sketch (.sig) path. + assembly_sigs : dict + sample_id -> assembly sketch (.sig) path. + ksize : int + k-mer size; read and assembly sketches must share it. + + Returns + ------- + pd.DataFrame + Rows = read sample_ids (sorted), columns = assembly sample_ids (sorted). + Values = containment(reads_row, assembly_col) in [0, 1]. Index name is + 'query_id', matching load_containment_matrix. + """ + asm_ids = sorted(assembly_sigs) + asm_mh = {a: _load_minhash(assembly_sigs[a], ksize) for a in asm_ids} + + rows: dict[str, dict[str, float]] = {} + for q in sorted(read_sigs): + read_mh = _load_minhash(read_sigs[q], ksize) + rows[q] = {a: float(read_mh.contained_by(asm_mh[a])) for a in asm_ids} + + df = pd.DataFrame.from_dict(rows, orient="index") + df = df.reindex(columns=asm_ids) + df.index.name = "query_id" + return df + + +def write_containment_matrix(matrix: pd.DataFrame, path: Path | str) -> Path: + """Write a containment matrix to TSV (the format load_containment_matrix reads).""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + matrix.to_csv(path, sep="\t") + return path diff --git a/src/refrover/coverage.py b/src/refrover/coverage.py new file mode 100644 index 0000000..cc0c94d --- /dev/null +++ b/src/refrover/coverage.py @@ -0,0 +1,131 @@ +""" +CoverM wrapper: compute per-contig depth (mean + variance) from sorted BAMs. + +Internal coverage schema +------------------------ +Everything downstream of `load_coverage` uses one normalized DataFrame: + + index contig name + column 'length' contig length (int), when available + '{sample}_depth' mean depth for that sample + '{sample}_var' depth variance for that sample (paired with _depth) + +`load_coverage` parses CoverM's raw `{sample} {Method}` columns into this schema. +The MetaBAT2 formatter consumes the variance columns; the others use only the +means. Keeping mean and variance together is what lets MetaBAT2 emit a real +per-sample variance instead of a placeholder. +""" + +import subprocess +from pathlib import Path +import pandas as pd + +DEPTH_SUFFIX = "_depth" +VAR_SUFFIX = "_var" +DEFAULT_METHODS = ("length", "mean", "variance") + + +def run_coverm( + bams_dir: Path | str, + outdir: Path | str, + *, + threads: int = 8, + force: bool = False, + methods: tuple[str, ...] = DEFAULT_METHODS, + coverm_path: str = "coverm", +) -> Path: + """ + Run CoverM on all sorted BAMs in bams_dir. + + Writes the raw CoverM table to outdir/coverage.tsv. By default requests + `length mean variance` so the MetaBAT2 formatter has a real variance column. + Use `load_coverage` to read the result into the normalized schema. + + Returns the path to coverage.tsv. + """ + bams_dir = Path(bams_dir) + outdir = Path(outdir) + outdir.mkdir(parents=True, exist_ok=True) + + out_tsv = outdir / "coverage.tsv" + if out_tsv.exists() and not force: + return out_tsv + + bam_files = sorted(bams_dir.glob("*.sorted.bam")) + if not bam_files: + raise FileNotFoundError(f"No *.sorted.bam files found in {bams_dir}") + + cmd = [ + coverm_path, "contig", + "--bam-files", *[str(b) for b in bam_files], + "--methods", *methods, + "--threads", str(threads), + "--output-file", str(out_tsv), + ] + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"CoverM failed:\n{result.stderr}") + + return out_tsv + + +def _clean_sample(stoit: str) -> str: + """Strip CoverM's BAM-derived decorations to recover the sample name.""" + name = stoit.strip() + for suffix in (".bam", ".sorted", ".sort"): + if name.endswith(suffix): + name = name[: -len(suffix)] + return name.strip() + + +def normalize_coverage(df: pd.DataFrame) -> pd.DataFrame: + """ + Normalize a raw CoverM table into RefRover's internal coverage schema. + + CoverM names columns `{sample} {Method}` (e.g. `S1.sorted Mean`, + `S1.sorted Variance`, `S1.sorted Length`). This maps: + ` Mean` -> `{sample}_depth` + ` Variance` -> `{sample}_var` + ` Length` -> a single `length` column (CoverM repeats it per BAM) + + If the frame already looks normalized (has any `_depth` column), it is + returned unchanged — so re-reading an already-normalized table is a no-op. + """ + if any(c.endswith(DEPTH_SUFFIX) for c in df.columns): + return df + + means: dict[str, pd.Series] = {} + variances: dict[str, pd.Series] = {} + length: pd.Series | None = None + + for col in df.columns: + low = col.lower() + if low.endswith(" mean"): + sample = _clean_sample(col[: -len(" Mean")]) + means[f"{sample}{DEPTH_SUFFIX}"] = df[col] + elif low.endswith(" variance"): + sample = _clean_sample(col[: -len(" Variance")]) + variances[f"{sample}{VAR_SUFFIX}"] = df[col] + elif low.endswith(" length") or low == "length": + if length is None: + length = df[col] + + out = pd.DataFrame(index=df.index) + if length is not None: + out["length"] = length.astype(int) + for name, series in means.items(): + out[name] = series + for name, series in variances.items(): + out[name] = series + + if not means: + # Nothing matched the CoverM naming — return the original frame so the + # caller at least sees its data rather than an empty table. + return df + return out + + +def load_coverage(tsv_path: Path | str) -> pd.DataFrame: + """Load a CoverM coverage.tsv and normalize it to the internal schema.""" + df = pd.read_csv(tsv_path, sep="\t", index_col=0) + return normalize_coverage(df) diff --git a/src/refrover/formatters/__init__.py b/src/refrover/formatters/__init__.py new file mode 100644 index 0000000..65f5445 --- /dev/null +++ b/src/refrover/formatters/__init__.py @@ -0,0 +1,18 @@ +from .registry import BINNER_REGISTRY +from pathlib import Path +import pandas as pd + + +def format_for_binner( + coverage_df: pd.DataFrame, + binner: str, + outdir: Path | str, + force: bool = False, +): + """Dispatch coverage_df to the formatter for the requested binner.""" + outdir = Path(outdir) + if binner not in BINNER_REGISTRY: + raise ValueError( + f"Unknown binner '{binner}'. Available: {sorted(BINNER_REGISTRY)}" + ) + return BINNER_REGISTRY[binner](coverage_df, outdir, force=force) diff --git a/src/refrover/formatters/_columns.py b/src/refrover/formatters/_columns.py new file mode 100644 index 0000000..7e06024 --- /dev/null +++ b/src/refrover/formatters/_columns.py @@ -0,0 +1,29 @@ +""" +Shared helpers for reading RefRover's normalized coverage schema. + +Schema (see refrover.coverage): + 'length' optional contig length + '{sample}_depth' mean depth per sample + '{sample}_var' depth variance per sample (optional, paired with _depth) +""" + +import pandas as pd + +DEPTH_SUFFIX = "_depth" +VAR_SUFFIX = "_var" + + +def mean_columns(df: pd.DataFrame) -> list[str]: + """Return the per-sample mean-depth columns, in order.""" + return [c for c in df.columns if c.endswith(DEPTH_SUFFIX)] + + +def sample_of(mean_col: str) -> str: + """Recover the sample name from a '{sample}_depth' column.""" + return mean_col[: -len(DEPTH_SUFFIX)] if mean_col.endswith(DEPTH_SUFFIX) else mean_col + + +def var_column(df: pd.DataFrame, mean_col: str) -> str | None: + """Return the variance column paired with a mean column, or None if absent.""" + candidate = sample_of(mean_col) + VAR_SUFFIX + return candidate if candidate in df.columns else None diff --git a/src/refrover/formatters/concoct.py b/src/refrover/formatters/concoct.py new file mode 100644 index 0000000..277df70 --- /dev/null +++ b/src/refrover/formatters/concoct.py @@ -0,0 +1,15 @@ +"""CONCOCT formatter: contig x sample matrix, tab-separated, integer depths.""" + +from pathlib import Path +import pandas as pd + +from ._columns import mean_columns + + +def write_concoct(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) -> Path: + out = outdir / "coverage_table.tsv" + if out.exists() and not force: + return out + df_out = coverage_df[mean_columns(coverage_df)].round(0).astype(int) + df_out.to_csv(out, sep="\t") + return out diff --git a/src/refrover/formatters/generic.py b/src/refrover/formatters/generic.py new file mode 100644 index 0000000..7716ffd --- /dev/null +++ b/src/refrover/formatters/generic.py @@ -0,0 +1,21 @@ +"""Generic TSV formatter: contig, length, sample1_depth, sample2_depth, ...""" + +from pathlib import Path +import pandas as pd + +from ._columns import mean_columns + + +def write_generic(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) -> Path: + """ + Write a generic coverage TSV: contig index, optional 'length', and one + mean-depth column per sample. Variance columns are omitted. + + Output: coverage_generic.tsv — tab-separated. + """ + out = outdir / "coverage_generic.tsv" + if out.exists() and not force: + return out + cols = (["length"] if "length" in coverage_df.columns else []) + mean_columns(coverage_df) + coverage_df[cols].to_csv(out, sep="\t") + return out diff --git a/src/refrover/formatters/maxbin2.py b/src/refrover/formatters/maxbin2.py new file mode 100644 index 0000000..50924b8 --- /dev/null +++ b/src/refrover/formatters/maxbin2.py @@ -0,0 +1,18 @@ +"""MaxBin2 formatter: one .abund file per sample, single depth column.""" + +from pathlib import Path +import pandas as pd + +from ._columns import mean_columns, sample_of + + +def write_maxbin2(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) -> list[Path]: + out_paths = [] + for col in mean_columns(coverage_df): + sample = sample_of(col) + out = outdir / f"{sample}.abund" + out_paths.append(out) + if out.exists() and not force: + continue + coverage_df[[col]].to_csv(out, sep="\t", header=False) + return out_paths diff --git a/src/refrover/formatters/metabat2.py b/src/refrover/formatters/metabat2.py new file mode 100644 index 0000000..f5ba836 --- /dev/null +++ b/src/refrover/formatters/metabat2.py @@ -0,0 +1,40 @@ +""" +MetaBAT2 formatter: jgi_summarize_bam_contig_depths format. + +Columns: contigName, contigLen, totalAvgDepth, sample1, sample1-var, sample2, sample2-var, ... + +The per-sample variance is taken from the coverage table's '{sample}_var' column +when present (CoverM run with `--methods variance`). If a sample has no paired +variance column, its variance falls back to 0.0. +""" + +from pathlib import Path +import pandas as pd + +from ._columns import mean_columns, sample_of, var_column + + +def write_metabat2(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) -> Path: + """Write depth.txt in MetaBAT2 jgi_summarize_bam_contig_depths format.""" + out = outdir / "depth.txt" + if out.exists() and not force: + return out + + depth_cols = mean_columns(coverage_df) + rows = [] + for contig, row in coverage_df.iterrows(): + rec = { + "contigName": contig, + "contigLen": int(row.get("length", 0)), + "totalAvgDepth": row[depth_cols].mean(), + } + for col in depth_cols: + sample = sample_of(col) + vcol = var_column(coverage_df, col) + rec[sample] = row[col] + rec[f"{sample}-var"] = row[vcol] if vcol is not None else 0.0 + rows.append(rec) + + df_out = pd.DataFrame(rows) + df_out.to_csv(out, sep="\t", index=False) + return out diff --git a/src/refrover/formatters/registry.py b/src/refrover/formatters/registry.py new file mode 100644 index 0000000..217c506 --- /dev/null +++ b/src/refrover/formatters/registry.py @@ -0,0 +1,15 @@ +"""Binner formatter registry.""" + +from .generic import write_generic +from .metabat2 import write_metabat2 +from .semibin2 import write_semibin2 +from .maxbin2 import write_maxbin2 +from .concoct import write_concoct + +BINNER_REGISTRY = { + "generic": write_generic, + "metabat2": write_metabat2, + "semibin2": write_semibin2, + "maxbin2": write_maxbin2, + "concoct": write_concoct, +} diff --git a/src/refrover/formatters/semibin2.py b/src/refrover/formatters/semibin2.py new file mode 100644 index 0000000..5c39d84 --- /dev/null +++ b/src/refrover/formatters/semibin2.py @@ -0,0 +1,14 @@ +"""SemiBin2 formatter: contig x sample depth matrix, no variance column.""" + +from pathlib import Path +import pandas as pd + +from ._columns import mean_columns + + +def write_semibin2(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) -> Path: + out = outdir / "coverage_metabinner.tsv" + if out.exists() and not force: + return out + coverage_df[mean_columns(coverage_df)].to_csv(out, sep="\t") + return out diff --git a/src/refrover/io.py b/src/refrover/io.py new file mode 100644 index 0000000..72a2a8d --- /dev/null +++ b/src/refrover/io.py @@ -0,0 +1,50 @@ +from pathlib import Path +import pandas as pd + +REQUIRED_COLUMNS = {"sample_id", "assembly"} +OPTIONAL_COLUMNS = {"r1", "r2", "long_reads"} + + +def read_manifest(path: Path | str) -> pd.DataFrame: + """ + Read and validate a RefRover sample manifest (TSV). + + Required columns: sample_id, assembly + At least one of r1 or long_reads must be present per row. + r2 requires r1; absent r2 = single-end short reads. + """ + df = pd.read_csv(path, sep="\t", dtype=str) + df.columns = df.columns.str.strip().str.lower() + + missing = REQUIRED_COLUMNS - set(df.columns) + if missing: + raise ValueError(f"Manifest missing required columns: {sorted(missing)}") + + if df["sample_id"].duplicated().any(): + dupes = df.loc[df["sample_id"].duplicated(keep=False), "sample_id"].unique() + raise ValueError(f"Duplicate sample_id values: {sorted(dupes)}") + + has_r1 = "r1" in df.columns + has_long = "long_reads" in df.columns + if not has_r1 and not has_long: + raise ValueError("Manifest must have at least one of: r1, long_reads") + + for _, row in df.iterrows(): + sid = row["sample_id"] + r1 = row.get("r1") if has_r1 else None + long = row.get("long_reads") if has_long else None + r2 = row.get("r2") if "r2" in df.columns else None + + if pd.isna(r1) and pd.isna(long): + raise ValueError( + f"Sample '{sid}': at least one of r1 or long_reads must be provided" + ) + if not pd.isna(r2) and pd.isna(r1): + raise ValueError(f"Sample '{sid}': r2 provided but r1 is missing") + + return df.reset_index(drop=True) + + +def write_assignments(assignments: pd.DataFrame, path: Path | str) -> None: + """Write prototype assignment table produced by a selector.""" + assignments.to_csv(path, sep="\t", index=False) diff --git a/src/refrover/pipeline.py b/src/refrover/pipeline.py new file mode 100644 index 0000000..72cb08f --- /dev/null +++ b/src/refrover/pipeline.py @@ -0,0 +1,102 @@ +""" +RefRoverPipeline: orchestrates the full sketch → select → align → coverage → format pipeline. +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional +import pandas as pd + +from refrover.selectors.base import BaseSelector +from refrover.selectors.containment import ContainmentSelector + + +@dataclass +class PipelineResults: + coverage_tables: dict[str, Path] = field(default_factory=dict) + assignments: pd.DataFrame = field(default_factory=pd.DataFrame) + + +class RefRoverPipeline: + def __init__( + self, + manifest: pd.DataFrame, + selector: BaseSelector, + binners: list[str], + threads: int = 8, + outdir: Path = Path("refrover_out"), + force: bool = False, + containment_matrix: Optional[pd.DataFrame] = None, + ): + self.manifest = manifest + self.selector = selector + self.binners = binners + self.threads = threads + self.outdir = Path(outdir) + self.force = force + self.containment_matrix = containment_matrix + + if isinstance(selector, ContainmentSelector) and containment_matrix is None: + raise ValueError( + "ContainmentSelector requires a containment_matrix to be passed " + "to RefRoverPipeline." + ) + + def run(self) -> PipelineResults: + from refrover.sketch import sketch_assemblies, compare_sketches + from refrover.similarity import load_similarity_matrix + from refrover.io import write_assignments + from refrover.align import run_alignment + from refrover.coverage import run_coverm, load_coverage + from refrover.formatters import format_for_binner + + outdir = self.outdir + sketch_dir = outdir / "sketches" + bam_dir = outdir / "bams" + cov_dir = outdir / "coverage" + fmt_dir = outdir / "formatted" + + # 1-2. Build the selection matrix. + # Containment selectors use a precomputed reads-vs-assembly containment + # matrix; Jaccard selectors sketch the assemblies and compare them. + if self.containment_matrix is not None: + sim_matrix = self.containment_matrix + else: + sig_paths = sketch_assemblies( + self.manifest["assembly"].tolist(), + outdir=sketch_dir, + threads=self.threads, + force=self.force, + ) + sim_csv = outdir / "similarity.csv" + compare_sketches(sig_paths, output_csv=sim_csv, force=self.force) + sim_matrix = load_similarity_matrix(sim_csv) + + # 3. Prototype selection + rows = [] + for _, row in self.manifest.iterrows(): + sid = row["sample_id"] + prototypes = self.selector.select(sim_matrix, sid) + rows.append({"sample_id": sid, "prototype_ids": ",".join(prototypes), + "n_prototypes": len(prototypes)}) + assignments = pd.DataFrame(rows) + write_assignments(assignments, outdir / "assignments.tsv") + + # 4. Align + run_alignment( + self.manifest, assignments, + outdir=bam_dir, threads=self.threads, force=self.force, + ) + + # 5. Coverage + cov_tsv = run_coverm(bam_dir, outdir=cov_dir, threads=self.threads, force=self.force) + coverage_df = load_coverage(cov_tsv) + + # 6. Format + fmt_dir.mkdir(parents=True, exist_ok=True) + coverage_tables = {} + for binner in self.binners + ["generic"]: + out = format_for_binner(coverage_df, binner=binner, outdir=fmt_dir, force=self.force) + coverage_tables[binner] = out if isinstance(out, Path) else out[0] + + return PipelineResults(coverage_tables=coverage_tables, assignments=assignments) diff --git a/src/refrover/selectors/__init__.py b/src/refrover/selectors/__init__.py new file mode 100644 index 0000000..e198154 --- /dev/null +++ b/src/refrover/selectors/__init__.py @@ -0,0 +1,36 @@ +from .base import BaseSelector +from .random import RandomSelector +from .maxmin import MaxMinSelector +from .kmedoids import KMedoidsSelector +from .archetype import ArchetypeSelector +from .greedy_var import GreedyVarSelector +from .containment import ContainmentSelector +from .feedback import FeedbackSelector + +SELECTOR_REGISTRY: dict[str, type[BaseSelector]] = { + "random": RandomSelector, + "maxmin": MaxMinSelector, + "kmedoids": KMedoidsSelector, + "archetype": ArchetypeSelector, + "greedy_var": GreedyVarSelector, + "containment": ContainmentSelector, + "feedback": FeedbackSelector, +} + +# Selectors that operate on a reads-vs-assembly containment matrix instead of +# the assembly-vs-assembly Jaccard matrix. The CLI/pipeline use this to decide +# which matrix to feed the selector. +CONTAINMENT_SELECTORS: set[str] = {"containment"} + +__all__ = [ + "BaseSelector", + "RandomSelector", + "MaxMinSelector", + "KMedoidsSelector", + "ArchetypeSelector", + "GreedyVarSelector", + "ContainmentSelector", + "FeedbackSelector", + "SELECTOR_REGISTRY", + "CONTAINMENT_SELECTORS", +] diff --git a/src/refrover/selectors/archetype.py b/src/refrover/selectors/archetype.py new file mode 100644 index 0000000..61c885c --- /dev/null +++ b/src/refrover/selectors/archetype.py @@ -0,0 +1,62 @@ +import numpy as np +import pandas as pd +from .base import BaseSelector + + +class ArchetypeSelector(BaseSelector): + """ + Archetype analysis in Jaccard similarity space. + + Treats each sample's row in the similarity submatrix as a feature vector, + then finds the k archetypes (extreme points spanning the convex hull). + Each archetype is mapped back to the nearest actual sample. + + Archetypes span the diversity space rather than summarising it, making + them preferable to medoids for differential coverage: you want assemblies + that together cover all the genomic variation present in the dataset, not + assemblies that are typical representatives. + + Requires the `archetypes` package (pip install archetypes). + """ + + def __init__(self, k: int, min_jaccard: float = 0.1, random_state: int | None = 42): + super().__init__(k, min_jaccard) + self.random_state = random_state + + def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + try: + from archetypes import AA + except ImportError as exc: + raise ImportError( + "ArchetypeSelector requires the 'archetypes' package: pip install archetypes" + ) from exc + + candidates = self._candidates(sim_matrix, query_id) + if len(candidates) <= self.k: + return self._trim_to_k(candidates) + + X = sim_matrix.loc[candidates, candidates].to_numpy().astype(float) + k = min(self.k, len(candidates)) + + aa = AA(n_archetypes=k, random_state=self.random_state) + aa.fit(X) + + # Map each archetype to its nearest candidate (unique assignment) + selected, used = [], set() + for arch in aa.archetypes_: + dists = np.linalg.norm(X - arch, axis=1) + for idx in np.argsort(dists): + cid = candidates[idx] + if cid not in used: + selected.append(cid) + used.add(cid) + break + + # Ensure query is included — if not, replace the archetype whose nearest + # sample is closest to query + if query_id not in selected: + sim_row = sim_matrix.loc[query_id, selected].to_numpy() + most_replaceable = selected[int(np.argmax(sim_row))] + selected[selected.index(most_replaceable)] = query_id + + return selected diff --git a/src/refrover/selectors/base.py b/src/refrover/selectors/base.py new file mode 100644 index 0000000..2350e6c --- /dev/null +++ b/src/refrover/selectors/base.py @@ -0,0 +1,45 @@ +from abc import ABC, abstractmethod +import warnings +import pandas as pd + + +class BaseSelector(ABC): + def __init__(self, k: int, min_jaccard: float = 0.1): + if k < 1: + raise ValueError(f"k must be >= 1, got {k}") + if not 0.0 <= min_jaccard <= 1.0: + raise ValueError(f"min_jaccard must be in [0, 1], got {min_jaccard}") + self.k = k + self.min_jaccard = min_jaccard + + @abstractmethod + def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + """ + Select up to k prototype IDs for query_id. + + sim_matrix: symmetric pairwise Jaccard similarity indexed by sample ID, + with 1.0 on the diagonal. + query_id: sample whose reads will be aligned to the returned prototypes. + returns: list of selected prototype IDs, length <= k. + """ + + def _candidates(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + if query_id not in sim_matrix.index: + raise KeyError(f"query_id '{query_id}' not found in similarity matrix") + row = sim_matrix.loc[query_id] + candidates = row[row >= self.min_jaccard].index.tolist() + if not candidates: + raise ValueError( + f"No candidates for '{query_id}' with min_jaccard={self.min_jaccard}. " + "Lower min_jaccard or check that the similarity matrix is correct." + ) + return candidates + + def _trim_to_k(self, candidates: list[str]) -> list[str]: + if len(candidates) < self.k: + warnings.warn( + f"Only {len(candidates)} candidates available (k={self.k}). " + "Returning all candidates.", + stacklevel=3, + ) + return candidates[: self.k] diff --git a/src/refrover/selectors/containment.py b/src/refrover/selectors/containment.py new file mode 100644 index 0000000..fd8afb4 --- /dev/null +++ b/src/refrover/selectors/containment.py @@ -0,0 +1,147 @@ +""" +ContainmentSelector: prototype selection driven by cross-sample containment. + +Unlike the Jaccard-based selectors, this one operates on a reads-vs-assembly +*containment* matrix rather than an assembly-vs-assembly Jaccard matrix. +containment[i, j] = fraction of sample i's read k-mers found in assembly j — +a direct measure of how well sample i's reads will map to assembly j, which is +the quantity differential-coverage binning actually cares about. + +Identifier model +---------------- +A prototype is identified by the sample_id of the sample whose assembly it is. +The containment matrix is therefore square and labelled by sample_id on both +axes: rows are query (read) samples, columns are candidate assemblies. The +manifest resolves each returned sample_id to its assembly FASTA at align time +(see refrover.align._resolve_fastas). This is the same contract every other +selector follows, which lets the pipeline treat all selectors polymorphically. + +Two modes +--------- +unweighted (default) + MaxMin greedy on containment-profile distance (1 - Pearson correlation of + the two assemblies' cross-sample containment columns). Selects prototypes + that are maximally diverse in read-mapping space. + +weighted (requires `weights`) + Each candidate's MaxMin score is multiplied by a per-assembly weight + (a pd.Series indexed by sample_id). Use this to up-weight assemblies + predicted to carry more differential signal. +""" + +import warnings +from typing import Optional + +import numpy as np +import pandas as pd + +from .base import BaseSelector + + +class ContainmentSelector(BaseSelector): + """ + Prototype selection using cross-sample read containment. + + Parameters + ---------- + k : int + Number of prototypes to select. + min_containment : float + Minimum containment(query_reads, prototype_assembly) required. + Assemblies below this threshold are excluded — reads won't map well. + weights : pd.Series or None + Per-assembly weight Series (index = sample_id). If provided, candidate + MaxMin scores are multiplied by weight[sample_id] (default 1.0 for any + assembly not in the Series). + min_jaccard : float or None + Alias accepted so the selector can be constructed uniformly from + SELECTOR_REGISTRY with a `min_jaccard=` keyword. When given it overrides + `min_containment` (the threshold semantics are identical: a minimum + similarity floor on the query row). + """ + + def __init__( + self, + k: int, + min_containment: float = 0.05, + weights: Optional[pd.Series] = None, + min_jaccard: Optional[float] = None, + ): + threshold = min_containment if min_jaccard is None else min_jaccard + super().__init__(k, min_jaccard=threshold) + # Expose under the domain-appropriate name as well. + self.min_containment = threshold + self.weights = weights + + def select(self, containment_matrix: pd.DataFrame, query_id: str) -> list[str]: + """ + Select up to k prototype sample_ids for query_id. + + Parameters + ---------- + containment_matrix : pd.DataFrame + Square, labelled by sample_id on both axes. Rows = query (read) + samples, columns = candidate assemblies. Values = containment. + query_id : str + Row index of the query sample. + + Returns + ------- + list[str] + Ordered prototype sample_ids, length <= k. The query's own assembly + is first when it clears the threshold (containment 1.0 by definition). + """ + if query_id not in containment_matrix.index: + raise KeyError(f"query_id '{query_id}' not found in containment matrix rows") + + query_row = containment_matrix.loc[query_id] + candidates = query_row[query_row >= self.min_containment].index.tolist() + if not candidates: + raise ValueError( + f"No assemblies meet min_containment={self.min_containment} for " + f"'{query_id}'. Lower min_containment or check the containment matrix." + ) + + # A prototype is identified by sample_id; the query's own assembly is the + # column labelled query_id. Pin it first when present. + if query_id in candidates: + candidates = [query_id] + [c for c in candidates if c != query_id] + else: + candidates = sorted(candidates, key=lambda c: -float(query_row[c])) + + if len(candidates) < self.k: + warnings.warn( + f"Only {len(candidates)} candidates available (k={self.k}). " + "Returning all candidates.", + UserWarning, + stacklevel=2, + ) + if len(candidates) <= self.k: + return candidates + + # Precompute the candidate-vs-candidate containment-profile distance ONCE. + # dist(j, k) = 1 - Pearson correlation of the two assemblies' columns. + corr = containment_matrix[candidates].corr().to_numpy() + dist = 1.0 - np.nan_to_num(corr, nan=0.0) + np.fill_diagonal(dist, 0.0) + + weights = None + if self.weights is not None: + weights = np.array([float(self.weights.get(c, 1.0)) for c in candidates]) + + # MaxMin greedy, seeded with the query's own assembly (index 0). + selected = [0] + remaining = list(range(1, len(candidates))) + min_dist = dist[0, remaining].copy() + + while len(selected) < self.k and remaining: + scores = min_dist if weights is None else min_dist * weights[remaining] + best_pos = int(np.argmax(scores)) + best = remaining[best_pos] + selected.append(best) + remaining.pop(best_pos) + min_dist = np.delete(min_dist, best_pos) + if remaining: + min_dist = np.minimum(min_dist, dist[best, remaining]) + + return [candidates[i] for i in selected] diff --git a/src/refrover/selectors/feedback.py b/src/refrover/selectors/feedback.py new file mode 100644 index 0000000..251abca --- /dev/null +++ b/src/refrover/selectors/feedback.py @@ -0,0 +1,20 @@ +import pandas as pd +from .base import BaseSelector + + +class FeedbackSelector(BaseSelector): + """ + Coverage-feedback selector: map reads → measure actual variance → iterate. + + Selects an initial set of prototypes, aligns reads, computes real per-contig + coverage variance, then refines the selection to maximise observed variance. + + Most accurate selector; most expensive (requires alignment per iteration). + Not yet implemented — depends on align.py + coverage.py being wired up. + """ + + def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + raise NotImplementedError( + "FeedbackSelector requires alignment + CoverM; not yet implemented. " + "Use archetype or greedy_var instead." + ) diff --git a/src/refrover/selectors/greedy_var.py b/src/refrover/selectors/greedy_var.py new file mode 100644 index 0000000..a2714d6 --- /dev/null +++ b/src/refrover/selectors/greedy_var.py @@ -0,0 +1,42 @@ +import pandas as pd +from .base import BaseSelector + + +class GreedyVarSelector(BaseSelector): + """ + Greedy coverage-variance maximization. + + For each candidate assembly j, estimates the marginal contribution to + cross-sample coverage variance using two factors: + 1. var_score: variance of j's similarity to all other samples — proxy + for how much j's coverage will differ across samples. + 2. diversity: minimum distance from j to already-selected prototypes — + penalizes adding a prototype redundant with one already chosen. + + Score = var_score * diversity. Query is always the first prototype. + """ + + def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + candidates = self._candidates(sim_matrix, query_id) + if len(candidates) <= self.k: + return self._trim_to_k(candidates) + + all_ids = sim_matrix.columns.tolist() + selected = [query_id] + remaining = [c for c in candidates if c != query_id] + + # Pre-compute per-candidate variance scores (these don't change) + var_scores = {c: float(sim_matrix.loc[c, all_ids].var()) for c in remaining} + + while len(selected) < self.k and remaining: + best, best_score = None, -1.0 + for c in remaining: + min_dist = min(1.0 - float(sim_matrix.loc[c, s]) for s in selected) + score = var_scores[c] * min_dist + if score > best_score: + best_score = score + best = c + selected.append(best) + remaining.remove(best) + + return selected diff --git a/src/refrover/selectors/kmedoids.py b/src/refrover/selectors/kmedoids.py new file mode 100644 index 0000000..9879fd4 --- /dev/null +++ b/src/refrover/selectors/kmedoids.py @@ -0,0 +1,76 @@ +import numpy as np +import pandas as pd +from .base import BaseSelector + + +def _kmedoids(dist: np.ndarray, k: int, rng: np.random.Generator) -> list[int]: + """ + PAM-style k-medoids on a precomputed distance matrix. + Initialises with greedy furthest-first, then runs swap improvements. + Returns indices of the k medoids. + """ + n = len(dist) + k = min(k, n) + + # Greedy furthest-first initialisation + first = int(rng.integers(n)) + medoids = [first] + while len(medoids) < k: + min_dists = dist[:, medoids].min(axis=1) + medoids.append(int(np.argmax(min_dists))) + + def total_cost(meds): + return dist[:, meds].min(axis=1).sum() + + # Swap phase: try replacing each medoid with each non-medoid + improved = True + while improved: + improved = False + for i in range(len(medoids)): + current_cost = total_cost(medoids) + for cand in range(n): + if cand in medoids: + continue + trial = medoids.copy() + trial[i] = cand + if total_cost(trial) < current_cost - 1e-10: + medoids[i] = cand + improved = True + break + if improved: + break + + return medoids + + +class KMedoidsSelector(BaseSelector): + """ + k-medoids clustering in Jaccard distance space; returns the k medoids. + Medoids are real samples (not centroids), making them valid alignment + references. + """ + + def __init__(self, k: int, min_jaccard: float = 0.1, random_state: int | None = 42): + super().__init__(k, min_jaccard) + self.random_state = random_state + + def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + candidates = self._candidates(sim_matrix, query_id) + if len(candidates) <= self.k: + return self._trim_to_k(candidates) + + dist = 1.0 - sim_matrix.loc[candidates, candidates].to_numpy() + np.fill_diagonal(dist, 0.0) + + rng = np.random.default_rng(self.random_state) + medoid_indices = _kmedoids(dist, self.k, rng) + + # Ensure query is in the result — swap the medoid most similar to query + # for query itself if query wasn't selected + query_idx = candidates.index(query_id) + if query_idx not in medoid_indices: + sim_to_query = sim_matrix.loc[query_id, candidates].to_numpy() + closest = min(medoid_indices, key=lambda i: -sim_to_query[i]) + medoid_indices[medoid_indices.index(closest)] = query_idx + + return [candidates[i] for i in medoid_indices] diff --git a/src/refrover/selectors/maxmin.py b/src/refrover/selectors/maxmin.py new file mode 100644 index 0000000..2d5ed06 --- /dev/null +++ b/src/refrover/selectors/maxmin.py @@ -0,0 +1,44 @@ +import numpy as np +import pandas as pd +from .base import BaseSelector + + +class MaxMinSelector(BaseSelector): + """ + Greedy MaxMin: iteratively selects the candidate furthest from all + already-selected prototypes, maximising the minimum pairwise distance + in Jaccard space. + + Initialises with the query's own assembly (Jaccard = 1.0 to itself), + so the query is always the first prototype. + """ + + def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + candidates = self._candidates(sim_matrix, query_id) + if len(candidates) <= self.k: + return self._trim_to_k(candidates) + + dist = 1.0 - sim_matrix.loc[candidates, candidates].to_numpy() + n = len(candidates) + id_to_idx = {sid: i for i, sid in enumerate(candidates)} + + # Seed with query itself + selected = [id_to_idx[query_id]] + remaining = list(range(n)) + remaining.remove(id_to_idx[query_id]) + + # min_dist_to_selected[i] = distance from candidate i to its nearest selected prototype + min_dist = dist[np.array(remaining), :][:, selected].min(axis=1) + + while len(selected) < self.k and remaining: + best_pos = int(np.argmax(min_dist)) + best_idx = remaining[best_pos] + selected.append(best_idx) + remaining.pop(best_pos) + min_dist = np.delete(min_dist, best_pos) + + if remaining: + new_dists = dist[np.array(remaining), best_idx] + min_dist = np.minimum(min_dist, new_dists) + + return [candidates[i] for i in selected] diff --git a/src/refrover/selectors/random.py b/src/refrover/selectors/random.py new file mode 100644 index 0000000..f57050d --- /dev/null +++ b/src/refrover/selectors/random.py @@ -0,0 +1,18 @@ +import random +import pandas as pd +from .base import BaseSelector + + +class RandomSelector(BaseSelector): + """Baseline: random k prototypes within Jaccard threshold.""" + + def __init__(self, k: int, min_jaccard: float = 0.1, seed: int | None = None): + super().__init__(k, min_jaccard) + self.seed = seed + + def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]: + candidates = self._candidates(sim_matrix, query_id) + if len(candidates) <= self.k: + return self._trim_to_k(candidates) + rng = random.Random(self.seed) + return rng.sample(candidates, self.k) diff --git a/src/refrover/similarity.py b/src/refrover/similarity.py new file mode 100644 index 0000000..3227265 --- /dev/null +++ b/src/refrover/similarity.py @@ -0,0 +1,108 @@ +""" +Load and operate on pairwise Jaccard similarity matrices. + +The primary input format is the CSV written by `sourmash compare --csv`. +""" + +import numpy as np +import pandas as pd +from pathlib import Path + + +def load_similarity_matrix(csv_path: Path | str) -> pd.DataFrame: + """ + Load a sourmash compare CSV into a labelled similarity DataFrame. + + sourmash compare --csv writes: + - Row 0: comma-separated sample labels (from sig names or filenames) + - Rows 1..n: the matrix rows, one per sample + + Returns a square DataFrame with sample IDs as both index and columns. + """ + csv_path = Path(csv_path) + df = pd.read_csv(csv_path, index_col=0) + # The first column after the index is also a label column in sourmash output; + # ensure the matrix is square and symmetric. + if df.shape[0] != df.shape[1]: + raise ValueError( + f"Similarity matrix is not square: {df.shape}. " + "Expected output from `sourmash compare --csv`." + ) + df.index.name = None + df.columns.name = None + return df + + +def jaccard_to_distance(sim_matrix: pd.DataFrame) -> pd.DataFrame: + """Convert a Jaccard similarity matrix to a distance matrix (1 - Jaccard).""" + return 1.0 - sim_matrix + + +def load_containment_matrix(tsv_path: Path | str) -> pd.DataFrame: + """ + Load a cross-sample containment matrix (TSV) for ContainmentSelector. + + Expected layout (as written by compute_containment.py): + - First column: query sample_id (the row index). + - Remaining columns: one per candidate assembly, labelled by sample_id. + - Values: containment(reads_row, assembly_col) in [0, 1]. + + The matrix is square and labelled by sample_id on both axes, but need not be + symmetric (containment is directional). Returns a labelled DataFrame. + """ + tsv_path = Path(tsv_path) + df = pd.read_csv(tsv_path, sep="\t", index_col=0) + df.index.name = None + df.columns.name = None + return df + + +def filter_by_jaccard( + sim_matrix: pd.DataFrame, + query_id: str, + min_jaccard: float, +) -> list[str]: + """ + Return sample IDs with Jaccard(query, candidate) >= min_jaccard. + + The query itself is always included (diagonal = 1.0). + Raises KeyError if query_id is not in the matrix. + """ + if query_id not in sim_matrix.index: + raise KeyError(f"query_id '{query_id}' not found in similarity matrix") + row = sim_matrix.loc[query_id] + return row[row >= min_jaccard].index.tolist() + + +def matrix_from_sigs( + sig_paths: "list[Path]", + *, + ksize: int = 31, +) -> pd.DataFrame: + """ + Compute a pairwise Jaccard similarity matrix directly from sourmash .sig files, + without calling the sourmash CLI. + + Uses the sourmash Python API for speed (no subprocess overhead). + Loads all sketches into memory then computes all n*(n-1)/2 pairs. + + Returns a labelled DataFrame (index and columns = sig names). + """ + import sourmash + + sigs = [] + for p in sig_paths: + sig = next(sourmash.load_file_as_signatures(str(p), ksize=ksize)) + sigs.append(sig) + + n = len(sigs) + names = [s.name or Path(sig_paths[i]).stem for i, s in enumerate(sigs)] + mat = np.eye(n, dtype=np.float64) + + for i in range(n): + for j in range(i + 1, n): + j_val = sigs[i].minhash.jaccard(sigs[j].minhash) + mat[i, j] = j_val + mat[j, i] = j_val + + return pd.DataFrame(mat, index=names, columns=names) diff --git a/src/refrover/sketch.py b/src/refrover/sketch.py new file mode 100644 index 0000000..5bf5216 --- /dev/null +++ b/src/refrover/sketch.py @@ -0,0 +1,96 @@ +""" +Wrappers around sourmash for sketching assemblies and computing pairwise similarity. +""" + +import subprocess +import sys +from pathlib import Path +from typing import Sequence + + +def sketch_assemblies( + fasta_paths: Sequence[Path | str], + outdir: Path | str, + *, + ksize: int = 31, + scaled: int = 1000, + threads: int = 1, + force: bool = False, + sourmash_path: str = "sourmash", +) -> list[Path]: + """ + Sketch each FASTA with sourmash sketch dna. + + Returns a list of .sig paths (one per input), in the same order. + Skips existing .sig files unless force=True. + """ + outdir = Path(outdir) + outdir.mkdir(parents=True, exist_ok=True) + + to_sketch: list[tuple[Path, Path]] = [] + all_sigs: list[Path] = [] + + for fa in fasta_paths: + fa = Path(fa) + sig_path = outdir / (fa.stem + ".sig") + all_sigs.append(sig_path) + if sig_path.exists() and not force: + continue + to_sketch.append((fa, sig_path)) + + if not to_sketch: + return all_sigs + + # sourmash sketch dna processes multiple files in one call + fasta_list = [str(fa) for fa, _ in to_sketch] + sig_list = [str(sig) for _, sig in to_sketch] + + cmd = [ + sourmash_path, "sketch", "dna", + f"--param-string=k={ksize},scaled={scaled}", + "--output-dir", str(outdir), + "-p", f"k={ksize},scaled={scaled}", + ] + fasta_list + + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"sourmash sketch failed:\n{result.stderr}" + ) + + return all_sigs + + +def compare_sketches( + sig_paths: Sequence[Path | str], + output_csv: Path | str, + *, + ksize: int = 31, + force: bool = False, + sourmash_path: str = "sourmash", +) -> Path: + """ + Run sourmash compare on a list of .sig files to get a pairwise CSV matrix. + + Returns the path to the output CSV. + Raises RuntimeError on failure. + """ + output_csv = Path(output_csv) + if output_csv.exists() and not force: + return output_csv + + output_csv.parent.mkdir(parents=True, exist_ok=True) + + cmd = [ + sourmash_path, "compare", + "--csv", str(output_csv), + f"--ksize={ksize}", + ] + [str(p) for p in sig_paths] + + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"sourmash compare failed:\n{result.stderr}" + ) + + return output_csv diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..703ad19 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,83 @@ +import numpy as np +import pandas as pd +import pytest +from pathlib import Path + + +def _make_sim(matrix: np.ndarray, ids: list[str]) -> pd.DataFrame: + np.fill_diagonal(matrix, 1.0) + sym = (matrix + matrix.T) / 2 + np.fill_diagonal(sym, 1.0) + return pd.DataFrame(sym, index=ids, columns=ids) + + +@pytest.fixture +def clustered_sim(): + """ + 12 samples in 3 clusters of 4. + Within-cluster Jaccard ~ 0.6-0.8, cross-cluster ~ 0.05. + """ + rng = np.random.default_rng(42) + n = 12 + ids = [f"s{i:02d}" for i in range(n)] + m = np.full((n, n), 0.05) + for start in (0, 4, 8): + for i in range(start, start + 4): + for j in range(start, start + 4): + if i != j: + m[i, j] = 0.6 + 0.2 * rng.random() + return _make_sim(m, ids) + + +@pytest.fixture +def uniform_sim(): + """10 samples, all pairs roughly equidistant (Jaccard 0.1-0.4).""" + rng = np.random.default_rng(0) + n = 10 + ids = [f"s{i:02d}" for i in range(n)] + m = rng.uniform(0.1, 0.4, (n, n)) + return _make_sim(m, ids) + + +@pytest.fixture +def sparse_sim(): + """8 samples, most pairs below min_jaccard=0.1 (stress-tests fallback).""" + rng = np.random.default_rng(7) + n = 8 + ids = [f"s{i:02d}" for i in range(n)] + m = rng.uniform(0.01, 0.09, (n, n)) + # s00 has a few valid neighbours + m[0, 1] = m[1, 0] = 0.3 + m[0, 2] = m[2, 0] = 0.2 + return _make_sim(m, ids) + + +@pytest.fixture +def coverage_df(): + """ + Minimal coverage DataFrame in RefRover's normalized schema: 5 contigs x + 3 samples, with a 'length' column, a '{sample}_depth' mean column, and a + paired '{sample}_var' variance column per sample. + """ + contigs = [f"contig_{i}" for i in range(5)] + data = { + "length": [1000, 2000, 500, 3000, 1500], + "s1_depth": [10.2, 0.0, 5.5, 22.1, 8.8], + "s2_depth": [0.0, 15.3, 6.1, 18.9, 0.0], + "s3_depth": [7.7, 12.0, 0.0, 25.4, 3.3], + "s1_var": [2.1, 0.0, 1.2, 4.4, 1.8], + "s2_var": [0.0, 3.3, 1.5, 3.9, 0.0], + "s3_var": [1.7, 2.4, 0.0, 5.1, 0.6], + } + return pd.DataFrame(data, index=contigs) + + +@pytest.fixture +def tiny_manifest(tmp_path): + """Write a minimal valid manifest TSV and return its path.""" + content = "sample_id\tassembly\tr1\n" + content += "S1\tassemblies/S1.fasta\treads/S1_R1.fastq.gz\n" + content += "S2\tassemblies/S2.fasta\treads/S2_R1.fastq.gz\n" + p = tmp_path / "manifest.tsv" + p.write_text(content) + return p diff --git a/tests/test_adaptive_k.py b/tests/test_adaptive_k.py new file mode 100644 index 0000000..ed6c35e --- /dev/null +++ b/tests/test_adaptive_k.py @@ -0,0 +1,111 @@ +import pytest +import numpy as np +import pandas as pd +from refrover.adaptive_k import ( + estimate_k, _scree_elbow, _similarity_gap, _saturation_curve, + _containment_saturation, +) + + +@pytest.fixture +def three_cluster_sim(): + """ + 15 samples, 3 tight clusters of 5. + Within-cluster Jaccard ~ 0.7, cross-cluster ~ 0.05. + The intrinsic dimensionality is ~3 → expect k_hat around 3. + """ + rng = np.random.default_rng(0) + n = 15 + ids = [f"s{i:02d}" for i in range(n)] + m = np.full((n, n), 0.05) + for start in (0, 5, 10): + for i in range(start, start + 5): + for j in range(start, start + 5): + if i != j: + m[i, j] = 0.65 + 0.1 * rng.random() + np.fill_diagonal(m, 1.0) + sym = (m + m.T) / 2 + np.fill_diagonal(sym, 1.0) + return pd.DataFrame(sym, index=ids, columns=ids) + + +def test_estimate_k_returns_int(three_cluster_sim): + k = estimate_k(three_cluster_sim, "s00", min_jaccard=0.01) + assert isinstance(k, int) + + +def test_estimate_k_respects_k_min(three_cluster_sim): + k = estimate_k(three_cluster_sim, "s00", min_jaccard=0.01, k_min=4) + assert k >= 4 + + +def test_estimate_k_respects_k_max(three_cluster_sim): + k = estimate_k(three_cluster_sim, "s00", min_jaccard=0.01, k_max=5) + assert k <= 5 + + +def test_estimate_k_unknown_query(three_cluster_sim): + with pytest.raises(KeyError): + estimate_k(three_cluster_sim, "does_not_exist") + + +def test_all_methods_return_valid_k(three_cluster_sim): + for method in ("scree_elbow", "similarity_gap", "saturation_curve", + "containment_saturation"): + k = estimate_k(three_cluster_sim, "s00", min_jaccard=0.01, + method=method, k_min=2, k_max=10) + assert 2 <= k <= 10, f"{method} returned k={k} outside [2, 10]" + + +def test_containment_saturation_finds_low_dim(three_cluster_sim): + """3-cluster data: the greedy residual-variance elbow should be small.""" + ids = three_cluster_sim.columns.tolist() + k = _containment_saturation(three_cluster_sim, ids, k_min=1, k_max=12) + assert 1 <= k <= 6 + + +def test_containment_saturation_monotone_curve_threshold(): + """ + A matrix with one dominant variance axis and a long tail of tiny ones + should saturate after very few prototypes. + """ + rng = np.random.default_rng(1) + n = 20 + base = rng.normal(size=n) + cols = {"big0": base * 5.0, "big1": rng.normal(size=n) * 4.0} + for i in range(8): + cols[f"tiny{i}"] = base * 5.0 + 1e-4 * rng.normal(size=n) # near-duplicates + df = pd.DataFrame(cols, index=[f"s{i}" for i in range(n)]) + k = _containment_saturation(df, list(df.columns), k_min=1, k_max=10) + assert k <= 3 # only ~2 real axes of variation + + +def test_similarity_gap_finds_natural_break(): + """ + Construct a similarity row with a clear gap: 3 high-sim candidates, + then a large drop, then low-sim candidates. + similarity_gap should return k=3. + """ + ids = [f"s{i}" for i in range(8)] + sims = [1.0, 0.75, 0.70, 0.68, # top 4 — tight group + 0.15, 0.12, 0.11, 0.10] # bottom 4 — after big drop + row = pd.Series(dict(zip(ids, sims))) + candidates = ids + + k = _similarity_gap(row, candidates, k_min=2, k_max=8) + assert k == 4 # gap is between index 3 (0.68) and 4 (0.15) + + +def test_saturation_curve_stops_at_elbow(three_cluster_sim): + """With 3 clusters, saturation should stop adding prototypes at ~3.""" + ids = three_cluster_sim.columns.tolist() + k = _saturation_curve(three_cluster_sim, "s00", ids, k_min=2, k_max=10) + # After the first 3 prototypes (one per cluster) diversity gain collapses + assert k <= 6 # generous upper bound + + +def test_scree_elbow_clustered(three_cluster_sim): + """3-cluster data has intrinsic dim ~3; scree elbow should be near that.""" + ids = three_cluster_sim.columns.tolist() + k = _scree_elbow(three_cluster_sim, ids, k_min=2, k_max=10) + assert 2 <= k <= 6 diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py new file mode 100644 index 0000000..e8ea468 --- /dev/null +++ b/tests/test_benchmark.py @@ -0,0 +1,96 @@ +"""Tests for the Tier-1 variance-explained selector ranking.""" + +import numpy as np +import pandas as pd +import pytest + +from refrover.benchmark import unique_variance_explained, rank_selectors + + +@pytest.fixture +def feature_matrix(): + """6 samples × 4 feature columns with structure (not full rank in practice).""" + rng = np.random.default_rng(0) + cols = ["a", "b", "c", "d"] + data = rng.normal(size=(6, 4)) + # make 'd' a near-duplicate of 'a' (correlated/redundant) + data[:, 3] = data[:, 0] + 1e-3 * rng.normal(size=6) + return pd.DataFrame(data, index=[f"s{i}" for i in range(6)], columns=cols) + + +def test_all_columns_explain_everything(feature_matrix): + ve = unique_variance_explained(feature_matrix, list(feature_matrix.columns)) + assert ve == pytest.approx(1.0, abs=1e-9) + + +def test_empty_selection_explains_nothing(feature_matrix): + assert unique_variance_explained(feature_matrix, []) == 0.0 + + +def test_value_in_unit_interval(feature_matrix): + ve = unique_variance_explained(feature_matrix, ["a", "b"]) + assert 0.0 <= ve <= 1.0 + + +def test_redundant_column_adds_little(feature_matrix): + """'d' duplicates 'a', so adding it explains almost no new variance.""" + base = unique_variance_explained(feature_matrix, ["a"]) + plus_dupe = unique_variance_explained(feature_matrix, ["a", "d"]) + plus_new = unique_variance_explained(feature_matrix, ["a", "b"]) + assert plus_dupe - base < plus_new - base + + +def test_more_columns_never_explain_less(feature_matrix): + """Orthogonal projection: a superset never explains less (nested sets).""" + small = unique_variance_explained(feature_matrix, ["a", "b"]) + big = unique_variance_explained(feature_matrix, ["a", "b", "c"]) + assert big >= small - 1e-9 + + +# ── rank_selectors ───────────────────────────────────────────────────────────── + +def test_rank_returns_expected_columns(clustered_sim): + ranking = rank_selectors( + clustered_sim, selectors=["random", "maxmin"], k_values=[3, 4], + min_similarity=0.1, + ) + assert list(ranking.columns) == [ + "selector", "k", "mean_var_explained", "median_var_explained", "n_samples", + ] + assert set(ranking["selector"]) == {"random", "maxmin"} + assert set(ranking["k"]) == {3, 4} + + +def test_rank_maxmin_monotone_in_k(clustered_sim): + """MaxMin is greedily nested, so more prototypes explain at least as much.""" + ranking = rank_selectors( + clustered_sim, selectors=["maxmin"], k_values=[3, 4], min_similarity=0.1, + ) + by_k = ranking.set_index("k")["mean_var_explained"] + assert by_k[4] >= by_k[3] - 1e-9 + + +def test_rank_skips_feedback_with_warning(clustered_sim): + with pytest.warns(UserWarning): + ranking = rank_selectors( + clustered_sim, selectors=["maxmin", "feedback"], k_values=[3], + min_similarity=0.1, + ) + assert set(ranking["selector"]) == {"maxmin"} + + +def test_rank_is_reproducible(clustered_sim): + """The random baseline is seeded, so the ranking is deterministic.""" + a = rank_selectors(clustered_sim, selectors=["random"], k_values=[3], min_similarity=0.1) + b = rank_selectors(clustered_sim, selectors=["random"], k_values=[3], min_similarity=0.1) + pd.testing.assert_frame_equal(a, b) + + +def test_rank_on_containment_matrix(clustered_sim): + """All selectors, including containment, run on a single shared matrix.""" + ranking = rank_selectors( + clustered_sim, selectors=["maxmin", "containment"], k_values=[3], + min_similarity=0.1, + ) + assert set(ranking["selector"]) == {"maxmin", "containment"} + assert (ranking["mean_var_explained"] >= 0).all() diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..29bfd46 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,90 @@ +"""CLI integration tests for the containment selector wiring.""" + +import numpy as np +import pandas as pd +import pytest +from click.testing import CliRunner + +from refrover.cli import main + + +@pytest.fixture +def containment_files(tmp_path): + """ + Write a 4-sample containment matrix TSV and a matching manifest. + Two groups: (s0, s1) and (s2, s3); within-group containment high, + cross-group low. + """ + ids = ["s0", "s1", "s2", "s3"] + mat = np.array([ + [1.00, 0.40, 0.05, 0.04], + [0.42, 1.00, 0.06, 0.05], + [0.05, 0.04, 1.00, 0.45], + [0.06, 0.05, 0.43, 1.00], + ]) + cont = pd.DataFrame(mat, index=ids, columns=ids) + cont.index.name = "query_id" + cont_path = tmp_path / "containment.tsv" + cont.to_csv(cont_path, sep="\t") + + manifest = pd.DataFrame({ + "sample_id": ids, + "assembly": [f"assemblies/{s}.fasta" for s in ids], + "r1": [f"reads/{s}_R1.fastq.gz" for s in ids], + }) + man_path = tmp_path / "manifest.tsv" + manifest.to_csv(man_path, sep="\t", index=False) + + return cont_path, man_path, tmp_path + + +def test_select_containment_writes_assignments(containment_files): + cont_path, man_path, tmp_path = containment_files + outdir = tmp_path / "out" + runner = CliRunner() + result = runner.invoke(main, [ + "select", + "--selector", "containment", + "--containment-matrix", str(cont_path), + "--manifest", str(man_path), + "--k", "2", + "--min-jaccard", "0.05", + "--outdir", str(outdir), + ]) + assert result.exit_code == 0, result.output + + assignments = pd.read_csv(outdir / "assignments.tsv", sep="\t") + assert set(assignments["sample_id"]) == {"s0", "s1", "s2", "s3"} + assert (assignments["n_prototypes"] <= 2).all() + # Each sample's own assembly should be the first prototype. + for _, row in assignments.iterrows(): + first = str(row["prototype_ids"]).split(",")[0] + assert first == row["sample_id"] + + +def test_select_containment_requires_matrix(containment_files): + _, man_path, tmp_path = containment_files + runner = CliRunner() + result = runner.invoke(main, [ + "select", + "--selector", "containment", + "--manifest", str(man_path), + "--k", "2", + "--outdir", str(tmp_path / "out2"), + ]) + assert result.exit_code != 0 + assert "requires --containment-matrix" in result.output + + +def test_select_jaccard_requires_sketches(containment_files): + _, man_path, tmp_path = containment_files + runner = CliRunner() + result = runner.invoke(main, [ + "select", + "--selector", "maxmin", + "--manifest", str(man_path), + "--k", "2", + "--outdir", str(tmp_path / "out3"), + ]) + assert result.exit_code != 0 + assert "requires --sketches" in result.output diff --git a/tests/test_containment.py b/tests/test_containment.py new file mode 100644 index 0000000..d7e01fc --- /dev/null +++ b/tests/test_containment.py @@ -0,0 +1,102 @@ +"""Tests for the package containment-matrix computation.""" + +from pathlib import Path + +import pytest + +from refrover.containment import ( + compute_containment_matrix, + sigs_by_sample, + write_containment_matrix, +) +from refrover.similarity import load_containment_matrix + + +def _write_sig(hashes, name, path, ksize=31): + """Build a sourmash signature from raw hashes and save it to path.""" + from sourmash import MinHash, SourmashSignature + from sourmash.sourmash_args import SaveSignaturesToLocation + + mh = MinHash(n=0, ksize=ksize, scaled=1) + for h in hashes: + mh.add_hash(h) + sig = SourmashSignature(mh, name=name) + with SaveSignaturesToLocation(str(path)) as save: + save.add(sig) + + +@pytest.fixture +def sketches(tmp_path): + """ + Two samples rA, rB with read and assembly sketches. + reads rA = {1,2,3,4}, reads rB = {5,6,7,8} + asm rA = {1,2,3,4}, asm rB = {3,4,5,6} + Expected containment (reads row vs assembly col): + rA->rA 1.0 rA->rB 0.5 + rB->rA 0.0 rB->rB 0.5 + """ + reads_dir = tmp_path / "reads" + asm_dir = tmp_path / "asm" + reads_dir.mkdir() + asm_dir.mkdir() + + _write_sig([1, 2, 3, 4], "rA", reads_dir / "rA.sig") + _write_sig([5, 6, 7, 8], "rB", reads_dir / "rB.sig") + _write_sig([1, 2, 3, 4], "rA", asm_dir / "rA.sig") + _write_sig([3, 4, 5, 6], "rB", asm_dir / "rB.sig") + + return reads_dir, asm_dir + + +def test_sigs_by_sample_maps_stems(sketches): + reads_dir, asm_dir = sketches + rmap = sigs_by_sample(reads_dir) + assert set(rmap) == {"rA", "rB"} + assert rmap["rA"].name == "rA.sig" + + +def test_containment_values(sketches): + reads_dir, asm_dir = sketches + matrix = compute_containment_matrix( + sigs_by_sample(reads_dir), sigs_by_sample(asm_dir) + ) + assert matrix.loc["rA", "rA"] == pytest.approx(1.0) + assert matrix.loc["rA", "rB"] == pytest.approx(0.5) + assert matrix.loc["rB", "rA"] == pytest.approx(0.0) + assert matrix.loc["rB", "rB"] == pytest.approx(0.5) + + +def test_matrix_shape_and_labels(sketches): + reads_dir, asm_dir = sketches + matrix = compute_containment_matrix( + sigs_by_sample(reads_dir), sigs_by_sample(asm_dir) + ) + assert matrix.shape == (2, 2) + assert list(matrix.index) == ["rA", "rB"] + assert list(matrix.columns) == ["rA", "rB"] + assert matrix.index.name == "query_id" + + +def test_write_and_reload_roundtrip(sketches, tmp_path): + reads_dir, asm_dir = sketches + matrix = compute_containment_matrix( + sigs_by_sample(reads_dir), sigs_by_sample(asm_dir) + ) + out = write_containment_matrix(matrix, tmp_path / "containment_matrix.tsv") + reloaded = load_containment_matrix(out) + assert reloaded.loc["rA", "rB"] == pytest.approx(0.5) + assert list(reloaded.columns) == ["rA", "rB"] + + +def test_feeds_containment_selector(sketches): + """The computed matrix is directly consumable by ContainmentSelector.""" + from refrover.selectors import ContainmentSelector + + reads_dir, asm_dir = sketches + matrix = compute_containment_matrix( + sigs_by_sample(reads_dir), sigs_by_sample(asm_dir) + ) + sel = ContainmentSelector(k=2, min_containment=0.01) + result = sel.select(matrix, "rA") + assert result[0] == "rA" # own assembly first (containment 1.0) + assert set(result) <= {"rA", "rB"} diff --git a/tests/test_coverage.py b/tests/test_coverage.py new file mode 100644 index 0000000..84baf90 --- /dev/null +++ b/tests/test_coverage.py @@ -0,0 +1,70 @@ +"""Tests for CoverM output normalization into RefRover's internal schema.""" + +import pandas as pd +import pytest + +from refrover.coverage import normalize_coverage, load_coverage + + +def _raw_coverm_frame(): + """Mimic a raw `coverm contig --methods length mean variance` table. + + CoverM names columns `{stoit} {Method}` and repeats Length per BAM. + """ + return pd.DataFrame( + { + "S1.sorted Length": [1000, 2000, 500], + "S1.sorted Mean": [10.2, 0.0, 5.5], + "S1.sorted Variance": [2.1, 0.0, 1.2], + "S2.sorted Length": [1000, 2000, 500], + "S2.sorted Mean": [0.0, 15.3, 6.1], + "S2.sorted Variance": [0.0, 3.3, 1.5], + }, + index=pd.Index(["contig_0", "contig_1", "contig_2"], name="Contig"), + ) + + +def test_normalize_maps_methods_to_schema(): + out = normalize_coverage(_raw_coverm_frame()) + assert list(out.columns) == ["length", "S1_depth", "S2_depth", "S1_var", "S2_var"] + assert out.loc["contig_0", "S1_depth"] == 10.2 + assert out.loc["contig_1", "S2_var"] == 3.3 + + +def test_normalize_collapses_repeated_length(): + out = normalize_coverage(_raw_coverm_frame()) + assert (out["length"] == [1000, 2000, 500]).all() + assert out["length"].dtype.kind == "i" # integer length + + +def test_normalize_handles_missing_length(): + raw = _raw_coverm_frame().drop(columns=["S1.sorted Length", "S2.sorted Length"]) + out = normalize_coverage(raw) + assert "length" not in out.columns + assert list(out.columns) == ["S1_depth", "S2_depth", "S1_var", "S2_var"] + + +def test_normalize_passthrough_when_already_normalized(coverage_df): + """An already-normalized frame (has _depth columns) is returned unchanged.""" + out = normalize_coverage(coverage_df) + assert out is coverage_df + + +def test_load_coverage_normalizes_from_disk(tmp_path): + raw = _raw_coverm_frame() + p = tmp_path / "coverage.tsv" + raw.to_csv(p, sep="\t") + out = load_coverage(p) + assert "S1_depth" in out.columns + assert "S1_var" in out.columns + assert out.loc["contig_2", "S1_depth"] == 5.5 + + +def test_normalize_mean_without_variance(): + """Mean-only CoverM output still normalizes (no _var columns produced).""" + raw = pd.DataFrame( + {"S1.sorted Mean": [3.0, 4.0], "S2.sorted Mean": [1.0, 2.0]}, + index=pd.Index(["c0", "c1"], name="Contig"), + ) + out = normalize_coverage(raw) + assert list(out.columns) == ["S1_depth", "S2_depth"] diff --git a/tests/test_formatters/__init__.py b/tests/test_formatters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_formatters/test_generic.py b/tests/test_formatters/test_generic.py new file mode 100644 index 0000000..55eb3c3 --- /dev/null +++ b/tests/test_formatters/test_generic.py @@ -0,0 +1,32 @@ +import pandas as pd +import pytest +from refrover.formatters import format_for_binner + + +def test_generic_writes_tsv(coverage_df, tmp_path): + out = format_for_binner(coverage_df, binner="generic", outdir=tmp_path) + assert out.exists() + df = pd.read_csv(out, sep="\t", index_col=0) + assert list(df.columns) == ["length", "s1_depth", "s2_depth", "s3_depth"] + assert len(df) == 5 + + +def test_generic_skips_existing(coverage_df, tmp_path): + out1 = format_for_binner(coverage_df, binner="generic", outdir=tmp_path) + mtime1 = out1.stat().st_mtime + out2 = format_for_binner(coverage_df, binner="generic", outdir=tmp_path) + assert out1 == out2 + assert out2.stat().st_mtime == mtime1 # file untouched + + +def test_generic_force_overwrites(coverage_df, tmp_path): + out1 = format_for_binner(coverage_df, binner="generic", outdir=tmp_path) + mtime1 = out1.stat().st_mtime + import time; time.sleep(0.01) + out2 = format_for_binner(coverage_df, binner="generic", outdir=tmp_path, force=True) + assert out2.stat().st_mtime >= mtime1 + + +def test_unknown_binner_raises(coverage_df, tmp_path): + with pytest.raises(ValueError, match="Unknown binner"): + format_for_binner(coverage_df, binner="not_a_binner", outdir=tmp_path) diff --git a/tests/test_formatters/test_metabat2.py b/tests/test_formatters/test_metabat2.py new file mode 100644 index 0000000..9067b46 --- /dev/null +++ b/tests/test_formatters/test_metabat2.py @@ -0,0 +1,53 @@ +import pandas as pd +from refrover.formatters import format_for_binner + + +def test_metabat2_filename(coverage_df, tmp_path): + out = format_for_binner(coverage_df, binner="metabat2", outdir=tmp_path) + assert out.name == "depth.txt" + + +def test_metabat2_required_columns(coverage_df, tmp_path): + out = format_for_binner(coverage_df, binner="metabat2", outdir=tmp_path) + df = pd.read_csv(out, sep="\t") + assert "contigName" in df.columns + assert "contigLen" in df.columns + assert "totalAvgDepth" in df.columns + + +def test_metabat2_variance_columns(coverage_df, tmp_path): + out = format_for_binner(coverage_df, binner="metabat2", outdir=tmp_path) + df = pd.read_csv(out, sep="\t") + # Each sample has a paired -var column carrying the real variance from the + # coverage table (not a zero placeholder). + for sample in ["s1", "s2", "s3"]: + assert sample in df.columns + assert f"{sample}-var" in df.columns + pd.testing.assert_series_equal( + df[f"{sample}-var"].reset_index(drop=True), + coverage_df[f"{sample}_var"].reset_index(drop=True), + check_names=False, + ) + + +def test_metabat2_variance_falls_back_to_zero_when_absent(coverage_df, tmp_path): + """A coverage table with no _var columns yields zero variance (back-compat).""" + means_only = coverage_df[["length", "s1_depth", "s2_depth", "s3_depth"]] + out = format_for_binner(means_only, binner="metabat2", outdir=tmp_path) + df = pd.read_csv(out, sep="\t") + for sample in ["s1", "s2", "s3"]: + assert (df[f"{sample}-var"] == 0.0).all() + + +def test_metabat2_total_avg_depth(coverage_df, tmp_path): + out = format_for_binner(coverage_df, binner="metabat2", outdir=tmp_path) + df = pd.read_csv(out, sep="\t") + # totalAvgDepth for first contig: mean of 10.2, 0.0, 7.7 = 5.9666... + expected = (10.2 + 0.0 + 7.7) / 3 + assert abs(df.loc[0, "totalAvgDepth"] - expected) < 1e-4 + + +def test_metabat2_row_count(coverage_df, tmp_path): + out = format_for_binner(coverage_df, binner="metabat2", outdir=tmp_path) + df = pd.read_csv(out, sep="\t") + assert len(df) == len(coverage_df) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py new file mode 100644 index 0000000..dc4c280 --- /dev/null +++ b/tests/test_pipeline.py @@ -0,0 +1,228 @@ +""" +End-to-end pipeline tests using mocked external tools. + +The pipeline calls sourmash, bwa-mem2/minimap2, samtools, and coverm via +subprocess. These tests patch those calls so the suite runs without any +external binaries installed. +""" + +import pandas as pd +import pytest +from pathlib import Path +from unittest.mock import patch, MagicMock + + +# ── io tests ────────────────────────────────────────────────────────────────── + +def test_read_manifest_valid(tiny_manifest): + from refrover.io import read_manifest + df = read_manifest(tiny_manifest) + assert list(df["sample_id"]) == ["S1", "S2"] + + +def test_read_manifest_missing_column(tmp_path): + from refrover.io import read_manifest + p = tmp_path / "bad.tsv" + p.write_text("sample_id\tr1\nS1\treads/S1.fastq.gz\n") + with pytest.raises(ValueError, match="missing required columns"): + read_manifest(p) + + +def test_read_manifest_duplicate_ids(tmp_path): + from refrover.io import read_manifest + p = tmp_path / "dup.tsv" + p.write_text("sample_id\tassembly\tr1\n" + "S1\ta.fa\tr1.fq\n" + "S1\tb.fa\tr2.fq\n") + with pytest.raises(ValueError, match="Duplicate"): + read_manifest(p) + + +def test_read_manifest_r2_without_r1(tmp_path): + from refrover.io import read_manifest + p = tmp_path / "bad_r2.tsv" + p.write_text("sample_id\tassembly\tr2\n" + "S1\ta.fa\treads/r2.fq\n") + with pytest.raises(ValueError, match="r1"): + read_manifest(p) + + +def test_write_read_assignments_roundtrip(tmp_path): + from refrover.io import write_assignments + df = pd.DataFrame({ + "sample_id": ["S1", "S2"], + "prototype_ids": ["S1,S2", "S1"], + "n_prototypes": [2, 1], + }) + out = tmp_path / "assignments.tsv" + write_assignments(df, out) + df2 = pd.read_csv(out, sep="\t") + assert list(df2["sample_id"]) == ["S1", "S2"] + assert df2.loc[0, "n_prototypes"] == 2 + + +# ── similarity tests ────────────────────────────────────────────────────────── + +def test_load_similarity_matrix(tmp_path): + from refrover.similarity import load_similarity_matrix + # Write a minimal sourmash compare CSV + csv = tmp_path / "sim.csv" + csv.write_text(",S1,S2,S3\nS1,1.0,0.5,0.2\nS2,0.5,1.0,0.3\nS3,0.2,0.3,1.0\n") + df = load_similarity_matrix(csv) + assert df.shape == (3, 3) + assert df.loc["S1", "S2"] == pytest.approx(0.5) + + +def test_load_similarity_matrix_not_square(tmp_path): + from refrover.similarity import load_similarity_matrix + csv = tmp_path / "bad.csv" + csv.write_text(",S1,S2\nS1,1.0,0.5\n") + with pytest.raises(ValueError, match="not square"): + load_similarity_matrix(csv) + + +def test_filter_by_jaccard(clustered_sim): + from refrover.similarity import filter_by_jaccard + result = filter_by_jaccard(clustered_sim, "s00", min_jaccard=0.5) + # Only within-cluster samples (0.6-0.8) should pass; cross-cluster (0.05) should not + assert "s00" in result + for r in result: + assert clustered_sim.loc["s00", r] >= 0.5 + + +def test_filter_by_jaccard_unknown_query(clustered_sim): + from refrover.similarity import filter_by_jaccard + with pytest.raises(KeyError): + filter_by_jaccard(clustered_sim, "does_not_exist", min_jaccard=0.1) + + +# ── sketch subprocess mock ──────────────────────────────────────────────────── + +def test_sketch_assemblies_skips_existing(tmp_path): + from refrover.sketch import sketch_assemblies + + # Pre-create the expected .sig file + fa = tmp_path / "asm.fasta" + fa.write_text(">c1\nACGT\n") + sig = tmp_path / "asm.sig" + sig.write_text("fake") + + with patch("subprocess.run") as mock_run: + result = sketch_assemblies([fa], outdir=tmp_path) + mock_run.assert_not_called() # should skip because sig exists + assert result == [sig] + + +def test_sketch_assemblies_calls_sourmash(tmp_path): + from refrover.sketch import sketch_assemblies + + fa = tmp_path / "asm.fasta" + fa.write_text(">c1\nACGT\n") + + mock_result = MagicMock() + mock_result.returncode = 0 + with patch("subprocess.run", return_value=mock_result) as mock_run: + sketch_assemblies([fa], outdir=tmp_path) + assert mock_run.called + cmd = mock_run.call_args[0][0] + assert "sourmash" in cmd[0] + assert "sketch" in cmd + + +# ── formatter integration ───────────────────────────────────────────────────── + +def test_all_binners_produce_output(coverage_df, tmp_path): + from refrover.formatters import format_for_binner, BINNER_REGISTRY + for binner in BINNER_REGISTRY: + result = format_for_binner(coverage_df, binner=binner, outdir=tmp_path) + if isinstance(result, list): + assert all(p.exists() for p in result) + else: + assert result.exists() + + +# ── align tests ─────────────────────────────────────────────────────────────── + +def _make_mock_run(returncode=0): + m = MagicMock() + m.returncode = returncode + return m + + +def test_align_skips_existing_bam(tmp_path): + from refrover.align import run_alignment + + manifest = pd.DataFrame({ + "sample_id": ["S1"], + "assembly": [str(tmp_path / "S1.fasta")], + "r1": [str(tmp_path / "S1_R1.fastq.gz")], + }) + assignments = pd.DataFrame({ + "sample_id": ["S1"], + "prototype_ids": ["S1"], + }) + # Pre-create the BAM so it should be skipped + bam = tmp_path / "S1.sorted.bam" + bam.write_text("fake") + + with patch("subprocess.run") as mock_run: + result = run_alignment(manifest, assignments, outdir=tmp_path) + mock_run.assert_not_called() + assert result == [bam] + + +def test_concat_fastas(tmp_path): + from refrover.align import _concat_fastas + + fa1 = tmp_path / "a.fasta" + fa2 = tmp_path / "b.fasta" + fa1.write_text(">seq1\nACGT\n") + fa2.write_text(">seq2\nTTTT\n") + out = tmp_path / "combined.fasta" + + _concat_fastas([fa1, fa2], out) + content = out.read_text() + assert ">seq1" in content + assert ">seq2" in content + + +def test_align_bwa_calls_correct_subcommands(tmp_path): + from refrover.align import run_alignment + + fa = tmp_path / "S1.fasta" + fa.write_text(">c1\nACGT\n") + r1 = tmp_path / "S1_R1.fastq.gz" + r1.write_text("") + + manifest = pd.DataFrame({ + "sample_id": ["S1"], + "assembly": [str(fa)], + "r1": [str(r1)], + }) + assignments = pd.DataFrame({ + "sample_id": ["S1"], + "prototype_ids": ["S1"], + }) + + calls = [] + + def fake_run(cmd, **kwargs): + calls.append(cmd) + m = MagicMock() + m.returncode = 0 + # For samtools sort, create the output BAM so pipeline doesn't error + if "sort" in cmd: + out_idx = cmd.index("-o") + 1 + Path(cmd[out_idx]).write_text("fake_bam") + return m + + with patch("subprocess.run", side_effect=fake_run), \ + patch("builtins.open", side_effect=open): + try: + run_alignment(manifest, assignments, outdir=tmp_path, aligner="bwa-mem2") + except Exception: + pass # may fail on SAM write details; we just check subprocess calls + + cmd_names = [c[0] for c in calls if c] + assert any("bwa-mem2" in str(c) or "bwa" in str(c) for c in cmd_names), \ + f"Expected bwa call, got: {cmd_names}" diff --git a/tests/test_selectors/__init__.py b/tests/test_selectors/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_selectors/test_archetype.py b/tests/test_selectors/test_archetype.py new file mode 100644 index 0000000..c05582d --- /dev/null +++ b/tests/test_selectors/test_archetype.py @@ -0,0 +1,39 @@ +import pytest +from refrover.selectors import ArchetypeSelector + + +def test_returns_k_results(clustered_sim): + sel = ArchetypeSelector(k=3) + assert len(sel.select(clustered_sim, "s00")) == 3 + + +def test_query_always_included(clustered_sim): + sel = ArchetypeSelector(k=3, min_jaccard=0.0) + assert "s08" in sel.select(clustered_sim, "s08") + + +def test_no_duplicates(clustered_sim): + sel = ArchetypeSelector(k=4, min_jaccard=0.0) + result = sel.select(clustered_sim, "s00") + assert len(result) == len(set(result)) + + +def test_spans_clusters(clustered_sim): + """Archetypes should select samples from all 3 distinct clusters.""" + sel = ArchetypeSelector(k=3, min_jaccard=0.0) + result = sel.select(clustered_sim, "s00") + cluster_of = lambda sid: int(sid[1:]) // 4 + assert len({cluster_of(r) for r in result}) == 3 + + +def test_valid_sample_ids(clustered_sim): + sel = ArchetypeSelector(k=3) + result = sel.select(clustered_sim, "s00") + assert all(r in clustered_sim.columns for r in result) + + +def test_fewer_candidates_than_k(sparse_sim): + sel = ArchetypeSelector(k=5, min_jaccard=0.1) + with pytest.warns(UserWarning, match="candidates available"): + result = sel.select(sparse_sim, "s00") + assert len(result) <= 5 diff --git a/tests/test_selectors/test_containment.py b/tests/test_selectors/test_containment.py new file mode 100644 index 0000000..aa40bbf --- /dev/null +++ b/tests/test_selectors/test_containment.py @@ -0,0 +1,127 @@ +import numpy as np +import pandas as pd +import pytest +from refrover.selectors.containment import ContainmentSelector +from refrover.selectors.base import BaseSelector +from refrover.selectors import SELECTOR_REGISTRY, CONTAINMENT_SELECTORS + + +@pytest.fixture +def containment_mat(): + """ + 6 samples x 6 assemblies. + Diagonal is 1.0 (each sample's own assembly). + Samples in the same 'group' share similar containment profiles. + Group A: s0, s1, s2 (assemblies a0-a2, contain each other's reads well) + Group B: s3, s4, s5 (assemblies a3-a5) + Cross-group containment is low (~0.05). + """ + rng = np.random.default_rng(42) + sample_ids = [f"s{i}" for i in range(6)] + asm_ids = [f"s{i}" for i in range(6)] # assembly IDs match sample IDs + + mat = np.full((6, 6), 0.04) + for grp in [(0, 1, 2), (3, 4, 5)]: + for i in grp: + for j in grp: + if i == j: + mat[i, j] = 1.0 + else: + mat[i, j] = 0.3 + 0.2 * rng.random() + + return pd.DataFrame(mat, index=sample_ids, columns=asm_ids) + + +def test_returns_k_results(containment_mat): + sel = ContainmentSelector(k=3, min_containment=0.05) + result = sel.select(containment_mat, "s0") + assert len(result) == 3 + + +def test_own_assembly_first(containment_mat): + sel = ContainmentSelector(k=3, min_containment=0.05) + result = sel.select(containment_mat, "s0") + assert result[0] == "s0" + + +def test_all_meet_min_containment(containment_mat): + sel = ContainmentSelector(k=3, min_containment=0.1) + result = sel.select(containment_mat, "s0") + for rid in result: + assert containment_mat.loc["s0", rid] >= 0.1 + + +def test_spans_groups_when_threshold_low(containment_mat): + """With min_containment=0.01 all assemblies are candidates; diversity should pull cross-group.""" + sel = ContainmentSelector(k=4, min_containment=0.01) + result = sel.select(containment_mat, "s0") + group_a = {"s0", "s1", "s2"} + from_b = [r for r in result if r not in group_a] + assert len(from_b) >= 1 + + +def test_unknown_query_raises(containment_mat): + sel = ContainmentSelector(k=3) + with pytest.raises(KeyError, match="not found"): + sel.select(containment_mat, "does_not_exist") + + +def test_fewer_than_k_warns(containment_mat): + # min_containment=0.5 leaves only the diagonal (1.0) in range + sel = ContainmentSelector(k=4, min_containment=0.5) + with pytest.warns(UserWarning, match="candidates available"): + result = sel.select(containment_mat, "s0") + assert len(result) <= 4 + + +def test_weighted_selection_favours_high_weight(containment_mat): + """Assembly s3 normally wouldn't be selected (cross-group), but with a high weight it should be.""" + weights = pd.Series({"s3": 10.0}, dtype=float) # strongly favour s3 + sel = ContainmentSelector(k=3, min_containment=0.01, weights=weights) + result = sel.select(containment_mat, "s0") + assert "s3" in result + + +def test_results_are_unique(containment_mat): + sel = ContainmentSelector(k=4, min_containment=0.05) + result = sel.select(containment_mat, "s0") + assert len(result) == len(set(result)) + + +# ── Registry / interface integration ────────────────────────────────────────── + +def test_is_base_selector_subclass(): + assert issubclass(ContainmentSelector, BaseSelector) + + +def test_registered_in_registry(): + assert SELECTOR_REGISTRY["containment"] is ContainmentSelector + assert "containment" in CONTAINMENT_SELECTORS + + +def test_constructible_via_registry_min_jaccard(containment_mat): + """The registry constructs selectors uniformly with k= and min_jaccard=.""" + sel_cls = SELECTOR_REGISTRY["containment"] + sel = sel_cls(k=3, min_jaccard=0.05) + assert sel.min_containment == 0.05 + result = sel.select(containment_mat, "s0") + assert result[0] == "s0" + assert len(result) == 3 + + +def test_min_jaccard_overrides_min_containment(): + sel = ContainmentSelector(k=3, min_containment=0.05, min_jaccard=0.2) + assert sel.min_containment == 0.2 + + +def test_base_validation_applies(): + """k < 1 should raise via BaseSelector validation.""" + with pytest.raises(ValueError, match="k must be"): + ContainmentSelector(k=0) + + +def test_returns_only_threshold_candidates(containment_mat): + """All returned prototypes clear the containment threshold for the query.""" + sel = ContainmentSelector(k=6, min_containment=0.05) + result = sel.select(containment_mat, "s0") + assert all(containment_mat.loc["s0", r] >= 0.05 for r in result) diff --git a/tests/test_selectors/test_greedy_var.py b/tests/test_selectors/test_greedy_var.py new file mode 100644 index 0000000..d0b3a43 --- /dev/null +++ b/tests/test_selectors/test_greedy_var.py @@ -0,0 +1,48 @@ +import pytest +from refrover.selectors import GreedyVarSelector + + +def test_returns_k_results(clustered_sim): + sel = GreedyVarSelector(k=4) + result = sel.select(clustered_sim, "s00") + assert len(result) == 4 + + +def test_query_always_first(clustered_sim): + sel = GreedyVarSelector(k=4) + result = sel.select(clustered_sim, "s00") + assert result[0] == "s00" + + +def test_all_within_threshold(clustered_sim): + sel = GreedyVarSelector(k=4, min_jaccard=0.05) + result = sel.select(clustered_sim, "s00") + for rid in result: + assert clustered_sim.loc["s00", rid] >= 0.05 + + +def test_spans_clusters(clustered_sim): + """With min_jaccard low enough to allow cross-cluster candidates, diverse selection expected.""" + # cross-cluster sim = 0.05, so min_jaccard must be < 0.05 to allow them as candidates + sel = GreedyVarSelector(k=4, min_jaccard=0.01) + result = sel.select(clustered_sim, "s00") + # cross-cluster assemblies have high variance in their similarity profiles + # (0.7 within their own cluster, 0.05 to others) → GreedyVarSelector should prefer them + cluster_a = {"s00", "s01", "s02", "s03"} + from_other_clusters = [r for r in result if r not in cluster_a] + assert len(from_other_clusters) >= 1, ( + "GreedyVarSelector should pick diverse prototypes spanning multiple clusters" + ) + + +def test_fewer_than_k_candidates_warns(sparse_sim): + sel = GreedyVarSelector(k=5, min_jaccard=0.1) + with pytest.warns(UserWarning, match="candidates available"): + result = sel.select(sparse_sim, "s00") + assert len(result) <= 5 + + +def test_results_are_unique(clustered_sim): + sel = GreedyVarSelector(k=5) + result = sel.select(clustered_sim, "s00") + assert len(result) == len(set(result)) diff --git a/tests/test_selectors/test_kmedoids.py b/tests/test_selectors/test_kmedoids.py new file mode 100644 index 0000000..a326f9b --- /dev/null +++ b/tests/test_selectors/test_kmedoids.py @@ -0,0 +1,33 @@ +import pytest +from refrover.selectors import KMedoidsSelector + + +def test_returns_k_results(clustered_sim): + sel = KMedoidsSelector(k=3) + assert len(sel.select(clustered_sim, "s00")) == 3 + + +def test_query_always_included(clustered_sim): + sel = KMedoidsSelector(k=3, min_jaccard=0.0) + result = sel.select(clustered_sim, "s04") + assert "s04" in result + + +def test_medoids_are_valid_ids(clustered_sim): + sel = KMedoidsSelector(k=4, min_jaccard=0.0) + result = sel.select(clustered_sim, "s00") + assert all(r in clustered_sim.columns for r in result) + assert len(result) == len(set(result)) + + +def test_reproduces_with_same_seed(clustered_sim): + sel = KMedoidsSelector(k=3, random_state=7) + assert sel.select(clustered_sim, "s00") == sel.select(clustered_sim, "s00") + + +def test_covers_clusters(clustered_sim): + """k-medoids should select one representative from each cluster.""" + sel = KMedoidsSelector(k=3, min_jaccard=0.0) + result = sel.select(clustered_sim, "s00") + cluster_of = lambda sid: int(sid[1:]) // 4 + assert len({cluster_of(r) for r in result}) == 3 diff --git a/tests/test_selectors/test_maxmin.py b/tests/test_selectors/test_maxmin.py new file mode 100644 index 0000000..9dca398 --- /dev/null +++ b/tests/test_selectors/test_maxmin.py @@ -0,0 +1,52 @@ +import numpy as np +import pandas as pd +import pytest +from refrover.selectors import MaxMinSelector + + +def test_query_is_first(clustered_sim): + sel = MaxMinSelector(k=3) + result = sel.select(clustered_sim, "s00") + assert result[0] == "s00" + + +def test_returns_k_results(clustered_sim): + sel = MaxMinSelector(k=4) + assert len(sel.select(clustered_sim, "s00")) == 4 + + +def test_selects_across_clusters(clustered_sim): + """With k=3 and 3 clusters, MaxMin should pick one from each cluster.""" + sel = MaxMinSelector(k=3, min_jaccard=0.0) + result = sel.select(clustered_sim, "s00") + cluster_of = lambda sid: int(sid[1:]) // 4 + clusters_hit = {cluster_of(r) for r in result} + assert len(clusters_hit) == 3 + + +def test_diversity_exceeds_random(uniform_sim): + """MaxMin selection should be more diverse than a random selection.""" + sel_mm = MaxMinSelector(k=5, min_jaccard=0.0) + from refrover.selectors import RandomSelector + sel_r = RandomSelector(k=5, min_jaccard=0.0, seed=0) + + def min_pairwise_dist(ids, sim): + vals = [1 - sim.loc[a, b] for i, a in enumerate(ids) for b in ids[i+1:]] + return min(vals) if vals else 0.0 + + mm_div = min_pairwise_dist(sel_mm.select(uniform_sim, "s00"), uniform_sim) + r_div = min_pairwise_dist(sel_r.select(uniform_sim, "s00"), uniform_sim) + assert mm_div >= r_div + + +def test_all_within_threshold(clustered_sim): + sel = MaxMinSelector(k=3, min_jaccard=0.5) + result = sel.select(clustered_sim, "s00") + for rid in result: + assert clustered_sim.loc["s00", rid] >= 0.5 + + +def test_no_duplicates(clustered_sim): + sel = MaxMinSelector(k=4) + result = sel.select(clustered_sim, "s00") + assert len(result) == len(set(result)) diff --git a/tests/test_selectors/test_random.py b/tests/test_selectors/test_random.py new file mode 100644 index 0000000..eb9c2fd --- /dev/null +++ b/tests/test_selectors/test_random.py @@ -0,0 +1,40 @@ +import pytest +from refrover.selectors import RandomSelector + + +def test_returns_k_results(clustered_sim): + sel = RandomSelector(k=3, seed=0) + result = sel.select(clustered_sim, "s00") + assert len(result) == 3 + + +def test_all_within_threshold(clustered_sim): + sel = RandomSelector(k=4, min_jaccard=0.5, seed=1) + result = sel.select(clustered_sim, "s00") + for rid in result: + assert clustered_sim.loc["s00", rid] >= 0.5 + + +def test_fewer_than_k_candidates(sparse_sim): + sel = RandomSelector(k=5, min_jaccard=0.1, seed=0) + with pytest.warns(UserWarning, match="candidates available"): + result = sel.select(sparse_sim, "s00") + assert len(result) <= 5 + + +def test_unknown_query_raises(sparse_sim): + sel = RandomSelector(k=3, min_jaccard=0.1) + with pytest.raises(KeyError, match="not found"): + sel.select(sparse_sim, "does_not_exist") + + +def test_seed_reproducible(uniform_sim): + sel = RandomSelector(k=4, seed=99) + assert sel.select(uniform_sim, "s00") == sel.select(uniform_sim, "s00") + + +def test_results_are_valid_ids(clustered_sim): + sel = RandomSelector(k=3, seed=0) + result = sel.select(clustered_sim, "s00") + assert all(rid in clustered_sim.columns for rid in result) + assert len(result) == len(set(result))