Function pdb_filepath_to_protein

flowdock/utils/data_utils.py:794–988

Takes a PDB filepath and constructs a FDProtein object. WARNING: All non-standard residue types will be converted into UNK. All non-standard atoms will be ignored. All water residues will be ignored. All hetero residues will be ignored if `filter_out_hetero_residues` is `True`.

(
    pdb_filepath: str,
    model_id: int = 0,
    atom_occupancy_min_threshold: int = 0.5,
    filter_out_hetero_residues: bool = True,
    allow_insertion_code: bool = True,
    accept_only_valid_backbone_residues: bool = True,
    chain_id: Optional[Union[List[str], str]] = None,
    bounding_box: Optional[np.ndarray] = None,
    res_start: Optional[int] = None,
    res_end: Optional[int] = None,
)

@beartype
def pdb_filepath_to_protein(
    pdb_filepath: str,
    model_id: int = 0,
    atom_occupancy_min_threshold: int = 0.5,
    filter_out_hetero_residues: bool = True,
    allow_insertion_code: bool = True,
    accept_only_valid_backbone_residues: bool = True,
    chain_id: Optional[Union[List[str], str]] = None,
    bounding_box: Optional[np.ndarray] = None,
    res_start: Optional[int] = None,
    res_end: Optional[int] = None,
) -> FDProtein:
    """Takes a PDB filepath and constructs a FDProtein object.

    WARNING: All non-standard residue types will be converted into UNK. All
      non-standard atoms will be ignored. All water residues will be ignored.
      All hetero residues will be ignored if `filter_out_hetero_residues` is `True`.
      All residues without valid positions for their N, Ca, C, and O atoms will be ignored.

    Adapted from: https://github.com/aqlaboratory/openfold and https://github.com/zrqiao/NeuralPLexer

    :param pdb_filepath: The filepath to the PDB file to parse.
    :param model_id: The model number to parse.
    :param atom_occupancy_min_threshold: The minimum occupancy threshold for atoms.
    :param filter_out_hetero_residues: If True, then hetero residues will be ignored.
    :param allow_insertion_code: If True, residues with insertion codes are parsed.
    :param accept_only_valid_backbone_residues: If True, only residues with valid N, Ca, C, and O atoms are parsed.
    :param chain_id: If `chain_id` is specified (e.g., `[A, B]`), then only those chains
        are parsed. Otherwise, all chains are parsed.
    :param bounding_box: If provided, only chains with backbone intersecting with
        the box are parsed.
    :param res_start: If provided, only residues with index >= res_start are parsed.
    :param res_end: If provided, only residues with index <= res_end are parsed.
    :return: A new `FDProtein` parsed from the PDB contents.
    """
    assert pdb_filepath.endswith(".pdb"), f"Invalid file extension: {pdb_filepath}"
    assert os.path.exists(pdb_filepath), f"File not found: {pdb_filepath}"

    with open(pdb_filepath) as pdb_fh:
        pdb_str = pdb_fh.read()

    pdb_fh = io.StringIO(pdb_str)
    parser = PDBParser(QUIET=True)
    structure = parser.get_structure("none", pdb_fh)
    models = list(structure.get_models())
    if not (0 <= model_id < len(models)):
        raise ValueError(f"Model ID {model_id} is out of range")
    model = models[model_id]
    if isinstance(chain_id, str):
        chain_id = [chain_id]

    seq = []
    atom_positions = []
    aatype = []
    atomtypes = []
    atom_mask = []
    residue_index = []
    chain_ids = []