Get the position of 14 atoms for each residue Args: protein: Protein object node_position: (num_atom, 3) Returns: node_position14: (num_residue, 14, 3) mask14: (num_atom,) indicate whether the atom is in the 14 atoms
(protein, node_position=None)
| 791 | |
| 792 | |
| 793 | def get_atom14_position(protein, node_position=None): |
| 794 | """ |
| 795 | Get the position of 14 atoms for each residue |
| 796 | Args: |
| 797 | protein: Protein object |
| 798 | node_position: (num_atom, 3) |
| 799 | |
| 800 | Returns: |
| 801 | node_position14: (num_residue, 14, 3) |
| 802 | mask14: (num_atom,) indicate whether the atom is in the 14 atoms |
| 803 | """ |
| 804 | if node_position is None: |
| 805 | node_position = protein.node_position |
| 806 | atom14index = restype_atom14_index_map.to(protein.device)[ |
| 807 | protein.residue_type[protein.atom2residue], protein.atom_name |
| 808 | ] # (num_atom, ) |
| 809 | node_position14 = torch.zeros((protein.num_residue, 14, 3), dtype=torch.float, device=protein.device) |
| 810 | mask14 = atom14index != -1 # (num_atom, ) |
| 811 | node_position14[protein.atom2residue[mask14], atom14index[mask14], :] = node_position[mask14] |
| 812 | return node_position14, mask14 |
| 813 | |
| 814 | |
| 815 | def get_atom37_position(protein, node_position=None, return_nan=True): |
no outgoing calls
no test coverage detected