Get atom position for each chi torsion angles of each residue. Args: protein: Protein object. node_position: (num_atom, 3) tensor, atom position. Returns: chi_atom_position: (num_residue, 4, 4, 3) tensor, atom position for each chi torsion angles of each residu
(protein, node_position=None)
| 681 | |
| 682 | |
| 683 | def get_chi_atom_position(protein, node_position=None): |
| 684 | """ |
| 685 | Get atom position for each chi torsion angles of each residue. |
| 686 | |
| 687 | Args: |
| 688 | protein: Protein object. |
| 689 | node_position: (num_atom, 3) tensor, atom position. |
| 690 | |
| 691 | Returns: |
| 692 | chi_atom_position: (num_residue, 4, 4, 3) tensor, atom position for each chi torsion angles of each residue. |
| 693 | `Nan` indicates that the atom does not exist. |
| 694 | """ |
| 695 | if node_position is None: |
| 696 | node_position = protein.node_position |
| 697 | node_position37 = get_atom37_position(protein, node_position) |
| 698 | chi_atom37_index = chi_atom37_index_map.to(protein.device)[protein.residue_type] # (num_residue, 4, 4) 0~36 |
| 699 | chi_atom37_mask = chi_atom37_index == -1 |
| 700 | chi_atom37_index[chi_atom37_mask] = 0 |
| 701 | chi_atom37_index = chi_atom37_index.flatten(-2, -1) # (num_residue, 16) |
| 702 | chi_atom_position = torch.gather(node_position37, -2, |
| 703 | chi_atom37_index[:, :, None].expand(-1, -1, 3)) # (num_residue, 16, 3) |
| 704 | chi_atom_position = chi_atom_position.view(-1, 4, 4, 3) # (num_residue, 4, 4, 3) |
| 705 | return chi_atom_position |
| 706 | |
| 707 | |
| 708 | @torch.no_grad() |
no test coverage detected