keep_chains.py

from Bio import PDB

# input_pdb_path = '/Users/muhsenalzzaqry/Desktop/diffdockmm/1kee.pdb'
# output_pdb_path = '/Users/muhsenalzzaqry/Desktop/diffdockmm/1kee_e_f.pdb'
# # chain_names = ['Q','B','A']  # Replace with the list of chain names you want to keep
# chain_names = ['E','F'] 

# keep_chains(input_pdb_path, output_pdb_path, chain_names)


def keep_chains_cif(input_cif_path, output_cif_path, chain_names):
    parser = PDB.MMCIFParser(QUIET=True)
    io = PDB.MMCIFIO()
    
    # Parse the input CIF file
    structure = parser.get_structure('structure', input_cif_path)
    
    class ChainKeeper(PDB.Select):
        def __init__(self, chain_names):
            self.chain_names = set(chain_names)
        
        def accept_chain(self, chain):
            return chain.id in self.chain_names

    # Save the structure with only the specified chains
    io.set_structure(structure)
    io.save(output_cif_path, ChainKeeper(chain_names))