rarelens/ml/tests/test_train.py

from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from rarelens_ml.train import label, read_vep_tab

HEADER = [
    "Uploaded_variation", "Location", "Allele", "Consequence", "IMPACT", "SYMBOL",
    "gnomADe_AF", "CLIN_SIG", "CADD_PHRED", "am_pathogenicity",
]


def write_vep_tab(path: Path, rows: list[list[str]]) -> Path:
    lines = [
        "## ENSEMBL VARIANT EFFECT PREDICTOR v113.0",
        "## Column descriptions:",
        "#" + "\t".join(HEADER),
        *("\t".join(r) for r in rows),
    ]
    path.write_text("\n".join(lines) + "\n")
    return path


@pytest.mark.parametrize(
    ("clin_sig", "expected"),
    [
        # VEP writes lowercase, comma-separated terms from co-located ClinVar records.
        ("pathogenic", 1),
        ("pathogenic,likely_pathogenic", 1),
        ("likely_benign", 0),
        ("benign,likely_benign", 0),
        # ClinVar VCF CLNSIG spelling must keep working too.
        ("Pathogenic/Likely_pathogenic", 1),
        ("Benign", 0),
        ("uncertain_significance", None),
        ("pathogenic,benign", None),  # conflicting evidence is not a label
        ("-", None),
        ("", None),
        (np.nan, None),
    ],
)
def test_label(clin_sig: object, expected: int | None) -> None:
    assert label(clin_sig) == expected


def test_read_vep_tab_uses_the_hash_header_and_keeps_dashes(tmp_path: Path) -> None:
    tsv = write_vep_tab(
        tmp_path / "x.vep.tsv",
        [["22_1_A_G", "22:1", "G", "missense_variant", "MODERATE", "TBX1", "-", "pathogenic", "28", "0.9"]],
    )
    df = read_vep_tab(tsv)
    assert list(df.columns) == HEADER
    assert df.loc[0, "gnomADe_AF"] == "-"
    assert df.loc[0, "CLIN_SIG"] == "pathogenic"


def test_load_returns_raw_serving_columns_and_labels(tmp_path: Path) -> None:
    from rarelens_ml.features import RAW_COLUMNS
    from rarelens_ml.train import load

    tsv = write_vep_tab(
        tmp_path / "x.vep.tsv",
        [
            ["a", "22:1", "G", "missense_variant", "MODERATE", "TBX1", "0.0001", "pathogenic", "28", "0.9"],
            ["b", "22:2", "A", "synonymous_variant", "LOW", "CHEK2", "0.12", "benign", "3", "-"],
            ["c", "22:3", "T", "intron_variant", "MODIFIER", "CHEK2", "0.3", "uncertain_significance", "1", "-"],
        ],
    )
    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:
    """The registered model must take the raw columns serving sends and return P(pathogenic)."""
    import mlflow

    from rarelens_ml.train import fit, log_and_register

    rng = np.random.default_rng(0)
    n = 400
    impact = rng.choice(["HIGH", "MODERATE", "LOW", "MODIFIER"], n)
    y = pd.Series(((impact == "HIGH") | (rng.random(n) < 0.1)).astype(int))
    X = pd.DataFrame({
        "impact": impact,
        "consequence": rng.choice(["stop_gained", "missense_variant", "intron_variant"], n),
        "gnomad_af": rng.random(n).round(4).astype(str),   # strings, as read from the DB
        "cadd_phred": (rng.random(n) * 40).round(1).astype(str),
        "am_pathogenicity": "-",
    })

    mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db")
    mlflow.set_experiment("test")
    clf = fit(X, y)
    version = log_and_register(clf, model_name="rarelens-test", alias="production")

    model = mlflow.pyfunc.load_model("models:/rarelens-test@production")
    scores = np.asarray(model.predict(X.head(50)))
    assert version == "1"
    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