recovery/scripts/run_baselines.py

#!/usr/bin/env python3
"""Benchmark leakage-auditing character n-gram ridge baselines."""

from __future__ import annotations

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

import numpy as np
from scipy.sparse import hstack
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import Ridge

from mitointeract_recovery.metrics import regression_metrics

ALPHAS = (0.1, 1.0, 10.0, 100.0)


def read_jsonl(path: Path) -> list[dict]:
    with path.open() as handle:
        return [json.loads(line) for line in handle if line.strip()]


def read_manifest(path: Path) -> dict[str, str]:
    return {row["pair_id"]: row["split"] for row in read_jsonl(path)}


def partition(rows: list[dict], manifest: dict[str, str]) -> dict[str, list[dict]]:
    result = {"train": [], "validation": [], "test": []}
    for row in rows:
        result[manifest[row["pair_id"]]].append(row)
    return result


def targets(rows: list[dict], target_key: str) -> np.ndarray:
    return np.asarray([row[target_key] for row in rows], dtype=np.float64)


def select_ridge(
    train_x,
    train_y: np.ndarray,
    validation_x,
    validation_y: np.ndarray,
) -> tuple[Ridge, float, list[dict]]:
    trials = []
    best = None
    for alpha in ALPHAS:
        model = Ridge(alpha=alpha, solver="lsqr", tol=1e-4)
        model.fit(train_x, train_y)
        predictions = model.predict(validation_x)
        metrics = regression_metrics(validation_y, predictions)
        trials.append({"alpha": alpha, "metrics": metrics})
        if best is None or metrics["rmse"] < best[0]:
            best = (metrics["rmse"], model, alpha)
    return best[1], best[2], trials


def evaluate_feature_set(
    name: str,
    train_x,
    validation_x,
    test_x,
    train_y: np.ndarray,
    validation_y: np.ndarray,
    test_y: np.ndarray,
) -> dict:
    started = time.monotonic()
    model, alpha, trials = select_ridge(train_x, train_y, validation_x, validation_y)
    return {
        "name": name,
        "selected_alpha": alpha,
        "validation_trials": trials,
        "validation": regression_metrics(validation_y, model.predict(validation_x)),
        "test": regression_metrics(test_y, model.predict(test_x)),
        "fit_and_eval_seconds": time.monotonic() - started,
    }


def benchmark_split(rows: list[dict], manifest_path: Path, target_key: str) -> dict:
    manifest = read_manifest(manifest_path)
    splits = partition(rows, manifest)
    train_y = targets(splits["train"], target_key)
    validation_y = targets(splits["validation"], target_key)
    test_y = targets(splits["test"], target_key)

    mean = float(train_y.mean())
    result = {
        "rows": {name: len(values) for name, values in splits.items()},
        "mean_baseline": {
            "prediction": mean,
            "validation": regression_metrics(
                validation_y, np.full_like(validation_y, mean)
            ),
            "test": regression_metrics(test_y, np.full_like(test_y, mean)),
        },
    }

    protein_vectorizer = TfidfVectorizer(
        analyzer="char",
        ngram_range=(3, 3),
        lowercase=False,
        min_df=2,
        max_features=4096,
        sublinear_tf=True,
        dtype=np.float32,
    )
    ligand_vectorizer = TfidfVectorizer(
        analyzer="char",
        ngram_range=(2, 5),
        lowercase=False,
        min_df=2,
        max_features=4096,
        sublinear_tf=True,
        dtype=np.float32,
    )

    protein_train = protein_vectorizer.fit_transform(
        [row["sequence"] for row in splits["train"]]
    )
    protein_validation = protein_vectorizer.transform(
        [row["sequence"] for row in splits["validation"]]
    )
    protein_test = protein_vectorizer.transform(
        [row["sequence"] for row in splits["test"]]
    )
    ligand_train = ligand_vectorizer.fit_transform(
        [row["smiles"] for row in splits["train"]]
    )
    ligand_validation = ligand_vectorizer.transform(
        [row["smiles"] for row in splits["validation"]]
    )
    ligand_test = ligand_vectorizer.transform([row["smiles"] for row in splits["test"]])

    result["feature_dimensions"] = {
        "protein": protein_train.shape[1],
        "ligand": ligand_train.shape[1],
    }
    result["models"] = [
        evaluate_feature_set(
            "protein_char3_ridge",
            protein_train,
            protein_validation,
            protein_test,
            train_y,
            validation_y,
            test_y,
        ),
        evaluate_feature_set(
            "ligand_char2_5_ridge",
            ligand_train,
            ligand_validation,
            ligand_test,
            train_y,
            validation_y,
            test_y,
        ),
        evaluate_feature_set(
            "combined_char_ridge",
            hstack([protein_train, ligand_train], format="csr"),
            hstack([protein_validation, ligand_validation], format="csr"),
            hstack([protein_test, ligand_test], format="csr"),
            train_y,
            validation_y,
            test_y,
        ),
    ]
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", type=Path, default=Path("artifacts/dev-10k"))
    parser.add_argument("--target-key", default="paffinity")
    parser.add_argument("--target-name", default="pAffinity")
    parser.add_argument(
        "--output", type=Path, default=Path("artifacts/dev-10k/baselines.json")
    )
    args = parser.parse_args()
    rows = read_jsonl(args.data_dir / "sample.jsonl")

    started = time.monotonic()
    report = {
        "sample_rows": len(rows),
        "target": args.target_name,
        "packages": {
            package: importlib.metadata.version(package)
            for package in ("numpy", "scipy", "scikit-learn")
        },
        "splits": {},
    }
    for manifest_path in sorted(args.data_dir.glob("split-*.jsonl")):
        split_name = manifest_path.stem.removeprefix("split-")
        report["splits"][split_name] = benchmark_split(
            rows, manifest_path, args.target_key
        )
    report["total_seconds"] = time.monotonic() - started

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()