(protein, chi_id)
| 772 | |
| 773 | @torch.no_grad() |
| 774 | def remove_by_chi(protein, chi_id): |
| 775 | new_protein = protein.clone() |
| 776 | mask_attrs = ['chi_1pi_periodic_mask', 'chi_2pi_periodic_mask', 'chi_mask'] |
| 777 | for attr in mask_attrs: |
| 778 | if hasattr(new_protein, attr): |
| 779 | getattr(new_protein, attr)[:, :chi_id] = 0 |
| 780 | getattr(new_protein, attr)[:, chi_id + 1:] = 0 |
| 781 | |
| 782 | if chi_id == 3: |
| 783 | return new_protein |
| 784 | else: |
| 785 | chi_atom14_index = chi_atom14_index_map.to(new_protein.device)[new_protein.residue_type] |
| 786 | atom_4 = chi_atom14_index[:, chi_id + 1, -1] |
| 787 | atom_mask = (new_protein.atom14index >= atom_4[new_protein.atom2residue]) & ( |
| 788 | atom_4[new_protein.atom2residue] != -1) |
| 789 | new_protein = new_protein.subgraph(~atom_mask) |
| 790 | return new_protein |
| 791 | |
| 792 | |
| 793 | def get_atom14_position(protein, node_position=None): |
nothing calls this directly
no outgoing calls
no test coverage detected