oncodsl / tests /test_h1.py
govindbalki's picture
Upload folder using huggingface_hub
0fff343 verified
Raw
History Blame Contribute Delete
2.43 kB
"""Smoke test for the H1 verification fixture (synthetic; no real-data dependency)."""
from __future__ import annotations
import numpy as np
import pandas as pd
from dsl import Cohort
from validate.h1 import IMMUNE_GENES, MMR_GENES, run_h1
def _synthetic_named_cohort(n_per_group: int = 60, seed: int = 0) -> Cohort:
"""A synthetic cohort where MSI-H tumours have low MMR and high immune expression."""
rng = np.random.default_rng(seed)
other_genes = [f"GENE_{i}" for i in range(20)]
feature_cols = MMR_GENES + IMMUNE_GENES + other_genes
rows = []
sample_ids = []
msi = []
stages = []
ages = []
for i in range(n_per_group):
row = rng.normal(loc=5.0, scale=0.5, size=len(feature_cols))
for g in MMR_GENES:
row[feature_cols.index(g)] = rng.normal(2.0, 0.3)
for g in IMMUNE_GENES:
row[feature_cols.index(g)] = rng.normal(8.0, 0.3)
rows.append(row)
sample_ids.append(f"H{i}")
msi.append("MSI-H")
stages.append(rng.choice(["I", "II", "III", "IV"]))
ages.append(float(rng.integers(40, 85)))
for i in range(n_per_group):
row = rng.normal(loc=5.0, scale=0.5, size=len(feature_cols))
for g in MMR_GENES:
row[feature_cols.index(g)] = rng.normal(8.0, 0.3)
for g in IMMUNE_GENES:
row[feature_cols.index(g)] = rng.normal(2.0, 0.3)
rows.append(row)
sample_ids.append(f"S{i}")
msi.append("MSS")
stages.append(rng.choice(["I", "II", "III", "IV"]))
ages.append(float(rng.integers(40, 85)))
idx = pd.Index(sample_ids, name="sample_id")
expr = pd.DataFrame(rows, index=idx, columns=feature_cols)
clinical = pd.DataFrame({"stage": stages, "age": ages}, index=idx)
labels = pd.DataFrame({"msi_status": msi, "tmb": [np.nan] * len(idx)}, index=idx)
return Cohort(expression=expr, clinical=clinical, labels=labels)
def test_h1_recovers_known_direction_on_synthetic_data():
cohort = _synthetic_named_cohort()
res = run_h1(cohort)
assert res.mmr_separates_correct_direction # MSI-H low on MMR
assert res.immune_separates_correct_direction # MSI-H high on immune
# With a clean synthetic signal Fit should crush it.
assert res.fit.auroc > 0.95
assert res.fit.balanced_acc > 0.9
# n_used is the full synthetic cohort (no missing stage/age).
assert res.n_used == 120