Converts a `Protein` instance to a PDB string. Adapted from: https://github.com/jasonkyuyim/se3_diffusion :param prot: The protein to convert to PDB. :param model: The model number to use. :param add_end: Whether to add an `END` record. :param add_endmdl: Whether to add an `END
(prot: Union[OFProtein, FDProtein], model=1, add_end=True, add_endmdl=True)
| 61 | |
| 62 | @beartype |
| 63 | def to_pdb(prot: Union[OFProtein, FDProtein], model=1, add_end=True, add_endmdl=True) -> str: |
| 64 | """Converts a `Protein` instance to a PDB string. |
| 65 | |
| 66 | Adapted from: https://github.com/jasonkyuyim/se3_diffusion |
| 67 | |
| 68 | :param prot: The protein to convert to PDB. |
| 69 | :param model: The model number to use. |
| 70 | :param add_end: Whether to add an `END` record. |
| 71 | :param add_endmdl: Whether to add an `ENDMDL` record. |
| 72 | :return: PDB string. |
| 73 | """ |
| 74 | restypes = residue_constants.restypes + ["X"] |
| 75 | atom_types = residue_constants.atom_types |
| 76 | |
| 77 | pdb_lines = [] |
| 78 | |
| 79 | atom_mask = prot.atom_mask |
| 80 | aatype = prot.aatype |
| 81 | atom_positions = prot.atom_positions |
| 82 | residue_index = prot.residue_index.astype(int) |
| 83 | chain_index = prot.chain_index.astype(int) |
| 84 | b_factors = prot.b_factors |
| 85 | |
| 86 | if np.any(aatype > residue_constants.restype_num): |
| 87 | raise ValueError("Invalid aatypes.") |
| 88 | |
| 89 | # construct a mapping from chain integer indices to chain ID strings |
| 90 | chain_ids = {} |
| 91 | for i in np.unique(chain_index): # NOTE: `np.unique` gives sorted output |
| 92 | if i >= PDB_MAX_CHAINS: |
| 93 | raise ValueError(f"The PDB format supports at most {PDB_MAX_CHAINS} chains.") |
| 94 | chain_ids[i] = PDB_CHAIN_IDS[i] |
| 95 | |
| 96 | pdb_lines.append(f"MODEL {model}") |
| 97 | atom_index = 1 |
| 98 | last_chain_index = chain_index[0] |
| 99 | # add all atom sites |
| 100 | for i in range(aatype.shape[0]): |
| 101 | # close the previous chain if in a multichain PDB |
| 102 | if last_chain_index != chain_index[i]: |
| 103 | pdb_lines.append( |
| 104 | _chain_end( |
| 105 | atom_index, |
| 106 | res_1to3(restypes, aatype[i - 1]), |
| 107 | chain_ids[chain_index[i - 1]], |
| 108 | residue_index[i - 1], |
| 109 | ) |
| 110 | ) |
| 111 | last_chain_index = chain_index[i] |
| 112 | atom_index += 1 # NOTE: atom index increases at the `TER` symbol |
| 113 | |
| 114 | res_name_3 = res_1to3(restypes, aatype[i]) |
| 115 | for atom_name, pos, mask, b_factor in zip( |
| 116 | atom_types, atom_positions[i], atom_mask[i], b_factors[i] |
| 117 | ): |
| 118 | if mask < 0.5: |
| 119 | continue |
| 120 |
no test coverage detected