Skip to content

Commit 7834124

Browse files
mikesshmikessh
committed
wip
1 parent 889f298 commit 7834124

4 files changed

Lines changed: 311 additions & 20 deletions

File tree

‎benchmarks/covid19_associations_benchmark.py‎

Lines changed: 76 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,11 @@
1010
import json
1111
import os
1212
import time
13+
import tracemalloc
1314
from pathlib import Path
1415

1516
import pandas as pd
17+
from sklearn.metrics import roc_auc_score
1618

1719
from mir.biomarkers.associations import AssociationParams, associate_clonotype_metadata, build_public_clonotype_panel
1820
from mir.common.filter import filter_functional
@@ -32,6 +34,41 @@ def _env_int(name: str, default: int) -> int:
3234
return max(1, value)
3335

3436

37+
def _reference_file(dataset_root: Path) -> Path | None:
38+
candidates = [
39+
dataset_root / "covid19_biomarker_clonotypes.csv",
40+
dataset_root / "covid_associated_clonotypes.csv",
41+
]
42+
for candidate in candidates:
43+
if candidate.exists():
44+
return candidate
45+
return None
46+
47+
48+
def _reference_cdr3_set(path: Path) -> set[str]:
49+
df = pd.read_csv(path)
50+
for col in ("cdr3", "junction_aa", "sequence"):
51+
if col in df.columns:
52+
return {str(x) for x in df[col].dropna().astype(str)}
53+
return set()
54+
55+
56+
def _sample_biomarker_scores(samples: list[SampleRepertoire], biomarker_cdr3: set[str]) -> pd.DataFrame:
57+
rows = []
58+
for sample in samples:
59+
rep = sample.get_locus("TRB")
60+
seqs = {str(c.junction_aa) for c in rep.clonotypes if c.junction_aa}
61+
score = float(len(seqs & biomarker_cdr3))
62+
rows.append(
63+
{
64+
"sample_id": sample.sample_id,
65+
"covid": 1 if str(sample.sample_metadata.get("COVID_status", "")) == "COVID" else 0,
66+
"score": score,
67+
}
68+
)
69+
return pd.DataFrame(rows)
70+
71+
3572
def main() -> int:
3673
dataset_root = ensure_airr_covid19()
3774
metadata = pd.read_csv(dataset_root / "metadata_trb_min100000.tsv", sep="\t", dtype={"donor_id": "string"}, low_memory=False)
@@ -45,11 +82,14 @@ def main() -> int:
4582
samples: list[SampleRepertoire] = []
4683

4784
t0 = time.perf_counter()
85+
tracemalloc.start()
4886
for _, row in metadata.sort_values(["COVID_status", "sample_id"]).head(max_samples).iterrows():
4987
path = Path(dataset_root) / str(row["file_name"])
5088
if not path.exists():
5189
continue
52-
clones = parser.parse(str(path))
90+
clones = [c for c in parser.parse(str(path)) if str(c.locus).upper() == "TRB"]
91+
if not clones:
92+
continue
5393
rep = filter_functional(LocusRepertoire(clonotypes=clones, locus="TRB", repertoire_id=str(row["sample_id"])))
5494
if rep.clonotype_count == 0:
5595
continue
@@ -62,8 +102,11 @@ def main() -> int:
62102
)
63103

64104
load_s = time.perf_counter() - t0
105+
_, load_peak = tracemalloc.get_traced_memory()
106+
tracemalloc.stop()
65107
targets = build_public_clonotype_panel(samples, locus="TRB", min_sample_fraction=min_fraction)[:max_targets]
66108

109+
tracemalloc.start()
67110
t1 = time.perf_counter()
68111
fisher_res = associate_clonotype_metadata(
69112
samples,
@@ -73,7 +116,10 @@ def main() -> int:
73116
params=AssociationParams(test="fisher", count_mode="sample", match_mode="none"),
74117
)
75118
fisher_s = time.perf_counter() - t1
119+
_, fisher_peak = tracemalloc.get_traced_memory()
120+
tracemalloc.stop()
76121

122+
tracemalloc.start()
77123
t2 = time.perf_counter()
78124
depth_res = associate_clonotype_metadata(
79125
samples,
@@ -83,17 +129,45 @@ def main() -> int:
83129
params=AssociationParams(test="depth_glm", count_mode="rearrangement", match_mode="none"),
84130
)
85131
depth_s = time.perf_counter() - t2
132+
_, depth_peak = tracemalloc.get_traced_memory()
133+
tracemalloc.stop()
134+
135+
fisher_df = fisher_res.table.to_pandas().sort_values(["q_value", "p_value"]).reset_index(drop=True)
136+
positive_hits = fisher_df[(fisher_df["odds_ratio"].fillna(0.0) > 1.0) & (fisher_df["q_value"] < 0.2)]
137+
if positive_hits.empty:
138+
positive_hits = fisher_df.head(30)
139+
biomarker_set = set(positive_hits["junction_aa"].astype(str))
140+
score_df = _sample_biomarker_scores(samples, biomarker_set)
141+
auc = float("nan")
142+
if score_df["covid"].nunique() == 2 and score_df["score"].nunique() > 1:
143+
auc = float(roc_auc_score(score_df["covid"], score_df["score"]))
144+
145+
ref = _reference_file(Path(dataset_root))
146+
ref_overlap_top100 = None
147+
ref_overlap_biomarkers = None
148+
if ref is not None:
149+
ref_set = _reference_cdr3_set(ref)
150+
top100 = set(fisher_df.head(100)["junction_aa"].astype(str))
151+
ref_overlap_top100 = len(top100 & ref_set)
152+
ref_overlap_biomarkers = len(biomarker_set & ref_set)
86153

