Protein structure representation.
| 33 | |
| 34 | @dataclasses.dataclass(frozen=True) |
| 35 | class Protein: |
| 36 | """Protein structure representation.""" |
| 37 | |
| 38 | # Cartesian coordinates of atoms in angstroms. The atom types correspond to |
| 39 | # residue_constants.atom_types, i.e. the first three are N, CA, CB. |
| 40 | atom_positions: np.ndarray # [num_res, num_atom_type, 3] |
| 41 | |
| 42 | # Amino-acid type for each residue represented as an integer between 0 and |
| 43 | # 20, where 20 is 'X'. |
| 44 | aatype: np.ndarray # [num_res] |
| 45 | |
| 46 | # Binary float mask to indicate presence of a particular atom. 1.0 if an atom |
| 47 | # is present and 0.0 if not. This should be used for loss masking. |
| 48 | atom_mask: np.ndarray # [num_res, num_atom_type] |
| 49 | |
| 50 | # Residue index as used in PDB. It is not necessarily continuous or 0-indexed. |
| 51 | residue_index: np.ndarray # [num_res] |
| 52 | |
| 53 | # 0-indexed number corresponding to the chain in the protein that this residue |
| 54 | # belongs to. |
| 55 | chain_index: np.ndarray # [num_res] |
| 56 | |
| 57 | # B-factors, or temperature factors, of each residue (in sq. angstroms units), |
| 58 | # representing the displacement of the residue from its ground truth mean |
| 59 | # value. |
| 60 | b_factors: np.ndarray # [num_res, num_atom_type] |
| 61 | |
| 62 | def __post_init__(self): |
| 63 | if len(np.unique(self.chain_index)) > PDB_MAX_CHAINS: |
| 64 | raise ValueError( |
| 65 | f'Cannot build an instance with more than {PDB_MAX_CHAINS} chains ' |
| 66 | 'because these cannot be written to PDB format.') |
| 67 | |
| 68 | def to_dict(self): |
| 69 | return dataclasses.asdict(self) |
| 70 | |
| 71 | |
| 72 | def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein: |