feat(ml): train a real model, and report the number that matters rather than the flattering one

"Variants are unscored" was accurate: nothing was ever trained, so a quarter of every rank was
dead weight and the UI leaked a connection error at the reader.

- scripts/make-training-set.sh derives a training table from ClinVar directly. ClinVar already
  carries the molecular consequence, the gene and an allele frequency, which is the feature set
  serving sends, so this avoids running VEP over hundreds of thousands of variants. 2-star
  records only.
- train.py now holds out whole genes (GroupShuffleSplit). docs/data.md had said to do this since
  the data pass; the code was still doing a random split, which is the leak Grimm 2015 describes.
- evaluate() reports missense on its own. On the last run: AUROC 0.986 over 74,239 held-out
  variants, but 0.872 over the 13,553 missense ones, and the docs say plainly why even that is
  flattered — within missense the only live feature is allele frequency, and ClinVar's benign
  calls often use allele frequency as evidence (ACMG BA1/BS1), so the feature partly caused the
  label.
- the 503 now names what is missing (model@alias via tracking URI) and leaves the exception in
  the server log instead of the UI.
- make training-set / make train; the 58 MB table is gitignored.

Verified end to end: model registered as v2, the simulated NF2 case scores 0.999 on the planted
variant, and it now ranks 1.00 with all four components live.

Tests: api 77, ml 22, loader 16, web 32; ruff, mypy, svelte-check clean.
This commit is contained in:
Kemal Yaylali
2026-09-12 09:13:54 +01:00
parent 3ab404ebe5
commit 197975cc42
9 changed files with 235 additions and 14 deletions
+64 -6
View File
@@ -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}")
+65 -1
View File
@@ -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