def load_structure(fpath, chain=None): """ Args: fpath: filepath to either pdb or cif file chain: the chain id Returns: biotite.structure.AtomArray """ if fpath.endswith('cif'): with open(fpath) as fin: pdbxf = pdbx.PDBxFile.read(fin) structure = pdbx.get_structure(pdbxf, model=1) elif fpath.endswith('pdb'): with open(fpath) as fin: pdbf = pdb.PDBFile.read(fin) structure = pdb.get_structure(pdbf, model=1) issolvent = filter_solvent(structure) structure = structure[~issolvent] chains = get_chains(structure) print(f'Found {len(chains)} chains:', chains, '\n') if len(chains) == 0: raise ValueError('No chains found in the input file.') if chain is None: chain = chains[0] if chain not in chains: raise ValueError(f'Chain {chain} not found in input file') structure = structure[structure.chain_id == chain] print(f'Loaded chain {chain}\n') return structure
def test_chain_iter(array): n = 0 for chain in struc.get_chains(array): n += 1 assert isinstance(array, struc.AtomArray) assert n == 6
def test_get_chains(array): assert struc.get_chains(array).tolist() == ["A", "B", "C", "D", "E", "F"]