rarelens/ml/rarelens_ml/train.py

"""Train a pathogenicity classifier on ClinVar labels ((likely) pathogenic vs (likely) benign).

Label leakage warning: CLIN_SIG must never be a feature. This is a learning exercise, not a clinical model.
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
import mlflow
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 GroupShuffleSplit

from rarelens_ml.features import RAW_COLUMNS, build
from rarelens_ml.model import PathogenicityModel

PACKAGE_DIR = Path(__file__).resolve().parent
MODEL_NAME = "rarelens-pathogenicity"
PARAMS = {
    "n_estimators": 400, "learning_rate": 0.05, "num_leaves": 31, "class_weight": "balanced",
    "verbose": -1,
}

POS = {"pathogenic", "likely_pathogenic"}
NEG = {"benign", "likely_benign"}
# VEP --tab column -> raw feature column (am_pathogenicity already matches).
VEP_TO_RAW = {
    "IMPACT": "impact", "Consequence": "consequence", "gnomADe_AF": "gnomad_af",
    "CADD_PHRED": "cadd_phred",
}


def label(clin_sig: object) -> int | None:
    """1 / 0 when every ClinVar term agrees, None for VUS, conflicts and missing values.

    Accepts VEP's lowercase comma-separated form ("pathogenic,likely_pathogenic") and ClinVar's
    CLNSIG form ("Pathogenic/Likely_pathogenic").
    """
    if not isinstance(clin_sig, str):
        return None
    terms = {t for t in re.split(r"[,&/|]", clin_sig.strip().lower()) if t and t != "-"}
    if terms and terms <= POS:
        return 1
    if terms and terms <= NEG:
        return 0
    return None


def read_vep_tab(path: str | Path) -> pd.DataFrame:
    """Read VEP --tab output as strings, keeping "-" (VEP's missing marker) verbatim.

    Skips the "##" preamble by position instead of comment="#", which would also cut any value
    containing "#".
    """
    with open(path) as fh:
        for n, line in enumerate(fh):
            if line.startswith("#Uploaded_variation"):
                break
        else:
            raise ValueError(f"{path}: no #Uploaded_variation header; is this VEP --tab output?")
    df = pd.read_csv(path, sep="\t", skiprows=n, dtype=str, keep_default_na=False)
    return df.rename(columns={"#Uploaded_variation": "Uploaded_variation"})


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)


def log_and_register(clf: lgb.LGBMClassifier, model_name: str, alias: str) -> str:
    """Log the pyfunc, register it and point `alias` at the new version. Returns the version."""
    info = mlflow.pyfunc.log_model(
        name="model",
        python_model=PathogenicityModel(clf),
        code_paths=[str(PACKAGE_DIR)],
        registered_model_name=model_name,
        pip_requirements=[
            f"lightgbm=={lgb.__version__}",
            f"pandas=={pd.__version__}",
            f"scikit-learn=={sklearn.__version__}",
        ],
    )
    version = str(info.registered_model_version)
    MlflowClient().set_registered_model_alias(model_name, alias, version)
    return version


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--tsv", required=True)
    p.add_argument("--register", action="store_true",
                   help="register the model and move the alias to the new version")
    p.add_argument("--alias", default="production")
    a = p.parse_args()

    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]
        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}")
        else:
            mlflow.pyfunc.log_model(name="model", python_model=PathogenicityModel(clf),
                                    code_paths=[str(PACKAGE_DIR)])


if __name__ == "__main__":
    main()