Script to make a video animation of adaptyv round 2 binders bound to EGFR

adaptyv_round_2_animation.py

"""
Uses https://github.com/hgbrian/biomodals
Makes a video like https://www.youtube.com/watch?v=DJkY-Tkg8RE
Top 100 PDBs already folded here: https://drive.google.com/drive/folders/10JoYaVgFyJ_y3iPB9-JyryjpJFpbObKo?usp=sharing
Adaptyv round 2 fasta from https://raw.githubusercontent.com/agitter/adaptyvbio-egfr/refs/heads/main/round2/screening/adaptyv-egfr-design-competition-round2-400-seqs-2024-11-07.fasta
"""
from pathlib import Path
from subprocess import run
import prody as pd

import sys
num = int(sys.argv[1])

EGFR = "LEEKKVCQGTSNKLTQLGTFEDHFLSLQRMFNNCEVVLGNLEITYVQRNYDLSFLKTIQEVAGYVLIALNTVERIPLENLQIIRGNMYYENSYALAVLSNYDANKTGLKELPMRNLQEILHGAVRFSNNPALCNVESIQWRDIVSSDFLSNMSMDFQNHLGSCQKCDPSCPNGSCWGAGEENCQKLTKIICAQQCSGRCRGKSPSDCCHNQCAAGCTGPRESDCLVCRKFRDEATCKDTCPPLMLYNPTTYQMDVNPEGKYSFGATCVKKCPRNYVVTDHGSCVRACGADSYEMEEDGVRKCKKCEGPCRKVCNGIGIGEFKDSLSINATNIKHFKNCTSISGDLHILPVAFRGDSFTHTPPLDPQELDILKTVKEITGFLLIQAWPENRTDLHAFENLEIIRGRTKQHGQFSLAVVSLNITSLGLRSLKEISDGDVIISGNKNLCYANTINWKKLFGTSGQKTKIISNRGENSCKATGQVCHALCSPEGCWGPEPRDCVSCRNVSRGRECVDKCNLLEGEPREFVENSECIQCHPECLPQAMNITCTGRGPDNCIQCAHYIDGPHCVKTCPAGVMGENNTLVWKYADAGHVCHLCHPNCTYGCTGPGLEGCPTNGPKIPS"

fasta = [(l.partition("\n")[0], l.partition("\n")[2].replace("\n",""))
         for l in open("adaptyv_round_2_results.faa").read().split(">") if l]

color_strs = ["0.8,0.8,0.6,0.8,0.6,0.8", "0.8,0.8,0.6,0.6,0.8,0.8", "0.8,0.8,0.6,0.85,0.7,0.5"]

for n, (seq_id, seq) in enumerate(fasta):
    if n != num:
        continue
    
    seq_id = f"{n:03d}_{seq_id.replace('.', '_').replace(' ','_').replace(',','_').replace('=','_')}"

    fasta_ = f"/tmp/{seq_id}.fasta"
    with open(fasta_, 'w') as out:
        out.write(f">{seq_id}\n{EGFR}:{seq}\n")

    if not list(Path("ad_anim").glob(f"{seq_id}*.pdb")):
        print(f"modal run modal_alphafold --input-fasta {fasta_} --out-dir ad_anim")
        run(f"modal run modal_alphafold --input-fasta {fasta_} --out-dir ad_anim", shell=True)
        run(f"cd ad_anim && unzip -n {seq_id}.result.zip", shell=True)

    pdb_path = str(next(Path("ad_anim").glob(f"{seq_id}*.pdb")))
    print(pdb_path)

    # Calculate transformation
    mobile = pd.parsePDB(pdb_path)
    target = pd.parsePDB('ad_anim/EGFR_binder_reference.pdb')
    mobile_chainA = mobile.select('chain A')
    target_chainA = target.select('chain A')
    transformation = pd.calcTransformation(mobile_chainA, target_chainA)
    transformed = pd.applyTransformation(transformation, mobile)

    # Write aligned structure
    pd.writePDB(f'/tmp/{seq_id}_aligned.pdb', transformed)

    angle = n * 15
    protein_color = color_strs[n%len(color_strs)]

    run(f"modal run modal_pdb2png.py --input-pdb /tmp/{seq_id}_aligned.pdb --out-dir aa_pngs"
        f" --protein-color {protein_color}"
        f" --width 1600 --height 1200 --protein-rotate '{angle}-{angle+15},0,90,20'", shell=True)

    rank = int(seq_id.partition("_")[0]) + 1 if n < 100 else ""
    sname = seq_id.partition("_")[2]

    for png_path in Path("aa_pngs").glob(f"{seq_id}_aligned_*[!text].png"):
        run(f'convert {png_path}'
            ' -gravity southwest -font Courier -pointsize 72 -fill white'
            f' -annotate +30+100 "rank: {rank}\n{sname}"'
            f' {Path(png_path).with_suffix("")}.text.png', shell=True)

run('ffmpeg -framerate 30 -pattern_type glob -i "aa_pngs/*text.png" -c:v libx264 -pix_fmt yuv420p output.mp4', shell=True)