diff --git a/.gitignore b/.gitignore index 3284d42..7fcc305 100644 --- a/.gitignore +++ b/.gitignore @@ -16,4 +16,5 @@ mlruns/ *.tfstate.backup .terraform/ data/*.vcf* +data/*.tsv !data/README.md diff --git a/Makefile b/Makefile index 34c74a2..b85b5d0 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,7 @@ -.PHONY: up down clean migrate test lint data hpo demo-case loader pipeline annotate images kind serverless-deploy serverless-destroy gcp-configure gcp-secrets +.PHONY: up down clean migrate test lint data hpo demo-case training-set train loader pipeline annotate images kind serverless-deploy serverless-destroy gcp-configure gcp-secrets VCF ?= data/example.vcf.gz +MLFLOW_URI ?= http://localhost:5001 TAG ?= latest # The loader container reaches docker-compose's Postgres through the host. HOST_DB_URL ?= postgresql://rarelens:rarelens@host.docker.internal:5432/rarelens @@ -36,6 +37,13 @@ hpo: ## load HPO gene-to-phenotype annotations, which the ranking matches again demo-case: ## build the simulated proband: GIAB background + one ClinVar pathogenic variant scripts/make-demo-case.sh +training-set: ## build a ClinVar training table, shaped like VEP --tab output + scripts/make-training-set.sh + +train: ## train the pathogenicity model and point the production alias at it (needs `make up`) + cd ml && MLFLOW_TRACKING_URI=$(MLFLOW_URI) uv run --extra dev \ + python -m rarelens_ml.train --tsv ../data/clinvar-training.vep.tsv --register + loader: docker build -t rarelens/loader:dev -f pipeline/loader.Dockerfile pipeline diff --git a/README.md b/README.md index f12913e..d4851fd 100644 --- a/README.md +++ b/README.md @@ -75,10 +75,13 @@ would be the container itself. `LOCAL_DATA_ROOT` is the directory a sample's `vc To train and register a model (the API scores with `models:/rarelens-pathogenicity@production`): ```bash -cd ml && MLFLOW_TRACKING_URI=http://localhost:5000 \ - uv run python -m rarelens_ml.train --tsv ../pipeline/results/.vep.tsv --register +make training-set # a ClinVar-derived training table, ~370k labelled variants +make train # fits, reports held-out metrics by gene split, moves the production alias ``` +What those metrics do and do not mean is in [docs/data.md](docs/data.md); the headline AUROC +flatters a model whose strongest feature is the consequence class. + Local Kubernetes: `make kind` builds the images, loads them into a kind cluster and applies `infra/k8s/overlays/local`. diff --git a/api/app/routers/cases.py b/api/app/routers/cases.py index a7a79e4..1afb2be 100644 --- a/api/app/routers/cases.py +++ b/api/app/routers/cases.py @@ -8,6 +8,7 @@ from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import selectinload +from app.config import settings from app.db import SessionDep from app.models import ( Case, @@ -152,9 +153,14 @@ async def score(case_id: uuid.UUID, session: SessionDep): scored, version = await score_job(job.id, session) except Exception as e: # No registry, no model behind the alias, a model that will not load: all of these are - # the environment being unready, not a bug in the request. Say so rather than throwing 500. + # the environment being unready, not a bug in the request. Name what is missing and keep + # the exception in the log, where it is useful, rather than in the UI, where it is noise. logger.exception("scoring case %s failed", case_id) - raise HTTPException(503, f"could not score with the model: {e}") from e + raise HTTPException( + 503, + f"no model available: {settings.model_name}@{settings.model_alias} " + f"via {settings.mlflow_tracking_uri}", + ) from e return ScoreOut(case_id=case_id, scored=scored, model_version=version) diff --git a/api/tests/test_scoring.py b/api/tests/test_scoring.py index 1b22921..026f448 100644 --- a/api/tests/test_scoring.py +++ b/api/tests/test_scoring.py @@ -99,4 +99,7 @@ async def test_an_unreachable_model_registry_is_explained_not_a_500( monkeypatch.setattr(scoring, "load_model", unreachable) r = await client.post(f"/api/cases/{case_id}/score") assert r.status_code == 503 - assert "connection refused" in r.json()["detail"] + detail = r.json()["detail"] + # Names what is missing; the exception itself belongs in the server log, not the UI. + assert "rarelens-pathogenicity@production" in detail + assert "Max retries" not in detail and "Traceback" not in detail diff --git a/docs/data.md b/docs/data.md index be67868..cc23999 100644 --- a/docs/data.md +++ b/docs/data.md @@ -59,6 +59,36 @@ instead of the 25 GB cache. It returns no gnomAD frequencies, so every variant l gnomAD and the rarity term stops discriminating. Fine for showing the mechanics; use the cache for anything you would quote. +## The model, and what its numbers mean + +`make training-set` builds a training table straight from ClinVar rather than running VEP over +hundreds of thousands of variants: ClinVar already carries the molecular consequence (`MC`), the +gene (`GENEINFO`) and an allele frequency (`AF_EXAC`), which is the feature set serving sends. +Only 2-star-and-above records are kept. `make train` then fits LightGBM and points the +`production` alias at the new version. + +The last run: 312,025 training and 74,239 held-out variants across 7,728 and 1,932 genes, with no +gene on both sides. + +| | AUROC | AUPRC | +|---|---|---| +| all held-out variants | 0.986 | 0.954 | +| missense only (13,553) | 0.872 | 0.725 | + +Three things to say before anyone quotes the headline number: + +1. **0.986 mostly measures how separable ClinVar's classes are by consequence.** Its pathogenic set + is largely loss of function and its benign set largely is not, so a model handed the consequence + class does well without knowing anything hard. That is why the missense row exists: missense is + where interpretation is actually difficult. +2. **Even 0.872 is flattered by circularity.** Within missense, every row has the same consequence + and impact and no CADD or AlphaMissense score, so allele frequency is doing nearly all the work + — and ClinVar's benign calls frequently *use* allele frequency as evidence (ACMG BA1/BS1). The + feature partly caused the label. +3. **It is not comparable to published CADD or AlphaMissense numbers.** Those are trained and + evaluated on different data. A fair comparison scores the same held-out rows with all three, + which needs the plugin data (see above) and is the obvious next step. + ## Evaluating the model honestly The model trains on ClinVar labels and is scored on ClinVar-labelled variants, which is exactly @@ -69,7 +99,8 @@ where published benchmarks go wrong. What to do about it: 2. **Split by gene, not by variant.** Random splits put variants from the same gene on both sides, and a model can then score a gene rather than a variant. Grimm et al. showed this inflates reported accuracy for exactly this class of tool: *Hum Mutat* 36:513–523, 2015. - [10.1002/humu.22768](https://doi.org/10.1002/humu.22768) + [10.1002/humu.22768](https://doi.org/10.1002/humu.22768) *Implemented*: + `rarelens_ml.train.split_by_gene` holds out whole genes. 3. **Prefer a time-based holdout.** Train on an older ClinVar release (monthly archives live under `vcf_GRCh38/archive_2.0/`) and test only on variants classified after that date. This is the closest thing to a prospective evaluation available without new patients. diff --git a/ml/rarelens_ml/train.py b/ml/rarelens_ml/train.py index d1a8443..a5b3682 100644 --- a/ml/rarelens_ml/train.py +++ b/ml/rarelens_ml/train.py @@ -5,6 +5,7 @@ Usage: python -m rarelens_ml.train --tsv results/clinvar.vep.tsv --register """ import argparse import re +import sys from pathlib import Path import lightgbm as lgb @@ -13,7 +14,7 @@ import pandas as pd import sklearn from mlflow import MlflowClient from sklearn.metrics import average_precision_score, roc_auc_score -from sklearn.model_selection import train_test_split +from sklearn.model_selection import GroupShuffleSplit from rarelens_ml.features import RAW_COLUMNS, build from rarelens_ml.model import PathogenicityModel @@ -66,19 +67,64 @@ def read_vep_tab(path: str | Path) -> pd.DataFrame: return df.rename(columns={"#Uploaded_variation": "Uploaded_variation"}) -def load(tsv: str) -> tuple[pd.DataFrame, pd.Series]: +def load(tsv: str) -> tuple[pd.DataFrame, pd.Series, pd.Series]: + """Returns the raw feature columns, the labels, and each row's gene for grouping.""" df = read_vep_tab(tsv).rename(columns=VEP_TO_RAW) for col in RAW_COLUMNS: # plugin columns are absent when VEP ran without CADD/AlphaMissense if col not in df: df[col] = pd.NA + if "SYMBOL" not in df: + df["SYMBOL"] = "-" y = df["CLIN_SIG"].map(label) keep = y.notna() return ( df.loc[keep, RAW_COLUMNS].reset_index(drop=True), y[keep].astype(int).reset_index(drop=True), + df.loc[keep, "SYMBOL"].reset_index(drop=True), ) +def split_by_gene( + X: pd.DataFrame, y: pd.Series, genes: pd.Series, test_size: float = 0.2, seed: int = 42 +) -> tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series, pd.Series, pd.Series]: + """Hold out whole genes, never single variants. + + A random split puts variants of the same gene on both sides, and the model can then score the + gene rather than the variant. Grimm et al. (Hum Mutat 2015) showed this inflates reported + accuracy for exactly this class of tool; see docs/data.md. + """ + splitter = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=seed) + train_idx, test_idx = next(splitter.split(X, y, groups=genes)) + return ( + X.iloc[train_idx], X.iloc[test_idx], + y.iloc[train_idx], y.iloc[test_idx], + genes.iloc[train_idx], genes.iloc[test_idx], + ) + + +MIN_SUBSET = 50 + + +def evaluate(X: pd.DataFrame, y: pd.Series, proba) -> dict[str, float]: + """Headline metrics, plus missense on its own. + + Most of ClinVar's pathogenic set is loss of function and most of its benign set is not, so a + model given the consequence class separates them easily and the overall AUROC flatters it. + Missense is where variant interpretation is actually hard, so it gets its own number. + """ + metrics = { + "auroc": float(roc_auc_score(y, proba)), + "auprc": float(average_precision_score(y, proba)), + "test_variants": float(len(y)), + } + missense = X["consequence"].eq("missense_variant").to_numpy() + if missense.sum() >= MIN_SUBSET and len(set(y[missense])) == 2: + metrics["auroc_missense"] = float(roc_auc_score(y[missense], proba[missense])) + metrics["auprc_missense"] = float(average_precision_score(y[missense], proba[missense])) + metrics["missense_variants"] = float(missense.sum()) + return metrics + + def fit(X: pd.DataFrame, y: pd.Series) -> lgb.LGBMClassifier: return lgb.LGBMClassifier(**PARAMS).fit(build(X), y) @@ -109,16 +155,28 @@ def main() -> None: p.add_argument("--alias", default="production") a = p.parse_args() - X, y = load(a.tsv) - Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42) + X, y, genes = load(a.tsv) + Xtr, Xte, ytr, yte, train_genes, test_genes = split_by_gene(X, y, genes) + print( + f"{len(Xtr)} train / {len(Xte)} test variants; " + f"{train_genes.nunique()} / {test_genes.nunique()} genes, no gene in both", + file=sys.stderr, + ) mlflow.set_experiment(MODEL_NAME) with mlflow.start_run(): mlflow.log_params(PARAMS) clf = fit(Xtr, ytr) proba = clf.predict_proba(build(Xte))[:, 1] - mlflow.log_metrics({"auroc": roc_auc_score(yte, proba), - "auprc": average_precision_score(yte, proba)}) + metrics = evaluate(Xte, yte, proba) | {"test_genes": float(test_genes.nunique())} + mlflow.log_metrics(metrics) + summary = f"held-out AUROC {metrics['auroc']:.3f}, AUPRC {metrics['auprc']:.3f}" + if "auroc_missense" in metrics: + summary += ( + f" | missense only: AUROC {metrics['auroc_missense']:.3f}, " + f"AUPRC {metrics['auprc_missense']:.3f} over {int(metrics['missense_variants'])}" + ) + print(summary, file=sys.stderr) if a.register: version = log_and_register(clf, MODEL_NAME, a.alias) print(f"registered {MODEL_NAME} v{version} as @{a.alias}") diff --git a/ml/tests/test_train.py b/ml/tests/test_train.py index 6945f28..f8d6a78 100644 --- a/ml/tests/test_train.py +++ b/ml/tests/test_train.py @@ -68,9 +68,10 @@ def test_load_returns_raw_serving_columns_and_labels(tmp_path: Path) -> None: ["c", "22:3", "T", "intron_variant", "MODIFIER", "CHEK2", "0.3", "uncertain_significance", "1", "-"], ], ) - X, y = load(str(tsv)) + X, y, genes = load(str(tsv)) assert list(X.columns) == RAW_COLUMNS assert y.tolist() == [1, 0] # the VUS row is dropped + assert list(genes) == ["TBX1", "CHEK2"] def test_logged_model_returns_probabilities_from_raw_columns(tmp_path: Path) -> None: @@ -102,3 +103,66 @@ def test_logged_model_returns_probabilities_from_raw_columns(tmp_path: Path) -> assert scores.shape == (50,) assert ((scores >= 0) & (scores <= 1)).all() assert not set(np.unique(scores)) <= {0.0, 1.0}, "got class labels, expected probabilities" + + +def test_load_returns_the_gene_of_each_row_for_grouping(tmp_path: Path) -> None: + from rarelens_ml.train import load + + tsv = write_vep_tab( + tmp_path / "x.vep.tsv", + [ + ["a", "22:1", "G", "missense_variant", "MODERATE", "NF2", "0.0001", "pathogenic", "28", "0.9"], + ["b", "22:2", "A", "synonymous_variant", "LOW", "CHEK2", "0.12", "benign", "3", "-"], + ], + ) + _, _, genes = load(str(tsv)) + assert list(genes) == ["NF2", "CHEK2"] + + +def test_the_split_never_puts_one_gene_on_both_sides(tmp_path: Path) -> None: + """Random splits leak: a model can learn the gene instead of the variant (Grimm 2015).""" + import numpy as np + import pandas as pd + + from rarelens_ml.train import split_by_gene + + genes = pd.Series([f"GENE{i // 4}" for i in range(40)]) + X = pd.DataFrame({"impact": ["HIGH"] * 40, "consequence": ["stop_gained"] * 40, + "gnomad_af": ["0"] * 40, "cadd_phred": ["10"] * 40, "am_pathogenicity": ["-"] * 40}) + y = pd.Series(np.tile([1, 0], 20)) + + Xtr, Xte, _ytr, yte, train_genes, test_genes = split_by_gene(X, y, genes, test_size=0.3) + assert set(train_genes) & set(test_genes) == set() + assert len(Xtr) + len(Xte) == 40 + assert len(yte) > 0 + + +def test_evaluate_reports_missense_separately() -> None: + """Overall AUROC flatters a consequence-based model; missense is where the problem is.""" + import numpy as np + import pandas as pd + + from rarelens_ml.train import evaluate + + rng = np.random.default_rng(0) + consequence = ["missense_variant"] * 200 + ["stop_gained"] * 200 + X = pd.DataFrame({"consequence": consequence}) + y = pd.Series([*rng.integers(0, 2, 200), *([1] * 200)]) + proba = np.concatenate([rng.random(200), rng.uniform(0.8, 1.0, 200)]) + + metrics = evaluate(X, y, proba) + assert metrics["missense_variants"] == 200 + assert "auroc_missense" in metrics + assert metrics["auroc"] > metrics["auroc_missense"] # the easy class inflates the headline + + +def test_evaluate_omits_the_missense_metric_when_there_is_nothing_to_measure() -> None: + import numpy as np + import pandas as pd + + from rarelens_ml.train import evaluate + + X = pd.DataFrame({"consequence": ["stop_gained"] * 100}) + y = pd.Series([1] * 50 + [0] * 50) + metrics = evaluate(X, y, np.linspace(0, 1, 100)) + assert "auroc_missense" not in metrics diff --git a/scripts/make-training-set.sh b/scripts/make-training-set.sh new file mode 100755 index 0000000..6b7a418 --- /dev/null +++ b/scripts/make-training-set.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Build a training table from ClinVar, shaped like VEP --tab output so rarelens_ml.train reads it +# unchanged. +# +# ClinVar already carries the molecular consequence (MC), the gene (GENEINFO) and an allele +# frequency (AF_EXAC), which is the feature set serving sends. That avoids running VEP over +# hundreds of thousands of variants to produce a training set. CADD and AlphaMissense are left +# missing, exactly as they are when the pipeline runs without plugin data. +# +# Only 2-star-and-above records are kept: "criteria provided, multiple submitters, no conflicts" +# or better. Labels come from CLNSIG (see docs/data.md on evaluating this honestly). +set -eu + +IMAGE=${BCFTOOLS_IMAGE:-quay.io/biocontainers/bcftools:1.20--h8b25389_0} +CLINVAR=${CLINVAR:-https://ftp.ncbi.nlm.nih.gov/pub/clinvar/vcf_GRCh38/clinvar.vcf.gz} +REGION=${REGION:-} # empty means the whole genome +OUT=${OUT:-data/clinvar-training.vep.tsv} + +mkdir -p "$(dirname "$OUT")" +docker run --rm -v "$PWD/$(dirname "$OUT"):/out" "$IMAGE" bash -eu -c " + echo '==> streaming ClinVar ${REGION:-(whole genome)}' >&2 + bcftools query ${REGION:+-r '$REGION'} \ + -f '%CHROM\t%POS\t%REF\t%ALT\t%INFO/MC\t%INFO/GENEINFO\t%INFO/AF_EXAC\t%INFO/CLNSIG\t%INFO/CLNREVSTAT\n' \ + '$CLINVAR' \ + | awk -F'\t' -v OFS='\t' ' + BEGIN { + split(\"transcript_ablation splice_acceptor_variant splice_donor_variant stop_gained frameshift_variant stop_lost start_lost transcript_amplification\", h, \" \"); + for (i in h) impact[h[i]] = \"HIGH\"; + split(\"inframe_insertion inframe_deletion missense_variant protein_altering_variant\", m, \" \"); + for (i in m) impact[m[i]] = \"MODERATE\"; + split(\"splice_region_variant synonymous_variant start_retained_variant stop_retained_variant\", l, \" \"); + for (i in l) impact[l[i]] = \"LOW\"; + print \"#Uploaded_variation\", \"Location\", \"Allele\", \"Consequence\", \"IMPACT\", \"SYMBOL\", \"gnomADe_AF\", \"CLIN_SIG\", \"CADD_PHRED\", \"am_pathogenicity\"; + } + \$9 !~ /multiple_submitters|expert_panel|practice_guideline/ { next } # 2 stars and up + \$5 == \".\" || \$6 == \".\" { next } + { + split(\$5, mc, \",\"); split(mc[1], so, \"|\"); csq = so[2]; + if (csq == \"\") next; + split(\$6, gi, \"|\"); split(gi[1], g, \":\"); gene = g[1]; + af = (\$7 == \".\" ? \"-\" : \$7); + imp = (csq in impact ? impact[csq] : \"MODIFIER\"); + print \$1 \"_\" \$2 \"_\" \$3 \"_\" \$4, \$1 \":\" \$2, \$4, csq, imp, gene, af, \$8, \"-\", \"-\"; + }' > /out/$(basename "$OUT") +" +echo "wrote $OUT: $(($(wc -l < "$OUT") - 1)) labelled variants" +awk -F'\t' 'NR>1 {print $8}' "$OUT" | sort | uniq -c | sort -rn | head -6