fixed residues ligand_mpnn using prody/mda

fixed_res_lmpnn.py

import numpy as np
from typing import List, Set, Tuple, Optional

def detect_binding_interface_mda(pdb_file: str, 
                                 protease_chain: str, 
                                 peptide_chain: str, 
                                 distance: float = 8.0,
                                 immutable_positions: Optional[str] = None) -> Tuple[List[int], List[int], str]:
    import MDAnalysis as mda
    
    immuts = set()
    if immutable_positions:
        immuts = set(map(int, immutable_positions.strip().split(',')))
    
    u = mda.Universe(pdb_file)
    ligand = u.select_atoms(f"chainID {peptide_chain} and backbone")
    nearby = u.select_atoms(f"(around {distance} group ligand) and not chainID {peptide_chain}", 
                           ligand=ligand)
    protein_bb = u.select_atoms(f"chainID {protease_chain} and backbone")
    binding_site = nearby & protein_bb
    
    if immutable_positions:
        site = [f"resnum {x}" for x in immutable_positions.strip().split(',')]
        exclude = u.select_atoms(f"chainID {protease_chain} and ({' or '.join(site)})")
        binding_site = binding_site - exclude
    
    design_ids_set = {int(atom.resid) for atom in binding_site}
    design_ids = sorted(design_ids_set)
    
    protein_ids_set = {int(atom.resid) for atom in protein_bb}
    fixed_ids_set = protein_ids_set - design_ids_set
    fixed_ids_set = fixed_ids_set | immuts
    fixed_ids = sorted(fixed_ids_set)
    
    ligand_ids = {int(atom.resid) for atom in ligand}
    fix_protein = design_string(fixed_ids, protease_chain)
    fix_ligand = design_string(ligand_ids, peptide_chain)
    fix_string = f"{fix_protein},{fix_ligand}"
    
    return design_ids, fixed_ids, fix_string

def detect_binding_interface_prody(pdb_file: str,
                                   protease_chain: str,
                                   peptide_chain: str,
                                   distance: float = 8.0,
                                   immutable_positions: Optional[str] = None) -> Tuple[List[int], List[int], str]:
    import prody as pd
    
    immuts = set()
    if immutable_positions:
        immuts = set(map(int, immutable_positions.strip().split(',')))
    
    structure = pd.parsePDB(pdb_file)
    peptide = structure.select(f"chain {peptide_chain} and backbone")
    protein = structure.select(f"chain {protease_chain}")
    
    contacts = pd.Contacts(protein)
    contacts.select(distance, peptide)
    
    binding_residues = set()
    for contact in contacts:
        res1, res2 = contact.getResindices()
        if contact.getAtoms()[0].getChid() == protease_chain:
            binding_residues.add(res1 + 1)
        if contact.getAtoms()[1].getChid() == protease_chain:
            binding_residues.add(res2 + 1)
    
    if immutable_positions:
        binding_residues = binding_residues - immuts
    
    design_ids = sorted(binding_residues)
    
    protein_residues = set()
    for residue in protein.iterResidues():
        protein_residues.add(residue.getResnum())
    
    fixed_ids_set = protein_residues - binding_residues
    fixed_ids_set = fixed_ids_set | immuts
    fixed_ids = sorted(fixed_ids_set)
    
    peptide_residues = set()
    for residue in peptide.iterResidues():
        peptide_residues.add(residue.getResnum())
    
    fix_protein = design_string(fixed_ids, protease_chain)
    fix_ligand = design_string(sorted(peptide_residues), peptide_chain)
    fix_string = f"{fix_protein},{fix_ligand}"
    
    return design_ids, fixed_ids, fix_string

def design_string(residue_ids: List[int], chain_id: str) -> str:
    if not residue_ids:
        return ""
    ranges = []
    start = residue_ids[0]
    end = start
    
    for i in range(1, len(residue_ids)):
        if residue_ids[i] == end + 1:
            end = residue_ids[i]
        else:
            if start == end:
                ranges.append(f"{chain_id}{start}")
            else:
                ranges.append(f"{chain_id}{start}-{end}")
            start = residue_ids[i]
            end = start
    
    if start == end:
        ranges.append(f"{chain_id}{start}")
    else:
        ranges.append(f"{chain_id}{start}-{end}")
    
    return ",".join(ranges)