From 83c22c149f834e81bd679fb1c218ce149f179238 Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Tue, 9 Jun 2026 22:13:59 -0400 Subject: [PATCH 01/10] Add package scaffold, four selector implementations, and tests - pyproject.toml: installable package with entry point refrover.cli:main - src/refrover/selectors/: BaseSelector ABC, RandomSelector, MaxMinSelector, KMedoidsSelector (pure-numpy PAM), ArchetypeSelector (archetypes package) - src/refrover/io.py: manifest read/validate - tests/: 23 passing unit tests covering correctness, diversity, and edge cases - .gitignore: excludes data/, *.sig, *.zip, build artifacts --- .gitignore | 32 +++++++++++ pyproject.toml | 30 ++++++++++ src/refrover/__init__.py | 6 ++ src/refrover/io.py | 50 +++++++++++++++++ src/refrover/selectors/__init__.py | 21 +++++++ src/refrover/selectors/archetype.py | 62 +++++++++++++++++++++ src/refrover/selectors/base.py | 45 +++++++++++++++ src/refrover/selectors/kmedoids.py | 76 ++++++++++++++++++++++++++ src/refrover/selectors/maxmin.py | 44 +++++++++++++++ src/refrover/selectors/random.py | 18 ++++++ tests/__init__.py | 0 tests/conftest.py | 51 +++++++++++++++++ tests/test_selectors/__init__.py | 0 tests/test_selectors/test_archetype.py | 39 +++++++++++++ tests/test_selectors/test_kmedoids.py | 33 +++++++++++ tests/test_selectors/test_maxmin.py | 52 ++++++++++++++++++ tests/test_selectors/test_random.py | 40 ++++++++++++++ 17 files changed, 599 insertions(+) create mode 100644 .gitignore create mode 100644 pyproject.toml create mode 100644 src/refrover/__init__.py create mode 100644 src/refrover/io.py create mode 100644 src/refrover/selectors/__init__.py create mode 100644 src/refrover/selectors/archetype.py create mode 100644 src/refrover/selectors/base.py create mode 100644 src/refrover/selectors/kmedoids.py create mode 100644 src/refrover/selectors/maxmin.py create mode 100644 src/refrover/selectors/random.py create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/test_selectors/__init__.py create mode 100644 tests/test_selectors/test_archetype.py create mode 100644 tests/test_selectors/test_kmedoids.py create mode 100644 tests/test_selectors/test_maxmin.py create mode 100644 tests/test_selectors/test_random.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..330ce21 --- /dev/null +++ b/.gitignore @@ -0,0 +1,32 @@ +# 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 +data/ +*.zip +*.sbt.zip +*.sig + +# macOS +.DS_Store + +# Editors +.vscode/ +.idea/ 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/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/selectors/__init__.py b/src/refrover/selectors/__init__.py new file mode 100644 index 0000000..6a1fd37 --- /dev/null +++ b/src/refrover/selectors/__init__.py @@ -0,0 +1,21 @@ +from .base import BaseSelector +from .random import RandomSelector +from .maxmin import MaxMinSelector +from .kmedoids import KMedoidsSelector +from .archetype import ArchetypeSelector + +SELECTOR_REGISTRY: dict[str, type[BaseSelector]] = { + "random": RandomSelector, + "maxmin": MaxMinSelector, + "kmedoids": KMedoidsSelector, + "archetype": ArchetypeSelector, +} + +__all__ = [ + "BaseSelector", + "RandomSelector", + "MaxMinSelector", + "KMedoidsSelector", + "ArchetypeSelector", + "SELECTOR_REGISTRY", +] 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/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/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..e6028fb --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,51 @@ +import numpy as np +import pandas as pd +import pytest + + +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) 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_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)) From 1ad2c205a7754f161dc732910555640cebc64f73 Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 06:45:55 -0400 Subject: [PATCH 02/10] Add full pipeline modules, three new selectors, and formatter suite - cli.py: Click entry point with sketch/select/align/coverage/format/run/benchmark subcommands - sketch.py: sourmash sketch + compare wrappers - similarity.py: load/filter Jaccard similarity matrices; direct Python API computation - align.py: BWA-MEM2, BWA, minimap2 wrappers; handles paired/single/long/hybrid reads - coverage.py: CoverM wrapper for per-contig depth - pipeline.py: RefRoverPipeline orchestrator - benchmark.py: stub (requires CheckM2) - selectors/greedy_var.py: greedy coverage-variance maximization selector - selectors/containment.py: read-containment-based selector with optional species weighting - selectors/feedback.py: stub (requires alignment loop) - formatters/: MetaBAT2, SemiBin2, MaxBin2, CONCOCT, generic TSV with registry dispatch - 61 tests passing; added coverage_df and tiny_manifest fixtures --- src/refrover/align.py | 275 +++++++++++++++++++++++ src/refrover/benchmark.py | 34 +++ src/refrover/cli.py | 237 +++++++++++++++++++ src/refrover/coverage.py | 55 +++++ src/refrover/formatters/__init__.py | 18 ++ src/refrover/formatters/concoct.py | 14 ++ src/refrover/formatters/generic.py | 20 ++ src/refrover/formatters/maxbin2.py | 16 ++ src/refrover/formatters/metabat2.py | 33 +++ src/refrover/formatters/registry.py | 15 ++ src/refrover/formatters/semibin2.py | 13 ++ src/refrover/pipeline.py | 89 ++++++++ src/refrover/selectors/__init__.py | 6 + src/refrover/selectors/containment.py | 165 ++++++++++++++ src/refrover/selectors/feedback.py | 20 ++ src/refrover/selectors/greedy_var.py | 42 ++++ src/refrover/similarity.py | 89 ++++++++ src/refrover/sketch.py | 96 ++++++++ tests/conftest.py | 28 +++ tests/test_formatters/__init__.py | 0 tests/test_formatters/test_generic.py | 32 +++ tests/test_formatters/test_metabat2.py | 39 ++++ tests/test_pipeline.py | 228 +++++++++++++++++++ tests/test_selectors/test_containment.py | 86 +++++++ tests/test_selectors/test_greedy_var.py | 48 ++++ 25 files changed, 1698 insertions(+) create mode 100644 src/refrover/align.py create mode 100644 src/refrover/benchmark.py create mode 100644 src/refrover/cli.py create mode 100644 src/refrover/coverage.py create mode 100644 src/refrover/formatters/__init__.py create mode 100644 src/refrover/formatters/concoct.py create mode 100644 src/refrover/formatters/generic.py create mode 100644 src/refrover/formatters/maxbin2.py create mode 100644 src/refrover/formatters/metabat2.py create mode 100644 src/refrover/formatters/registry.py create mode 100644 src/refrover/formatters/semibin2.py create mode 100644 src/refrover/pipeline.py create mode 100644 src/refrover/selectors/containment.py create mode 100644 src/refrover/selectors/feedback.py create mode 100644 src/refrover/selectors/greedy_var.py create mode 100644 src/refrover/similarity.py create mode 100644 src/refrover/sketch.py create mode 100644 tests/test_formatters/__init__.py create mode 100644 tests/test_formatters/test_generic.py create mode 100644 tests/test_formatters/test_metabat2.py create mode 100644 tests/test_pipeline.py create mode 100644 tests/test_selectors/test_containment.py create mode 100644 tests/test_selectors/test_greedy_var.py 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..7765150 --- /dev/null +++ b/src/refrover/benchmark.py @@ -0,0 +1,34 @@ +""" +Benchmark selector strategies against each other. + +Runs each (selector, k) combination, evaluates MAG yield via CheckM2, +and writes a comparison table. +""" + +from pathlib import Path +import pandas as pd + + +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: + """ + For each (selector, k) combination: + 1. Run RefRoverPipeline with that selector + 2. Run a binner (MetaBAT2) on the coverage output + 3. Evaluate bins with CheckM2 + + Returns a summary DataFrame with columns: + selector, k, n_mags_passing, completeness_mean, contamination_mean, + cpu_seconds, total_alignment_cpu_seconds + """ + raise NotImplementedError( + "Benchmark runner not yet implemented. " + "Requires CheckM2 installation and a labelled ground-truth community." + ) diff --git a/src/refrover/cli.py b/src/refrover/cli.py new file mode 100644 index 0000000..cd164cc --- /dev/null +++ b/src/refrover/cli.py @@ -0,0 +1,237 @@ +""" +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}") + + +# ── select ──────────────────────────────────────────────────────────────────── + +@main.command() +@click.option("--sketches", required=True, type=click.Path(exists=True), + help="Directory of assembly .sig files") +@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"])) +@click.option("--k", default=5, show_default=True, help="Number of prototypes per sample") +@click.option("--min-jaccard", default=0.1, show_default=True) +@click.option("--outdir", required=True, type=click.Path()) +@click.option("--force", is_flag=True) +def select(sketches, manifest, selector, k, 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 + from refrover.similarity import matrix_from_sigs + + 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) + + 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) + + sel_cls = SELECTOR_REGISTRY[selector] + sel = sel_cls(k=k, min_jaccard=min_jaccard) + + rows = [] + for _, row in df.iterrows(): + sid = row["sample_id"] + prototypes = sel.select(sim_matrix, sid) + rows.append({ + "sample_id": sid, + "prototype_ids": ",".join(prototypes), + "n_prototypes": len(prototypes), + }) + + 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 + import pandas as pd + + 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 = pd.read_csv(cov_tsv, sep="\t", index_col=0) + + 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"])) +@click.option("--k", default=5, show_default=True) +@click.option("--min-jaccard", default=0.1, show_default=True) +@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, 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 + + 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(",")] + + pipeline = RefRoverPipeline( + manifest=df, + selector=sel, + binners=binner_list, + threads=threads, + outdir=Path(outdir), + force=force, + ) + results = pipeline.run() + click.echo(f"Done. Coverage tables: {list(results.coverage_tables.keys())}") + + +# ── 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/coverage.py b/src/refrover/coverage.py new file mode 100644 index 0000000..c5bc61b --- /dev/null +++ b/src/refrover/coverage.py @@ -0,0 +1,55 @@ +""" +CoverM wrapper: compute per-contig mean depth from sorted BAMs. +""" + +import subprocess +from pathlib import Path +import pandas as pd + + +def run_coverm( + bams_dir: Path | str, + outdir: Path | str, + *, + threads: int = 8, + force: bool = False, + coverm_path: str = "coverm", +) -> Path: + """ + Run CoverM on all BAMs in bams_dir. + + Writes coverage.tsv to outdir with columns: + Contig Length sample1_depth sample2_depth ... + + Returns 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", "mean", + "--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 load_coverage(tsv_path: Path | str) -> pd.DataFrame: + """Load a CoverM coverage.tsv into a DataFrame indexed by contig name.""" + df = pd.read_csv(tsv_path, sep="\t", index_col=0) + return 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/concoct.py b/src/refrover/formatters/concoct.py new file mode 100644 index 0000000..3f79c37 --- /dev/null +++ b/src/refrover/formatters/concoct.py @@ -0,0 +1,14 @@ +"""CONCOCT formatter: contig x sample matrix, tab-separated, integer depths.""" + +from pathlib import Path +import pandas as pd + + +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 + depth_cols = [c for c in coverage_df.columns if c != "length"] + df_out = coverage_df[depth_cols].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..46d1f16 --- /dev/null +++ b/src/refrover/formatters/generic.py @@ -0,0 +1,20 @@ +"""Generic TSV formatter: contig, length, sample1_depth, sample2_depth, ...""" + +from pathlib import Path +import pandas as pd + + +def write_generic(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) -> Path: + """ + Write a generic coverage TSV. + + Input: DataFrame with index=contig, columns including 'length' and one depth + column per sample (e.g. 's001_depth'). + + Output: coverage_generic.tsv — same layout, tab-separated. + """ + out = outdir / "coverage_generic.tsv" + if out.exists() and not force: + return out + coverage_df.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..0db49b5 --- /dev/null +++ b/src/refrover/formatters/maxbin2.py @@ -0,0 +1,16 @@ +"""MaxBin2 formatter: one .abund file per sample, single depth column.""" + +from pathlib import Path +import pandas as pd + + +def write_maxbin2(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) -> list[Path]: + depth_cols = [c for c in coverage_df.columns if c != "length"] + out_paths = [] + for col in depth_cols: + out = outdir / f"{col}.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..a4749fd --- /dev/null +++ b/src/refrover/formatters/metabat2.py @@ -0,0 +1,33 @@ +""" +MetaBAT2 formatter: jgi_summarize_bam_contig_depths format. + +Columns: contigName, contigLen, totalAvgDepth, sample1, sample1-var, sample2, sample2-var, ... +Variance columns are set to 0 (RefRover does not compute variance from CoverM mean output). +""" + +from pathlib import Path +import pandas as pd + + +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 = [c for c in coverage_df.columns if c != "length"] + 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: + rec[col] = row[col] + rec[f"{col}-var"] = 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..2c0168f --- /dev/null +++ b/src/refrover/formatters/semibin2.py @@ -0,0 +1,13 @@ +"""SemiBin2 formatter: contig x sample depth matrix, no variance column.""" + +from pathlib import Path +import pandas as pd + + +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 + depth_cols = [c for c in coverage_df.columns if c != "length"] + coverage_df[depth_cols].to_csv(out, sep="\t") + return out diff --git a/src/refrover/pipeline.py b/src/refrover/pipeline.py new file mode 100644 index 0000000..8a7afda --- /dev/null +++ b/src/refrover/pipeline.py @@ -0,0 +1,89 @@ +""" +RefRoverPipeline: orchestrates the full sketch → select → align → coverage → format pipeline. +""" + +from dataclasses import dataclass, field +from pathlib import Path +import pandas as pd + +from refrover.selectors.base import BaseSelector + + +@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, + ): + self.manifest = manifest + self.selector = selector + self.binners = binners + self.threads = threads + self.outdir = Path(outdir) + self.force = force + + 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. Sketch + sig_paths = sketch_assemblies( + self.manifest["assembly"].tolist(), + outdir=sketch_dir, + threads=self.threads, + force=self.force, + ) + + # 2. Pairwise similarity + 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 index 6a1fd37..c72c2de 100644 --- a/src/refrover/selectors/__init__.py +++ b/src/refrover/selectors/__init__.py @@ -3,12 +3,16 @@ from .maxmin import MaxMinSelector from .kmedoids import KMedoidsSelector from .archetype import ArchetypeSelector +from .greedy_var import GreedyVarSelector +from .feedback import FeedbackSelector SELECTOR_REGISTRY: dict[str, type[BaseSelector]] = { "random": RandomSelector, "maxmin": MaxMinSelector, "kmedoids": KMedoidsSelector, "archetype": ArchetypeSelector, + "greedy_var": GreedyVarSelector, + "feedback": FeedbackSelector, } __all__ = [ @@ -17,5 +21,7 @@ "MaxMinSelector", "KMedoidsSelector", "ArchetypeSelector", + "GreedyVarSelector", + "FeedbackSelector", "SELECTOR_REGISTRY", ] diff --git a/src/refrover/selectors/containment.py b/src/refrover/selectors/containment.py new file mode 100644 index 0000000..548d9cc --- /dev/null +++ b/src/refrover/selectors/containment.py @@ -0,0 +1,165 @@ +""" +ContainmentSelector: prototype selection driven by cross-sample containment. + +Unlike the other selectors, this one does NOT use the assembly-vs-assembly +Jaccard similarity matrix. Instead it uses a reads-vs-assembly containment +matrix (containment[i, j] = fraction of sample i's read k-mers found in +assembly j), which directly measures how well sample i's reads will map to +assembly j. + +Two modes +--------- +unweighted (default) + Apply MaxMin greedy selection on the containment matrix, treating + containment as the similarity measure. Selects prototypes that are + maximally diverse in read-mapping space. + +weighted (requires species_weights) + Weight each candidate assembly j by its predicted differential coverage + signal, estimated from per-species abundance variance. Assemblies whose + dominant species vary across samples get higher weight. + + species_weights: pd.Series indexed by assembly_id, values in [0, 1]. + Can be derived from gtdb_species_stats.tsv (see explore_species.R): + weight_j = mean CV of species whose primary assembly is j, + restricted to mid-abundance species. + +Input matrix shape +------------------ +The containment matrix has: + - rows: read sample IDs (query samples) + - columns: assembly IDs (prototype candidates) + +Unlike BaseSelector which takes a square assembly-vs-assembly matrix, +ContainmentSelector takes a rectangular reads-vs-assemblies matrix. +The query_id must be a row index (read sample ID). +""" + +import warnings +import pandas as pd +import numpy as np +from typing import Optional + + +class ContainmentSelector: + """ + 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 = assembly_id). If provided, + candidate scores are multiplied by weight[assembly_id]. + """ + + def __init__( + self, + k: int, + min_containment: float = 0.05, + weights: Optional[pd.Series] = None, + ): + self.k = k + self.min_containment = min_containment + self.weights = weights + + def select(self, containment_matrix: pd.DataFrame, query_id: str) -> list[str]: + """ + Select k prototype assemblies for query_id. + + Parameters + ---------- + containment_matrix : pd.DataFrame + Rows = read sample IDs, columns = assembly IDs. + Values = containment(reads_i, assembly_j). + query_id : str + Row index of the query sample in containment_matrix. + + Returns + ------- + list[str] + Ordered list of prototype assembly IDs. The query sample's own + assembly is always first (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] + + # Filter candidates by minimum containment + candidates = query_row[query_row >= self.min_containment].index.tolist() + + # The query's own assembly should always be first if present + own_assembly = query_id # assumes assembly ID matches sample ID + if own_assembly in candidates: + candidates = [own_assembly] + [c for c in candidates if c != own_assembly] + elif candidates: + # Put the highest-containment assembly first + candidates = sorted(candidates, key=lambda c: -float(query_row[c])) + + if len(candidates) == 0: + raise ValueError( + f"No assemblies meet min_containment={self.min_containment} for {query_id}" + ) + + 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 + + # MaxMin greedy selection in containment space, with optional weighting + selected = [candidates[0]] + remaining = candidates[1:] + + # Track min containment-distance from selected set for each remaining candidate + # distance = 1 - containment(assembly_j, assembly_k), approximated via + # the cross-sample containment profiles: dist(j, k) = 1 - corr(col_j, col_k) + all_rows = containment_matrix.values # shape (n_samples, n_assemblies) + col_index = {c: containment_matrix.columns.get_loc(c) for c in candidates} + + def _profile(assembly_id): + return all_rows[:, col_index[assembly_id]] + + # Pairwise distance proxy: 1 - Pearson correlation between containment profiles + min_dist = np.array([ + 1.0 - float(np.corrcoef(_profile(selected[0]), _profile(c))[0, 1]) + for c in remaining + ]) + min_dist = np.nan_to_num(min_dist, nan=0.0) + + while len(selected) < self.k and remaining: + scores = min_dist.copy() + + # Apply species-level weights if provided + if self.weights is not None: + w = np.array([ + float(self.weights.get(c, 1.0)) for c in remaining + ]) + scores = scores * w + + best_idx = int(np.argmax(scores)) + best = remaining[best_idx] + selected.append(best) + + # Update min distances + new_dists = np.array([ + 1.0 - float(np.corrcoef(_profile(best), _profile(c))[0, 1]) + for c in remaining + ]) + new_dists = np.nan_to_num(new_dists, nan=0.0) + min_dist = np.minimum(min_dist, new_dists) + + remaining.pop(best_idx) + min_dist = np.delete(min_dist, best_idx) + + return 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/similarity.py b/src/refrover/similarity.py new file mode 100644 index 0000000..f619775 --- /dev/null +++ b/src/refrover/similarity.py @@ -0,0 +1,89 @@ +""" +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 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/conftest.py b/tests/conftest.py index e6028fb..a7a61be 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ 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: @@ -49,3 +50,30 @@ def sparse_sim(): 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: 5 contigs x 3 samples + length column. + Matches the output format from CoverM (index=contig, columns=length + depth 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], + } + 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_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..95092b9 --- /dev/null +++ b/tests/test_formatters/test_metabat2.py @@ -0,0 +1,39 @@ +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 depth column must have a paired -var column set to 0 + depth_cols = ["s1_depth", "s2_depth", "s3_depth"] + for col in depth_cols: + assert f"{col}-var" in df.columns + assert (df[f"{col}-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/test_containment.py b/tests/test_selectors/test_containment.py new file mode 100644 index 0000000..37028b3 --- /dev/null +++ b/tests/test_selectors/test_containment.py @@ -0,0 +1,86 @@ +import numpy as np +import pandas as pd +import pytest +from refrover.selectors.containment import ContainmentSelector + + +@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)) 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)) From c1a01dd0a5bbe530033006c20a7fe2d21fb4c2a7 Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 07:24:00 -0400 Subject: [PATCH 03/10] Add adaptive-k estimators (scree_elbow, similarity_gap, saturation_curve) and --k auto CLI flag --- src/refrover/adaptive_k.py | 200 +++++++++++++++++++++++++++++++++++++ src/refrover/cli.py | 21 +++- tests/test_adaptive_k.py | 84 ++++++++++++++++ 3 files changed, 302 insertions(+), 3 deletions(-) create mode 100644 src/refrover/adaptive_k.py create mode 100644 tests/test_adaptive_k.py diff --git a/src/refrover/adaptive_k.py b/src/refrover/adaptive_k.py new file mode 100644 index 0000000..36338c3 --- /dev/null +++ b/src/refrover/adaptive_k.py @@ -0,0 +1,200 @@ +""" +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. + +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"] + + +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) + 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) diff --git a/src/refrover/cli.py b/src/refrover/cli.py index cd164cc..78ee577 100644 --- a/src/refrover/cli.py +++ b/src/refrover/cli.py @@ -55,11 +55,14 @@ def sketch(manifest, outdir, ksize, scaled, threads, force): @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"])) -@click.option("--k", default=5, show_default=True, help="Number of prototypes per sample") +@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"])) @click.option("--min-jaccard", default=0.1, show_default=True) @click.option("--outdir", required=True, type=click.Path()) @click.option("--force", is_flag=True) -def select(sketches, manifest, selector, k, min_jaccard, outdir, force): +def select(sketches, 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 @@ -84,17 +87,29 @@ def select(sketches, manifest, selector, k, min_jaccard, outdir, force): 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] - sel = sel_cls(k=k, min_jaccard=min_jaccard) 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) diff --git a/tests/test_adaptive_k.py b/tests/test_adaptive_k.py new file mode 100644 index 0000000..37a6c42 --- /dev/null +++ b/tests/test_adaptive_k.py @@ -0,0 +1,84 @@ +import pytest +import numpy as np +import pandas as pd +from refrover.adaptive_k import estimate_k, _scree_elbow, _similarity_gap, _saturation_curve + + +@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"): + 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_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 From c72e7c15e38ca5411af42161b597cf26ae91d954 Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 12:37:21 -0400 Subject: [PATCH 04/10] Add GTDB gather pipeline, species-level analysis scripts, and assembly weight computation --- .gitignore | 10 +- data/mouse_rewilding/build_community_subdb.py | 203 ++++++++++++++++++ data/mouse_rewilding/compute_containment.py | 114 ++++++++++ .../compute_species_weights.py | 186 ++++++++++++++++ data/mouse_rewilding/explore_containment.R | 160 ++++++++++++++ data/mouse_rewilding/explore_species.R | 184 ++++++++++++++++ data/mouse_rewilding/run_gtdb_gather.py | 138 ++++++++++++ 7 files changed, 993 insertions(+), 2 deletions(-) create mode 100644 data/mouse_rewilding/build_community_subdb.py create mode 100644 data/mouse_rewilding/compute_containment.py create mode 100644 data/mouse_rewilding/compute_species_weights.py create mode 100644 data/mouse_rewilding/explore_containment.R create mode 100644 data/mouse_rewilding/explore_species.R create mode 100644 data/mouse_rewilding/run_gtdb_gather.py diff --git a/.gitignore b/.gitignore index 330ce21..4130f65 100644 --- a/.gitignore +++ b/.gitignore @@ -18,8 +18,14 @@ htmlcov/ # Type checking .mypy_cache/ -# Data and databases — large files not tracked -data/ +# 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 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/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_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() From a27838f4c560f7e6b35e4c4a06c6bb1c6d87f7bc Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 15:45:00 -0400 Subject: [PATCH 05/10] Add selector comparison and adaptive-k analysis scripts compare_selectors.py evaluates four prototype selection strategies on the 183-sample containment matrix using two metrics: - Naive: sum of selected column variances (biased toward correlated selectors) - Unique: orthogonal projection of the full 182-column matrix onto the selected set (removes double-counting of correlated assemblies) explore_selectors.R generates five diagnostic plots: variance-explained curves (both metrics side-by-side), per-sample distribution at k=5, selector agreement heatmap, adaptive-k distribution, and per-sample saturation curves. Key algorithm: the saturation curve uses greedy forward selection by residual variance (Gram-Schmidt incremental QR) rather than MaxMin order, which ensures a monotonically decreasing gains sequence necessary for threshold-based elbow detection. --- data/mouse_rewilding/compare_selectors.py | 339 ++++++++++++++++++++++ data/mouse_rewilding/explore_selectors.R | 178 ++++++++++++ 2 files changed, 517 insertions(+) create mode 100644 data/mouse_rewilding/compare_selectors.py create mode 100644 data/mouse_rewilding/explore_selectors.R 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/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") From 7b45c6fdc093df468f3c1881d3916ed1c62033b2 Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 19:36:36 -0400 Subject: [PATCH 06/10] Wire ContainmentSelector into the selector interface, CLI, and pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The containment-based selection validated on the rewilded-mouse dataset was implemented but unreachable: ContainmentSelector used a bespoke interface, was absent from SELECTOR_REGISTRY, and had no CLI or pipeline path. This makes it a first-class selector. Interface unification - ContainmentSelector now subclasses BaseSelector and returns prototype sample_ids, the same contract every selector follows (a prototype is the sample whose assembly it is; align resolves sample_id -> FASTA). This removes the implicit sample==assembly coupling in favour of the explicit manifest mapping that align already uses. - Accepts min_jaccard= as an alias for min_containment so the registry can construct it uniformly. Gains BaseSelector's k/threshold validation. Performance - Precompute the candidate correlation-distance matrix once instead of calling np.corrcoef pairwise inside the greedy loop. ~4 ms/sample on the 183-sample matrix (was O(k*n) corrcoef calls per query). Wiring - Registered as "containment" in SELECTOR_REGISTRY; added CONTAINMENT_SELECTORS so the CLI/pipeline know to feed it a containment matrix. - refrover select / refrover run accept --containment-matrix; error clearly when the required matrix (or --sketches for Jaccard selectors) is missing. - RefRoverPipeline accepts containment_matrix= and uses it for selection, skipping the sketch+compare step. - load_containment_matrix() added to similarity.py. Tests: 9 new (registry membership, BaseSelector conformance, registry construction, CLI select happy path + both missing-input errors). 78 pass. Descoped from this branch: the MetaBAT2 variance fix (it is a coverage-schema change touching load_coverage and all five formatters, not the few lines first estimated) — to follow as its own change. --- CLAUDE.md | 7 +- opus_suggestions.md | 257 +++++++++++++++++++++++ src/refrover/cli.py | 69 ++++-- src/refrover/pipeline.py | 37 ++-- src/refrover/selectors/__init__.py | 9 + src/refrover/selectors/containment.py | 172 +++++++-------- src/refrover/similarity.py | 19 ++ tests/test_cli.py | 90 ++++++++ tests/test_selectors/test_containment.py | 41 ++++ 9 files changed, 574 insertions(+), 127 deletions(-) create mode 100644 opus_suggestions.md create mode 100644 tests/test_cli.py diff --git a/CLAUDE.md b/CLAUDE.md index b271e1f..ece0d0e 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. 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/src/refrover/cli.py b/src/refrover/cli.py index 78ee577..8ab6e4f 100644 --- a/src/refrover/cli.py +++ b/src/refrover/cli.py @@ -50,24 +50,29 @@ def sketch(manifest, outdir, ksize, scaled, threads, force): # ── select ──────────────────────────────────────────────────────────────────── @main.command() -@click.option("--sketches", required=True, type=click.Path(exists=True), - help="Directory of assembly .sig files") +@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"])) + 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"])) -@click.option("--min-jaccard", default=0.1, show_default=True) +@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, manifest, selector, k, adaptive_k_method, min_jaccard, outdir, force): +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 - from refrover.similarity import matrix_from_sigs + 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) @@ -79,13 +84,26 @@ def select(sketches, manifest, selector, k, adaptive_k_method, min_jaccard, outd df = read_manifest(manifest) - 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) + # 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) @@ -185,26 +203,40 @@ def format(coverage_dir, binners, outdir, force): @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"])) + 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) +@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, aligner, binners, threads, outdir, force): +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 + 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, @@ -212,6 +244,7 @@ def run(manifest, selector, k, min_jaccard, aligner, binners, threads, outdir, f 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())}") diff --git a/src/refrover/pipeline.py b/src/refrover/pipeline.py index 8a7afda..72cb08f 100644 --- a/src/refrover/pipeline.py +++ b/src/refrover/pipeline.py @@ -4,9 +4,11 @@ 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 @@ -24,6 +26,7 @@ def __init__( threads: int = 8, outdir: Path = Path("refrover_out"), force: bool = False, + containment_matrix: Optional[pd.DataFrame] = None, ): self.manifest = manifest self.selector = selector @@ -31,6 +34,13 @@ def __init__( 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 @@ -46,18 +56,21 @@ def run(self) -> PipelineResults: cov_dir = outdir / "coverage" fmt_dir = outdir / "formatted" - # 1. Sketch - sig_paths = sketch_assemblies( - self.manifest["assembly"].tolist(), - outdir=sketch_dir, - threads=self.threads, - force=self.force, - ) - - # 2. Pairwise similarity - sim_csv = outdir / "similarity.csv" - compare_sketches(sig_paths, output_csv=sim_csv, force=self.force) - sim_matrix = load_similarity_matrix(sim_csv) + # 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 = [] diff --git a/src/refrover/selectors/__init__.py b/src/refrover/selectors/__init__.py index c72c2de..e198154 100644 --- a/src/refrover/selectors/__init__.py +++ b/src/refrover/selectors/__init__.py @@ -4,6 +4,7 @@ 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]] = { @@ -12,9 +13,15 @@ "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", @@ -22,6 +29,8 @@ "KMedoidsSelector", "ArchetypeSelector", "GreedyVarSelector", + "ContainmentSelector", "FeedbackSelector", "SELECTOR_REGISTRY", + "CONTAINMENT_SELECTORS", ] diff --git a/src/refrover/selectors/containment.py b/src/refrover/selectors/containment.py index 548d9cc..fd8afb4 100644 --- a/src/refrover/selectors/containment.py +++ b/src/refrover/selectors/containment.py @@ -1,47 +1,44 @@ """ ContainmentSelector: prototype selection driven by cross-sample containment. -Unlike the other selectors, this one does NOT use the assembly-vs-assembly -Jaccard similarity matrix. Instead it uses a reads-vs-assembly containment -matrix (containment[i, j] = fraction of sample i's read k-mers found in -assembly j), which directly measures how well sample i's reads will map to -assembly j. +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) - Apply MaxMin greedy selection on the containment matrix, treating - containment as the similarity measure. Selects prototypes that are - maximally diverse in read-mapping space. - -weighted (requires species_weights) - Weight each candidate assembly j by its predicted differential coverage - signal, estimated from per-species abundance variance. Assemblies whose - dominant species vary across samples get higher weight. - - species_weights: pd.Series indexed by assembly_id, values in [0, 1]. - Can be derived from gtdb_species_stats.tsv (see explore_species.R): - weight_j = mean CV of species whose primary assembly is j, - restricted to mid-abundance species. - -Input matrix shape ------------------- -The containment matrix has: - - rows: read sample IDs (query samples) - - columns: assembly IDs (prototype candidates) - -Unlike BaseSelector which takes a square assembly-vs-assembly matrix, -ContainmentSelector takes a rectangular reads-vs-assemblies matrix. -The query_id must be a row index (read sample ID). + 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 -import pandas as pd -import numpy as np from typing import Optional +import numpy as np +import pandas as pd + +from .base import BaseSelector -class ContainmentSelector: + +class ContainmentSelector(BaseSelector): """ Prototype selection using cross-sample read containment. @@ -53,8 +50,14 @@ class ContainmentSelector: 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 = assembly_id). If provided, - candidate scores are multiplied by weight[assembly_id]. + 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__( @@ -62,50 +65,50 @@ def __init__( k: int, min_containment: float = 0.05, weights: Optional[pd.Series] = None, + min_jaccard: Optional[float] = None, ): - self.k = k - self.min_containment = min_containment + 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 k prototype assemblies for query_id. + Select up to k prototype sample_ids for query_id. Parameters ---------- containment_matrix : pd.DataFrame - Rows = read sample IDs, columns = assembly IDs. - Values = containment(reads_i, assembly_j). + 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 in containment_matrix. + Row index of the query sample. Returns ------- list[str] - Ordered list of prototype assembly IDs. The query sample's own - assembly is always first (containment = 1.0 by definition). + 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] - - # Filter candidates by minimum containment candidates = query_row[query_row >= self.min_containment].index.tolist() - - # The query's own assembly should always be first if present - own_assembly = query_id # assumes assembly ID matches sample ID - if own_assembly in candidates: - candidates = [own_assembly] + [c for c in candidates if c != own_assembly] - elif candidates: - # Put the highest-containment assembly first - candidates = sorted(candidates, key=lambda c: -float(query_row[c])) - - if len(candidates) == 0: + if not candidates: raise ValueError( - f"No assemblies meet min_containment={self.min_containment} for {query_id}" + 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}). " @@ -113,53 +116,32 @@ def select(self, containment_matrix: pd.DataFrame, query_id: str) -> list[str]: UserWarning, stacklevel=2, ) - if len(candidates) <= self.k: return candidates - # MaxMin greedy selection in containment space, with optional weighting - selected = [candidates[0]] - remaining = candidates[1:] - - # Track min containment-distance from selected set for each remaining candidate - # distance = 1 - containment(assembly_j, assembly_k), approximated via - # the cross-sample containment profiles: dist(j, k) = 1 - corr(col_j, col_k) - all_rows = containment_matrix.values # shape (n_samples, n_assemblies) - col_index = {c: containment_matrix.columns.get_loc(c) for c in 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) - def _profile(assembly_id): - return all_rows[:, col_index[assembly_id]] + weights = None + if self.weights is not None: + weights = np.array([float(self.weights.get(c, 1.0)) for c in candidates]) - # Pairwise distance proxy: 1 - Pearson correlation between containment profiles - min_dist = np.array([ - 1.0 - float(np.corrcoef(_profile(selected[0]), _profile(c))[0, 1]) - for c in remaining - ]) - min_dist = np.nan_to_num(min_dist, nan=0.0) + # 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.copy() - - # Apply species-level weights if provided - if self.weights is not None: - w = np.array([ - float(self.weights.get(c, 1.0)) for c in remaining - ]) - scores = scores * w - - best_idx = int(np.argmax(scores)) - best = remaining[best_idx] + 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]) - # Update min distances - new_dists = np.array([ - 1.0 - float(np.corrcoef(_profile(best), _profile(c))[0, 1]) - for c in remaining - ]) - new_dists = np.nan_to_num(new_dists, nan=0.0) - min_dist = np.minimum(min_dist, new_dists) - - remaining.pop(best_idx) - min_dist = np.delete(min_dist, best_idx) - - return selected + return [candidates[i] for i in selected] diff --git a/src/refrover/similarity.py b/src/refrover/similarity.py index f619775..3227265 100644 --- a/src/refrover/similarity.py +++ b/src/refrover/similarity.py @@ -38,6 +38,25 @@ def jaccard_to_distance(sim_matrix: pd.DataFrame) -> pd.DataFrame: 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, 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_selectors/test_containment.py b/tests/test_selectors/test_containment.py index 37028b3..aa40bbf 100644 --- a/tests/test_selectors/test_containment.py +++ b/tests/test_selectors/test_containment.py @@ -2,6 +2,8 @@ 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 @@ -84,3 +86,42 @@ 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) From b19dbe5ed34c55bf22c1fcecede0680b43f6dfa8 Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 21:04:45 -0400 Subject: [PATCH 07/10] Emit real per-sample variance in MetaBAT2 output The MetaBAT2 formatter wrote every per-sample variance column as 0 because CoverM was only ever asked for mean depth, and the internal coverage schema had no place to carry variance. MetaBAT2's jgi format uses that column, so binning silently lost differential signal. Coverage schema now carries variance: length, {sample}_depth (mean), {sample}_var (variance) - run_coverm requests `--methods length mean variance` (configurable). - load_coverage normalizes CoverM's `{sample} {Method}` columns into the schema via normalize_coverage: ` Mean` -> _depth, ` Variance` -> _var, ` Length` -> a single collapsed `length`. Already-normalized frames pass through unchanged. - MetaBAT2 formatter emits the real paired variance, falling back to 0.0 only when a sample has no variance column (mean-only CoverM runs). - generic / semibin2 / maxbin2 / concoct select mean columns by suffix via a shared formatters/_columns helper, so variance columns never leak into the mean-only formats. MaxBin2 .abund files are now named by clean sample id. - The standalone `refrover format` command routes through load_coverage instead of a raw pd.read_csv, so it normalizes raw CoverM tables too. Tests: new test_coverage.py locks the CoverM-parsing behaviour (which can't be exercised without CoverM installed); test_metabat2 now asserts real variance passes through and that absent variance falls back to zero. 85 pass. --- src/refrover/cli.py | 4 +- src/refrover/coverage.py | 92 +++++++++++++++++++++++--- src/refrover/formatters/_columns.py | 29 ++++++++ src/refrover/formatters/concoct.py | 5 +- src/refrover/formatters/generic.py | 13 ++-- src/refrover/formatters/maxbin2.py | 8 ++- src/refrover/formatters/metabat2.py | 15 +++-- src/refrover/formatters/semibin2.py | 5 +- tests/conftest.py | 8 ++- tests/test_coverage.py | 70 ++++++++++++++++++++ tests/test_formatters/test_metabat2.py | 24 +++++-- 11 files changed, 239 insertions(+), 34 deletions(-) create mode 100644 src/refrover/formatters/_columns.py create mode 100644 tests/test_coverage.py diff --git a/src/refrover/cli.py b/src/refrover/cli.py index 8ab6e4f..4b3bda6 100644 --- a/src/refrover/cli.py +++ b/src/refrover/cli.py @@ -184,14 +184,14 @@ def coverage(bams, outdir, threads, force): def format(coverage_dir, binners, outdir, force): """Format coverage tables for downstream binners.""" from refrover.formatters import format_for_binner - import pandas as pd + 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 = pd.read_csv(cov_tsv, sep="\t", index_col=0) + df = load_coverage(cov_tsv) for binner in binner_list: out = format_for_binner(df, binner=binner, outdir=outdir, force=force) diff --git a/src/refrover/coverage.py b/src/refrover/coverage.py index c5bc61b..cc0c94d 100644 --- a/src/refrover/coverage.py +++ b/src/refrover/coverage.py @@ -1,11 +1,29 @@ """ -CoverM wrapper: compute per-contig mean depth from sorted BAMs. +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, @@ -13,15 +31,17 @@ def run_coverm( *, threads: int = 8, force: bool = False, + methods: tuple[str, ...] = DEFAULT_METHODS, coverm_path: str = "coverm", ) -> Path: """ - Run CoverM on all BAMs in bams_dir. + Run CoverM on all sorted BAMs in bams_dir. - Writes coverage.tsv to outdir with columns: - Contig Length sample1_depth sample2_depth ... + 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 path to coverage.tsv. + Returns the path to coverage.tsv. """ bams_dir = Path(bams_dir) outdir = Path(outdir) @@ -38,7 +58,7 @@ def run_coverm( cmd = [ coverm_path, "contig", "--bam-files", *[str(b) for b in bam_files], - "--methods", "mean", + "--methods", *methods, "--threads", str(threads), "--output-file", str(out_tsv), ] @@ -49,7 +69,63 @@ def run_coverm( 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 into a DataFrame indexed by contig name.""" + """Load a CoverM coverage.tsv and normalize it to the internal schema.""" df = pd.read_csv(tsv_path, sep="\t", index_col=0) - return df + return normalize_coverage(df) 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 index 3f79c37..277df70 100644 --- a/src/refrover/formatters/concoct.py +++ b/src/refrover/formatters/concoct.py @@ -3,12 +3,13 @@ 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 - depth_cols = [c for c in coverage_df.columns if c != "length"] - df_out = coverage_df[depth_cols].round(0).astype(int) + 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 index 46d1f16..7716ffd 100644 --- a/src/refrover/formatters/generic.py +++ b/src/refrover/formatters/generic.py @@ -3,18 +3,19 @@ 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. - - Input: DataFrame with index=contig, columns including 'length' and one depth - column per sample (e.g. 's001_depth'). + Write a generic coverage TSV: contig index, optional 'length', and one + mean-depth column per sample. Variance columns are omitted. - Output: coverage_generic.tsv — same layout, tab-separated. + Output: coverage_generic.tsv — tab-separated. """ out = outdir / "coverage_generic.tsv" if out.exists() and not force: return out - coverage_df.to_csv(out, sep="\t") + 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 index 0db49b5..50924b8 100644 --- a/src/refrover/formatters/maxbin2.py +++ b/src/refrover/formatters/maxbin2.py @@ -3,12 +3,14 @@ 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]: - depth_cols = [c for c in coverage_df.columns if c != "length"] out_paths = [] - for col in depth_cols: - out = outdir / f"{col}.abund" + 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 diff --git a/src/refrover/formatters/metabat2.py b/src/refrover/formatters/metabat2.py index a4749fd..f5ba836 100644 --- a/src/refrover/formatters/metabat2.py +++ b/src/refrover/formatters/metabat2.py @@ -2,12 +2,17 @@ MetaBAT2 formatter: jgi_summarize_bam_contig_depths format. Columns: contigName, contigLen, totalAvgDepth, sample1, sample1-var, sample2, sample2-var, ... -Variance columns are set to 0 (RefRover does not compute variance from CoverM mean output). + +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.""" @@ -15,7 +20,7 @@ def write_metabat2(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) if out.exists() and not force: return out - depth_cols = [c for c in coverage_df.columns if c != "length"] + depth_cols = mean_columns(coverage_df) rows = [] for contig, row in coverage_df.iterrows(): rec = { @@ -24,8 +29,10 @@ def write_metabat2(coverage_df: pd.DataFrame, outdir: Path, force: bool = False) "totalAvgDepth": row[depth_cols].mean(), } for col in depth_cols: - rec[col] = row[col] - rec[f"{col}-var"] = 0.0 + 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) diff --git a/src/refrover/formatters/semibin2.py b/src/refrover/formatters/semibin2.py index 2c0168f..5c39d84 100644 --- a/src/refrover/formatters/semibin2.py +++ b/src/refrover/formatters/semibin2.py @@ -3,11 +3,12 @@ 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 - depth_cols = [c for c in coverage_df.columns if c != "length"] - coverage_df[depth_cols].to_csv(out, sep="\t") + coverage_df[mean_columns(coverage_df)].to_csv(out, sep="\t") return out diff --git a/tests/conftest.py b/tests/conftest.py index a7a61be..703ad19 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -55,8 +55,9 @@ def sparse_sim(): @pytest.fixture def coverage_df(): """ - Minimal coverage DataFrame: 5 contigs x 3 samples + length column. - Matches the output format from CoverM (index=contig, columns=length + depth per sample). + 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 = { @@ -64,6 +65,9 @@ def coverage_df(): "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) 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/test_metabat2.py b/tests/test_formatters/test_metabat2.py index 95092b9..9067b46 100644 --- a/tests/test_formatters/test_metabat2.py +++ b/tests/test_formatters/test_metabat2.py @@ -18,11 +18,25 @@ def test_metabat2_required_columns(coverage_df, tmp_path): 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 depth column must have a paired -var column set to 0 - depth_cols = ["s1_depth", "s2_depth", "s3_depth"] - for col in depth_cols: - assert f"{col}-var" in df.columns - assert (df[f"{col}-var"] == 0.0).all() + # 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): From 801dd102d822997aff6cddeffb5bb29eb04841ed Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 22:19:49 -0400 Subject: [PATCH 08/10] Add refrover.containment: in-package containment-matrix computation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cross-sample containment matrix that ContainmentSelector depends on was only produced by an out-of-tree script (data/mouse_rewilding/compute_containment.py). This makes it a first-class package stage so containment selection is self-sufficient. - refrover/containment.py: compute_containment_matrix(read_sigs, assembly_sigs) computes containment(reads_i, assembly_j) via sourmash MinHash.contained_by, returning a sample x sample DataFrame (query_id index) ready for ContainmentSelector and load_containment_matrix. Plus sigs_by_sample() and write_containment_matrix() helpers. - refrover containment CLI subcommand: takes --read-sketches / --assembly-sketches directories, writes containment_matrix.tsv. Idempotent (skips existing unless --force). Uses the sourmash API rather than the analysis script's raw-JSON hash parsing — slower but correct across sketch parameters and maintainable in-tree. Tests: new test_containment.py builds real sourmash sketches in a tmp dir and checks exact (asymmetric) containment values, labelling, TSV round-trip through load_containment_matrix, and that the output feeds ContainmentSelector. 90 pass. --- CLAUDE.md | 9 +++- src/refrover/cli.py | 38 ++++++++++++++ src/refrover/containment.py | 78 +++++++++++++++++++++++++++ tests/test_containment.py | 102 ++++++++++++++++++++++++++++++++++++ 4 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 src/refrover/containment.py create mode 100644 tests/test_containment.py diff --git a/CLAUDE.md b/CLAUDE.md index ece0d0e..a4e4715 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -96,6 +96,12 @@ 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/ +# 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/ + # Benchmark selectors against each other on a dataset with known ground truth refrover benchmark \ --manifest samples.tsv \ @@ -178,7 +184,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/src/refrover/cli.py b/src/refrover/cli.py index 4b3bda6..b0bbbcf 100644 --- a/src/refrover/cli.py +++ b/src/refrover/cli.py @@ -47,6 +47,44 @@ def sketch(manifest, outdir, ksize, scaled, threads, 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() 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/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"} From e1fadcf9c521b3bd58f3f0e4e009bbd269bbdea8 Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 22:25:53 -0400 Subject: [PATCH 09/10] Implement Tier-1 selector ranking (variance-explained proxy) + rank-selectors CLI Turns benchmark.py from a stub into a working alignment-free selector evaluator, operationalising the data/mouse_rewilding selector-comparison analysis with the real package selectors. - unique_variance_explained(matrix, selected_ids): fraction of total cross-sample column variance the selected columns explain, via orthogonal (QR) projection of the full matrix onto the selected subspace. Correlated/redundant selections count once, so it rewards prototypes spanning independent axes of variation. - rank_selectors(matrix, selectors, k_values, ...): feeds the same matrix (Jaccard or containment) to every selector for an apples-to-apples ranking. Skips unavailable selectors (feedback; archetype without its package) with a warning instead of aborting; seeds the random baseline for reproducibility. - refrover rank-selectors CLI: takes --sketches or --containment-matrix, writes selector_ranking.tsv and prints the table. - run_benchmark (Tier 2, CheckM2 MAG yield) documented as the gold standard the Tier-1 proxy is meant to predict; still NotImplementedError. On the real 183-sample containment matrix the ranking reproduces the analysis: maxmin and containment lead, greedy_var trails, variance rises with k. Tests: test_benchmark.py covers the projection metric (all-columns=1.0, empty=0, redundancy adds little, nested monotonicity) and ranking (columns, k-monotonicity for nested maxmin, feedback skipped, reproducibility, containment on shared matrix). 100 pass. --- CLAUDE.md | 12 +++- src/refrover/benchmark.py | 147 ++++++++++++++++++++++++++++++++++---- src/refrover/cli.py | 48 +++++++++++++ tests/test_benchmark.py | 96 +++++++++++++++++++++++++ 4 files changed, 290 insertions(+), 13 deletions(-) create mode 100644 tests/test_benchmark.py diff --git a/CLAUDE.md b/CLAUDE.md index a4e4715..05fdfe7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -102,7 +102,15 @@ refrover select --containment-matrix cont/containment_matrix.tsv \ --manifest samples.tsv --selector containment --k 5 --min-jaccard 0.05 \ --outdir assignments/ -# Benchmark selectors against each other on a dataset with known ground truth +# 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 \ @@ -112,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. diff --git a/src/refrover/benchmark.py b/src/refrover/benchmark.py index 7765150..abfffac 100644 --- a/src/refrover/benchmark.py +++ b/src/refrover/benchmark.py @@ -1,13 +1,140 @@ """ Benchmark selector strategies against each other. -Runs each (selector, k) combination, evaluates MAG yield via CheckM2, -and writes a comparison table. +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, @@ -19,16 +146,12 @@ def run_benchmark( checkm2_db: Path | str | None = None, ) -> pd.DataFrame: """ - For each (selector, k) combination: - 1. Run RefRoverPipeline with that selector - 2. Run a binner (MetaBAT2) on the coverage output - 3. Evaluate bins with CheckM2 - - Returns a summary DataFrame with columns: - selector, k, n_mags_passing, completeness_mean, contamination_mean, - cpu_seconds, total_alignment_cpu_seconds + 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( - "Benchmark runner not yet implemented. " - "Requires CheckM2 installation and a labelled ground-truth community." + "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 index b0bbbcf..65a9efd 100644 --- a/src/refrover/cli.py +++ b/src/refrover/cli.py @@ -288,6 +288,54 @@ def run(manifest, selector, k, min_jaccard, containment_matrix, aligner, binners 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() 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() From ec85541a8759c021de71fba94ed70349e80a070c Mon Sep 17 00:00:00 2001 From: Daniel Sprockett Date: Wed, 10 Jun 2026 22:28:23 -0400 Subject: [PATCH 10/10] Add containment_saturation adaptive-k estimator (the validated method) The three existing estimators (scree_elbow, similarity_gap, saturation_curve) all run on a Jaccard matrix as a proxy. This adds the method validated on the rewilded-mouse data: greedy forward selection by residual cross-sample variance (incremental Gram-Schmidt), which gives a monotonically decreasing marginal-gain curve, with the elbow at the first k whose gain drops below 10% of the first. - _containment_saturation in adaptive_k.py; wired into estimate_k and the AdaptiveMethod literal. - refrover select --adaptive-k-method gains the containment_saturation choice (use with a containment matrix for the intended behaviour). On the real 183-sample containment matrix it reproduces the analysis exactly: median k=4 (182 samples at 4, one at 5), matching compare_selectors.py. Tests: containment_saturation added to the all-methods validity check, plus a low-dimensional-data test and a dominant-axis saturation test. 102 pass. --- src/refrover/adaptive_k.py | 76 +++++++++++++++++++++++++++++++++++++- src/refrover/cli.py | 3 +- tests/test_adaptive_k.py | 31 +++++++++++++++- 3 files changed, 106 insertions(+), 4 deletions(-) diff --git a/src/refrover/adaptive_k.py b/src/refrover/adaptive_k.py index 36338c3..0e1f0af 100644 --- a/src/refrover/adaptive_k.py +++ b/src/refrover/adaptive_k.py @@ -24,6 +24,15 @@ 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. """ @@ -34,7 +43,9 @@ from typing import Literal -AdaptiveMethod = Literal["scree_elbow", "similarity_gap", "saturation_curve"] +AdaptiveMethod = Literal[ + "scree_elbow", "similarity_gap", "saturation_curve", "containment_saturation" +] def estimate_k( @@ -77,6 +88,8 @@ def estimate_k( 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}") @@ -198,3 +211,64 @@ def _saturation_curve( 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/cli.py b/src/refrover/cli.py index 65a9efd..2a6f168 100644 --- a/src/refrover/cli.py +++ b/src/refrover/cli.py @@ -99,7 +99,8 @@ def containment(read_sketches, assembly_sketches, ksize, outdir, force): @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"])) + 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()) diff --git a/tests/test_adaptive_k.py b/tests/test_adaptive_k.py index 37a6c42..ed6c35e 100644 --- a/tests/test_adaptive_k.py +++ b/tests/test_adaptive_k.py @@ -1,7 +1,10 @@ import pytest import numpy as np import pandas as pd -from refrover.adaptive_k import estimate_k, _scree_elbow, _similarity_gap, _saturation_curve +from refrover.adaptive_k import ( + estimate_k, _scree_elbow, _similarity_gap, _saturation_curve, + _containment_saturation, +) @pytest.fixture @@ -47,12 +50,36 @@ def test_estimate_k_unknown_query(three_cluster_sim): def test_all_methods_return_valid_k(three_cluster_sim): - for method in ("scree_elbow", "similarity_gap", "saturation_curve"): + 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,