scripts/run_h2.py

"""Run the H2 GP engine end-to-end and persist artefacts.

This is the biology-aware orchestrator: it loads the processed cohort,
restricts to the usable MSI-H vs MSS samples, anonymises the expression
matrix via the airgap, then hands a binary y plus an opaque-ID matrix to
the engine. The engine never sees the cohort labels' string values or any
gene symbols.

Usage:
    python -m scripts.run_h2 [--seed 42] [--population 150] [--generations 30]

Outputs:
    data/processed/h2/evolution_log.json   (anonymised, opaque IDs only)
    data/processed/h2/result.json          (anonymised, opaque IDs only)
"""

from __future__ import annotations

import argparse
import json
import sys

from airgap import anonymise
from data_pipeline import schema
from dsl import Load
from engine import run_gp_pipeline
from validate.h1 import POSITIVE_LABEL, usable_msi_cohort


OUT_DIR = schema.PROCESSED_DIR / "h2"
EVOLUTION_PATH = OUT_DIR / "evolution_log.json"
RESULT_PATH = OUT_DIR / "result.json"


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Run the H2 GP discovery pipeline.")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--population", type=int, default=150)
    ap.add_argument("--generations", type=int, default=30)
    ap.add_argument("--prefilter-n", type=int, default=2000)
    ap.add_argument("--permutations", type=int, default=200)
    args = ap.parse_args(argv)

    print("=" * 75)
    print("OncoDSL H2  —  blind discovery via genetic programming")
    print("=" * 75)
    print(f"Seed             : {args.seed}")
    print(f"Population size  : {args.population}")
    print(f"Generations      : {args.generations}")
    print(f"Prefilter top N  : {args.prefilter_n}")
    print(f"Permutations     : {args.permutations}")
    print()

    print("[1/3] Loading + anonymising cohort ...")
    cohort = usable_msi_cohort(Load("processed"))
    matrix = anonymise(cohort.expression)
    y = (cohort.labels["msi_status"] == POSITIVE_LABEL).astype(int).values
    print(f"      {len(cohort.sample_ids)} samples, "
          f"{matrix.shape[1]} anonymised feature IDs "
          f"(MSI-H {int(y.sum())} / MSS {int((1 - y).sum())}).")

    print()
    print("[2/3] Running prefilter + GP + baseline + permutation null ...")
    evolution_log, result = run_gp_pipeline(
        matrix, y,
        seed=args.seed,
        prefilter_n=args.prefilter_n,
        population_size=args.population,
        n_generations=args.generations,
        n_permutations=args.permutations,
    )

    print()
    print("[3/3] Writing artefacts ...")
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    EVOLUTION_PATH.write_text(json.dumps(evolution_log, indent=2))
    RESULT_PATH.write_text(json.dumps(result, indent=2))
    print(f"      wrote {EVOLUTION_PATH}")
    print(f"      wrote {RESULT_PATH}")

    win = result["winning"]
    base = result["baseline"]
    perm = result["permutation_summary"]
    print()
    print("-" * 75)
    print("WINNER (anonymised)")
    print("-" * 75)
    print(f"  program        : {win['program_repr']}")
    print(f"  feature sets   : {win['feature_sets']}")
    print(f"  n_genes        : {len(win['gene_ids'])}")
    print(f"  CV fitness     : {win['cv_fitness']:.4f}")
    print(f"  HELD-OUT AUROC : {win['holdout_auroc']:.4f}")
    print(f"  Permutation p  : {win['permutation_p']:.4f}")
    print()
    print(f"BASELINE  ({len(base['gene_ids'])} top-prefilter genes)  "
          f"AUROC: {base['holdout_auroc']:.4f}")
    print(f"NULL      mean AUROC: {perm['null_auroc_mean']:.4f}, "
          f"p95: {perm['null_auroc_p95']:.4f}")
    print("-" * 75)
    print(
        "Start the API to serve these artefacts:\n"
        "  uvicorn api.app:app --reload\n"
        "Then open the H2 tab in the Streamlit viewer."
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())