Set the position of 14 atoms for each residue Args: protein: Protein object node_position14: (num_residue, 14, 3) Returns: protein: Protein object
(protein, node_position14)
| 838 | |
| 839 | @torch.no_grad() |
| 840 | def set_atom14_position(protein, node_position14): |
| 841 | """ |
| 842 | Set the position of 14 atoms for each residue |
| 843 | |
| 844 | Args: |
| 845 | protein: Protein object |
| 846 | node_position14: (num_residue, 14, 3) |
| 847 | |
| 848 | Returns: |
| 849 | protein: Protein object |
| 850 | """ |
| 851 | atom14index = restype_atom14_index_map[ |
| 852 | protein.residue_type[protein.atom2residue], protein.atom_name |
| 853 | ] # (num_atom, ) |
| 854 | mask14 = atom14index != -1 |
| 855 | protein.node_position[mask14] = node_position14[protein.atom2residue[mask14], atom14index[mask14]] |
| 856 | return protein |
| 857 | |
| 858 | |
| 859 | @torch.no_grad() |