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
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 train_test_split

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]:
    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
    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),
    )


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 = load(a.tsv)
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

    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)})
        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()