recovery/scripts/run_source_atlas_gate.py

#!/usr/bin/env python3
"""Run the CPU-only Source Atlas transfer gate without test evaluation."""

from __future__ import annotations

import argparse
import importlib.metadata
import json
import time
from pathlib import Path
from typing import Any

from mitointeract_recovery.chimera import read_manifest
from mitointeract_recovery.source_atlas import sha256_file
from mitointeract_recovery.source_atlas_gate import (
    atlas_input_hashes,
    build_entity_features,
    read_cluster_assignments,
    read_jsonl,
    run_source_atlas_gate,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--atlas-dir", type=Path, required=True)
    parser.add_argument("--clusters", type=Path, required=True)
    parser.add_argument("--sample", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--alpha", type=float, default=100.0)
    parser.add_argument("--outer-splits", type=int, default=5)
    parser.add_argument("--bootstrap-iterations", type=int, default=2000)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--minimum-rmse-improvement", type=float, default=0.05)
    parser.add_argument("--catastrophic-fold-tolerance", type=float, default=0.10)
    return parser.parse_args()


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    with path.open("w") as handle:
        for row in rows:
            handle.write(json.dumps(row, sort_keys=True) + "\n")


def main() -> None:
    args = parse_args()
    if args.output_dir.exists() and any(args.output_dir.iterdir()):
        raise FileExistsError(f"output directory is not empty: {args.output_dir}")
    for path in (
        args.atlas_dir / "records.jsonl",
        args.atlas_dir / "proteins.jsonl",
        args.atlas_dir / "ligands.jsonl",
        args.clusters,
        args.sample,
        args.manifest,
    ):
        if not path.is_file():
            raise FileNotFoundError(path)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    started = time.monotonic()
    atlas_records = read_jsonl(args.atlas_dir / "records.jsonl")
    proteins = read_jsonl(args.atlas_dir / "proteins.jsonl")
    ligands = read_jsonl(args.atlas_dir / "ligands.jsonl")
    primary_rows = read_jsonl(args.sample)
    manifest = read_manifest(args.manifest)
    if set(manifest) != {row["pair_id"] for row in primary_rows}:
        raise ValueError("primary manifest pair set must exactly match the sample")
    clusters = read_cluster_assignments(args.clusters)
    protein_vectors, ligand_vectors = build_entity_features(proteins, ligands)
    result = run_source_atlas_gate(
        atlas_records=atlas_records,
        primary_rows=primary_rows,
        primary_manifest=manifest,
        cluster_assignments=clusters,
        protein_vectors=protein_vectors,
        ligand_vectors=ligand_vectors,
        alpha=args.alpha,
        outer_splits=args.outer_splits,
        bootstrap_iterations=args.bootstrap_iterations,
        seed=args.seed,
        minimum_rmse_improvement=args.minimum_rmse_improvement,
        catastrophic_fold_tolerance=args.catastrophic_fold_tolerance,
    )

    prediction_data = result.pop("predictions")
    prediction_rows = []
    for index, row in enumerate(prediction_data["rows"]):
        prediction_rows.append(
            {
                "observation_id": row["observation_id"],
                "pair_id": row["pair_id"],
                "protein_id": row["protein_id"],
                "cluster_id": str(prediction_data["groups"][index]),
                "source_partition": manifest[row["pair_id"]],
                "target_pkd": float(prediction_data["targets"][index]),
                "outer_fold": int(prediction_data["fold_assignments"][index]),
                "kd_only_prediction": float(prediction_data["kd_only"][index]),
                "multitask_prediction": float(prediction_data["multitask"][index]),
            }
        )
    predictions_path = args.output_dir / "development-predictions.jsonl"
    write_jsonl(predictions_path, prediction_rows)

    report = {
        "status": "development_only_no_test_evaluation",
        "decision": (
            "advance_to_neural_pretraining"
            if result["gate"]["passed"]
            else "reject_neural_pretraining_not_justified"
        ),
        "protocol": {
            "description": (
                "fixed-alpha task-balanced shared-plus-task-specific Ridge transfer gate"
            ),
            "tasks": ["Kd", "Ki", "IC50", "EC50"],
            "features": {
                "protein": "log length plus 20 amino-acid fractions",
                "ligand": "512-bit radius-2 Morgan fingerprint plus eight descriptors",
                "architecture": "shared feature block plus one task-specific feature block",
            },
            "alpha": args.alpha,
            "outer_splits": args.outer_splits,
            "grouping": "joint MMseqs2 identity 0.5 coverage 0.8 clusters",
            "test_cluster_policy": (
                "all Atlas and primary development rows in clusters containing any "
                "benchmark test protein are excluded"
            ),
            "test_evaluations": 0,
            "seed": args.seed,
        },
        **result,
        "inputs": atlas_input_hashes(args.atlas_dir, args.clusters, args.sample, args.manifest),
        "packages": {
            package: importlib.metadata.version(package)
            for package in ("numpy", "rdkit", "scikit-learn", "scipy")
        },
        "artifacts": {
            "predictions": predictions_path.name,
            "predictions_sha256": sha256_file(predictions_path),
        },
        "elapsed_seconds": time.monotonic() - started,
    }
    report_path = args.output_dir / "report.json"
    report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
    print(
        json.dumps(
            {
                "decision": report["decision"],
                "development_rows": report["development_rows"],
                "safe_atlas_records": report["safe_atlas_records"],
                "kd_only_rmse": report["kd_only"]["metrics"]["rmse"],
                "multitask_rmse": report["multitask"]["metrics"]["rmse"],
                "rmse_improvement": report["gate"]["actual_rmse_improvement"],
                "bootstrap_ci": [
                    report["bootstrap"]["ci_2_5"],
                    report["bootstrap"]["ci_97_5"],
                ],
                "gate_passed": report["gate"]["passed"],
                "test_evaluations": report["test_evaluations"],
                "elapsed_seconds": report["elapsed_seconds"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()