87154
out = {
88155
"dataset_root": str(dataset_root),
89156
"samples": len(samples),
90157
"targets": len(targets),
91158
"load_seconds": load_s,
159+
"load_peak_mib": float(load_peak / (1024 ** 2)),
92160
"fisher_seconds": fisher_s,
161+
"fisher_peak_mib": float(fisher_peak / (1024 ** 2)),
93162
"depth_glm_seconds": depth_s,
163+
"depth_glm_peak_mib": float(depth_peak / (1024 ** 2)),
94164
"fisher_rows": int(fisher_res.table.height),
95165
"depth_rows": int(depth_res.table.height),
96-
"reference_csv_exists": bool((Path(dataset_root) / "covid_associated_clonotypes.csv").exists()),
166+
"biomarker_count": int(len(biomarker_set)),
167+
"separation_auc": auc,
168+
"reference_csv": str(ref) if ref is not None else None,
169+
"reference_overlap_top100": ref_overlap_top100,
170+
"reference_overlap_biomarkers": ref_overlap_biomarkers,
97171
}
98172

99173
print(json.dumps(out, indent=2, sort_keys=True))

‎benchmarks/covid19_reference_split_benchmark.py‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
"""Benchmark alpha/beta separation using split reference biomarker clonotype file.
2+
3+
This benchmark does not infer biomarkers; it scores samples by counting presence
4+
of reference biomarker clonotypes from ``covid_associated_clonotypes.csv`` split
5+
by chain and reports runtime, memory footprint, and AUC.
6+
7+
Usage:
8+
source .venv/bin/activate.fish
9+
/Users/mikesh/vcs/code/mirpy/.venv/bin/python benchmarks/covid19_reference_split_benchmark.py
10+
"""
11+
12+
from __future__ import annotations
13+
14+
import json
15+
import os
16+
import time
17+
import tracemalloc
18+
from pathlib import Path
19+
20+
import pandas as pd
21+
from sklearn.metrics import roc_auc_score
22+
23+
from mir.common.filter import filter_functional
24+
from mir.common.parser import ClonotypeTableParser
25+
from mir.common.repertoire import LocusRepertoire
26+
from mir.utils.notebook_assets import ensure_airr_covid19
27+
28+
29+
def _env_int(name: str, default: int) -> int:
30+
raw = os.getenv(name)
31+
if raw is None:
32+
return default
33+
try:
34+
value = int(raw)
35+
except ValueError:
36+
return default
37+
return max(1, value)
38+
39+
40+
def _load_reference(dataset_root: Path) -> dict[str, set[str]]:
41+
ref_path = dataset_root / "covid_associated_clonotypes.csv"
42+
if not ref_path.exists():
43+
raise FileNotFoundError(f"Reference file is missing: {ref_path}")
44+
45+
ref = pd.read_csv(ref_path)
46+
required = {"cdr3", "chain", "has_covid_association"}
47+
missing = sorted(required - set(ref.columns))
48+
if missing:
49+
raise ValueError(f"Reference file missing required columns: {missing}")
50+
51+
ref = ref[ref["has_covid_association"] == True].copy()
52+
chain_to_locus = {"alpha": "TRA", "beta": "TRB"}
53+
out: dict[str, set[str]] = {}
54+
for chain, locus in chain_to_locus.items():
55+
out[locus] = set(ref[ref["chain"] == chain]["cdr3"].astype(str))
56+
return out
57+
58+
59+
def _score_chain(
60+
*,
61+
dataset_root: Path,
62+
metadata: pd.DataFrame,
63+
locus: str,
64+
biomarkers: set[str],
65+
max_samples: int,
66+
) -> dict[str, float | int | str]:
67+
parser = ClonotypeTableParser()
68+
selected = metadata[
69+
(metadata["locus"].astype(str).str.upper() == locus)
70+
& (metadata["COVID_status"].isin(["COVID", "healthy"]))
71+
].copy()
72+
73+
if "is_bad_reseq" in selected.columns:
74+
bad_mask = selected["is_bad_reseq"].fillna("").astype(str).str.strip().str.lower().isin({"1", "true", "yes"})
75+
selected = selected[~bad_mask].copy()
76+
77+
selected = selected.sort_values(["sample_id"])
78+
half = max_samples // 2
79+
covid_part = selected[selected["COVID_status"] == "COVID"].head(half)
80+
healthy_part = selected[selected["COVID_status"] == "healthy"].head(half)
81+
selected = pd.concat([covid_part, healthy_part], axis=0).sort_values(["COVID_status", "sample_id"]).reset_index(drop=True)
82+
83+
labels: list[int] = []
84+
scores: list[float] = []
85+
n_loaded = 0
86+
87+
tracemalloc.start()
88+
t0 = time.perf_counter()
89+
90+
for _, row in selected.iterrows():
91+
sample_path = dataset_root / str(row["file_name"])
92+
if not sample_path.exists():
93+
continue
94+
clones = [c for c in parser.parse(str(sample_path)) if str(c.locus).upper() == locus]
95+
if not clones:
96+
continue
97+
98+
rep = filter_functional(LocusRepertoire(clonotypes=clones, locus=locus, repertoire_id=str(row["sample_id"])))
99+
if rep.clonotype_count == 0:
100+
continue
101+
102+
seqs = {str(c.junction_aa) for c in rep.clonotypes if c.junction_aa}
103+
score = float(len(seqs & biomarkers))
104+
105+
labels.append(1 if str(row["COVID_status"]) == "COVID" else 0)
106+
scores.append(score)
107+
n_loaded += 1
108+
109+
elapsed = time.perf_counter() - t0
110+
_, peak = tracemalloc.get_traced_memory()
111+
tracemalloc.stop()
112+
113+
auc = float("nan")
114+
if len(set(labels)) == 2 and len(set(scores)) > 1:
115+
auc = float(roc_auc_score(labels, scores))
116+
117+
return {
118+
"locus": locus,
119+
"samples_loaded": int(n_loaded),
120+
"covid_loaded": int(sum(labels)),
121+
"healthy_loaded": int(len(labels) - sum(labels)),
122+
"biomarker_count": int(len(biomarkers)),
123+
"elapsed_seconds": float(elapsed),
124+
"peak_mib": float(peak / (1024 ** 2)),
125+
"auc": float(auc),
126+
"score_mean": float(sum(scores) / len(scores)) if scores else 0.0,
127+
}
128+
129+
130+
def main() -> int:
131+
dataset_root = ensure_airr_covid19()
132+
metadata = pd.read_csv(dataset_root / "metadata.tsv", sep="\t", dtype={"donor_id": "string"}, low_memory=False)
133+
134+
max_samples = _env_int("MIRPY_COVID_REF_BENCH_SAMPLES", 160)
135+
refs = _load_reference(dataset_root)
136+
137+
out = {
138+
"dataset_root": str(dataset_root),
139+
"max_samples_per_chain": int(max_samples),
140+
"results": [],
141+
}
142+
143+
for locus in ("TRA", "TRB"):
144+
out["results"].append(
145+
_score_chain(
146+
dataset_root=dataset_root,
147+
metadata=metadata,
148+
locus=locus,
149+
biomarkers=refs.get(locus, set()),
150+
max_samples=max_samples,
151+
)
152+
)
153+
154+
print(json.dumps(out, indent=2, sort_keys=True))
155+
return 0
156+
157+
158+
if __name__ == "__main__":
159+
raise SystemExit(main())

