(protein, chi_id=None)
| 657 | |
| 658 | |
| 659 | def get_chi_mask(protein, chi_id=None): |
| 660 | chi_atom14_index = chi_atom14_index_map.to(protein.device)[protein.residue_type] # (num_residue, 4, 4) 0~13 |
| 661 | chi_atom14_mask = chi_atom14_index != -1 |
| 662 | chi_mask = chi_atom14_mask.all(dim=-1) # (num_residue, 4) |
| 663 | |
| 664 | chi_atom_position = get_chi_atom_position(protein) # (num_residue, 4, 4, 3) |
| 665 | has_atom_mask = ~torch.isnan(chi_atom_position).any(dim=-1).any(dim=-1) # (num_residue, 4) |
| 666 | chi_mask = chi_mask & has_atom_mask |
| 667 | |
| 668 | if chi_id is not None: |
| 669 | chi_mask[:, :chi_id] = False |
| 670 | chi_mask[:, chi_id + 1:] = False |
| 671 | return chi_mask |
| 672 | |
| 673 | |
| 674 | def get_chis(protein, node_position=None): |
no test coverage detected