(protein, chi_id=None)
| 707 | |
| 708 | @torch.no_grad() |
| 709 | def get_chi_periodic_mask(protein, chi_id=None): |
| 710 | chi_mask = get_chi_mask(protein) |
| 711 | chi_1pi_periodic_mask = torch.tensor(chi_pi_periodic).to(protein.device)[protein.residue_type] # [num_residue, 4] |
| 712 | chi_2pi_periodic_mask = ~chi_1pi_periodic_mask # [num_residue, 4] |
| 713 | chi_1pi_periodic_mask = torch.logical_and(chi_mask, chi_1pi_periodic_mask) |
| 714 | chi_2pi_periodic_mask = torch.logical_and(chi_mask, chi_2pi_periodic_mask) |
| 715 | for mask in [chi_1pi_periodic_mask, chi_2pi_periodic_mask]: |
| 716 | if chi_id is not None: |
| 717 | mask[:, :chi_id] = False |
| 718 | mask[:, chi_id + 1:] = False |
| 719 | return chi_1pi_periodic_mask, chi_2pi_periodic_mask |
| 720 | |
| 721 | |
| 722 | @torch.no_grad() |
nothing calls this directly
no test coverage detected