‎notebooks/covid19_biomarkers.ipynb‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -518,19 +518,28 @@
518518
"candidate_path = derived_dir / \"covid_associated_clonotypes_candidates.csv\"\n",
519519
"fisher_df.to_csv(candidate_path, index=False)\n",
520520
"\n",
521-
"reference_path = dataset_root / \"covid_associated_clonotypes.csv\"\n",
521+
"reference_candidates = [\n",
522+
" dataset_root / \"covid19_biomarker_clonotypes.csv\",\n",
523+
" dataset_root / \"covid_associated_clonotypes.csv\",\n",
524+
"]\n",
525+
"reference_path = next((p for p in reference_candidates if p.exists()), None)\n",
526+
"\n",
522527
"concordance = {\n",
523528
" \"candidate_path\": str(candidate_path),\n",
524-
" \"reference_path\": str(reference_path),\n",
525-
" \"reference_exists\": reference_path.exists(),\n",
529+
" \"reference_path\": str(reference_path) if reference_path is not None else None,\n",
530+
" \"reference_exists\": reference_path is not None,\n",
526531
"}\n",
527532
"\n",
528-
"if reference_path.exists():\n",
533+
"if reference_path is not None:\n",
529534
" ref = pd.read_csv(reference_path)\n",
530-
" if \"junction_aa\" in ref.columns:\n",
535+
" ref_set = None\n",
536+
" for col in (\"junction_aa\", \"cdr3\", \"sequence\"):\n",
537+
" if col in ref.columns:\n",
538+
" ref_set = set(ref[col].dropna().astype(str))\n",
539+
" break\n",
540+
" if ref_set is not None:\n",
531541
" top_n = 200\n",
532542
" cand_top = fisher_df.head(top_n)[\"junction_aa\"].astype(str)\n",
533-
" ref_set = set(ref[\"junction_aa\"].astype(str))\n",
534543
" overlap = int(cand_top.isin(ref_set).sum())\n",
535544
" concordance.update(\n",
536545
" {\n",

0 commit comments

Comments
 (0)