filter_top_designs.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Filter top AlphaFold designs by minimum number of grafted epitopes.

Reads the af2_metrics_summary CSV generated by post_alphafold_analysis.py
and copies PDB files that have at least MIN_EPITOPES grafted into a
dedicated output folder.

Usage:
    python filter_top_designs.py \
        --summary  AF2_analysis/af2_metrics_summary_*.csv \
        --pdb_dir  AlphaFold_results/ \
        --output   top_designs/ \
        --min_epitopes 3
"""

import os
import sys
import csv
import glob
import shutil
import argparse
from datetime import datetime


def get_args():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "--summary", required=True,
        help="Path to af2_metrics_summary_*.csv (wildcards accepted)."
    )
    parser.add_argument(
        "--pdb_dir", required=True,
        help="Directory containing the AlphaFold PDB files."
    )
    parser.add_argument(
        "--output", default="top_designs/",
        help="Output directory for selected PDBs (default: top_designs/)."
    )
    parser.add_argument(
        "--min_epitopes", type=int, default=3,
        help="Minimum number of grafted epitopes required (default: 3)."
    )
    return parser.parse_args()


def load_summary(csv_path):
    """Loads the AF2 metrics summary CSV into a list of dicts."""
    rows = []
    with open(csv_path, newline="") as f:
        for row in csv.DictReader(f):
            rows.append(row)
    return rows


def count_epitopes(epitopes_found_str):
    """
    Counts the number of epitopes from the 'epitopes_found' column.
    Value is semicolon-separated, e.g. 'frag1;frag3;frag5' -> 3
    Empty string -> 0
    """
    if not epitopes_found_str or epitopes_found_str.strip() == "":
        return 0
    return len([e for e in epitopes_found_str.split(";") if e.strip()])


def main():
    args = get_args()

    # ── resolve summary CSV (support wildcards) ───────────────────
    csv_matches = sorted(glob.glob(args.summary))
    if not csv_matches:
        print(f"Error: No summary CSV found at: {args.summary}")
        sys.exit(1)
    csv_path = csv_matches[-1]   # use most recent if multiple
    print(f"Summary CSV  : {csv_path}")
    print(f"PDB dir      : {args.pdb_dir}")
    print(f"Min epitopes : {args.min_epitopes}")

    rows = load_summary(csv_path)
    print(f"Models loaded: {len(rows)}\n")

    # ── filter ────────────────────────────────────────────────────
    passing = []
    failing = []

    for row in rows:
        n_ep = count_epitopes(row.get("epitopes_found", ""))
        row["n_epitopes_found"] = n_ep
        if n_ep >= args.min_epitopes:
            passing.append(row)
        else:
            failing.append(row)

    # Sort passing by global_mean_plddt descending (best confidence first)
    passing.sort(key=lambda r: float(r.get("global_mean_plddt") or 0), reverse=True)

    print(f"{'='*55}")
    print(f"Passing (>= {args.min_epitopes} epitopes) : {len(passing)}")
    print(f"Failing                    : {len(failing)}")
    print(f"{'='*55}\n")

    if not passing:
        print("No models passed the filter. Exiting.")
        sys.exit(0)

    # ── copy PDBs ─────────────────────────────────────────────────
    os.makedirs(args.output, exist_ok=True)
    copied  = 0
    missing = 0

    print(f"{'Rank':<5} {'Model':<45} {'Epitopes':>8} {'pLDDT':>7}")
    print("-" * 70)

    for rank, row in enumerate(passing, start=1):
        model_name = row["model_name"]
        n_ep       = row["n_epitopes_found"]
        plddt      = float(row.get("global_mean_plddt") or 0)
        ep_names   = row.get("epitopes_found", "")

        # Try to find the PDB file
        pdb_src = os.path.join(args.pdb_dir, f"{model_name}.pdb")
        if not os.path.exists(pdb_src):
            # AlphaFold sometimes adds suffixes — try glob
            matches = glob.glob(os.path.join(args.pdb_dir, f"{model_name}*.pdb"))
            pdb_src = matches[0] if matches else None

        print(f"  #{rank:<3} {model_name:<43} {n_ep:>8}  {plddt:>6.1f}"
              f"  [{ep_names}]")

        if pdb_src and os.path.exists(pdb_src):
            dst = os.path.join(args.output, os.path.basename(pdb_src))
            shutil.copy2(pdb_src, dst)
            copied += 1
        else:
            print(f"       ⚠ PDB not found: {model_name}.pdb")
            missing += 1

    # ── save filtered CSV ─────────────────────────────────────────
    timestamp    = datetime.now().strftime("%Y%m%d_%H%M%S")
    filtered_csv = os.path.join(args.output, f"filtered_min{args.min_epitopes}ep_{timestamp}.csv")

    if rows:
        all_cols = list(rows[0].keys()) + ["n_epitopes_found"]
        # deduplicate while preserving order
        seen = set()
        cols = [c for c in all_cols if not (c in seen or seen.add(c))]

        with open(filtered_csv, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=cols + ["passed_filter"])
            writer.writeheader()
            for row in passing:
                row["passed_filter"] = "YES"
                writer.writerow({c: row.get(c, "") for c in cols + ["passed_filter"]})
            for row in failing:
                row["passed_filter"] = "NO"
                writer.writerow({c: row.get(c, "") for c in cols + ["passed_filter"]})

    print(f"\n{'='*55}")
    print(f"PDBs copied   : {copied}")
    print(f"PDBs missing  : {missing}")
    print(f"Output dir    : {args.output}")
    print(f"Filtered CSV  : {filtered_csv}")
    print(f"{'='*55}")


if __name__ == "__main__":
    main()