#!/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()