| |
| """Create a sequence-similarity-disjoint protein split using MMseqs2.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import shutil |
| import subprocess |
| from collections import Counter |
| from pathlib import Path |
|
|
| from mitointeract_recovery.splits import assign_grouped_splits, assert_group_disjoint |
|
|
|
|
| def read_rows(path: Path) -> list[dict]: |
| with path.open() as handle: |
| return [json.loads(line) for line in handle if line.strip()] |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--data-dir", type=Path, default=Path("artifacts/dev-10k")) |
| parser.add_argument("--min-seq-id", type=float, default=0.5) |
| parser.add_argument("--coverage", type=float, default=0.8) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--threads", type=int, default=8) |
| args = parser.parse_args() |
|
|
| if not shutil.which("mmseqs"): |
| raise RuntimeError("MMseqs2 is required; install the mmseqs2 package") |
| rows = read_rows(args.data_dir / "sample.jsonl") |
| audit_path = args.data_dir / "audit.json" |
| audit = json.loads(audit_path.read_text()) |
| split_ratios = audit.get( |
| "split_ratios", {"train": 0.8, "validation": 0.1, "test": 0.1} |
| ) |
| proteins = sorted({row["protein_id"]: row["sequence"] for row in rows}.items()) |
| suffix = ( |
| f"identity-{int(args.min_seq_id * 100)}-coverage-{int(args.coverage * 100)}" |
| ) |
| work = args.data_dir / f"mmseqs-{suffix}" |
| work.mkdir(parents=True, exist_ok=True) |
| fasta = work / "proteins.fasta" |
| with fasta.open("w") as handle: |
| for protein_id, sequence in proteins: |
| handle.write(f">{protein_id}\n{sequence}\n") |
|
|
| prefix = work / "clusters" |
| tmp = work / "tmp" |
| command = [ |
| "mmseqs", |
| "easy-cluster", |
| str(fasta), |
| str(prefix), |
| str(tmp), |
| "--min-seq-id", |
| str(args.min_seq_id), |
| "-c", |
| str(args.coverage), |
| "--cov-mode", |
| "0", |
| "--alignment-mode", |
| "3", |
| "--threads", |
| str(args.threads), |
| ] |
| cluster_tsv = prefix.with_name(prefix.name + "_cluster.tsv") |
| if not cluster_tsv.exists(): |
| subprocess.run(command, check=True) |
| protein_to_cluster: dict[str, str] = {} |
| with cluster_tsv.open() as handle: |
| for line in handle: |
| representative, member = line.rstrip("\n").split("\t") |
| protein_to_cluster[member] = representative |
| missing = {protein_id for protein_id, _ in proteins} - protein_to_cluster.keys() |
| if missing: |
| raise AssertionError(f"MMseqs2 omitted {len(missing)} proteins") |
|
|
| enriched = [ |
| dict(row, protein_similarity_cluster=protein_to_cluster[row["protein_id"]]) |
| for row in rows |
| ] |
| assignments = assign_grouped_splits( |
| enriched, |
| group_key="protein_similarity_cluster", |
| seed=args.seed, |
| ratios=split_ratios, |
| ) |
| assert_group_disjoint(enriched, assignments, group_key="protein_similarity_cluster") |
| manifest_name = f"split-protein_similarity_{int(args.min_seq_id * 100)}.jsonl" |
| manifest = args.data_dir / manifest_name |
| with manifest.open("w") as handle: |
| for row in rows: |
| handle.write( |
| json.dumps( |
| {"pair_id": row["pair_id"], "split": assignments[row["pair_id"]]} |
| ) |
| + "\n" |
| ) |
|
|
| cluster_sizes = Counter(protein_to_cluster.values()) |
| row_counts = Counter(assignments.values()) |
| report = { |
| "tool": "MMseqs2", |
| "command": command, |
| "min_sequence_identity": args.min_seq_id, |
| "minimum_alignment_coverage": args.coverage, |
| "coverage_mode": "both query and target", |
| "alignment_mode": "actual identical residues / aligned columns", |
| "unique_proteins": len(proteins), |
| "clusters": len(cluster_sizes), |
| "largest_cluster_proteins": max(cluster_sizes.values()), |
| "split_rows": dict(row_counts), |
| "manifest": manifest_name, |
| "protein_cluster_overlap": 0, |
| } |
| (args.data_dir / f"protein-clusters-{suffix}.json").write_text( |
| json.dumps(report, indent=2) + "\n" |
| ) |
| audit.setdefault("split_manifests", {})[ |
| f"protein_similarity_{int(args.min_seq_id * 100)}" |
| ] = manifest_name |
| audit.setdefault("external_split_audits", {})[ |
| f"protein_similarity_{int(args.min_seq_id * 100)}" |
| ] = report |
| audit_path.write_text(json.dumps(audit, indent=2) + "\n") |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
