Removes all non-amino-acid residues, including alt conformations

clean.py

from Bio.PDB import *
from glob import glob
import warnings
warnings.filterwarnings("ignore")

amino_acids = ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLU', 'GLN', 'GLY', 
               'HIS', 'ILE', 'LEU', 'LYS', 'MET', 'PHE', 'PRO', 'SER', 
               'THR', 'TRP', 'TYR', 'VAL']
amino_acids += ["3NK", "0CU", "IGP", "6NT", "6VP", "3NY"]
#amino_acids_keep = ["A" + aa for aa in amino_acids]

all_aas = set(amino_acids)
altlocs = {" ", "A"}

class ResidueSelect(Select):
    def accept_residue(self, residue):
        return (residue.get_resname() in all_aas)
    def accept_atom(self, atom):
        #print(atom.get_full_id())
        return (atom.element != "H" and (atom.get_altloc() in altlocs))
    def accept_chain(self, chain):
        return chain.get_id() == "A"
        
res_select = ResidueSelect()
parser = PDBParser()

if __name__ == "__main__":
    pdbs = [input("Write name of pdb: ")] #glob("*.pdb")
    for pdb in pdbs: 
        struct = parser.get_structure(pdb[:-4], pdb)
        io = PDBIO()
        io.set_structure(struct)

        io.save(f'{pdb[:-4]}_clean.pdb', select=res_select)