Get the position of 37 atoms for each residue Args: protein: Protein object node_position: (num_atom, 3) return_nan: whether to return nan if the atom is not in the 37 atoms Returns: node_position37: (num_residue, 37, 3): nan if the atom is not in the 3
(protein, node_position=None, return_nan=True)
| 813 | |
| 814 | |
| 815 | def get_atom37_position(protein, node_position=None, return_nan=True): |
| 816 | """ |
| 817 | Get the position of 37 atoms for each residue |
| 818 | |
| 819 | Args: |
| 820 | protein: Protein object |
| 821 | node_position: (num_atom, 3) |
| 822 | return_nan: whether to return nan if the atom is not in the 37 atoms |
| 823 | |
| 824 | Returns: |
| 825 | node_position37: (num_residue, 37, 3): nan if the atom is not in the 37 atoms |
| 826 | """ |
| 827 | if node_position is None: |
| 828 | node_position = protein.node_position |
| 829 | |
| 830 | if return_nan: |
| 831 | node_position37 = torch.ones((protein.num_residue, 37, 3), dtype=torch.float, device=protein.device) * np.nan |
| 832 | else: |
| 833 | node_position37 = torch.zeros((protein.num_residue, 37, 3), dtype=torch.float, device=protein.device) |
| 834 | |
| 835 | node_position37[protein.atom2residue, protein.atom_name, :] = node_position |
| 836 | return node_position37 |
| 837 | |
| 838 | |
| 839 | @torch.no_grad() |
no outgoing calls
no test coverage detected