| 973 | |
| 974 | |
| 975 | def _get_symm_atoms(pos_per_residue, residue_type): |
| 976 | sym_pos_per_residue = pos_per_residue.clone() # [num_residues, 37, 3] |
| 977 | for i in range(2): |
| 978 | atom_to = res_sym_atom_posn.to(residue_type.device)[residue_type][:, i, 0] # [num_residues] |
| 979 | atom_from = res_sym_atom_posn.to(residue_type.device)[residue_type][:, i, 1] # [num_residues] |
| 980 | sym_pos_per_residue[torch.arange(len(pos_per_residue)), atom_to] = \ |
| 981 | pos_per_residue[torch.arange(len(pos_per_residue)), atom_from] |
| 982 | sym_pos_per_residue[torch.arange(len(pos_per_residue)), atom_from] = \ |
| 983 | pos_per_residue[torch.arange(len(pos_per_residue)), atom_to] |
| 984 | return sym_pos_per_residue |