recovery/scripts/prepare_sample.py

#!/usr/bin/env python3
"""Create a deterministic, deduplicated DTA development sample and split manifests."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import random
import statistics
from collections import Counter, defaultdict
from pathlib import Path

import pyarrow.parquet as pq

from mitointeract_recovery import micromolar_to_paffinity
from mitointeract_recovery.chemistry import canonicalize_smiles, scaffold_id, stable_id
from mitointeract_recovery.splits import (
    assign_grouped_splits,
    assign_random_splits,
    assert_group_disjoint,
)

COLUMNS = ["seq", "smiles", "smiles_can", "affinity_uM", "neg_log10_affinity_M"]


def stable_pair_id(sequence: str, smiles: str) -> str:
    return stable_id("pair", f"{sequence}\0{smiles}")


def sample_raw_rows(
    parquet_path: Path, candidate_count: int, seed: int
) -> tuple[list[dict], int]:
    parquet = pq.ParquetFile(parquet_path)
    total_rows = parquet.metadata.num_rows
    if candidate_count > total_rows:
        candidate_count = total_rows
    wanted = sorted(random.Random(seed).sample(range(total_rows), candidate_count))
    selected: list[dict] = []
    pointer = 0
    offset = 0

    for batch in parquet.iter_batches(batch_size=8192, columns=COLUMNS):
        batch_end = offset + batch.num_rows
        local_indices: list[int] = []
        while pointer < len(wanted) and wanted[pointer] < batch_end:
            local_indices.append(wanted[pointer] - offset)
            pointer += 1
        if local_indices:
            data = batch.to_pydict()
            for local in local_indices:
                selected.append({column: data[column][local] for column in COLUMNS})
        offset = batch_end
        if pointer == len(wanted):
            break

    if len(selected) != candidate_count:
        raise RuntimeError(
            f"requested {candidate_count} rows but selected {len(selected)}"
        )
    return selected, total_rows


def normalize_and_deduplicate(
    raw_rows: list[dict], sample_size: int, seed: int
) -> tuple[list[dict], Counter]:
    rejected = Counter()
    pairs: dict[str, dict] = {}
    targets: dict[str, list[float]] = defaultdict(list)

    for raw in raw_rows:
        sequence = "".join(str(raw.get("seq") or "").split()).upper()
        source_smiles = str(raw.get("smiles_can") or raw.get("smiles") or "").strip()
        if not sequence:
            rejected["empty_sequence"] += 1
            continue
        try:
            canonical = canonicalize_smiles(source_smiles)
        except ValueError:
            rejected["invalid_smiles"] += 1
            continue
        try:
            affinity_um = float(raw["affinity_uM"])
            published_paffinity = float(raw["neg_log10_affinity_M"])
            calculated_paffinity = micromolar_to_paffinity(affinity_um)
        except (TypeError, ValueError, OverflowError):
            rejected["invalid_affinity"] += 1
            continue
        if (
            not math.isfinite(published_paffinity)
            or abs(calculated_paffinity - published_paffinity) > 1e-4
        ):
            rejected["unit_mismatch"] += 1
            continue

        pair_id = stable_pair_id(sequence, canonical)
        if pair_id not in pairs:
            pairs[pair_id] = {
                "pair_id": pair_id,
                "protein_id": stable_id("protein", sequence),
                "ligand_id": stable_id("ligand", canonical),
                "scaffold_id": scaffold_id(canonical),
                "sequence": sequence,
                "smiles": canonical,
                "protein_length": len(sequence),
                "smiles_length": len(canonical),
            }
        targets[pair_id].append(published_paffinity)

    normalized: list[dict] = []
    for pair_id, row in pairs.items():
        values = targets[pair_id]
        normalized.append(
            {
                **row,
                "paffinity": statistics.median(values),
                "replicate_count": len(values),
                "replicate_paffinity_range": max(values) - min(values),
            }
        )

    normalized.sort(
        key=lambda row: hashlib.sha256(f"{seed}:{row['pair_id']}".encode()).hexdigest()
    )
    if len(normalized) < sample_size:
        raise RuntimeError(
            f"only {len(normalized)} valid unique pairs remained; increase --candidate-multiplier"
        )
    return normalized[:sample_size], rejected


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


def target_stats(rows: list[dict], assignments: dict[str, str]) -> dict:
    result = {}
    for split in ("train", "validation", "test"):
        values = [
            row["paffinity"] for row in rows if assignments[row["pair_id"]] == split
        ]
        result[split] = {
            "rows": len(values),
            "mean_paffinity": statistics.mean(values),
            "std_paffinity": statistics.stdev(values) if len(values) > 1 else 0.0,
            "min_paffinity": min(values),
            "max_paffinity": max(values),
        }
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--parquet", type=Path, required=True)
    parser.add_argument(
        "--dataset-revision",
        default="11e49b7ece33d62afd7f65bc05ce60ad37f9ba7b",
    )
    parser.add_argument("--sample-size", type=int, default=10_000)
    parser.add_argument("--candidate-multiplier", type=int, default=3)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/dev-10k"))
    args = parser.parse_args()
    if args.sample_size < 100:
        raise ValueError("sample size must be at least 100")

    raw_rows, source_rows = sample_raw_rows(
        args.parquet,
        args.sample_size * args.candidate_multiplier,
        args.seed,
    )
    rows, rejected = normalize_and_deduplicate(raw_rows, args.sample_size, args.seed)

    manifests = {
        "random_pair": assign_random_splits(rows, seed=args.seed),
        "cold_protein_exact": assign_grouped_splits(rows, "protein_id", seed=args.seed),
        "cold_ligand_scaffold": assign_grouped_splits(
            rows, "scaffold_id", seed=args.seed
        ),
    }
    assert_group_disjoint(rows, manifests["cold_protein_exact"], "protein_id")
    assert_group_disjoint(rows, manifests["cold_ligand_scaffold"], "scaffold_id")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    write_jsonl(args.output_dir / "sample.jsonl", rows)
    for name, assignments in manifests.items():
        write_jsonl(
            args.output_dir / f"split-{name}.jsonl",
            [
                {"pair_id": pair_id, "split": split}
                for pair_id, split in sorted(assignments.items())
            ],
        )

    with args.parquet.open("rb") as handle:
        source_sha256 = hashlib.file_digest(handle, "sha256").hexdigest()
    audit = {
        "source_parquet": str(args.parquet),
        "source_sha256": source_sha256,
        "dataset": "jglaser/binding_affinity",
        "dataset_revision": args.dataset_revision,
        "source_rows": source_rows,
        "candidate_rows": len(raw_rows),
        "sample_rows": len(rows),
        "seed": args.seed,
        "target": {"column": "neg_log10_affinity_M", "unit": "pAffinity"},
        "rejected": dict(rejected),
        "unique_proteins": len({row["protein_id"] for row in rows}),
        "unique_ligands": len({row["ligand_id"] for row in rows}),
        "unique_scaffold_groups": len({row["scaffold_id"] for row in rows}),
        "duplicate_measurements_in_sample": sum(
            row["replicate_count"] - 1 for row in rows
        ),
        "protein_length": {
            "min": min(row["protein_length"] for row in rows),
            "median": statistics.median(row["protein_length"] for row in rows),
            "max": max(row["protein_length"] for row in rows),
            "over_512": sum(row["protein_length"] > 512 for row in rows),
        },
        "splits": {
            name: target_stats(rows, assignments)
            for name, assignments in manifests.items()
        },
        "limitations": [
            "the source target can combine Ki, Kd, IC50, and EC50; pAffinity is not pKd",
            "assay type, relation operator, and source provenance are absent from the combined Parquet",
            "cold_protein_exact is entity-disjoint but not sequence-similarity-clustered",
            "scaffold groups use Bemis-Murcko with exact-molecule fallback for acyclic ligands",
            "this development sample is not an external biological benchmark",
        ],
    }
    (args.output_dir / "audit.json").write_text(json.dumps(audit, indent=2) + "\n")
    print(json.dumps(audit, indent=2))


if __name__ == "__main__":
    main()